GeoDataset: rtree -> geopandas

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.4.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:'
         ):