Repository: microsoft/torchgeo

Issue:
Sampling from custom object detection dataset
Hi all. Thanks a lot for this library... I am really looking forward to integrating it into my workflow.

I am trying to sample from my own multiclass object detection dataset. I have a set of large raster images (geotiffs) and polygons in a geojson. All the source data is in `EPSG:4326`. I am using torchgeo `v0.2.0`.

Here are my dataset classes:

```
class MyImageDataset(RasterDataset):
    
    filename_glob = "**/*.tif"


class MyObjectDataset(VectorDataset):
    
    filename_glob = "*.geojson"
```

Then I instantiate the dataset like this:

```
imd = MyImageDataset(root=MY_IMAGE_ROOT)
objd = MyObjectDataset(root=MY_VECTORFILE_ROOT)

dataset = imd & objd
```

I am trying to sample from the image dataset, but only in spots where there is an object in my vector dataset. So, I combine the two datasets as above, since I believe that an `IntersectionDataset` is what I want.

Then I set up my sampler and dataloader:

```
sampler = RandomGeoSampler(dataset, size=0.002, length=1)
dl = DataLoader(dataset, sampler=sampler, collate_fn=stack_samples)

for sample in dl:
    image = sample['image']
    ...
```

However, I keep getting this error: 

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [156], in <cell line: 15>()
     12 sampler = RandomGeoSampler(dataset, size=0.002, length=1)
     13 dl = DataLoader(dataset, sampler=sampler, collate_fn=stack_samples)
---> 15 for sample in dl:
     16     image = sample['image']

...

File ~/torchgeo_test/venv-torchgeo/lib/python3.8/site-packages/rasterio/features.py:321, in rasterize(shapes, out_shape, fill, out, transform, all_touched, merge_alg, default_value, dtype)
    318         warnings.warn('Invalid or empty shape {} at index {} will not be rasterized.'.format(geom, index), ShapeSkipWarning)
    320 if not valid_shapes:
--> 321     raise ValueError('No valid geometry objects found for rasterize')
    323 shape_values = np.array(shape_values)
    325 if not validate_dtype(shape_values, valid_dtypes):

ValueError: No valid geometry objects found for rasterize
```

I have experimented with different values for `size` but I did a sanity check by looking at a sample from just the raster dataset at this size and everything looks good. (Note that I am on `v0.2.0` and I don't think the `unit` option is available for specifying pixel units.)

To troubleshoot, I have limited my image dataset to a single tiff and my vector dataset to a geojson with only 100 objects, some (but not all) of which fall on the image in my image dataset. I have also focused on one specific ROI on that image, a bounding box of size 1000 m^2 where I know there are objects located. I still receive the errors above.

I can index into the raster dataset using that ROI bounding box and get back the correct imagery, and I can index into the vector dataset using that ROI bounding box and get back the correct mask for my objects. Also, I can sample from the `MyImageDataset` raster dataset, but I get the same error as above when attempting to sampling from the `MyObjectDataset` vector dataset.

What am I doing wrong? Feel free to point me to any examples of setting up this type of dataset that I may have missed. 



Discussion:
Could be related to #399
I saw that issue and read through it. In my case, shouldn't there be shapes that survive the filtering based on my other observations above? It certainly seems that no shapes are being returned, but I don't understand why. Any ideas on how I can test this other than the methods I tried above?

Resolving patch:
diff --git a/torchgeo/datasets/geo.py b/torchgeo/datasets/geo.py
--- a/torchgeo/datasets/geo.py
+++ b/torchgeo/datasets/geo.py
@@ -661,9 +661,14 @@ def __getitem__(self, query: BoundingBox) -> Dict[str, Any]:
         transform = rasterio.transform.from_bounds(
             query.minx, query.miny, query.maxx, query.maxy, width, height
         )
-        masks = rasterio.features.rasterize(
-            shapes, out_shape=(int(height), int(width)), transform=transform
-        )
+        if shapes:
+            masks = rasterio.features.rasterize(
+                shapes, out_shape=(int(height), int(width)), transform=transform
+            )
+        else:
+            # If no features are found in this query, return an empty mask
+            # with the default fill value and dtype used by rasterize
+            masks = np.zeros((int(height), int(width)), dtype=np.uint8)
 
         sample = {"mask": torch.tensor(masks), "crs": self.crs, "bbox": query}
 


Test patch:
diff --git a/tests/data/vector/data.py b/tests/data/vector/data.py
new file mode 100755
--- /dev/null
+++ b/tests/data/vector/data.py
@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+import json
+
+# Create an L shape:
+#
+# +--+
+# |  |
+# +--+--+
+# |  |  |
+# +--+--+
+#
+# This allows us to test queries:
+#
+# * within the L
+# * within the dataset bounding box but with no features
+# * outside the dataset bounding box
+
+geojson = {
+    "type": "FeatureCollection",
+    "crs": {"type": "name", "properties": {"name": "urn:ogc:def:crs:OGC:1.3:CRS84"}},
+    "features": [
+        {
+            "type": "Feature",
+            "properties": {},
+            "geometry": {
+                "type": "Polygon",
+                "coordinates": [
+                    [[0.0, 0.0], [0.0, 1.0], [1.0, 1.0], [1.0, 0.0], [0.0, 0.0]]
+                ],
+            },
+        },
+        {
+            "type": "Feature",
+            "properties": {},
+            "geometry": {
+                "type": "Polygon",
+                "coordinates": [
+                    [[1.0, 0.0], [1.0, 1.0], [2.0, 1.0], [2.0, 0.0], [1.0, 0.0]]
+                ],
+            },
+        },
+        {
+            "type": "Feature",
+            "properties": {},
+            "geometry": {
+                "type": "Polygon",
+                "coordinates": [
+                    [[0.0, 1.0], [0.0, 2.0], [1.0, 2.0], [1.0, 1.0], [0.0, 1.0]]
+                ],
+            },
+        },
+    ],
+}
+
+with open("vector.geojson", "w") as f:
+    json.dump(geojson, f)
diff --git a/tests/data/vector/vector.geojson b/tests/data/vector/vector.geojson
new file mode 100644
--- /dev/null
+++ b/tests/data/vector/vector.geojson
@@ -0,0 +1 @@
+{"type": "FeatureCollection", "crs": {"type": "name", "properties": {"name": "urn:ogc:def:crs:OGC:1.3:CRS84"}}, "features": [{"type": "Feature", "properties": {}, "geometry": {"type": "Polygon", "coordinates": [[[0.0, 0.0], [0.0, 1.0], [1.0, 1.0], [1.0, 0.0], [0.0, 0.0]]]}}, {"type": "Feature", "properties": {}, "geometry": {"type": "Polygon", "coordinates": [[[1.0, 0.0], [1.0, 1.0], [2.0, 1.0], [2.0, 0.0], [1.0, 0.0]]]}}, {"type": "Feature", "properties": {}, "geometry": {"type": "Polygon", "coordinates": [[[0.0, 1.0], [0.0, 2.0], [1.0, 2.0], [1.0, 1.0], [0.0, 1.0]]]}}]}
\ No newline at end of file
diff --git a/tests/datasets/test_geo.py b/tests/datasets/test_geo.py
--- a/tests/datasets/test_geo.py
+++ b/tests/datasets/test_geo.py
@@ -16,7 +16,6 @@
 from torchgeo.datasets import (
     NAIP,
     BoundingBox,
-    CanadianBuildingFootprints,
     GeoDataset,
     IntersectionDataset,
     RasterDataset,
@@ -44,6 +43,10 @@ def __getitem__(self, query: BoundingBox) -> Dict[str, BoundingBox]:
         return {"index": query}
 
 
+class CustomVectorDataset(VectorDataset):
+    filename_glob = "*.geojson"
+
+
 class CustomVisionDataset(VisionDataset):
     def __getitem__(self, index: int) -> Dict[str, int]:
         return {"index": index}
@@ -201,20 +204,25 @@ def test_plot_with_cmap(self, custom_dtype_ds: RasterDataset) -> None:
 
 
 class TestVectorDataset:
-    @pytest.fixture
-    def dataset(self) -> CanadianBuildingFootprints:
-        root = os.path.join("tests", "data", "cbf")
+    @pytest.fixture(scope="class")
+    def dataset(self) -> CustomVectorDataset:
+        root = os.path.join("tests", "data", "vector")
         transforms = nn.Identity()  # type: ignore[no-untyped-call]
-        return CanadianBuildingFootprints(root, res=0.1, transforms=transforms)
+        return CustomVectorDataset(root, res=0.1, transforms=transforms)
 
-    def test_getitem(self, dataset: CanadianBuildingFootprints) -> None:
+    def test_getitem(self, dataset: CustomVectorDataset) -> None:
         x = dataset[dataset.bounds]
         assert isinstance(x, dict)
         assert isinstance(x["crs"], CRS)
         assert isinstance(x["mask"], torch.Tensor)
 
-    def test_invalid_query(self, dataset: CanadianBuildingFootprints) -> None:
-        query = BoundingBox(2, 2, 2, 2, 2, 2)
+    def test_empty_shapes(self, dataset: CustomVectorDataset) -> None:
+        query = BoundingBox(1.1, 1.9, 1.1, 1.9, 0, 0)
+        x = dataset[query]
+        assert torch.equal(x["mask"], torch.zeros(7, 7, dtype=torch.uint8))
+
+    def test_invalid_query(self, dataset: CustomVectorDataset) -> None:
+        query = BoundingBox(3, 3, 3, 3, 0, 0)
         with pytest.raises(
             IndexError, match="query: .* not found in index with bounds:"
         ):