random_bbox_assignment: port to geopandas

Author: adamjstewart

tests/datasets/test_splits.py

@@ -2,11 +2,16 @@
 # Licensed under the MIT License.
 
 from collections.abc import Sequence
+from datetime import datetime
 from math import floor, isclose
 from typing import Any
 
+import pandas as pd
 import pytest
+import shapely
+from geopandas import GeoDataFrame
 from pyproj import CRS
+from shapely import Geometry
 
 from torchgeo.datasets import (
     BoundingBox,
@@ -18,13 +23,12 @@
     time_series_split,
 )
 
+MINT = datetime(2025, 4, 24)
+MAXT = datetime(2025, 4, 25)
 
-def total_area(dataset: GeoDataset) -> float:
-    total_area = 0.0
-    for hit in dataset.index.intersection(dataset.index.bounds, objects=True):
-        total_area += BoundingBox(*hit.bounds).area
 
-    return total_area
+def total_area(dataset: GeoDataset) -> float:
+    return dataset.index.geometry.area.sum()
 
 
 def no_overlap(ds1: GeoDataset, ds2: GeoDataset) -> bool:
@@ -39,20 +43,20 @@ def no_overlap(ds1: GeoDataset, ds2: GeoDataset) -> bool:
 class CustomGeoDataset(GeoDataset):
     def __init__(
         self,
-        items: list[tuple[BoundingBox, str]] = [(BoundingBox(0, 1, 0, 1, 0, 40), '')],
-        crs: CRS = CRS.from_epsg(3005),
-        res: tuple[float, float] = (1, 1),
+        index: pd.IntervalIndex | None = None,
+        geometry: Sequence[Geometry] = [shapely.box(0, 0, 1, 1)],
     ) -> None:
-        super().__init__()
-        for box, content in items:
-            self.index.insert(0, tuple(box), content)
-        self._crs = crs
-        self.res = res
+        if index is None:
+            intervals = [(MINT, MAXT)] * len(geometry)
+            index = pd.IntervalIndex.from_tuples(
+                intervals, closed='both', name='datetime'
+            )
+        crs = CRS.from_epsg(3005)
+        self.index = GeoDataFrame(index=index, geometry=geometry, crs=crs)
+        self.res = (1, 1)
 
     def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
-        hits = self.index.intersection(tuple(query), objects=True)
-        hit = next(iter(hits))
-        return {'content': hit.object}
+        return {'index': query}
 
 
 @pytest.mark.parametrize(
@@ -67,14 +71,13 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
 def test_random_bbox_assignment(
     lengths: Sequence[int | float], expected_lengths: Sequence[int]
 ) -> None:
-    ds = CustomGeoDataset(
-        [
-            (BoundingBox(0, 1, 0, 1, 0, 0), 'a'),
-            (BoundingBox(1, 2, 0, 1, 0, 0), 'b'),
-            (BoundingBox(2, 3, 0, 1, 0, 0), 'c'),
-            (BoundingBox(3, 4, 0, 1, 0, 0), 'd'),
-        ]
-    )
+    geometry = [
+        shapely.box(0, 0, 1, 1),
+        shapely.box(1, 0, 2, 1),
+        shapely.box(2, 0, 3, 1),
+        shapely.box(3, 0, 4, 1),
+    ]
+    ds = CustomGeoDataset(geometry=geometry)
 
     train_ds, val_ds, test_ds = random_bbox_assignment(ds, lengths)
 
@@ -94,7 +97,6 @@ def test_random_bbox_assignment(
     # Test __getitem__
     x = train_ds[train_ds.bounds]
     assert isinstance(x, dict)
-    assert isinstance(x['content'], str)
 
 
 def test_random_bbox_assignment_invalid_inputs() -> None:
@@ -110,14 +112,13 @@ def test_random_bbox_assignment_invalid_inputs() -> None:
 
 
 def test_random_bbox_splitting() -> None:
-    ds = CustomGeoDataset(
-        [
-            (BoundingBox(0, 1, 0, 1, 0, 0), 'a'),
-            (BoundingBox(1, 2, 0, 1, 0, 0), 'b'),
-            (BoundingBox(2, 3, 0, 1, 0, 0), 'c'),
-            (BoundingBox(3, 4, 0, 1, 0, 0), 'd'),
-        ]
-    )
+    geometry = [
+        shapely.box(0, 0, 1, 1),
+        shapely.box(1, 0, 2, 1),
+        shapely.box(2, 0, 3, 1),
+        shapely.box(3, 0, 4, 1),
+    ]
+    ds = CustomGeoDataset(geometry=geometry)
 
     ds_area = total_area(ds)
 
@@ -145,7 +146,6 @@ def test_random_bbox_splitting() -> None:
     # Test __get_item__
     x = train_ds[train_ds.bounds]
     assert isinstance(x, dict)
-    assert isinstance(x['content'], str)
 
     # Test invalid input fractions
     with pytest.raises(ValueError, match='Sum of input fractions must equal 1.'):
@@ -157,12 +157,8 @@ def test_random_bbox_splitting() -> None:
 
 
 def test_random_grid_cell_assignment() -> None:
-    ds = CustomGeoDataset(
-        [
-            (BoundingBox(0, 12, 0, 12, 0, 0), 'a'),
-            (BoundingBox(12, 24, 0, 12, 0, 0), 'b'),
-        ]
-    )
+    geometry = [shapely.box(0, 0, 12, 12), shapely.box(12, 0, 24, 12)]
+    ds = CustomGeoDataset(geometry=geometry)
 
     train_ds, val_ds, test_ds = random_grid_cell_assignment(
         ds, fractions=[1 / 2, 1 / 4, 1 / 4], grid_size=5
@@ -185,7 +181,6 @@ def test_random_grid_cell_assignment() -> None:
     # Test __get_item__
     x = train_ds[train_ds.bounds]
     assert isinstance(x, dict)
-    assert isinstance(x['content'], str)
 
     # Test invalid input fractions
     with pytest.raises(ValueError, match='Sum of input fractions must equal 1.'):
@@ -199,21 +194,20 @@ def test_random_grid_cell_assignment() -> None:
 
 
 def test_roi_split() -> None:
-    ds = CustomGeoDataset(
-        [
-            (BoundingBox(0, 1, 0, 1, 0, 0), 'a'),
-            (BoundingBox(1, 2, 0, 1, 0, 0), 'b'),
-            (BoundingBox(2, 3, 0, 1, 0, 0), 'c'),
-            (BoundingBox(3, 4, 0, 1, 0, 0), 'd'),
-        ]
-    )
+    geometry = [
+        shapely.box(0, 0, 1, 1),
+        shapely.box(1, 0, 2, 1),
+        shapely.box(2, 0, 3, 1),
+        shapely.box(3, 0, 4, 1),
+    ]
+    ds = CustomGeoDataset(geometry=geometry)
 
     train_ds, val_ds, test_ds = roi_split(
         ds,
         rois=[
-            BoundingBox(0, 2, 0, 1, 0, 0),
-            BoundingBox(2, 3.5, 0, 1, 0, 0),
-            BoundingBox(3.5, 4, 0, 1, 0, 0),
+            BoundingBox(0, 2, 0, 1, MINT, MAXT),
+            BoundingBox(2, 3.5, 0, 1, MINT, MAXT),
+            BoundingBox(3.5, 4, 0, 1, MINT, MAXT),
         ],
     )
 
@@ -234,7 +228,6 @@ def test_roi_split() -> None:
     # Test __get_item__
     x = train_ds[train_ds.bounds]
     assert isinstance(x, dict)
-    assert isinstance(x['content'], str)
 
     # Test invalid input rois
     with pytest.raises(ValueError, match="ROIs in input rois can't overlap."):
@@ -257,14 +250,23 @@ def test_roi_split() -> None:
 def test_time_series_split(
     lengths: Sequence[tuple[int, int] | int | float], expected_lengths: Sequence[int]
 ) -> None:
-    ds = CustomGeoDataset(
+    geometry = [
+        shapely.box(0, 0, 1, 1),
+        shapely.box(0, 0, 1, 1),
+        shapely.box(0, 0, 1, 1),
+        shapely.box(0, 0, 1, 1),
+    ]
+    index = pd.IntervalIndex.from_tuples(
         [
-            (BoundingBox(0, 1, 0, 1, 0, 10), 'a'),
-            (BoundingBox(0, 1, 0, 1, 10, 20), 'b'),
-            (BoundingBox(0, 1, 0, 1, 20, 30), 'c'),
-            (BoundingBox(0, 1, 0, 1, 30, 40), 'd'),
-        ]
+            (datetime(2025, 4, 25), datetime(2025, 4, 26)),
+            (datetime(2025, 4, 26), datetime(2025, 4, 27)),
+            (datetime(2025, 4, 27), datetime(2025, 4, 28)),
+            (datetime(2025, 4, 28), datetime(2025, 4, 29)),
+        ],
+        closed='both',
+        name='datetime',
     )
+    ds = CustomGeoDataset(index, geometry)
 
     train_ds, val_ds, test_ds = time_series_split(ds, lengths)
 
@@ -284,7 +286,6 @@ def test_time_series_split(
     # Test __get_item__
     x = train_ds[train_ds.bounds]
     assert isinstance(x, dict)
-    assert isinstance(x['content'], str)
 
 
 def test_time_series_split_invalid_input() -> None:

torchgeo/datasets/splits.py

@@ -3,6 +3,7 @@
 
 """Dataset splitting utilities."""
 
+import itertools
 from collections.abc import Sequence
 from copy import deepcopy
 from itertools import accumulate
@@ -71,23 +72,12 @@ def random_bbox_assignment(
         lengths = _fractions_to_lengths(lengths, len(dataset))
     lengths = cast(Sequence[int], lengths)
 
-    hits = list(dataset.index.intersection(dataset.index.bounds, objects=True))
-
-    hits = [hits[i] for i in randperm(sum(lengths), generator=generator)]
-
-    new_indexes = [
-        Index(interleaved=False, properties=Property(dimension=3)) for _ in lengths
-    ]
-
-    for i, length in enumerate(lengths):
-        for j in range(length):
-            hit = hits.pop()
-            new_indexes[i].insert(j, hit.bounds, hit.object)
+    indices = randperm(sum(lengths), generator=generator)
 
     new_datasets = []
-    for index in new_indexes:
+    for offset, length in zip(itertools.accumulate(lengths), lengths):
         ds = deepcopy(dataset)
-        ds.index = index
+        ds.index = dataset.index.iloc[indices[offset - length : offset]]
         new_datasets.append(ds)
 
     return new_datasets