add label after transform

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

Author: CarlosGomes98

terratorch/datasets/m_bigearthnet.py

@@ -82,12 +82,12 @@ def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
         labels_tensor = torch.tensor(labels_vector, dtype=torch.float)
 
         output =  {
-            "image": image,
-            "label": labels_tensor
+            "image": image
         }
 
         output = self.transform(**output)
 
+        output["label"] = labels_tensor
         return output
 
     def _validate_bands(self, bands: Sequence[str]) -> None:

terratorch/datasets/m_brick_kiln.py

@@ -81,10 +81,12 @@ def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
             attr_dict = pickle.loads(ast.literal_eval(h5file.attrs["pickle"]))
             class_index = attr_dict["label"]
 
-        output = {"image": image.astype(np.float32), "label": class_index}
+        output = {"image": image.astype(np.float32)}
 
         output = self.transform(**output)
 
+        output["label"] = class_index
+
         return output
 
     def __len__(self):

terratorch/datasets/m_eurosat.py

@@ -86,10 +86,12 @@ def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
         label_class = self.id_to_class[image_id]
         label_index = list(self.label_map.keys()).index(label_class)
 
-        output = {"image": image.astype(np.float32), "label": label_index}
+        output = {"image": image.astype(np.float32)}
 
         output = self.transform(**output)
 
+        output["label"] = label_index
+
         return output
 
     def _validate_bands(self, bands: Sequence[str]) -> None:

terratorch/datasets/m_forestnet.py

@@ -74,10 +74,12 @@ def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
             attr_dict = pickle.loads(ast.literal_eval(h5file.attrs["pickle"]))
             class_index = attr_dict["label"]
 
-        output = {"image": image.astype(np.float32), "label": class_index}
+        output = {"image": image.astype(np.float32)}
 
         output = self.transform(**output)
 
+        output["label"] = class_index
+
         return output
 
     def _validate_bands(self, bands: Sequence[str]) -> None:

terratorch/datasets/m_pv4ger.py

@@ -67,10 +67,12 @@ def __getitem__(self, index: int) -> dict[str, torch.Tensor]:
             attr_dict = pickle.loads(ast.literal_eval(h5file.attrs["pickle"]))
             class_index = attr_dict["label"]
 
-        output = {"image": image.astype(np.float32), "label": class_index}
+        output = {"image": image.astype(np.float32)}
 
         output = self.transform(**output)
 
+        output["label"] = class_index
+
         return output
 
     def _validate_bands(self, bands: Sequence[str]) -> None: