GridGeoSampler: change stride of last patch to sample entire ROI

tests/samplers/test_single.py

@@ -171,9 +171,13 @@ def sampler(self, dataset: CustomGeoDataset, request: SubRequest) -> GridGeoSamp
 
     def test_iter(self, sampler: GridGeoSampler) -> None:
         for query in sampler:
-            assert sampler.roi.minx <= query.minx <= query.maxx <= sampler.roi.maxx
-            assert sampler.roi.miny <= query.miny <= query.miny <= sampler.roi.maxy
-            assert sampler.roi.mint <= query.mint <= query.maxt <= sampler.roi.maxt
+            assert sampler.roi.minx <= query.minx
+            assert sampler.roi.miny <= query.miny
+            assert sampler.roi.mint <= query.mint
+            if query.maxx > sampler.roi.maxx:
+                assert (query.maxx - sampler.roi.maxx) < sampler.size[1]
+            if query.maxy > sampler.roi.maxy:
+                assert (query.maxy - sampler.roi.maxy) < sampler.size[0]
 
             assert math.isclose(query.maxx - query.minx, sampler.size[1])
             assert math.isclose(query.maxy - query.miny, sampler.size[0])
@@ -182,24 +186,74 @@ def test_iter(self, sampler: GridGeoSampler) -> None:
             )
 
     def test_len(self, sampler: GridGeoSampler) -> None:
-        rows = int((100 - sampler.size[0]) // sampler.stride[0]) + 1
-        cols = int((100 - sampler.size[1]) // sampler.stride[1]) + 1
-        length = rows * cols * 2
+        rows = math.ceil(
+            (100 - sampler.size[0] + sampler.stride[0]) / sampler.stride[0]
+        )
+        cols = math.ceil(
+            (100 - sampler.size[1] + sampler.stride[1]) / sampler.stride[1]
+        )
+        length = rows * cols * 2  # two items in dataset
         assert len(sampler) == length
 
+    def test_len_larger(self, sampler: GridGeoSampler) -> None:
+        entire_rows = (100 - sampler.size[0] + sampler.stride[0]) // sampler.stride[0]
+        entire_cols = (100 - sampler.size[1] + sampler.stride[1]) // sampler.stride[1]
+        leftover_row = (100 - sampler.size[0] + sampler.stride[0]) / sampler.stride[
+            0
+        ] - entire_rows
+        leftover_col = (100 - sampler.size[1] + sampler.stride[1]) / sampler.stride[
+            1
+        ] - entire_cols
+        assert (
+            len(sampler)
+            == (entire_rows + math.ceil(leftover_row))
+            * (entire_cols + math.ceil(leftover_col))
+            * 2
+        )
+
     def test_roi(self, dataset: CustomGeoDataset) -> None:
         roi = BoundingBox(0, 50, 200, 250, 400, 450)
         sampler = GridGeoSampler(dataset, 2, 1, roi=roi)
         for query in sampler:
             assert query in roi
 
     def test_small_area(self) -> None:
+        ds = CustomGeoDataset()
+        ds.index.insert(0, (0, 1, 0, 1, 0, 1))
+        sampler = GridGeoSampler(ds, 2, 10)
+        assert len(sampler) == 0
+
+    # TODO: skip patches with area=0 when two tiles are
+    #  side-by-side with an overlapping edge face.
+    def test_tiles_side_by_side(self) -> None:
         ds = CustomGeoDataset()
         ds.index.insert(0, (0, 10, 0, 10, 0, 10))
-        ds.index.insert(1, (20, 21, 20, 21, 20, 21))
+        ds.index.insert(0, (0, 10, 10, 20, 0, 10))
         sampler = GridGeoSampler(ds, 2, 10)
-        for _ in sampler:
-            continue
+        for bbox in sampler:
+            assert bbox.area > 0
+
+    def test_equal_area(self) -> None:
+        ds = CustomGeoDataset()
+        ds.index.insert(0, (0, 10, 0, 10, 0, 10))
+        sampler = GridGeoSampler(ds, 10, 10, units=Units.CRS)
+        assert len(sampler) == 1
+        for bbox in sampler:
+            assert bbox == BoundingBox(
+                minx=0.0, maxx=10.0, miny=0.0, maxy=10.0, mint=0.0, maxt=10.0
+            )
+
+    def test_larger_area(self) -> None:
+        ds = CustomGeoDataset()
+        ds.index.insert(0, (0, 6, 0, 5, 0, 10))
+        sampler = GridGeoSampler(ds, 5, 5, units=Units.CRS)
+        assert len(sampler) == 2
+        assert list(sampler)[0] == BoundingBox(
+            minx=0.0, maxx=5.0, miny=0.0, maxy=5.0, mint=0.0, maxt=10.0
+        )
+        assert list(sampler)[1] == BoundingBox(
+            minx=1.0, maxx=6.0, miny=0.0, maxy=5.0, mint=0.0, maxt=10.0
+        )
 
     @pytest.mark.slow
     @pytest.mark.parametrize("num_workers", [0, 1, 2])

torchgeo/samplers/single.py

@@ -4,6 +4,7 @@
 """TorchGeo samplers."""
 
 import abc
+import math
 from typing import Callable, Iterable, Iterator, Optional, Tuple, Union
 
 import torch
@@ -200,17 +201,24 @@ def __init__(
         for hit in self.index.intersection(tuple(self.roi), objects=True):
             bounds = BoundingBox(*hit.bounds)
             if (
-                bounds.maxx - bounds.minx > self.size[1]
-                and bounds.maxy - bounds.miny > self.size[0]
+                bounds.maxx - bounds.minx >= self.size[1]
+                and bounds.maxy - bounds.miny >= self.size[0]
             ):
                 self.hits.append(hit)
 
         self.length: int = 0
         for hit in self.hits:
             bounds = BoundingBox(*hit.bounds)
 
-            rows = int((bounds.maxy - bounds.miny - self.size[0]) // self.stride[0]) + 1
-            cols = int((bounds.maxx - bounds.minx - self.size[1]) // self.stride[1]) + 1
+            # last patch samples outside the bounds
+            rows = math.ceil(
+                (bounds.maxy - bounds.miny - self.size[0] + self.stride[0])
+                / self.stride[0]
+            )
+            cols = math.ceil(
+                (bounds.maxx - bounds.minx - self.size[1] + self.stride[1])
+                / self.stride[1]
+            )
             self.length += rows * cols
 
     def __iter__(self) -> Iterator[BoundingBox]:
@@ -223,8 +231,14 @@ def __iter__(self) -> Iterator[BoundingBox]:
         for hit in self.hits:
             bounds = BoundingBox(*hit.bounds)
 
-            rows = int((bounds.maxy - bounds.miny - self.size[0]) // self.stride[0]) + 1
-            cols = int((bounds.maxx - bounds.minx - self.size[1]) // self.stride[1]) + 1
+            rows = math.ceil(
+                (bounds.maxy - bounds.miny - self.size[0] + self.stride[0])
+                / self.stride[0]
+            )
+            cols = math.ceil(
+                (bounds.maxx - bounds.minx - self.size[1] + self.stride[1])
+                / self.stride[1]
+            )
 
             mint = bounds.mint
             maxt = bounds.maxt
@@ -233,11 +247,17 @@ def __iter__(self) -> Iterator[BoundingBox]:
             for i in range(rows):
                 miny = bounds.miny + i * self.stride[0]
                 maxy = miny + self.size[0]
+                if maxy > bounds.maxy:
+                    maxy = bounds.maxy
+                    miny = bounds.maxy - self.size[0]
 
                 # For each column...
                 for j in range(cols):
                     minx = bounds.minx + j * self.stride[1]
                     maxx = minx + self.size[1]
+                    if maxx > bounds.maxx:
+                        maxx = bounds.maxx
+                        minx = bounds.maxx - self.size[1]
 
                     yield BoundingBox(minx, maxx, miny, maxy, mint, maxt)