Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add ChangeDetectionTask #2422

Draft
wants to merge 18 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions tests/conf/oscd.yaml
Original file line number Diff line number Diff line change
@@ -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'
82 changes: 0 additions & 82 deletions tests/datamodules/test_oscd.py

This file was deleted.

12 changes: 4 additions & 8 deletions tests/datasets/test_oscd.py
Original file line number Diff line number Diff line change
Expand Up @@ -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':
Expand Down
7 changes: 4 additions & 3 deletions tests/trainers/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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


Expand Down
169 changes: 169 additions & 0 deletions tests/trainers/test_change.py
Original file line number Diff line number Diff line change
@@ -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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How hard would it be to do late fusion, so pass each image through the encoder separately, then concatenate them, then pass them through the decoder? This would make it easier to use pre-trained models.

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's definitely possible, although I think we would need a custom Unet implementation in torchgeo/models to do this. It would simplify using the pretrained weights but is late fusion a common enough approach that many people would find this useful?

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)
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Remind me what [6] means here?

Copy link
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Number of input channels (2 3-channel images stacked)

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()
]
)
6 changes: 4 additions & 2 deletions torchgeo/datamodules/oscd.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down
7 changes: 4 additions & 3 deletions torchgeo/datasets/oscd.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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')
Expand Down
3 changes: 3 additions & 0 deletions torchgeo/trainers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -14,8 +15,10 @@
from .simclr import SimCLRTask

__all__ = (
# Supervised
'BYOLTask',
'BaseTask',
'ChangeDetectionTask',
'ClassificationTask',
'IOBenchTask',
'MoCoTask',
Expand Down
Loading
Loading