Metadata-Version: 2.4
Name: rslearn
Version: 0.0.21
Summary: A library for developing remote sensing datasets and models
Author: OlmoEarth Team
License:                                  Apache License
                                   Version 2.0, January 2004
                                http://www.apache.org/licenses/
        
           TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
        
           1. Definitions.
        
              "License" shall mean the terms and conditions for use, reproduction,
              and distribution as defined by Sections 1 through 9 of this document.
        
              "Licensor" shall mean the copyright owner or entity authorized by
              the copyright owner that is granting the License.
        
              "Legal Entity" shall mean the union of the acting entity and all
              other entities that control, are controlled by, or are under common
              control with that entity. For the purposes of this definition,
              "control" means (i) the power, direct or indirect, to cause the
              direction or management of such entity, whether by contract or
              otherwise, or (ii) ownership of fifty percent (50%) or more of the
              outstanding shares, or (iii) beneficial ownership of such entity.
        
              "You" (or "Your") shall mean an individual or Legal Entity
              exercising permissions granted by this License.
        
              "Source" form shall mean the preferred form for making modifications,
              including but not limited to software source code, documentation
              source, and configuration files.
        
              "Object" form shall mean any form resulting from mechanical
              transformation or translation of a Source form, including but
              not limited to compiled object code, generated documentation,
              and conversions to other media types.
        
              "Work" shall mean the work of authorship, whether in Source or
              Object form, made available under the License, as indicated by a
              copyright notice that is included in or attached to the work
              (an example is provided in the Appendix below).
        
              "Derivative Works" shall mean any work, whether in Source or Object
              form, that is based on (or derived from) the Work and for which the
              editorial revisions, annotations, elaborations, or other modifications
              represent, as a whole, an original work of authorship. For the purposes
              of this License, Derivative Works shall not include works that remain
              separable from, or merely link (or bind by name) to the interfaces of,
              the Work and Derivative Works thereof.
        
              "Contribution" shall mean any work of authorship, including
              the original version of the Work and any modifications or additions
              to that Work or Derivative Works thereof, that is intentionally
              submitted to Licensor for inclusion in the Work by the copyright owner
              or by an individual or Legal Entity authorized to submit on behalf of
              the copyright owner. For the purposes of this definition, "submitted"
              means any form of electronic, verbal, or written communication sent
              to the Licensor or its representatives, including but not limited to
              communication on electronic mailing lists, source code control systems,
              and issue tracking systems that are managed by, or on behalf of, the
              Licensor for the purpose of discussing and improving the Work, but
              excluding communication that is conspicuously marked or otherwise
              designated in writing by the copyright owner as "Not a Contribution."
        
              "Contributor" shall mean Licensor and any individual or Legal Entity
              on behalf of whom a Contribution has been received by Licensor and
              subsequently incorporated within the Work.
        
           2. Grant of Copyright License. Subject to the terms and conditions of
              this License, each Contributor hereby grants to You a perpetual,
              worldwide, non-exclusive, no-charge, royalty-free, irrevocable
              copyright license to reproduce, prepare Derivative Works of,
              publicly display, publicly perform, sublicense, and distribute the
              Work and such Derivative Works in Source or Object form.
        
           3. Grant of Patent License. Subject to the terms and conditions of
              this License, each Contributor hereby grants to You a perpetual,
              worldwide, non-exclusive, no-charge, royalty-free, irrevocable
              (except as stated in this section) patent license to make, have made,
              use, offer to sell, sell, import, and otherwise transfer the Work,
              where such license applies only to those patent claims licensable
              by such Contributor that are necessarily infringed by their
              Contribution(s) alone or by combination of their Contribution(s)
              with the Work to which such Contribution(s) was submitted. If You
              institute patent litigation against any entity (including a
              cross-claim or counterclaim in a lawsuit) alleging that the Work
              or a Contribution incorporated within the Work constitutes direct
              or contributory patent infringement, then any patent licenses
              granted to You under this License for that Work shall terminate
              as of the date such litigation is filed.
        
           4. Redistribution. You may reproduce and distribute copies of the
              Work or Derivative Works thereof in any medium, with or without
              modifications, and in Source or Object form, provided that You
              meet the following conditions:
        
              (a) You must give any other recipients of the Work or
                  Derivative Works a copy of this License; and
        
              (b) You must cause any modified files to carry prominent notices
                  stating that You changed the files; and
        
              (c) You must retain, in the Source form of any Derivative Works
                  that You distribute, all copyright, patent, trademark, and
                  attribution notices from the Source form of the Work,
                  excluding those notices that do not pertain to any part of
                  the Derivative Works; and
        
              (d) If the Work includes a "NOTICE" text file as part of its
                  distribution, then any Derivative Works that You distribute must
                  include a readable copy of the attribution notices contained
                  within such NOTICE file, excluding those notices that do not
                  pertain to any part of the Derivative Works, in at least one
                  of the following places: within a NOTICE text file distributed
                  as part of the Derivative Works; within the Source form or
                  documentation, if provided along with the Derivative Works; or,
                  within a display generated by the Derivative Works, if and
                  wherever such third-party notices normally appear. The contents
                  of the NOTICE file are for informational purposes only and
                  do not modify the License. You may add Your own attribution
                  notices within Derivative Works that You distribute, alongside
                  or as an addendum to the NOTICE text from the Work, provided
                  that such additional attribution notices cannot be construed
                  as modifying the License.
        
              You may add Your own copyright statement to Your modifications and
              may provide additional or different license terms and conditions
              for use, reproduction, or distribution of Your modifications, or
              for any such Derivative Works as a whole, provided Your use,
              reproduction, and distribution of the Work otherwise complies with
              the conditions stated in this License.
        
           5. Submission of Contributions. Unless You explicitly state otherwise,
              any Contribution intentionally submitted for inclusion in the Work
              by You to the Licensor shall be under the terms and conditions of
              this License, without any additional terms or conditions.
              Notwithstanding the above, nothing herein shall supersede or modify
              the terms of any separate license agreement you may have executed
              with Licensor regarding such Contributions.
        
           6. Trademarks. This License does not grant permission to use the trade
              names, trademarks, service marks, or product names of the Licensor,
              except as required for reasonable and customary use in describing the
              origin of the Work and reproducing the content of the NOTICE file.
        
           7. Disclaimer of Warranty. Unless required by applicable law or
              agreed to in writing, Licensor provides the Work (and each
              Contributor provides its Contributions) on an "AS IS" BASIS,
              WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
              implied, including, without limitation, any warranties or conditions
              of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
              PARTICULAR PURPOSE. You are solely responsible for determining the
              appropriateness of using or redistributing the Work and assume any
              risks associated with Your exercise of permissions under this License.
        
           8. Limitation of Liability. In no event and under no legal theory,
              whether in tort (including negligence), contract, or otherwise,
              unless required by applicable law (such as deliberate and grossly
              negligent acts) or agreed to in writing, shall any Contributor be
              liable to You for damages, including any direct, indirect, special,
              incidental, or consequential damages of any character arising as a
              result of this License or out of the use or inability to use the
              Work (including but not limited to damages for loss of goodwill,
              work stoppage, computer failure or malfunction, or any and all
              other commercial damages or losses), even if such Contributor
              has been advised of the possibility of such damages.
        
           9. Accepting Warranty or Additional Liability. While redistributing
              the Work or Derivative Works thereof, You may choose to offer,
              and charge a fee for, acceptance of support, warranty, indemnity,
              or other liability obligations and/or rights consistent with this
              License. However, in accepting such obligations, You may act only
              on Your own behalf and on Your sole responsibility, not on behalf
              of any other Contributor, and only if You agree to indemnify,
              defend, and hold each Contributor harmless for any liability
              incurred by, or claims asserted against, such Contributor by reason
              of your accepting any such warranty or additional liability.
        
           END OF TERMS AND CONDITIONS
        
           APPENDIX: How to apply the Apache License to your work.
        
              To apply the Apache License to your work, attach the following
              boilerplate notice, with the fields enclosed by brackets "[]"
              replaced with your own identifying information. (Don't include
              the brackets!)  The text should be enclosed in the appropriate
              comment syntax for the file format. We also recommend that a
              file or class name and description of purpose be included on the
              same "printed page" as the copyright notice for easier
              identification within third-party archives.
        
           Copyright 2024 Allen Institute for AI
        
           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.
        
Project-URL: homepage, https://github.com/allenai/rslearn
Project-URL: issues, https://github.com/allenai/rslearn/issues
Project-URL: repository, https://github.com/allenai/rslearn
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: boto3>=1.39
Requires-Dist: fiona>=1.10
Requires-Dist: fsspec>=2025.10.0
Requires-Dist: jsonargparse>=4.35.0
Requires-Dist: lightning>=2.5.1.post0
Requires-Dist: Pillow>=11.3
Requires-Dist: pyproj>=3.7
Requires-Dist: python-dateutil>=2.9
Requires-Dist: pytimeparse>=1.1
Requires-Dist: rasterio>=1.4
Requires-Dist: shapely>=2.1
Requires-Dist: torch>=2.7.0
Requires-Dist: torchvision>=0.22.0
Requires-Dist: tqdm>=4.67
Requires-Dist: universal_pathlib>=0.2.6
Provides-Extra: extra
Requires-Dist: accelerate>=1.10; extra == "extra"
Requires-Dist: cdsapi>=0.7.6; extra == "extra"
Requires-Dist: earthdaily[platform]>=1.0.7; extra == "extra"
Requires-Dist: earthengine-api>=1.6.3; extra == "extra"
Requires-Dist: einops>=0.8; extra == "extra"
Requires-Dist: fsspec[gcs,s3]; extra == "extra"
Requires-Dist: google-cloud-bigquery>=3.35; extra == "extra"
Requires-Dist: google-cloud-storage>=2.18; extra == "extra"
Requires-Dist: huggingface_hub>=0.34.4; extra == "extra"
Requires-Dist: netCDF4>=1.7.2; extra == "extra"
Requires-Dist: osmium>=4.0.2; extra == "extra"
Requires-Dist: planet>=3.1; extra == "extra"
Requires-Dist: planetary_computer>=1.0; extra == "extra"
Requires-Dist: pycocotools>=2.0; extra == "extra"
Requires-Dist: pystac_client>=0.9; extra == "extra"
Requires-Dist: rtree>=1.4; extra == "extra"
Requires-Dist: termcolor>=3.0; extra == "extra"
Requires-Dist: satlaspretrain_models>=0.3; extra == "extra"
Requires-Dist: scipy>=1.16; extra == "extra"
Requires-Dist: terratorch>=1.0.2; extra == "extra"
Requires-Dist: transformers>=4.55; extra == "extra"
Requires-Dist: wandb>=0.21; extra == "extra"
Requires-Dist: timm>=0.9.7; extra == "extra"
Provides-Extra: dev
Requires-Dist: interrogate>=1.7.0; extra == "dev"
Requires-Dist: mypy<2,>=1.17.1; extra == "dev"
Requires-Dist: pre-commit>=4.3.0; extra == "dev"
Requires-Dist: pytest>=8.0; extra == "dev"
Requires-Dist: pytest_httpserver; extra == "dev"
Requires-Dist: ruff>=0.12.9; extra == "dev"
Requires-Dist: pytest-dotenv; extra == "dev"
Requires-Dist: pytest-xdist; extra == "dev"
Dynamic: license-file

Overview
--------

rslearn is a library and tool for developing remote sensing datasets and models.

rslearn helps with:

1. Developing remote sensing datasets, starting with defining spatiotemporal windows
   (roughly equivalent to training examples) that should be annotated.
2. Importing raster and vector data from various online or local data sources into the
   dataset.
3. Annotating new categories of vector data (like points, polygons, and classification
   labels) using integrated web-based labeling apps.
4. Fine-tuning remote sensing foundation models on these datasets.
5. Applying models on new locations and times.


Quick links:
- [CoreConcepts](docs/CoreConcepts.md) summarizes key concepts in rslearn, including
  datasets, windows, layers, and data sources.
- [Examples](docs/Examples.md) contains more examples, including customizing different
  stages of rslearn with additional code.
- [DatasetConfig](docs/DatasetConfig.md) documents the dataset configuration file.
- [ModelConfig](docs/ModelConfig.md) documents the model configuration file.


Setup
-----

rslearn requires Python 3.10+ (Python 3.12 is recommended).

```
git clone https://github.com/allenai/rslearn.git
cd rslearn
pip install .[extra]
```


Supported Data Sources
----------------------

rslearn supports ingesting raster and vector data from the following data sources. Even
if you don't plan to train models within rslearn, you can still use it to easily
download, crop, and re-project data based on spatiotemporal rectangles (windows) that
you define. See [Examples](docs/Examples.md) and [DatasetConfig](docs/DatasetConfig.md)
for how to setup these data sources.

- Sentinel-1
- Sentinel-2 L1C and L2A
- Landsat 8/9 OLI-TIRS
- National Agriculture Imagery Program
- OpenStreetMap
- Xyz (Slippy) Tiles (e.g., Mapbox tiles)
- Planet Labs (PlanetScope, SkySat)
- ESA WorldCover 2021

rslearn can also be used to easily mosaic, crop, and re-project any sets of local
raster and vector files you may have.


Example Usage
-------------

This is an example of building a remote sensing dataset, and then training a model
on that dataset, using rslearn. Specifically, we will train a model that inputs
Sentinel-2 images and predicts land cover through a semantic segmentation task.

Let's start by defining a region of interest and obtaining Sentinel-2 images. Create a
directory `/path/to/dataset` and corresponding configuration file at
`/path/to/dataset/config.json` as follows:

```json
{
    "layers": {
        "sentinel2": {
            "type": "raster",
            "band_sets": [{
                "dtype": "uint8",
                "bands": ["R", "G", "B"]
            }],
            "data_source": {
                "class_path": "rslearn.data_sources.gcp_public_data.Sentinel2",
                "init_args": {
                    "index_cache_dir": "cache/sentinel2/",
                    "sort_by": "cloud_cover",
                    "use_rtree_index": false
                }
            }
        }
    }
}
```

Here, we have initialized an empty dataset and defined a raster layer called
`sentinel2`. Because it specifies a data source, it will be populated automatically. In
particular, the data will be sourced from a
[public Google Cloud Storage bucket containing Sentinel-2 imagery](https://cloud.google.com/storage/docs/public-datasets/sentinel-2).
The `sort_by` option sorts scenes in ascending order by cloud cover, so we will end up
choosing the scenes with minimal cloud cover.

Next, let's create our spatiotemporal windows. These will correspond to training
examples.

```
export DATASET_PATH=/path/to/dataset
rslearn dataset add_windows --root $DATASET_PATH --group default --utm --resolution 10 --grid_size 128 --src_crs EPSG:4326 --box=-122.6901,47.2079,-121.4955,47.9403 --start 2024-06-01T00:00:00+00:00 --end 2024-08-01T00:00:00+00:00 --name seattle
```

This creates windows along a 128x128 grid in the specified projection (i.e.,
appropriate UTM zone for the location with 10 m/pixel resolution) covering the
specified bounding box, which is centered at Seattle.

We can now obtain the Sentinel-2 images by running prepare, ingest, and materialize.

* Prepare: lookup items (in this case, Sentinel-2 scenes) in the data source that match with the spatiotemporal windows we created.
* Ingest: retrieve those items. This step populates the `tiles` directory within the dataset.
* Materialize: crop/mosaic the items to align with the windows. This populates the `layers` folder in each window directory.

```
rslearn dataset prepare --root $DATASET_PATH --workers 32 --batch-size 8
rslearn dataset ingest --root $DATASET_PATH --workers 32 --no-use-initial-job --jobs-per-process 1
rslearn dataset materialize --root $DATASET_PATH --workers 32 --no-use-initial-job
```

For ingestion, you may need to reduce the number of workers depending on the available
memory on your system.

You should now be able to open the GeoTIFF images. Let's find the window that
corresponds to downtown Seattle:

```python
import shapely
from rslearn.const import WGS84_PROJECTION
from rslearn.dataset import Dataset
from rslearn.utils import Projection, STGeometry
from upath import UPath

# Define longitude and latitude for downtown Seattle.
downtown_seattle = shapely.Point(-122.333, 47.606)

# Iterate over the windows and find the closest one.
dataset = Dataset(path=UPath("/path/to/dataset"))
best_window_name = None
best_distance = None
for window in dataset.load_windows(workers=32):
    shp = window.get_geometry().to_projection(WGS84_PROJECTION).shp
    distance = shp.distance(downtown_seattle)
    if best_distance is None or distance < best_distance:
        best_window_name = window.name
        best_distance = distance

print(best_window_name)
```

It should be `seattle_54912_-527360`, so let's open it in qgis (or your favorite GIS
software):

```
qgis $DATASET_PATH/windows/default/seattle_54912_-527360/layers/sentinel2/R_G_B/geotiff.tif
```


### Adding Land Cover Labels

Before we can train a land cover prediction model, we need labels. Here, we will use
the ESA WorldCover land cover map as labels.

Start by downloading the WorldCover data from https://worldcover2021.esa.int

```
wget https://worldcover2021.esa.int/data/archive/ESA_WorldCover_10m_2021_v200_60deg_macrotile_N30W180.zip
mkdir world_cover_tifs
unzip ESA_WorldCover_10m_2021_v200_60deg_macrotile_N30W180.zip -d world_cover_tifs/
```

It would require some work to write a script to re-project and crop these GeoTIFFs so
that they align with the windows we have previously defined (and the Sentinel-2 images
we have already ingested). We can use the LocalFiles data source to have rslearn
automate this process. Update the dataset `config.json` with a new layer:

```jsonc
"layers": {
    "sentinel2": {
        # ...
    },
    "worldcover": {
        "type": "raster",
        "band_sets": [{
            "dtype": "uint8",
            "bands": ["B1"]
        }],
        "resampling_method": "nearest",
        "data_source": {
            "class_path": "rslearn.data_sources.local_files.LocalFiles",
            "init_args": {
                "src_dir": "file:///path/to/world_cover_tifs/"
            }
        }
    }
},
# ...
```

Repeat the materialize process so we populate the data for this new layer:

```
rslearn dataset prepare --root $DATASET_PATH --workers 32 --batch-size 8
rslearn dataset ingest --root $DATASET_PATH --workers 32 --no-use-initial-job --jobs-per-process 1
rslearn dataset materialize --root $DATASET_PATH --workers 32 --no-use-initial-job
```

We can visualize both the GeoTIFFs together in qgis:

```
qgis $DATASET_PATH/windows/default/seattle_54912_-527360/layers/*/*/geotiff.tif
```


### Training a Model

Create a model configuration file `land_cover_model.yaml`:

```yaml
model:
  class_path: rslearn.train.lightning_module.RslearnLightningModule
  init_args:
    # This part defines the model architecture.
    # Essentially we apply the SatlasPretrain Sentinel-2 backbone with a UNet decoder
    # that terminates at a segmentation prediction head.
    # The backbone outputs four feature maps at different scales, and the UNet uses
    # these to compute a feature map at the input scale.
    # Finally the segmentation head applies per-pixel softmax to compute the land
    # cover class.
    model:
      class_path: rslearn.models.singletask.SingleTaskModel
      init_args:
        encoder:
          - class_path: rslearn.models.satlaspretrain.SatlasPretrain
            init_args:
              model_identifier: "Sentinel2_SwinB_SI_RGB"
        decoder:
          - class_path: rslearn.models.unet.UNetDecoder
            init_args:
              in_channels: [[4, 128], [8, 256], [16, 512], [32, 1024]]
              # We use 101 classes because the WorldCover classes are 10, 20, 30, 40
              # 50, 60, 70, 80, 90, 95, 100.
              # We could process the GeoTIFFs to collapse them to 0-10 (the 11 actual
              # classes) but the model will quickly learn that the intermediate
              # values are never used.
              out_channels: 101
              conv_layers_per_resolution: 2
          - class_path: rslearn.train.tasks.segmentation.SegmentationHead
    # Remaining parameters in RslearnLightningModule define different aspects of the
    # training process like initial learning rate.
    lr: 0.0001
data:
  class_path: rslearn.train.data_module.RslearnDataModule
  init_args:
    path: ${DATASET_PATH}
    # This defines the layers that should be read for each window.
    # The key ("image" / "targets") is what the data will be called in the model,
    # while the layers option specifies which layers will be read.
    inputs:
      image:
        data_type: "raster"
        layers: ["sentinel2"]
        bands: ["R", "G", "B"]
        passthrough: true
      targets:
        data_type: "raster"
        layers: ["worldcover"]
        bands: ["B1"]
        is_target: true
    task:
      # Train for semantic segmentation.
      # The remap option is only used when visualizing outputs during testing.
      class_path: rslearn.train.tasks.segmentation.SegmentationTask
      init_args:
        num_classes: 101
        remap_values: [[0, 1], [0, 255]]
    batch_size: 8
    num_workers: 32
    # These define different options for different phases/splits, like training,
    # validation, and testing.
    # Here we use the same transform across splits except training where we add a
    # flipping augmentation.
    # For now we are using the same windows for training and validation.
    default_config:
      transforms:
        - class_path: rslearn.train.transforms.normalize.Normalize
          init_args:
            mean: 0
            std: 255
    train_config:
      transforms:
        - class_path: rslearn.train.transforms.normalize.Normalize
          init_args:
            mean: 0
            std: 255
        - class_path: rslearn.train.transforms.flip.Flip
          init_args:
            image_selectors: ["image", "target/classes", "target/valid"]
      groups: ["default"]
    val_config:
      groups: ["default"]
    test_config:
      groups: ["default"]
    predict_config:
      groups: ["predict"]
      load_all_patches: true
      skip_targets: true
      patch_size: 512
trainer:
  max_epochs: 10
  callbacks:
    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
      init_args:
        save_top_k: 1
        save_last: true
        monitor: val_accuracy
        mode: max
        dirpath: ./land_cover_model_checkpoints/
```

Now we can train the model:

```
rslearn model fit --config land_cover_model.yaml
```


### Apply the Model

Let's apply the model on Portland, OR (you can change it to Portland, ME if you like).
We start by defining a new window around Portland. This time, instead of creating
windows along a grid, we just create one big window. This is because we are just going
to run the prediction over the whole window rather than use different windows as
different training examples.

```
rslearn dataset add_windows --root $DATASET_PATH --group predict --utm --resolution 10 --src_crs EPSG:4326 --box=-122.712,45.477,-122.621,45.549 --start 2024-06-01T00:00:00+00:00 --end 2024-08-01T00:00:00+00:00 --name portland
rslearn dataset prepare --root $DATASET_PATH --workers 32 --batch-size 8
rslearn dataset ingest --root $DATASET_PATH --workers 32 --no-use-initial-job --jobs-per-process 1
rslearn dataset materialize --root $DATASET_PATH --workers 32 --no-use-initial-job
```

We also need to add an RslearnPredictionWriter to the trainer callbacks in the model
configuration file, as it will handle writing the outputs from the model to a GeoTIFF.

```yaml
trainer:
  callbacks:
    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
      ...
    - class_path: rslearn.train.prediction_writer.RslearnWriter
      init_args:
        # We need to include this argument, but it will be overridden with the dataset
        # path from data.init_args.path.
        path: placeholder
        output_layer: output
```

Because of our `predict_config`, when we run `model predict` it will apply the model on
windows in the "predict" group, which is where we added the Portland window.

And it will be written in a new output_layer called "output". But we have to update the
dataset configuration so it specifies the layer:

```jsonc
"layers": {
    "sentinel2": {
        # ...
    },
    "worldcover": {
        # ...
    },
    "output": {
        "type": "raster",
        "band_sets": [{
            "dtype": "uint8",
            "bands": ["output"]
        }]
    }
},
```

Now we can apply the model:

```
# Find model checkpoint in lightning_logs dir.
ls lightning_logs/*/checkpoints/last.ckpt
rslearn model predict --config land_cover_model.yaml --ckpt_path land_cover_model_checkpoints/last.ckpt
```

And visualize the Sentinel-2 image and output in qgis:

```
qgis $DATASET_PATH/windows/predict/portland/layers/*/*/geotiff.tif
```


### Defining Train and Validation Splits

We can visualize the logged metrics using Tensorboard:

```
tensorboard --logdir=lightning_logs/
```

However, because our training and validation data are identical, the validation metrics
are not meaningful.

There are two suggested ways to split windows into different subsets:

1. Assign windows to different groups.
2. Use different key-value pairs in the windows' options dicts for different splits.

We will use the second approach. The script below sets a "split" key in the options
dict (which is stored in each window's `metadata.json` file) to "train" or "val"
based on the SHA-256 hash of the window name.

```python
import hashlib
import tqdm
from rslearn.dataset import Dataset, Window
from upath import UPath

ds_path = UPath("/path/to/dataset/")
dataset = Dataset(ds_path)
windows = dataset.load_windows(show_progress=True, workers=32)
for window in tqdm.tqdm(windows):
    if hashlib.sha256(window.name.encode()).hexdigest()[0] in ["0", "1"]:
        split = "val"
    else:
        split = "train"
    if "split" in window.options and window.options["split"] == split:
        continue
    window.options["split"] = split
    window.save()
```

Now we can update the model configuration file to use these splits:

```yaml
default_config:
  transforms:
    - class_path: rslearn.train.transforms.normalize.Normalize
      init_args:
        mean: 0
        std: 255
train_config:
  transforms:
    - class_path: rslearn.train.transforms.normalize.Normalize
      init_args:
        mean: 0
        std: 255
    - class_path: rslearn.train.transforms.flip.Flip
      init_args:
        image_selectors: ["image", "target/classes", "target/valid"]
  groups: ["default"]
  tags:
    split: train
val_config:
  groups: ["default"]
  tags:
    split: val
test_config:
  groups: ["default"]
  tags:
    split: val
predict_config:
  groups: ["predict"]
  load_all_patches: true
  skip_targets: true
  patch_size: 512
```

The `tags` option that we are adding here tells rslearn to only load windows with a
matching key and value in the window options.

Previously when we run `model fit`, it should show the same number of windows for
training and validation:

```
got 4752 examples in split train
got 4752 examples in split val
```

With the updates, it should show different numbers like this:

```
got 4167 examples in split train
got 585 examples in split val
```


### Visualizing with `model test`

We can visualize the ground truth labels and model predictions in the test set using
the `model test` command:

```
mkdir ./vis
rslearn model test --config land_cover_model.yaml --ckpt_path land_cover_model_checkpoints/last.ckpt --model.init_args.visualize_dir=./vis/
```

This will produce PNGs in the vis directory. The visualizations are produced by the
`Task.visualize` function, so we could customize the visualization by subclassing
SegmentationTask and overriding the visualize function.


### Checkpoint and Logging Management

Above, we needed to configure the checkpoint directory in the model config (the
`dirpath` option under `lightning.pytorch.callbacks.ModelCheckpoint`), and explicitly
specify the checkpoint path when applying the model. Additionally, metrics are logged
to the local filesystem and not well organized.

We can instead let rslearn automatically manage checkpoints, along with logging to
Weights & Biases. To do so, we add project_name, run_name, and management_dir options
to the model config. The project_name corresponds to the W&B project, and the run name
corresponds to the W&B name. The management_dir is a directory to store project data;
rslearn determines a per-project directory at `{management_dir}/{project_name}/{run_name}/`
and uses it to store checkpoints.

```yaml
model:
  # ...
data:
  # ...
trainer:
  # ...
project_name: land_cover_model
run_name: version_00
# This sets the option via the MANAGEMENT_DIR environment variable.
management_dir: ${MANAGEMENT_DIR}
```

Now, set the `MANAGEMENT_DIR` environment variable and run `model fit`:

```
export MANAGEMENT_DIR=./project_data
rslearn model fit --config land_cover_model.yaml
```

The training and validation loss and accuracy metric should now be logged to W&B. The
accuracy metric is provided by SegmentationTask, and additional metrics can be enabled
by passing the relevant init_args to the task, e.g. mean IoU and F1:

```yaml
      class_path: rslearn.train.tasks.segmentation.SegmentationTask
      init_args:
        num_classes: 101
        remap_values: [[0, 1], [0, 255]]
        enable_miou_metric: true
        enable_f1_metric: true
```

When calling `model test` and `model predict` with management_dir set, rslearn will
automatically load the best checkpoint from the project directory, or raise an error if
no existing checkpoint exists. This behavior can be overridden with the
`--load_checkpoint_mode` and `--load_checkpoint_required` options (see `--help` for
details). Logging will be enabled during fit but not test/predict, and this can also
be overridden, using `--log_mode`.


### Inputting Multiple Sentinel-2 Images

Currently our model inputs a single Sentinel-2 image. However, for most tasks where
labels are not expected to change from week to week, we find that accuracy can be
significantly improved by inputting multiple images, regardless of the pre-trained
model used. Multiple images makes the model more resilient to clouds and image
artifacts, and allows the model to synthesize information across different views that
may come from different seasons or weather conditions.

We first update our dataset configuration to obtain three images, by customizing the
query_config section. This can replace the sentinel2 layer:

```jsonc
"layers": {
    "sentinel2_multi": {
        "type": "raster",
        "band_sets": [{
            "dtype": "uint8",
            "bands": ["R", "G", "B"]
        }],
        "data_source": {
            "class_path": "rslearn.data_sources.gcp_public_data.Sentinel2",
            "init_args": {
              "index_cache_dir": "cache/sentinel2/",
              "sort_by": "cloud_cover",
              "use_rtree_index": false
            },
            "query_config": {
                "max_matches": 3
            }
        }
    },
    "worldcover": {
        # ...
    },
    "output": {
        # ...
    }
}
```

Repeat the steps from earlier to prepare, ingest, and materialize the dataset.

Now we update our model configuration file. First, we modify the model architecture to
be able to input an image time series. We use the SimpleTimeSeries model, which takes
an encoder that expects a single-image input, and applies that encoder on each image in
the time series. It then applies max temporal pooling to combine the per-image feature
maps extracted by the encoder.

Image time series in rslearn are currently stored as [T*C, H, W] tensors. So we pass
the `image_channels` to SimpleTimeSeries so it knows how to slice up the tensor to
recover the per-timestep images.

```yaml
model:
  class_path: rslearn.train.lightning_module.RslearnLightningModule
  init_args:
    model:
      class_path: rslearn.models.singletask.SingleTaskModel
      init_args:
        encoder:
          - class_path: rslearn.models.simple_time_series.SimpleTimeSeries
            init_args:
              encoder:
                class_path: rslearn.models.satlaspretrain.SatlasPretrain
                init_args:
                  model_identifier: "Sentinel2_SwinB_SI_RGB"
              image_channels: 3
        decoder:
          # ...
```

Next, we update the data module section so that the dataset loads the image time series
rather than a single image. The `load_all_layers` option tells the dataset to stack the
rasters from all of the layers specified, and also to ignore windows where any of those
layers are missing.

```yaml
data:
  class_path: rslearn.train.data_module.RslearnDataModule
  init_args:
    path: # ...
    inputs:
      image:
        data_type: "raster"
        layers: ["sentinel2_multi", "sentinel2_multi.1", "sentinel2_multi.2"]
        bands: ["R", "G", "B"]
        passthrough: true
        load_all_layers: true
      targets:
        # ...
```

Now we can train an updated model:

```
rslearn model fit --config land_cover_model.yaml
```


Contact
-------

For questions and suggestions, please open an issue on GitHub.
