SemanticSegmentationTask: add class-wise metrics

.pre-commit-config.yaml

@@ -30,7 +30,7 @@ repos:
           - pytest>=6.1.2
           - scikit-image>=0.22.0
           - torch>=2.6
-          - torchmetrics>=0.10
+          - torchmetrics>=1.1.1
           - torchvision>=0.18
         exclude: (build|data|dist|logo|logs|output)/
   - repo: https://github.com/pre-commit/mirrors-prettier

pyproject.toml

@@ -74,8 +74,8 @@ dependencies = [
     "timm>=0.4.12",
     # torch 1.13+ required by torchvision
     "torch>=1.13",
-    # torchmetrics 0.10+ required for binary/multiclass/multilabel classification metrics
-    "torchmetrics>=0.10",
+    # torchmetrics 1.1.1+ required for average argument to MeanAveragePrecision
+    "torchmetrics>=1.1.1",
     # torchvision 0.14+ required for torchvision.models.swin_v2_b
     "torchvision>=0.14",
 ]

requirements/min-reqs.old

@@ -18,7 +18,7 @@ segmentation-models-pytorch==0.2.0
 shapely==1.8.0
 timm==0.4.12
 torch==1.13.0
-torchmetrics==0.10.0
+torchmetrics==1.1.1
 torchvision==0.14.0
 
 # datasets

tests/datamodules/test_fair1m.py

@@ -3,7 +3,6 @@
 
 import os
 
-import matplotlib.pyplot as plt
 import pytest
 
 from torchgeo.datamodules import FAIR1MDataModule
@@ -26,17 +25,10 @@ def test_val_dataloader(self, datamodule: FAIR1MDataModule) -> None:
         datamodule.setup('validate')
         next(iter(datamodule.val_dataloader()))
 
+    def test_test_dataloader(self, datamodule: FAIR1MDataModule) -> None:
+        datamodule.setup('test')
+        next(iter(datamodule.test_dataloader()))
+
     def test_predict_dataloader(self, datamodule: FAIR1MDataModule) -> None:
         datamodule.setup('predict')
         next(iter(datamodule.predict_dataloader()))
-
-    def test_plot(self, datamodule: FAIR1MDataModule) -> None:
-        datamodule.setup('validate')
-        batch = next(iter(datamodule.val_dataloader()))
-        sample = {
-            'image': batch['image'][0],
-            'boxes': batch['boxes'][0],
-            'label': batch['label'][0],
-        }
-        datamodule.plot(sample)
-        plt.close()

tests/datamodules/test_geo.py

@@ -3,12 +3,10 @@
 
 from typing import Any
 
-import matplotlib.pyplot as plt
 import pytest
 import torch
 from _pytest.fixtures import SubRequest
 from lightning.pytorch import Trainer
-from matplotlib.figure import Figure
 from rasterio.crs import CRS
 from torch import Tensor
 
@@ -31,12 +29,9 @@ def __init__(
         self.res = 1
 
     def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
-        image = torch.arange(3 * 2 * 2, dtype=torch.float).view(3, 2, 2)
+        image = torch.arange(3 * 2 * 2).view(3, 2, 2)
         return {'image': image, 'crs': CRS.from_epsg(4326), 'bounds': query}
 
-    def plot(self, *args: Any, **kwargs: Any) -> Figure:
-        return plt.figure()
-
 
 class CustomGeoDataModule(GeoDataModule):
     def __init__(self) -> None:
@@ -68,14 +63,11 @@ def __init__(
         self.length = length
 
     def __getitem__(self, index: int) -> dict[str, Tensor]:
-        return {'image': torch.arange(3 * 2 * 2, dtype=torch.float).view(3, 2, 2)}
+        return {'image': torch.arange(3 * 2 * 2).view(3, 2, 2)}
 
     def __len__(self) -> int:
         return self.length
 
-    def plot(self, *args: Any, **kwargs: Any) -> Figure:
-        return plt.figure()
-
 
 class CustomNonGeoDataModule(NonGeoDataModule):
     def __init__(self) -> None:
@@ -133,11 +125,6 @@ def test_predict(self, datamodule: CustomGeoDataModule) -> None:
         batch = datamodule.transfer_batch_to_device(batch, torch.device('cpu'), 1)
         batch = datamodule.on_after_batch_transfer(batch, 0)
 
-    def test_plot(self, datamodule: CustomGeoDataModule) -> None:
-        datamodule.setup('validate')
-        datamodule.plot()
-        plt.close()
-
     def test_no_datasets(self) -> None:
         dm = CustomGeoDataModule()
         msg = r'CustomGeoDataModule\.setup must define one of '
@@ -235,11 +222,6 @@ def test_predict(self, datamodule: CustomNonGeoDataModule) -> None:
         batch = next(iter(datamodule.predict_dataloader()))
         batch = datamodule.on_after_batch_transfer(batch, 0)
 
-    def test_plot(self, datamodule: CustomNonGeoDataModule) -> None:
-        datamodule.setup('validate')
-        datamodule.plot()
-        plt.close()
-
     def test_no_datasets(self) -> None:
         dm = CustomNonGeoDataModule()
         msg = r'CustomNonGeoDataModule\.setup must define one of '

tests/datamodules/test_usavars.py

@@ -3,12 +3,10 @@
 
 import os
 
-import matplotlib.pyplot as plt
 import pytest
 from _pytest.fixtures import SubRequest
 
 from torchgeo.datamodules import USAVarsDataModule
-from torchgeo.datasets import unbind_samples
 
 
 class TestUSAVarsDataModule:
@@ -41,10 +39,3 @@ def test_test_dataloader(self, datamodule: USAVarsDataModule) -> None:
         assert len(datamodule.test_dataloader()) == 1
         batch = next(iter(datamodule.test_dataloader()))
         assert batch['image'].shape[0] == datamodule.batch_size
-
-    def test_plot(self, datamodule: USAVarsDataModule) -> None:
-        datamodule.setup('validate')
-        batch = next(iter(datamodule.val_dataloader()))
-        sample = unbind_samples(batch)[0]
-        datamodule.plot(sample)
-        plt.close()

tests/datamodules/test_xview2.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import os
+
+import pytest
+
+from torchgeo.datamodules import XView2DataModule
+
+
+class TestXView2DataModule:
+    @pytest.fixture
+    def datamodule(self) -> XView2DataModule:
+        root = os.path.join('tests', 'data', 'xview2')
+        batch_size = 1
+        num_workers = 0
+        dm = XView2DataModule(
+            root=root, batch_size=batch_size, num_workers=num_workers, val_split_pct=0.5
+        )
+        dm.prepare_data()
+        return dm
+
+    def test_train_dataloader(self, datamodule: XView2DataModule) -> None:
+        datamodule.setup('fit')
+        next(iter(datamodule.train_dataloader()))
+
+    def test_val_dataloader(self, datamodule: XView2DataModule) -> None:
+        datamodule.setup('validate')
+        next(iter(datamodule.val_dataloader()))
+
+    def test_test_dataloader(self, datamodule: XView2DataModule) -> None:
+        datamodule.setup('test')
+        next(iter(datamodule.test_dataloader()))
+
+    def test_predict_dataloader(self, datamodule: XView2DataModule) -> None:
+        datamodule.setup('predict')
+        next(iter(datamodule.predict_dataloader()))

tests/trainers/test_classification.py

@@ -9,6 +9,7 @@
 import timm
 import torch
 import torch.nn as nn
+import torchvision
 from lightning.pytorch import Trainer
 from pytest import MonkeyPatch
 from torch.nn.modules import Module
@@ -19,7 +20,7 @@
     EuroSATDataModule,
     MisconfigurationException,
 )
-from torchgeo.datasets import BigEarthNet, EuroSAT, RGBBandsMissingError
+from torchgeo.datasets import BigEarthNet, EuroSAT
 from torchgeo.main import main
 from torchgeo.models import ResNet18_Weights
 from torchgeo.trainers import ClassificationTask, MultiLabelClassificationTask
@@ -55,12 +56,9 @@
     return ClassificationTestModel(**kwargs)
 
 
-def plot(*args: Any, **kwargs: Any) -> None:
-    return None
-
-
-def plot_missing_bands(*args: Any, **kwargs: Any) -> None:
-    raise RGBBandsMissingError()
+def load(url: str, *args: Any, **kwargs: Any) -> dict[str, Any]:
+    state_dict: dict[str, Any] = torch.load(url)
+    return state_dict
 
 
 class TestClassificationTask:
@@ -103,13 +101,13 @@
             '1',
         ]
 
-        main(['fit', *args])
+        main(['fit'] + args)
         try:
-            main(['test', *args])
+            main(['test'] + args)
         except MisconfigurationException:
             pass
         try:
-            main(['predict', *args])
+            main(['predict'] + args)
         except MisconfigurationException:
             pass
 
@@ -119,11 +117,7 @@
 
     @pytest.fixture
     def mocked_weights(
-        self,
-        tmp_path: Path,
-        monkeypatch: MonkeyPatch,
-        weights: WeightsEnum,
-        load_state_dict_from_url: None,
+        self, tmp_path: Path, monkeypatch: MonkeyPatch, weights: WeightsEnum
     ) -> WeightsEnum:
         path = tmp_path / f'{weights}.pth'
         model = timm.create_model(  # type: ignore[attr-defined]
@@ -134,6 +128,7 @@
             monkeypatch.setattr(weights.value, 'url', str(path))
         except AttributeError:
             monkeypatch.setattr(weights, 'url', str(path))
+        monkeypatch.setattr(torchvision.models._api, 'load_state_dict_from_url', load)
         return weights
 
     def test_weight_file(self, checkpoint: str) -> None:
@@ -183,34 +178,6 @@
         with pytest.raises(ValueError, match=match):
             ClassificationTask(model='resnet18', loss='invalid_loss')
 
-    def test_no_plot_method(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
-        monkeypatch.setattr(EuroSATDataModule, 'plot', plot)
-        datamodule = EuroSATDataModule(
-            root='tests/data/eurosat', batch_size=1, num_workers=0
-        )
-        model = ClassificationTask(model='resnet18', in_channels=13, num_classes=10)
-        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(EuroSATDataModule, 'plot', plot_missing_bands)
-        datamodule = EuroSATDataModule(
-            root='tests/data/eurosat', batch_size=1, num_workers=0
-        )
-        model = ClassificationTask(model='resnet18', in_channels=13, num_classes=10)
-        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_predict(self, fast_dev_run: bool) -> None:
         datamodule = PredictClassificationDataModule(
             root='tests/data/eurosat', batch_size=1, num_workers=0
@@ -237,7 +204,7 @@
 
 class TestMultiLabelClassificationTask:
     @pytest.mark.parametrize(
-        'name', ['bigearthnet_all', 'bigearthnet_s1', 'bigearthnet_s2', 'treesatai']
+        'name', ['bigearthnet_all', 'bigearthnet_s1', 'bigearthnet_s2']
     )
     def test_trainer(
         self, monkeypatch: MonkeyPatch, name: str, fast_dev_run: bool
@@ -259,13 +226,13 @@
             '1',
         ]
 
-        main(['fit', *args])
+        main(['fit'] + args)
         try:
-            main(['test', *args])
+            main(['test'] + args)
         except MisconfigurationException:
             pass
         try:
-            main(['predict', *args])
+            main(['predict'] + args)
         except MisconfigurationException:
             pass
 
@@ -274,38 +241,6 @@
         with pytest.raises(ValueError, match=match):
             MultiLabelClassificationTask(model='resnet18', loss='invalid_loss')
 
-    def test_no_plot_method(self, monkeypatch: MonkeyPatch, fast_dev_run: bool) -> None:
-        monkeypatch.setattr(BigEarthNetDataModule, 'plot', plot)
-        datamodule = BigEarthNetDataModule(
-            root='tests/data/bigearthnet', batch_size=1, num_workers=0
-        )
-        model = MultiLabelClassificationTask(
-            model='resnet18', in_channels=14, num_classes=19, loss='bce'
-        )
-        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(BigEarthNetDataModule, 'plot', plot_missing_bands)
-        datamodule = BigEarthNetDataModule(
-            root='tests/data/bigearthnet', batch_size=1, num_workers=0
-        )
-        model = MultiLabelClassificationTask(
-            model='resnet18', in_channels=14, num_classes=19, loss='bce'
-        )
-        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_predict(self, fast_dev_run: bool) -> None:
         datamodule = PredictMultiLabelClassificationDataModule(
             root='tests/data/bigearthnet', batch_size=1, num_workers=0