Adding SEN12MS training task

requirements.txt

@@ -18,6 +18,7 @@ rarfile
 rasterio>=1.0
 rtree>=0.5.0
 scikit-learn
+segmentation-models-pytorch
 setuptools>=42
 sphinx
 torch>=1.7

setup.cfg

@@ -64,3 +64,4 @@ tests =
 train =
     pytorch-lightning
     scikit-learn
+    segmentation-models-pytorch

spack.yaml

@@ -21,6 +21,7 @@ spack:
   - "py-rasterio@0.3:"
   - "py-rtree@0.5.0:"
   - py-scikit-learn
+  - py-segmentation-models-pytorch
   - "py-setuptools@42:"
   - py-sphinx
   - "py-torch@1.7:"

tests/datasets/test_sen12ms.py

@@ -18,20 +18,18 @@ def dataset(
         self, monkeypatch: Generator[MonkeyPatch, None, None], request: SubRequest
     ) -> SEN12MS:
         md5s = [
-            "3079d1c5038fa101ec2072657f2cb1ab",
-            "f11487a4b2e641b64ed80a031c4d121d",
-            "299691b948b37028398d4506d0195c6d",
-            "76e6847b10ee9323ce022508721e2c6c",
-            "dfbe57486455c31ae6f4d243186a8da5",
-            "8d0aae3b12d420cab2feff5035400cbf",
-            "f524074dcd90b9a770031cbfec50db71",
-            "5256cf09bd2a0ec44bdff78f28e6653d",
-            "b85b1641971444c87fedbc7134c437ac",
-            "af28777ee277e3f9577c10a3c6d952eb",
-            "44d18ee9efeb83f921b3b7aa6d511bbf",
-            "00e18016c6af1e55528c535d9b06c35a",
-            "02d5128ac1fc2bf8762091b4f319762d",
-            "02d5128ac1fc2bf8762091b4f319762d",
+            "7f14be13d3f62c09b4dd5b4d55c97fd6",
+            "48182d44b375360381f36d432956b225",
+            "96cf1b8405d4149c6fe61ad7100bd65d",
+            "ba8e7e10fba9eea6900ddc530c86025a",
+            "7ba7c51f2fb3a2074b7bbd3e24f9d70d",
+            "280c9be2d1e13e663824dccd85e1e42f",
+            "a5284baf48534d4bc77acb1b103ff16c",
+            "c6b176fed0cdd5033cb1835506e40ee4",
+            "adc672746b79be4c4edc8b1a564e3ff4",
+            "194fab4a4e067a0452824c4e39f61b77",
+            "7899c0c36c884ae8c991ab8518b0d177",
+            "ccfee543d4351bcc5aa68729e8cc795c",
         ]
 
         monkeypatch.setattr(SEN12MS, "md5s", md5s)  # type: ignore[attr-defined]

tests/test_train.py

@@ -4,6 +4,8 @@
 import sys
 from pathlib import Path
 
+import pytest
+
 
 def test_help() -> None:
     args = [sys.executable, "train.py", "--help"]
@@ -83,3 +85,28 @@ def test_overwrite_experiment_dir(tmp_path: Path) -> None:
         b"The experiment directory, .*, already exists, we might overwrite data in it!",
         ps.stdout,
     )
+
+
+@pytest.mark.parametrize("task", ["cyclone", "sen12ms"])
+def test_tasks(task: str, tmp_path: Path) -> None:
+    experiment_name = "test"
+    output_dir = tmp_path / "output"
+    data_dir = os.path.join("tests", "data")
+    log_dir = tmp_path / "logs"
+    args = [
+        sys.executable,
+        "train.py",
+        "--experiment_name",
+        experiment_name,
+        "--output_dir",
+        str(output_dir),
+        "--data_dir",
+        data_dir,
+        "--log_dir",
+        str(log_dir),
+        "--fast_dev_run",
+        "1",
+        "--task",
+        task,
+    ]
+    subprocess.run(args, check=True)

torchgeo/datasets/sen12ms.py

@@ -136,7 +136,7 @@ def __init__(
                 raise RuntimeError("Dataset not found or corrupted.")
 
         with open(os.path.join(self.root, self.base_folder, split + "_list.txt")) as f:
-            self.ids = f.readlines()
+            self.ids = [line.rstrip() for line in f.readlines()]
 
     def __getitem__(self, index: int) -> Dict[str, Tensor]:
         """Return an index within the dataset.
@@ -147,7 +147,7 @@ def __getitem__(self, index: int) -> Dict[str, Tensor]:
         Returns:
             data and label at that index
         """
-        filename = self.ids[index].rstrip()
+        filename = self.ids[index]
 
         lc = self._load_raster(filename, "lc")
         s1 = self._load_raster(filename, "s1")

torchgeo/trainers/__init__.py

@@ -1,8 +1,14 @@
 """TorchGeo trainers."""
 
 from .cyclone import CycloneDataModule, CycloneSimpleRegressionTask
+from .sen12ms import SEN12MSDataModule, SEN12MSSegmentationTask
 
-__all__ = ("CycloneDataModule", "CycloneSimpleRegressionTask")
+__all__ = (
+    "CycloneDataModule",
+    "CycloneSimpleRegressionTask",
+    "SEN12MSDataModule",
+    "SEN12MSSegmentationTask",
+)
 
 # https://stackoverflow.com/questions/40018681
 for module in __all__:

torchgeo/trainers/cyclone.py

@@ -90,16 +90,14 @@ def configure_optimizers(self) -> Dict[str, Any]:
         """Initialize the optimizer and learning rate scheduler."""
         optimizer = torch.optim.Adam(
             self.model.parameters(),
-            lr=self.hparams["learning_rate"],  # type: ignore[index]
+            lr=self.hparams["learning_rate"],
         )
         return {
             "optimizer": optimizer,
             "lr_scheduler": {
                 "scheduler": ReduceLROnPlateau(
                     optimizer,
-                    patience=self.hparams[
-                        "learning_rate_schedule_patience"
-                    ],  # type: ignore[index]
+                    patience=self.hparams["learning_rate_schedule_patience"],
                 ),
                 "monitor": "val_loss",
             },