Fix for classification & regression tests

microsoft · Jul 14, 2024 · 96f6b48 · 96f6b48
1 parent 3ced3de
commit 96f6b48
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 6 deletions.
diff --git a/torchgeo/datasets/cyclone.py b/torchgeo/datasets/cyclone.py
@@ -96,17 +96,19 @@ def __getitem__(self, index: int) -> dict[str, Any]:
         Returns:
             data, labels, field ids, and metadata at that index
         """
-        sample = {
+        features = {
             'relative_time': torch.tensor(self.features.iat[index, 2]),
             'ocean': torch.tensor(self.features.iat[index, 3]),
             'label': torch.tensor(self.labels.iat[index, 1]),
         }
 
         image_id = self.labels.iat[index, 0]
-        sample['image'] = self._load_image(image_id)
+        sample = {'image': self._load_image(image_id)}
+        sample['label'] = features['label']
 
         if self.transforms is not None:
             sample = self.transforms(sample)
+            sample.update({x: features[x] for x in features if x != 'label'})
 
         return sample
 

diff --git a/torchgeo/datasets/inria.py b/torchgeo/datasets/inria.py
@@ -217,7 +217,10 @@ def plot(
         show_predictions = 'prediction' in sample
 
         if show_mask:
-            mask = sample['mask'].numpy()
+            mask = sample['mask']
+            if mask.ndim == 3 and mask.shape[0] == 1:
+                mask = mask.squeeze(0)
+            mask = mask.numpy()
             ncols += 1
 
         if show_predictions:

diff --git a/torchgeo/datasets/quakeset.py b/torchgeo/datasets/quakeset.py
@@ -113,10 +113,11 @@ def __getitem__(self, index: int) -> dict[str, Tensor]:
         label = torch.tensor(self.data[index]['label'])
         magnitude = torch.tensor(self.data[index]['magnitude'])
 
-        sample = {'image': image, 'label': label, 'magnitude': magnitude}
+        sample = {'image': image, 'label': label}
 
         if self.transforms is not None:
             sample = self.transforms(sample)
+            sample['magnitude'] = magnitude
 
         return sample
 

diff --git a/torchgeo/datasets/skippd.py b/torchgeo/datasets/skippd.py
@@ -144,10 +144,12 @@ def __getitem__(self, index: int) -> dict[str, str | Tensor]:
             data and label at that index
         """
         sample: dict[str, str | Tensor] = {'image': self._load_image(index)}
-        sample.update(self._load_features(index))
+        features = self._load_features(index)
+        sample['label'] = features['label']
 
         if self.transforms is not None:
             sample = self.transforms(sample)
+            sample.update({x: features[x] for x in features if x != 'label'})
 
         return sample
 

diff --git a/torchgeo/datasets/sustainbench_crop_yield.py b/torchgeo/datasets/sustainbench_crop_yield.py
@@ -149,10 +149,17 @@ def __getitem__(self, index: int) -> dict[str, Tensor]:
             data and label at that index
         """
         sample: dict[str, Tensor] = {'image': self.images[index]}
-        sample.update(self.features[index])
+        sample['label'] = self.features[index]['label']
 
         if self.transforms is not None:
             sample = self.transforms(sample)
+            sample.update(
+                {
+                    x: self.features[index][x]
+                    for x in self.features[index]
+                    if x != 'label'
+                }
+            )
 
         return sample