Add ChangeDetectionTask

tests/conf/oscd.yaml

@@ -0,0 +1,15 @@
+model:
+  class_path: ChangeDetectionTask
+  init_args:
+    loss: 'bce'
+    model: 'unet'
+    backbone: 'resnet18'
+    in_channels: 13
+data:
+  class_path: OSCDDataModule
+  init_args:
+    batch_size: 2
+    patch_size: 16
+    val_split_pct: 0.5
+  dict_kwargs:
+    root: 'tests/data/oscd'

tests/datasets/test_oscd.py

@@ -66,19 +66,15 @@ def dataset(
     def test_getitem(self, dataset: OSCD) -> None:
         x = dataset[0]
         assert isinstance(x, dict)
-        assert isinstance(x['image1'], torch.Tensor)
-        assert x['image1'].ndim == 3
-        assert isinstance(x['image2'], torch.Tensor)
-        assert x['image2'].ndim == 3
+        assert isinstance(x['image'], torch.Tensor)
+        assert x['image'].ndim == 4
         assert isinstance(x['mask'], torch.Tensor)
         assert x['mask'].ndim == 2
 
         if dataset.bands == OSCD.rgb_bands:
-            assert x['image1'].shape[0] == 3
-            assert x['image2'].shape[0] == 3
+            assert x['image'].shape[1] == 3
         else:
-            assert x['image1'].shape[0] == 13
-            assert x['image2'].shape[0] == 13
+            assert x['image'].shape[1] == 13
 
     def test_len(self, dataset: OSCD) -> None:
         if dataset.split == 'train':

tests/trainers/conftest.py

@@ -6,8 +6,8 @@
 from pathlib import Path
 
 import pytest
+import timm
 import torch
-import torchvision
 from _pytest.fixtures import SubRequest
 from torch import Tensor
 from torch.nn.modules import Module
@@ -22,8 +22,9 @@ def fast_dev_run(request: SubRequest) -> bool:
 
 
 @pytest.fixture(scope='package')
-def model() -> Module:
-    model: Module = torchvision.models.resnet18(weights=None)
+def model(request: SubRequest) -> Module:
+    in_channels = getattr(request, 'param', 3)
+    model: Module = timm.create_model('resnet18', in_chans=in_channels)
     return model
 
 

tests/trainers/test_change.py

@@ -0,0 +1,169 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+from pathlib import Path
+from typing import Any, cast
+
+import pytest
+import segmentation_models_pytorch as smp
+import timm
+import torch
+import torch.nn as nn
+from pytest import MonkeyPatch
+from torch.nn.modules import Module
+from torchvision.models._api import WeightsEnum
+
+from torchgeo.datamodules import MisconfigurationException
+from torchgeo.main import main
+from torchgeo.models import ResNet18_Weights
+from torchgeo.trainers import ChangeDetectionTask
+
+
+class ChangeDetectionTestModel(Module):
+    def __init__(self, in_channels: int = 3, classes: int = 3, **kwargs: Any) -> None:
+        super().__init__()
+        self.conv1 = nn.Conv2d(
+            in_channels=in_channels, out_channels=classes, kernel_size=1, padding=0
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return cast(torch.Tensor, self.conv1(x))
+
+
+def create_model(**kwargs: Any) -> Module:
+    return ChangeDetectionTestModel(**kwargs)
+
+
+class TestChangeDetectionTask:
+    @pytest.mark.parametrize('name', ['oscd'])
+    def test_trainer(
+        self, monkeypatch: MonkeyPatch, name: str, fast_dev_run: bool
+    ) -> None:
+        config = os.path.join('tests', 'conf', name + '.yaml')
+
+        monkeypatch.setattr(smp, 'Unet', create_model)
+
+        args = [
+            '--config',
+            config,
+            '--trainer.accelerator',
+            'cpu',
+            '--trainer.fast_dev_run',
+            str(fast_dev_run),
+            '--trainer.max_epochs',
+            '1',
+            '--trainer.log_every_n_steps',
+            '1',
+        ]
+
+        main(['fit', *args])
+        try:
+            main(['test', *args])
+        except MisconfigurationException:
+            pass
+        try:
+            main(['predict', *args])
+        except MisconfigurationException:
+            pass
+
+    @pytest.fixture
+    def weights(self) -> WeightsEnum:
+        return ResNet18_Weights.SENTINEL2_ALL_MOCO
+
+    @pytest.fixture
+    def mocked_weights(
+        self,
+        tmp_path: Path,
+        monkeypatch: MonkeyPatch,
+        weights: WeightsEnum,
+        load_state_dict_from_url: None,
+    ) -> WeightsEnum:
+        path = tmp_path / f'{weights}.pth'
+        # multiply in_chans by 2 since images are concatenated
+        model = timm.create_model(
+            weights.meta['model'], in_chans=weights.meta['in_chans'] * 2
+        )
+        torch.save(model.state_dict(), path)
+        try:
+            monkeypatch.setattr(weights.value, 'url', str(path))
+        except AttributeError:
+            monkeypatch.setattr(weights, 'url', str(path))
+        return weights
+
+    @pytest.mark.parametrize('model', [6], indirect=True)
+    def test_weight_file(self, checkpoint: str) -> None:
+        ChangeDetectionTask(backbone='resnet18', weights=checkpoint)
+
+    def test_weight_enum(self, mocked_weights: WeightsEnum) -> None:
+        ChangeDetectionTask(
+            backbone=mocked_weights.meta['model'],
+            weights=mocked_weights,
+            in_channels=mocked_weights.meta['in_chans'],
+        )
+
+    def test_weight_str(self, mocked_weights: WeightsEnum) -> None:
+        ChangeDetectionTask(
+            backbone=mocked_weights.meta['model'],
+            weights=str(mocked_weights),
+            in_channels=mocked_weights.meta['in_chans'],
+        )
+
+    @pytest.mark.slow
+    def test_weight_enum_download(self, weights: WeightsEnum) -> None:
+        ChangeDetectionTask(
+            backbone=weights.meta['model'],
+            weights=weights,
+            in_channels=weights.meta['in_chans'],
+        )
+
+    @pytest.mark.slow
+    def test_weight_str_download(self, weights: WeightsEnum) -> None:
+        ChangeDetectionTask(
+            backbone=weights.meta['model'],
+            weights=str(weights),
+            in_channels=weights.meta['in_chans'],
+        )
+
+    def test_invalid_model(self) -> None:
+        match = "Model type 'invalid_model' is not valid."
+        with pytest.raises(ValueError, match=match):
+            ChangeDetectionTask(model='invalid_model')
+
+    def test_invalid_loss(self) -> None:
+        match = "Loss type 'invalid_loss' is not valid."
+        with pytest.raises(ValueError, match=match):
+            ChangeDetectionTask(loss='invalid_loss')
+
+    @pytest.mark.parametrize('model_name', ['unet'])
+    @pytest.mark.parametrize(
+        'backbone', ['resnet18', 'mobilenet_v2', 'efficientnet-b0']
+    )
+    def test_freeze_backbone(self, model_name: str, backbone: str) -> None:
+        model = ChangeDetectionTask(
+            model=model_name, backbone=backbone, freeze_backbone=True
+        )
+        assert all(
+            [param.requires_grad is False for param in model.model.encoder.parameters()]
+        )
+        assert all([param.requires_grad for param in model.model.decoder.parameters()])
+        assert all(
+            [
+                param.requires_grad
+                for param in model.model.segmentation_head.parameters()
+            ]
+        )
+
+    @pytest.mark.parametrize('model_name', ['unet'])
+    def test_freeze_decoder(self, model_name: str) -> None:
+        model = ChangeDetectionTask(model=model_name, freeze_decoder=True)
+        assert all(
+            [param.requires_grad is False for param in model.model.decoder.parameters()]
+        )
+        assert all([param.requires_grad for param in model.model.encoder.parameters()])
+        assert all(
+            [
+                param.requires_grad
+                for param in model.model.segmentation_head.parameters()
+            ]
+        )

torchgeo/datamodules/oscd.py

@@ -85,8 +85,10 @@ def __init__(
         self.std = torch.tensor([STD[b] for b in self.bands])
 
         self.aug = K.AugmentationSequential(
-            K.Normalize(mean=self.mean, std=self.std),
-            _RandomNCrop(self.patch_size, batch_size),
+            K.VideoSequential(
+                K.Normalize(mean=self.mean, std=self.std),
+                _RandomNCrop(self.patch_size, batch_size),
+            ),
             data_keys=None,
             keepdim=True,
         )

torchgeo/datasets/oscd.py

@@ -150,7 +150,8 @@ def __getitem__(self, index: int) -> dict[str, Tensor]:
         image1 = self._load_image(files['images1'])
         image2 = self._load_image(files['images2'])
         mask = self._load_target(str(files['mask']))
-        sample = {'image1': image1, 'image2': image2, 'mask': mask}
+        image = torch.stack(tensors=[image1, image2], dim=0)
+        sample = {'image': image, 'mask': mask}
 
         if self.transforms is not None:
             sample = self.transforms(sample)
@@ -324,8 +325,8 @@ def get_masked(img: Tensor) -> 'np.typing.NDArray[np.uint8]':
             )
             return array
 
-        image1 = get_masked(sample['image1'])
-        image2 = get_masked(sample['image2'])
+        image1 = get_masked(sample['image'][0])
+        image2 = get_masked(sample['image'][1])
         fig, axs = plt.subplots(ncols=ncols, figsize=(ncols * 10, 10))
         axs[0].imshow(image1)
         axs[0].axis('off')

torchgeo/trainers/__init__.py

@@ -5,6 +5,7 @@
 
 from .base import BaseTask
 from .byol import BYOLTask
+from .change import ChangeDetectionTask
 from .classification import ClassificationTask, MultiLabelClassificationTask
 from .detection import ObjectDetectionTask
 from .iobench import IOBenchTask
@@ -14,8 +15,10 @@
 from .simclr import SimCLRTask
 
 __all__ = (
+    # Supervised
     'BYOLTask',
     'BaseTask',
+    'ChangeDetectionTask',
     'ClassificationTask',
     'IOBenchTask',
     'MoCoTask',