Merge branch 'main' into main

docs/api/datasets.rst

@@ -404,6 +404,11 @@ Rwanda Field Boundary
 
 .. autoclass:: RwandaFieldBoundary
 
+SatlasPretrain
+^^^^^^^^^^^^^^
+
+.. autoclass:: SatlasPretrain
+
 Seasonal Contrast
 ^^^^^^^^^^^^^^^^^
 

docs/api/datasets/non_geo_datasets.csv

@@ -11,7 +11,7 @@ Dataset,Task,Source,License,# Samples,# Classes,Size (px),Resolution (m),Bands
 `Kenya Crop Type`_,S,Sentinel-2,"CC-BY-SA-4.0","4,688",7,"3,035x2,016",10,MSI
 `DeepGlobe Land Cover`_,S,DigitalGlobe +Vivid,-,803,7,"2,448x2,448",0.5,RGB
 `DFC2022`_,S,Aerial,"CC-BY-4.0","3,981",15,"2,000x2,000",0.5,RGB
-`Digital Typhoon`_,"C, R",Himawari,"CC-BY-4.0","189,364",8,512,5000,Infrared 
+`Digital Typhoon`_,"C, R",Himawari,"CC-BY-4.0","189,364",8,512,5000,Infrared
 `ETCI2021 Flood Detection`_,S,Sentinel-1,-,"66,810",2,256x256,5--20,SAR
 `EuroSAT`_,C,Sentinel-2,"MIT","27,000",10,64x64,10,MSI
 `FAIR1M`_,OD,Gaofen/Google Earth,"CC-BY-NC-SA-3.0","15,000",37,"1,024x1,024",0.3--0.8,RGB
@@ -38,6 +38,7 @@ Dataset,Task,Source,License,# Samples,# Classes,Size (px),Resolution (m),Bands
 `ReforesTree`_,"OD, R",Aerial,"CC-BY-4.0",100,6,"4,000x4,000",0.02,RGB
 `RESISC45`_,C,Google Earth,-,"31,500",45,256x256,0.2--30,RGB
 `Rwanda Field Boundary`_,S,Planetscope,"NICFI AND CC-BY-4.0",70,2,256x256,4.7,RGB + NIR
+`SatlasPretrain`_,"C, R, S, I, OD","NAIP, Landsat, Sentinel",ESA AND CC0-1.0 AND ODbL-1.0 AND CC-BY-4.0,302M,137,512,0.6--30,"SAR, MSI"
 `Seasonal Contrast`_,T,Sentinel-2,"CC-BY-4.0",100K--1M,-,264x264,10,MSI
 `SeasoNet`_,S,Sentinel-2,"CC-BY-4.0","1,759,830",33,120x120,10,MSI
 `SEN12MS`_,S,"Sentinel-1/2, MODIS","CC-BY-4.0","180,662",33,256x256,10,"SAR, MSI"

experiments/torchgeo/benchmark.py

@@ -217,7 +217,7 @@ def main(args: argparse.Namespace) -> None:
 
     criterion = nn.CrossEntropyLoss()
     params = model.parameters()
-    optimizer = optim.SGD(params, lr=0.0001)  # type: ignore[attr-defined]
+    optimizer = optim.SGD(params, lr=0.0001)
 
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu', args.device)
     model = model.to(device)

requirements/required.txt

@@ -1,5 +1,5 @@
 # setup
-setuptools==75.1.0
+setuptools==75.2.0
 
 # install
 einops==0.8.0
@@ -17,6 +17,6 @@ rtree==1.3.0
 segmentation-models-pytorch==0.3.4
 shapely==2.0.6
 timm==0.9.7
-torch==2.4.1
+torch==2.5.0
 torchmetrics==1.4.3
-torchvision==0.19.1
+torchvision==0.20.0

requirements/style.txt

@@ -1,3 +1,3 @@
 # style
 mypy==1.12.0
-ruff==0.6.9
+ruff==0.7.0

tests/data/satlas/data.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import json
+import os
+import shutil
+
+from PIL import Image
+
+SIZE = 32
+landsat_size = {
+    'b1': SIZE // 2,
+    'b2': SIZE // 2,
+    'b3': SIZE // 2,
+    'b4': SIZE // 2,
+    'b5': SIZE // 2,
+    'b6': SIZE // 2,
+    'b7': SIZE // 2,
+    'b8': SIZE,
+    'b9': SIZE // 2,
+    'b10': SIZE // 2,
+    'b11': SIZE // 4,
+    'b12': SIZE // 4,
+}
+
+index = [[7149, 3246], [1234, 5678]]
+good_images = [
+    [7149, 3246, '2022-03'],
+    [1234, 5678, '2022-03'],
+    [7149, 3246, 'm_3808245_se_17_1_20110801'],
+    [1234, 5678, 'm_3808245_se_17_1_20110801'],
+    [7149, 3246, '2022-01'],
+    [1234, 5678, '2022-01'],
+    [7149, 3246, 'S2A_MSIL1C_20220309T032601_N0400_R018_T48RYR_20220309T060235'],
+    [1234, 5678, 'S2A_MSIL1C_20220309T032601_N0400_R018_T48RYR_20220309T060235'],
+]
+times = {
+    '2022-03': '2022-03-01T00:00:00+00:00',
+    'm_3808245_se_17_1_20110801': '2011-08-01T12:00:00+00:00',
+    '2022-01': '2022-01-01T00:00:00+00:00',
+    'S2A_MSIL1C_20220309T032601_N0400_R018_T48RYR_20220309T060235': '2022-03-09T06:02:35+00:00',
+}
+
+FILENAME_HIERARCHY = dict[str, 'FILENAME_HIERARCHY'] | list[str]
+filenames: FILENAME_HIERARCHY = {
+    'landsat': {'2022-03': list(f'b{i}' for i in range(1, 12))},
+    'naip': {'m_3808245_se_17_1_20110801': ['tci', 'ir']},
+    'sentinel1': {'2022-01': ['vh', 'vv']},
+    'sentinel2': {
+        'S2A_MSIL1C_20220309T032601_N0400_R018_T48RYR_20220309T060235': [
+            'tci',
+            'b05',
+            'b06',
+            'b07',
+            'b08',
+            'b11',
+            'b12',
+        ]
+    },
+}
+
+
+def create_files(path: str) -> None:
+    os.makedirs(path, exist_ok=True)
+    for col, row in index:
+        band = os.path.basename(path)
+        mode = 'RGB' if band == 'tci' else 'L'
+        size = SIZE
+        if 'landsat' in path:
+            size = landsat_size[band]
+        img = Image.new(mode, (size, size))
+        img.save(os.path.join(path, f'{col}_{row}.png'))
+
+
+def create_directory(directory: str, hierarchy: FILENAME_HIERARCHY) -> None:
+    if isinstance(hierarchy, dict):
+        # Recursive case
+        for key, value in hierarchy.items():
+            path = os.path.join(directory, key)
+            create_directory(path, value)
+    else:
+        # Base case
+        for value in hierarchy:
+            path = os.path.join(directory, value)
+            create_files(path)
+
+
+if __name__ == '__main__':
+    create_directory('.', filenames)
+
+    col, row = index[0]
+    path = os.path.join('static', f'{col}_{row}')
+    os.makedirs(path, exist_ok=True)
+    img = Image.new('L', (SIZE, SIZE))
+    img.save(os.path.join(path, 'land_cover.png'))
+
+    os.makedirs('metadata', exist_ok=True)
+    with open(os.path.join('metadata', 'train_lowres.json'), 'w') as f:
+        json.dump(index, f)
+
+    with open(os.path.join('metadata', 'good_images_lowres_all.json'), 'w') as f:
+        json.dump(good_images, f)
+
+    with open(os.path.join('metadata', 'image_times.json'), 'w') as f:
+        json.dump(times, f)
+
+    for path in os.listdir('.'):
+        if os.path.isdir(path):
+            shutil.make_archive(path, 'tar', '.', path)

tests/data/satlas/metadata/good_images_lowres_all.json

@@ -0,0 +1 @@
+[[7149, 3246, "2022-03"], [1234, 5678, "2022-03"], [7149, 3246, "m_3808245_se_17_1_20110801"], [1234, 5678, "m_3808245_se_17_1_20110801"], [7149, 3246, "2022-01"], [1234, 5678, "2022-01"], [7149, 3246, "S2A_MSIL1C_20220309T032601_N0400_R018_T48RYR_20220309T060235"], [1234, 5678, "S2A_MSIL1C_20220309T032601_N0400_R018_T48RYR_20220309T060235"]]

tests/data/satlas/metadata/image_times.json

@@ -0,0 +1 @@
+{"2022-03": "2022-03-01T00:00:00+00:00", "m_3808245_se_17_1_20110801": "2011-08-01T12:00:00+00:00", "2022-01": "2022-01-01T00:00:00+00:00", "S2A_MSIL1C_20220309T032601_N0400_R018_T48RYR_20220309T060235": "2022-03-09T06:02:35+00:00"}

tests/data/satlas/metadata/train_lowres.json

@@ -0,0 +1 @@
+[[7149, 3246], [1234, 5678]]

tests/datasets/test_satlas.py

@@ -0,0 +1,59 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+import shutil
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import pytest
+import torch.nn as nn
+from pytest import MonkeyPatch
+from torch import Tensor
+
+from torchgeo.datasets import DatasetNotFoundError, SatlasPretrain
+from torchgeo.datasets.utils import Executable
+
+
+class TestSatlasPretrain:
+    @pytest.fixture
+    def dataset(
+        self, aws: Executable, monkeypatch: MonkeyPatch, tmp_path: Path
+    ) -> SatlasPretrain:
+        url = os.path.join('tests', 'data', 'satlas', '')
+        monkeypatch.setattr(SatlasPretrain, 'url', url)
+        images = ('landsat', 'naip', 'sentinel1', 'sentinel2')
+        products = (*images, 'static', 'metadata')
+        tarballs = {product: (f'{product}.tar',) for product in products}
+        monkeypatch.setattr(SatlasPretrain, 'tarballs', tarballs)
+        transforms = nn.Identity()
+        return SatlasPretrain(
+            tmp_path, images=images, transforms=transforms, download=True
+        )
+
+    @pytest.mark.parametrize('index', [0, 1])
+    def test_getitem(self, dataset: SatlasPretrain, index: int) -> None:
+        x = dataset[index]
+        assert isinstance(x, dict)
+        for image in dataset.images:
+            assert isinstance(x[f'image_{image}'], Tensor)
+            assert isinstance(x[f'time_{image}'], Tensor)
+        for label in dataset.labels:
+            assert isinstance(x[f'mask_{label}'], Tensor)
+
+    def test_len(self, dataset: SatlasPretrain) -> None:
+        assert len(dataset) == 2
+
+    def test_already_downloaded(self, dataset: SatlasPretrain) -> None:
+        shutil.rmtree(os.path.join(dataset.root, 'landsat'))
+        SatlasPretrain(root=dataset.root, download=True)
+
+    def test_not_downloaded(self, tmp_path: Path) -> None:
+        with pytest.raises(DatasetNotFoundError, match='Dataset not found'):
+            SatlasPretrain(tmp_path)
+
+    def test_plot(self, dataset: SatlasPretrain) -> None:
+        x = dataset[0]
+        x['prediction_land_cover'] = x['mask_land_cover']
+        dataset.plot(x, suptitle='Test')
+        plt.close()

torchgeo/datasets/__init__.py

@@ -99,6 +99,7 @@
 from .reforestree import ReforesTree
 from .resisc45 import RESISC45
 from .rwanda_field_boundary import RwandaFieldBoundary
+from .satlas import SatlasPretrain
 from .seasonet import SeasoNet
 from .seco import SeasonalContrastS2
 from .sen12ms import SEN12MS
@@ -244,6 +245,7 @@
     'RESISC45',
     'ReforesTree',
     'RwandaFieldBoundary',
+    'SatlasPretrain',
     'SeasonalContrastS2',
     'SeasoNet',
     'SEN12MS',