Coverage fix

microsoft · Oct 12, 2021 · 1074d6b · 1074d6b
1 parent 7b111b4
commit 1074d6b
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/tests/datasets/test_spacenet.py b/tests/datasets/test_spacenet.py
@@ -176,7 +176,9 @@ def dataset(
         )
 
     def test_getitem(self, dataset: SpaceNet4) -> None:
-        x = dataset[0]
+        # Get image-label pair with empty label to
+        # enusre coverage
+        x = dataset[2]
         assert isinstance(x, dict)
         assert isinstance(x["image"], torch.Tensor)
         assert isinstance(x["mask"], torch.Tensor)

diff --git a/torchgeo/datasets/spacenet.py b/torchgeo/datasets/spacenet.py
@@ -14,6 +14,7 @@
 import rasterio as rio
 import torch
 from affine import Affine
+from fiona.errors import DriverError
 from rasterio.features import rasterize
 from torch import Tensor
 
@@ -157,7 +158,7 @@ def _load_mask(self, path: str, tfm: Affine, shape: Tuple[int, int]) -> Tensor:
         try:
             with fiona.open(path) as src:
                 labels = [feature["geometry"] for feature in src]
-        except Exception:
+        except DriverError:
             labels = []
 
         if not labels: