Fix empty geometries

microsoft · Aug 6, 2024 · 9641f42 · 9641f42
1 parent f07607a
commit 9641f42
Showing 1 changed file with 12 additions and 9 deletions.
diff --git a/torchgeo/datasets/spacenet.py b/torchgeo/datasets/spacenet.py
@@ -196,16 +196,19 @@ def _load_mask(
                 if feature['geometry']
             ]
 
-        mask_data = rasterize(
-            labels,
-            out_shape=shape,
-            fill=0,  # nodata value
-            transform=tfm,
-            all_touched=False,
-            dtype=np.int64,
-        )
+        if labels:
+            mask = rasterize(
+                labels,
+                out_shape=shape,
+                fill=0,  # nodata value
+                transform=tfm,
+                all_touched=False,
+                dtype=np.int64,
+            )
+        else:
+            mask = np.zeros(shape=shape)
 
-        return torch.from_numpy(mask_data)
+        return torch.from_numpy(mask)
 
     def __getitem__(self, index: int) -> dict[str, Tensor]:
         """Return an index within the dataset.