Merge pull request #65 from IBM/pin/albumentations

pin albumentations

pyproject.toml

@@ -39,7 +39,9 @@ dependencies = [
   "geobench>=1.0.0",
   "mlflow>=2.12.1",
   # broken due to https://github.com/Lightning-AI/pytorch-lightning/issues/19977
-  "lightning>=2, <=2.2.5"
+  "lightning>=2, <=2.2.5",
+  # see issue #64
+  "albumentations<=1.4.10"
 ]
 
 [project.optional-dependencies]

terratorch/datasets/generic_pixel_wise_dataset.py

@@ -155,14 +155,15 @@ def __getitem__(self, index: int) -> dict[str, Any]:
             "image": image.astype(np.float32) * self.constant_scale,
             "mask": self._load_file(self.segmentation_mask_files[index], nan_replace=self.no_label_replace).to_numpy()[
                 0
-            ],
-            "filename": self.image_files[index],
+            ]
         }
 
         if self.reduce_zero_label:
             output["mask"] -= 1
         if self.transform:
             output = self.transform(**output)
+        output["filename"] = self.image_files[index]
+
         return output
 
     def _load_file(self, path, nan_replace: int | float | None = None) -> xr.DataArray:

terratorch/datasets/generic_scalar_label_dataset.py

@@ -143,13 +143,12 @@ def __getitem__(self, index: int) -> dict[str, Any]:
 
         output = {
             "image": image.astype(np.float32) * self.constant_scale,
-            "label": label,
-            "filename": self.samples[index][
-                0
-            ],  # samples is an attribute of ImageFolder. Contains a tuple of (Path, Target)
+            "label": label,  # samples is an attribute of ImageFolder. Contains a tuple of (Path, Target)
         }
         if self.transforms:
             output = self.transforms(**output)
+        output["filename"] = self.image_files[index]
+
         return output
 
     def _load_file(self, path) -> xr.DataArray: