Model training and inference for geospatial remote sensing and Earth Observation in JAX.
Jeo provides a framework for training machine learning models for geospatial remote sensing and Earth Observation using JAX and Flax. It leverages tf.data for efficient data pipelines, with a focus on TensorFlow Datasets for scalable and reproducible input. While TensorFlow Datasets are preferred, other dataset loaders are also supported. The codebase is designed to run seamlessly on CPUs, GPUs, or Google Cloud TPU VMs.
This project is open-sourced to share research code and facilitate collaboration in geospatial and sustainability model development. Jeo integrates effectively with the GeeFlow library to construct large-scale geospatial datasets using Google Earth Engine (GEE). An example workflow is outlined below.
Projects and publications that used this codebase:
- Light-weight geospatial model for global deforestation attribution, by Anton Raichuk, Michelle Sims, Radost Stanimirova, and Maxim Neumann. Presented at the NeurIPS 2024 Workshop on Tackling Climate Change with Machine Learning in Vancouver, BC, Canada, Dec 2024.
- Planted: a dataset for planted forest identification from multi-satellite time series, by Luis Miguel Pazos-Outón, Cristina Nader Vasconcelos, Anton Raichuk, Anurag Arnab, Dan Morris, and Maxim Neumann. In proceedings of IGARSS 2024, Athens, Greece, Jul 2024.
- Global drivers of forest loss at 1 km resolution, by Michelle Sims, Radost Stanimirova, Anton Raichuk, Maxim Neumann, Jessica Richter, Forrest Follett, James MacCarthy, Kristine Lister, Christopher Randle, Lindsey Sloat, Elena Esipova, Jaelah Jupiter, Charlotte Stanton, Dan Morris, Christy Slay, Drew Purves, and Nancy Harris, Environmental Research Letters, Jun 2025.
- Not every tree is a forest: benchmarking forest types from remote sensing, by Yuchang Jiang and Maxim Neumann. In proceedings of IGARSS 2025, Brisbane, Australia, Aug 2025.
- Uncertainty evaluation of segmentation models for Earth observation, by Mélanie Rey, Andriy Mnih, Maxim Neumann, Matt Overlan, Drew Purves. ArXiv, Oct 2025.
- ForestCast: Forecasting Deforestation Risk at Scale with Deep Learning, by Matthew Overlan, Charlotte Stanton, Maxim Neumann, Arianna Manzini, Julia Haas, Michelangelo Conserva, Melanie Rey, Kira Prabhu, Youngin Shin, Kuan Lu, Drew Purves, EarthArXiv, Nov 2025.
- Natural forests of the world: a 2020 baseline for deforestation and degradation monitoring, by Maxim Neumann, Anton Raichuk, Yuchang Jiang, Mélanie Rey, Radost Stanimirova, Michelle J. Sims, Sarah Carter, Elizabeth Goldman, Keith Anderson, Petra Poklukar, Katelyn Tarrio, Myroslava Lesiv, Steffen Fritz, Nicholas Clinton, Charlotte Stanton, Dan Morris, Drew Purves. Nature Scientific Data, Nov 2025.
The first step is to checkout JEO and install relevant python dependencies in a virtual environment:
git clone https://github.com/google-deepmind/jeo
# Go into the code directory for the examples below.
cd jeo/jeo
# Install and activate a virtual environment
python -m venv .venv
source .venv/bin/activate
# Install JEO.
pip install -e ..Launching a quick local model training for just a few steps on the CPU:
python -m jeo.train --config configs/tests/tiny_bit.py:runlocal \
--workdir /tmp/jeo/demo_tiny_bitThis will start to train a bit model (which is a modified convolutional
neureal net (CNN) model based on ResNet, see jeo/models/bit.py) on
CIFAR-10 TFDS dataset.
Since we run this on a local CPU just as a demo, we appended the :runlocal
config arg above, which specifies within the config to run just for a few
training and evaluation steps. For more configuration details, see the config
file jeo/configs/tests/tiny_bit.py.
In a standard workflow, the given workdir will be used to save checkpoints and
potentially other artifacts, such as final evaluation metrics.
Alex Wilson, Anton Raichuk, Burak Ekim, Cristina Nader Vasconcelos, Evelyn Camacho Soberon, Giovanni Ciampi, Keith Anderson, Linus Scheibenreif, Luis Miguel Pazos-Outón, Matt Overlan, Maxim Neumann, Mélanie Rey, Michelangelo Conserva, Petra Poklukar, Vishal Batchu, Yuchang Jiang
To cite this repository:
@software{jeo2025:github,
author = {Jeo team},
title = {{JEO}: Model training and inference for geospatial remote sensing and
{E}arth {O}bservation in {JAX}},
url = {https://github.com/google-deepmind/jeo},
year = {2025}
}
Copyright 2024 DeepMind Technologies Limited
This code is licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0.
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an AS IS BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.
This is not an official Google product.
