Switch class -> label

tests/datasets/test_vhr10.py

@@ -41,7 +41,7 @@ def test_getitem(self, dataset: VHR10) -> None:
             assert isinstance(x, dict)
             assert isinstance(x['image'], torch.Tensor)
             if dataset.split == 'positive':
-                assert isinstance(x['class'], torch.Tensor)
+                assert isinstance(x['label'], torch.Tensor)
                 assert isinstance(x['bbox_xyxy'], torch.Tensor)
                 if 'mask' in x:
                     assert isinstance(x['mask'], torch.Tensor)
@@ -82,7 +82,7 @@ def test_plot(self, dataset: VHR10) -> None:
             scores = [0.7, 0.3, 0.7]
             for i in range(3):
                 x = dataset[i]
-                x['prediction_labels'] = x['class']
+                x['prediction_labels'] = x['label']
                 x['prediction_boxes'] = x['bbox_xyxy']
                 x['prediction_scores'] = torch.Tensor([scores[i]])
                 if 'mask' in x:

torchgeo/datamodules/utils.py

@@ -31,10 +31,10 @@ def collate_fn_detection(batch: list[dict[str, Tensor]]) -> dict[str, Any]:
     output: dict[str, Any] = {}
     output['image'] = torch.stack([sample['image'] for sample in batch])
     output['bbox_xyxy'] = [sample['bbox_xyxy'].float() for sample in batch]
-    if 'class' in batch[0].keys():
-        output['class'] = [sample['class'] for sample in batch]
+    if 'label' in batch[0].keys():
+        output['label'] = [sample['label'] for sample in batch]
     else:
-        output['class'] = [
+        output['label'] = [
             torch.tensor([1] * len(sample['bbox_xyxy'])) for sample in batch
         ]
 

torchgeo/datasets/vhr10.py

@@ -251,7 +251,7 @@ def __getitem__(self, index: int) -> dict[str, Any]:
             sample['class'] = sample['label']['labels']
             sample['bbox_xyxy'] = sample['label']['boxes']
             sample['mask'] = sample['label']['masks'].float()
-            del sample['label']
+            sample['label'] = sample.pop('class')
 
         if self.transforms is not None:
             sample = self.transforms(sample)
@@ -402,7 +402,7 @@ def plot(
 
         image = sample['image'].permute(1, 2, 0).numpy()
         boxes = sample['bbox_xyxy'].cpu().numpy()
-        labels = sample['class'].cpu().numpy()
+        labels = sample['label'].cpu().numpy()
         if 'mask' in sample:
             masks = [mask.squeeze().cpu().numpy() for mask in sample['mask']]
 

torchgeo/trainers/detection.py

@@ -240,7 +240,7 @@ def training_step(
         batch_size = x.shape[0]
         assert 'bbox_xyxy' in batch, 'bbox_xyxy is required for object detection.'
         y = [
-            {'boxes': batch['bbox_xyxy'][i], 'labels': batch['class'][i]}
+            {'boxes': batch['bbox_xyxy'][i], 'labels': batch['label'][i]}
             for i in range(batch_size)
         ]
         loss_dict = self(x, y)
@@ -262,7 +262,7 @@ def validation_step(
         batch_size = x.shape[0]
         assert 'bbox_xyxy' in batch, 'bbox_xyxy is required for object detection.'
         y = [
-            {'boxes': batch['bbox_xyxy'][i], 'labels': batch['class'][i]}
+            {'boxes': batch['bbox_xyxy'][i], 'labels': batch['label'][i]}
             for i in range(batch_size)
         ]
         y_hat = self(x)
@@ -317,7 +317,7 @@ def test_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> None
         batch_size = x.shape[0]
         assert 'bbox_xyxy' in batch, 'bbox_xyxy is required for object detection.'
         y = [
-            {'boxes': batch['bbox_xyxy'][i], 'labels': batch['class'][i]}
+            {'boxes': batch['bbox_xyxy'][i], 'labels': batch['label'][i]}
             for i in range(batch_size)
         ]
         y_hat = self(x)