GeoDataset: rtree -> geopandas (#2747)

* GeoDataset: rtree -> geopandas

* Update documented spatiotemporal indexing backend

* Re-add rtree dep until we are ready to remove

* EDDMapS: update __init__

* GBIF: update __init__

* iNaturalist: update __init__

* Fewer backwards-incompatible changes

* L7 Irish: pray for the best?

* MMFlood: pray for the best?

* Minor docstring change

* Intersection/Union

* Add missing import

* Raster/VectorDataset: update __getitem__

* GeoDataset: tests pass

* Intersection/UnionDataset: various fixes

* Use real datetime objects in tests

* RasterDataset: all tests now pass

* VectorDataset: all tests passing

* Fix remaining type hint issues in geo.py

* GlobBiomass: try to fix docs

* Fix intersection/union index reprojection

* Fix no intersection

* Line/point intersection now counts as intersection

* Silence geopandas/shapely warning

* rasterio.crs.CRS to pyproj.CRS

* sjoin -> overlay

* BoundingBox: convert from int to datetime

* Docs: link to pyproj.CRS

* VectorDataset: fix fiona/pyproj CRS compatibility

* CDL: fix tests

* CopernicusBench: fix tests

* EDDMapS: fix tests

* GBIF: fix getitem

* HySpecNet-11k: fix tests

* iNaturalist: fix tests

* IntersectionDataset: implement temporal intersection

* BoundingBox and shapely.box have different order

* Restore backwards compatibility for edge overlap

* Break backwards compatibility for point datasets

* Update util tests

* GeoSampler: port to geopandas

* RandomGeoSampler: port to geopandas

* Ruff

* GridGeoSampler: convert to geopandas

* PreChippedGeoSampler: port to geopandas

* BatchGeoSampler: port to geopandas

* RandomBatchGeoSampler: port to geopandas

* Simpler logic

* RandomGeoSampler: revert most backwards-incompatible changes

* Samplers: fix type hints

* random_bbox_assignment: port to geopandas

* random_bbox_splitting: port to geopandas

* random_grid_cell_assignment: port to geopandas

* roi_split: port to geopandas

* time_series_split: port to geopandas

* Splitters: fix type hints

* EDDMapS: fix tests

* GBIF: fix tests

* iNaturalist: fix tests

* Landsat/Sentinel: remove redundant length tests

* SSL4EO: fix tests

* SACT: port to geopandas

* L7 Irish: fix tests

* AgriFieldNet: port to geopandas

* Docs: correct link to pandas

* Datamodules: port tests to geopandas

* Ruff

* BoundingBox: update tests, type hints

* Tutorials: port to geopandas

* shapely.LineString did not exist in shapely 1.X

* shapely.Geometry did not exist in shapely 1.X

* Deps: bump shapely to 2+

* Tutorials: port to geopandas

* Open Buildings: port to geopandas

* Try geopandas 0.12.0

* Try geopandas 0.12.1

* Add geopandas min version

* GlobBiomass: port to geopandas

* LandCoverAI: port to geopandas

* MMFlood: port to geopandas

* EnviroAtlas: port to geopandas

* Chesapeake CVPR: port to geopandas

* mypy: fix all type hints

* Tutorials: port to geopandas

* L7 Irish: 1 + 1 != 3

* Fiona: switch from PROJ dict to WKT for greater compatibility

* Explicitly drop geom type conversions during intersection to silence warning

* Ruff

* IntersectionDataset: add more tests

* UnionDataset: increase code coverage

* random_bbox_splitting: 100% test coverage

* RandomGeoSampler: 100% test coverage

* index.sindex.query -> index.cx

Author: adamjstewart

docs/conf.py

@@ -113,10 +113,11 @@
     'kornia': ('https://kornia.readthedocs.io/en/stable/', None),
     'matplotlib': ('https://matplotlib.org/stable/', None),
     'numpy': ('https://numpy.org/doc/stable/', None),
+    'pandas': ('https://pandas.pydata.org/docs/', None),
+    'pyproj': ('https://pyproj4.github.io/pyproj/stable/', None),
     'python': ('https://docs.python.org/3', None),
     'lightning': ('https://lightning.ai/docs/pytorch/stable/', None),
     'rasterio': ('https://rasterio.readthedocs.io/en/stable/', None),
-    'rtree': ('https://rtree.readthedocs.io/en/stable/', None),
     'segmentation_models_pytorch': ('https://smp.readthedocs.io/en/stable/', None),
     'sklearn': ('https://scikit-learn.org/stable/', None),
     'timm': ('https://huggingface.co/docs/timm/main/en/', None),

docs/tutorials/custom_raster_dataset.ipynb

@@ -49,7 +49,7 @@
    "source": [
     "### GeoDataset\n",
     "\n",
-    "In TorchGeo, each `GeoDataset` uses an [R-tree](https://en.wikipedia.org/wiki/R-tree) to store the spatiotemporal bounding box of each file or data point. To simplify this process and reduce code duplication, we provide two subclasses of `GeoDataset`:\n",
+    "In TorchGeo, each `GeoDataset` uses a [GeoDataFrame](https://geopandas.org/en/stable/docs/reference/api/geopandas.GeoDataFrame.html) to store the spatiotemporal bounding box of each file or data point. To simplify this process and reduce code duplication, we provide two subclasses of `GeoDataset`:\n",
     "\n",
     "* `RasterDataset`: recursively search for raster files in a directory\n",
     "* `VectorDataset`: recursively search for vector files in a directory\n",

docs/tutorials/earth_surface_water.ipynb

@@ -492,10 +492,8 @@
     "    \"\"\"\n",
     "\n",
     "    # To avoid loading the entire dataset in memory, we will loop through each img\n",
-    "    # The filenames will be retrieved from the dataset's rtree index\n",
-    "    files = [\n",
-    "        item.object for item in dset.index.intersection(dset.index.bounds, objects=True)\n",
-    "    ]\n",
+    "    # The filenames will be retrieved from the dataset's GeoDataFrame index\n",
+    "    files = dset.index.filepath\n",
     "\n",
     "    # Resetting statistics\n",
     "    accum_mean = 0\n",

docs/tutorials/torchgeo.ipynb

@@ -70,8 +70,8 @@
    "source": [
     "import os\n",
     "import tempfile\n",
-    "from datetime import datetime\n",
     "\n",
+    "import pandas as pd\n",
     "from matplotlib import pyplot as plt\n",
     "from torch.utils.data import DataLoader\n",
     "\n",
@@ -273,11 +273,7 @@
    "source": [
     "### Spatiotemporal indexing\n",
     "\n",
-    "How did we do this? TorchGeo uses a data structure called an *R-tree* to store the spatiotemporal bounding box of every file in the dataset. \n",
-    "\n",
-    "![R-tree](https://raw.githubusercontent.com/davidmoten/davidmoten.github.io/master/resources/rtree-3d/plot2.png)\n",
-    "\n",
-    "TorchGeo extracts the spatial bounding box from the metadata of each file, and the timestamp from the filename. This geospatial and geotemporal metadata allows us to efficiently compute the intersection or union of two datasets. It also lets us quickly retrieve an image and corresponding mask for a particular location in space and time."
+    "How did we do this? TorchGeo uses a [GeoDataFrame](https://geopandas.org/en/stable/docs/reference/api/geopandas.GeoDataFrame.html) to store the spatiotemporal bounding box of every file in the dataset. TorchGeo extracts the spatial bounding box from the metadata of each file, and the timestamp from the filename. This geospatial and geotemporal metadata allows us to efficiently compute the intersection or union of two datasets. It also lets us quickly retrieve an image and corresponding mask for a particular location in space and time."
    ]
   },
   {
@@ -293,8 +289,8 @@
     "xmax = xmin + size * 30\n",
     "ymin = 4470000\n",
     "ymax = ymin + size * 30\n",
-    "tmin = datetime(2023, 1, 1).timestamp()\n",
-    "tmax = datetime(2023, 12, 31).timestamp()\n",
+    "tmin = pd.Timestamp(2023, 1, 1)\n",
+    "tmax = pd.Timestamp(2023, 12, 31)\n",
     "\n",
     "bbox = BoundingBox(xmin, xmax, ymin, ymax, tmin, tmax)\n",
     "sample = dataset[bbox]\n",

docs/user/metrics/features.csv

@@ -1,10 +1,10 @@
 Library,ML Backend,I/O Backend,Spatial Backend,Transform Backend,Datasets,Weights,CLI,GUI,Reprojection,STAC,Time Series
-`TorchGeo`_,PyTorch,"GDAL, h5py, laspy, NetCDF4, OpenCV, pandas, pillow, scipy, xarray",R-tree,Kornia,125,93,✅,❌,✅,❌,🚧
+`TorchGeo`_,PyTorch,"GDAL, h5py, laspy, NetCDF4, OpenCV, pandas, pillow, scipy, xarray",geopandas,Kornia,125,93,✅,❌,✅,❌,🚧
 `eo-learn`_,scikit-learn,"GDAL, OpenCV, pandas, scipy, Zarr",geopandas,numpy,0,0,❌,❌,✅,❌,🚧
 `Raster Vision`_,"PyTorch, TensorFlow*","GDAL, OpenCV, pandas, pillow, scipy, xarray",STAC,Albumentations,0,6,✅,❌,✅,✅,🚧
 `DeepForest`_,"PyTorch, TensorFlow*","GDAL, OpenCV, pandas, pillow, scipy",R-tree,Albumentations,0,4,❌,❌,❌,❌,❌
 `samgeo`_,PyTorch,"GDAL, OpenCV, pandas, xarray",geopandas,numpy,0,0,❌,✅,✅,❌,❌
-`TerraTorch`_,PyTorch,"GDAL, h5py, pandas, xarray",R-tree,Albumentations,27,1,✅,❌,✅,❌,🚧
+`TerraTorch`_,PyTorch,"GDAL, h5py, pandas, xarray",-,Albumentations,27,1,✅,❌,✅,❌,🚧
 `SITS`_,R Torch,GDAL,-,tidyverse,22,0,❌,❌,✅,✅,✅
 `srai`_,PyTorch,"pandas, polars, duckdb","geopandas, duckdb-spatial",-,0,0,❌,❌,❌,❌,❌
 `scikit-eo`_,"scikit-learn, TensorFlow","GDAL, pandas, scipy","geopandas",numpy,0,0,❌,❌,❌,❌,🚧

experiments/ssl4eo/sample_conus.py

@@ -11,7 +11,8 @@
 import fiona
 from rtree import index
 from sample_ssl4eo import create_bbox, km2deg
-from shapely.geometry import MultiPolygon, Point, shape
+from shapely import MultiPolygon, Point
+from shapely.geometry import shape
 from shapely.ops import unary_union
 from torchvision.datasets.utils import download_and_extract_archive
 from tqdm import tqdm

pyproject.toml

@@ -39,6 +39,8 @@ dependencies = [
     "einops>=0.3",
     # fiona 1.8.22+ required for Python 3.11 wheels
     "fiona>=1.8.22",
+    # geopandas 0.12.1+ required for Shapely 2+ support
+    "geopandas>=0.12.1",
     # kornia 0.7.4+ required for AugmentationSequential support for unknown keys
     "kornia>=0.7.4",
     # lightly 1.4.5+ required for LARS optimizer
@@ -63,12 +65,10 @@ dependencies = [
     # rasterio 1.4.0-1.4.2 lack support for merging WarpedVRT objects
     # https://github.com/rasterio/rasterio/issues/3196
     "rasterio>=1.3.3,!=1.4.0,!=1.4.1,!=1.4.2",
-    # rtree 1.0.1+ required for Python 3.11 wheels
-    "rtree>=1.0.1",
     # segmentation-models-pytorch 0.5+ required for new UnetDecoder API
     "segmentation-models-pytorch>=0.5",
-    # shapely 1.8.5+ required for Python 3.11 wheels
-    "shapely>=1.8.5",
+    # shapely 2+ required for shapely.Geometry
+    "shapely>=2",
     # timm 0.8+ required for timm.models.adapt_input_conv, 0.9.2 required by SMP
     "timm>=0.9.2",
     # torch 2+ required for Python 3.11 wheels
@@ -259,6 +259,8 @@ filterwarnings = [
     "ignore:Failing to pass a value to the 'type_params' parameter of 'typing.*' is deprecated:DeprecationWarning:pydantic",
     # https://github.com/Lightning-AI/pytorch-lightning/pull/20802
     "ignore::jsonargparse._deprecated.JsonargparseDeprecationWarning:lightning.pytorch.cli",
+    # https://github.com/geopandas/geopandas/pull/3453
+    "ignore:The 'shapely.geos' module is deprecated:DeprecationWarning:geopandas._compat",
 
     # Expected warnings
     # Lightning warns us about using num_workers=0, but it's faster on macOS

requirements/min-reqs.old

@@ -4,6 +4,7 @@ setuptools==77.0.1
 # install
 einops==0.3.0
 fiona==1.8.22
+geopandas==0.12.1
 kornia==0.7.4
 lightly==1.4.5
 lightning[pytorch-extra]==2.0.0
@@ -13,9 +14,8 @@ pandas==1.5.0
 pillow==9.2.0
 pyproj==3.4.0
 rasterio==1.3.11
-rtree==1.0.1
 segmentation-models-pytorch==0.5.0
-shapely==1.8.5
+shapely==2.0.0
 timm==0.9.2
 torch==2.0.0
 torchmetrics==1.2.0

requirements/required.txt

@@ -4,6 +4,7 @@ setuptools==80.7.0
 # install
 einops==0.8.1
 fiona==1.10.1
+geopandas==1.0.1
 kornia==0.8.1
 lightly==1.5.20
 lightning[pytorch-extra]==2.5.0.post0
@@ -13,7 +14,6 @@ pandas==2.2.3
 pillow==11.2.1
 pyproj==3.7.1
 rasterio==1.4.3
-rtree==1.4.0
 segmentation-models-pytorch==0.5.0
 shapely==2.0.7
 timm==1.0.15

tests/data/eurocrops/data.py

@@ -10,7 +10,8 @@
 
 import fiona
 from rasterio.crs import CRS
-from shapely.geometry import Polygon, mapping
+from shapely import Polygon
+from shapely.geometry import mapping
 
 # Size of example crop field polygon in projection units.
 # This is set to align with Sentinel-2 test data, which is a 128x128 image at 10

tests/data/openbuildings/data.py

@@ -10,7 +10,7 @@
 import os
 import shutil
 
-from shapely.geometry import Polygon
+from shapely import Polygon
 
 SIZE = 0.05
 

tests/datamodules/test_geo.py

@@ -4,12 +4,15 @@
 from typing import Any
 
 import matplotlib.pyplot as plt
+import pandas as pd
 import pytest
+import shapely
 import torch
 from _pytest.fixtures import SubRequest
+from geopandas import GeoDataFrame
 from lightning.pytorch import Trainer
 from matplotlib.figure import Figure
-from rasterio.crs import CRS
+from pyproj import CRS
 from torch import Tensor
 
 from torchgeo.datamodules import (
@@ -20,19 +23,23 @@
 from torchgeo.datasets import BoundingBox, GeoDataset, NonGeoDataset
 from torchgeo.samplers import RandomBatchGeoSampler, RandomGeoSampler
 
+MINT = pd.Timestamp(2025, 4, 24)
+MAXT = pd.Timestamp(2025, 4, 25)
+
 
 class CustomGeoDataset(GeoDataset):
     def __init__(
         self, split: str = 'train', length: int = 1, download: bool = False
     ) -> None:
-        super().__init__()
-        for i in range(length):
-            self.index.insert(i, (0, 1, 2, 3, 4, 5))
+        geometry = [shapely.box(0, 0, 1, 1)] * length
+        index = pd.IntervalIndex([pd.Interval(MINT, MAXT)] * length, name='datetime')
+        crs = CRS.from_epsg(4326)
+        self.index = GeoDataFrame(index=index, geometry=geometry, crs=crs)
         self.res = (1, 1)
 
     def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
         image = torch.arange(3 * 2 * 2, dtype=torch.float).view(3, 2, 2)
-        return {'image': image, 'crs': CRS.from_epsg(4326), 'bounds': query}
+        return {'image': image, 'crs': self.index.crs, 'bounds': query}
 
     def plot(self, *args: Any, **kwargs: Any) -> Figure:
         return plt.figure()

tests/datasets/test_agb_live_woody_density.py

@@ -8,8 +8,8 @@
 import pytest
 import torch
 import torch.nn as nn
+from pyproj import CRS
 from pytest import MonkeyPatch
-from rasterio.crs import CRS
 
 from torchgeo.datasets import (
     AbovegroundLiveWoodyBiomassDensity,

tests/datasets/test_agrifieldnet.py

@@ -5,11 +5,12 @@
 from pathlib import Path
 
 import matplotlib.pyplot as plt
+import pandas as pd
 import pytest
 import torch
 import torch.nn as nn
+from pyproj import CRS
 from pytest import MonkeyPatch
-from rasterio.crs import CRS
 
 from torchgeo.datasets import (
     AgriFieldNet,
@@ -69,7 +70,7 @@ def test_plot_prediction(self, dataset: AgriFieldNet) -> None:
         plt.close()
 
     def test_invalid_query(self, dataset: AgriFieldNet) -> None:
-        query = BoundingBox(0, 0, 0, 0, 0, 0)
+        query = BoundingBox(0, 0, 0, 0, pd.Timestamp.min, pd.Timestamp.min)
         with pytest.raises(
             IndexError, match='query: .* not found in index with bounds:'
         ):

tests/datasets/test_airphen.py

@@ -5,10 +5,11 @@
 from pathlib import Path
 
 import matplotlib.pyplot as plt
+import pandas as pd
 import pytest
 import torch
 import torch.nn as nn
-from rasterio.crs import CRS
+from pyproj import CRS
 
 from torchgeo.datasets import (
     Airphen,
@@ -55,7 +56,7 @@ def test_no_data(self, tmp_path: Path) -> None:
             Airphen(tmp_path)
 
     def test_invalid_query(self, dataset: Airphen) -> None:
-        query = BoundingBox(0, 0, 0, 0, 0, 0)
+        query = BoundingBox(0, 0, 0, 0, pd.Timestamp.min, pd.Timestamp.min)
         with pytest.raises(
             IndexError, match='query: .* not found in index with bounds:'
         ):

tests/datasets/test_astergdem.py

@@ -6,10 +6,11 @@
 from pathlib import Path
 
 import matplotlib.pyplot as plt
+import pandas as pd
 import pytest
 import torch
 import torch.nn as nn
-from rasterio.crs import CRS
+from pyproj import CRS
 
 from torchgeo.datasets import (
     AsterGDEM,
@@ -66,7 +67,7 @@ def test_plot_prediction(self, dataset: AsterGDEM) -> None:
         plt.close()
 
     def test_invalid_query(self, dataset: AsterGDEM) -> None:
-        query = BoundingBox(100, 100, 100, 100, 0, 0)
+        query = BoundingBox(100, 100, 100, 100, pd.Timestamp.min, pd.Timestamp.min)
         with pytest.raises(
             IndexError, match='query: .* not found in index with bounds:'
         ):

tests/datasets/test_cbf.py

@@ -5,11 +5,12 @@
 from pathlib import Path
 
 import matplotlib.pyplot as plt
+import pandas as pd
 import pytest
 import torch
 import torch.nn as nn
+from pyproj import CRS
 from pytest import MonkeyPatch
-from rasterio.crs import CRS
 
 from torchgeo.datasets import (
     BoundingBox,
@@ -76,7 +77,7 @@ def test_not_downloaded(self, tmp_path: Path) -> None:
             CanadianBuildingFootprints(tmp_path)
 
     def test_invalid_query(self, dataset: CanadianBuildingFootprints) -> None:
-        query = BoundingBox(2, 2, 2, 2, 2, 2)
+        query = BoundingBox(2, 2, 2, 2, pd.Timestamp.min, pd.Timestamp.min)
         with pytest.raises(
             IndexError, match='query: .* not found in index with bounds:'
         ):

tests/datasets/test_cdl.py

@@ -4,15 +4,15 @@
 import glob
 import os
 import shutil
-from datetime import datetime
 from pathlib import Path
 
 import matplotlib.pyplot as plt
+import pandas as pd
 import pytest
 import torch
 import torch.nn as nn
+from pyproj import CRS
 from pytest import MonkeyPatch
-from rasterio.crs import CRS
 
 from torchgeo.datasets import (
     CDL,
@@ -71,9 +71,9 @@ def test_or(self, dataset: CDL) -> None:
 
     def test_full_year(self, dataset: CDL) -> None:
         bbox = dataset.bounds
-        time = datetime(2023, 6, 1).timestamp()
+        time = pd.Timestamp(2023, 6, 1)
         query = BoundingBox(bbox.minx, bbox.maxx, bbox.miny, bbox.maxy, time, time)
-        next(dataset.index.intersection(tuple(query)))
+        dataset[query]
 
     def test_already_extracted(self, dataset: CDL) -> None:
         CDL(dataset.paths, years=[2023, 2022])
@@ -117,7 +117,7 @@ def test_not_downloaded(self, tmp_path: Path) -> None:
             CDL(tmp_path)
 
     def test_invalid_query(self, dataset: CDL) -> None:
-        query = BoundingBox(0, 0, 0, 0, 0, 0)
+        query = BoundingBox(0, 0, 0, 0, pd.Timestamp.min, pd.Timestamp.min)
         with pytest.raises(
             IndexError, match='query: .* not found in index with bounds:'
         ):

tests/datasets/test_chesapeake.py

@@ -6,12 +6,13 @@
 from pathlib import Path
 
 import matplotlib.pyplot as plt
+import pandas as pd
 import pytest
 import torch
 import torch.nn as nn
 from _pytest.fixtures import SubRequest
+from pyproj import CRS
 from pytest import MonkeyPatch
-from rasterio.crs import CRS
 
 from torchgeo.datasets import (
     BoundingBox,
@@ -83,7 +84,7 @@ def test_plot(self, dataset: ChesapeakeDC) -> None:
         plt.close()
 
     def test_invalid_query(self, dataset: ChesapeakeDC) -> None:
-        query = BoundingBox(0, 0, 0, 0, 0, 0)
+        query = BoundingBox(0, 0, 0, 0, pd.Timestamp.min, pd.Timestamp.min)
         with pytest.raises(
             IndexError, match='query: .* not found in index with bounds:'
         ):
@@ -191,7 +192,7 @@ def test_not_downloaded(self, tmp_path: Path) -> None:
             ChesapeakeCVPR(tmp_path, checksum=True)
 
     def test_out_of_bounds_query(self, dataset: ChesapeakeCVPR) -> None:
-        query = BoundingBox(0, 0, 0, 0, 0, 0)
+        query = BoundingBox(0, 0, 0, 0, pd.Timestamp.min, pd.Timestamp.min)
         with pytest.raises(
             IndexError, match='query: .* not found in index with bounds:'
         ):