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,236 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+from pathlib import Path
+from typing import Any, Literal
+
+import pytest
+import segmentation_models_pytorch as smp
+import timm
+import torch
+import torch.nn as nn
+from lightning.pytorch import Trainer
+from pytest import MonkeyPatch
+from torch.nn.modules import Module
+from torchvision.models._api import WeightsEnum
+
+from torchgeo.datamodules import MisconfigurationException, OSCDDataModule
+from torchgeo.datasets import OSCD, RGBBandsMissingError
+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:
+        x = self.conv1(x)
+        return x
+
+
+def create_model(**kwargs: Any) -> Module:
+    return ChangeDetectionTestModel(**kwargs)
+
+
+def plot(*args: Any, **kwargs: Any) -> None:
+    return None
+
+
+def plot_missing_bands(*args: Any, **kwargs: Any) -> None:
+    raise RGBBandsMissingError()
+
+
+class PredictChangeDetectionDataModule(OSCDDataModule):
+    def setup(self, stage: str) -> None:
+        self.predict_dataset = OSCD(**self.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
+
+    def test_predict(self, fast_dev_run: bool) -> None:
+        datamodule = PredictChangeDetectionDataModule(
+            root=os.path.join('tests', 'data', 'oscd'),
+            batch_size=2,
+            patch_size=32,
+            val_split_pct=0.5,
+            num_workers=0,
+        )
+        model = ChangeDetectionTask(backbone='resnet18', in_channels=13, model='unet')
+        trainer = Trainer(
+            accelerator='cpu',
+            fast_dev_run=fast_dev_run,
+            log_every_n_steps=1,
+            max_epochs=1,
+        )
+        trainer.predict(model=model, datamodule=datamodule)
+
+    @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'],
+        )
+
+    @pytest.mark.parametrize('model_name', ['unet', 'fcsiamdiff', 'fcsiamconc'])
+    @pytest.mark.parametrize(
+        'backbone', ['resnet18', 'mobilenet_v2', 'efficientnet-b0']
+    )
+    def test_freeze_backbone(
+        self, model_name: Literal['unet', 'fcsiamdiff', 'fcsiamconc'], 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', 'fcsiamdiff', 'fcsiamconc'])
+    def test_freeze_decoder(
+        self, model_name: Literal['unet', 'fcsiamdiff', 'fcsiamconc']
+    ) -> 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()
+            ]
+        )
+
+    @pytest.mark.parametrize('loss_fn', ['bce', 'jaccard', 'focal'])
+    def test_losses(self, loss_fn: Literal['bce', 'jaccard', 'focal']) -> None:
+        ChangeDetectionTask(loss=loss_fn)
+
+    def test_no_plot_method(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
+        monkeypatch.setattr(OSCDDataModule, 'plot', plot)
+        datamodule = OSCDDataModule(
+            root=os.path.join('tests', 'data', 'oscd'),
+            batch_size=2,
+            patch_size=32,
+            val_split_pct=0.5,
+            num_workers=0,
+        )
+        model = ChangeDetectionTask(backbone='resnet18', in_channels=13, model='unet')
+        trainer = Trainer(
+            accelerator='cpu',
+            fast_dev_run=fast_dev_run,
+            log_every_n_steps=1,
+            max_epochs=1,
+        )
+        trainer.validate(model=model, datamodule=datamodule)
+
+    def test_no_rgb(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
+        monkeypatch.setattr(OSCDDataModule, 'plot', plot_missing_bands)
+        datamodule = OSCDDataModule(
+            root=os.path.join('tests', 'data', 'oscd'),
+            batch_size=2,
+            patch_size=32,
+            val_split_pct=0.5,
+            num_workers=0,
+        )
+        model = ChangeDetectionTask(backbone='resnet18', in_channels=13, model='unet')
+        trainer = Trainer(
+            accelerator='cpu',
+            fast_dev_run=fast_dev_run,
+            log_every_n_steps=1,
+            max_epochs=1,
+        )
+        trainer.validate(model=model, datamodule=datamodule)