add filename only after transforms

Signed-off-by: Carlos Gomes <[email protected]>

Author: CarlosGomes98

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: