Add methods for extracting true footprint for sampling valid data only

torchgeo/datasets/agrifieldnet.py

@@ -7,7 +7,7 @@
 import pathlib
 import re
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, cast
+from typing import Any
 
 import matplotlib.pyplot as plt
 import torch
@@ -183,8 +183,7 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
         """
         assert isinstance(self.paths, str | pathlib.Path)
 
-        hits = self.index.intersection(tuple(query), objects=True)
-        filepaths = cast(list[Path], [hit.object for hit in hits])
+        filepaths = self.filespaths_intersecting_query(query)
 
         if not filepaths:
             raise IndexError(

torchgeo/datasets/chesapeake.py

@@ -8,7 +8,7 @@
 import pathlib
 import sys
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, cast
+from typing import Any
 
 import fiona
 import matplotlib.pyplot as plt
@@ -604,8 +604,7 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
         Raises:
             IndexError: if query is not found in the index
         """
-        hits = self.index.intersection(tuple(query), objects=True)
-        filepaths = cast(list[dict[str, str]], [hit.object for hit in hits])
+        filepaths = self.filespaths_intersecting_query(query)
 
         sample = {'image': [], 'mask': [], 'crs': self.crs, 'bounds': query}
 

torchgeo/datasets/enviroatlas.py

@@ -6,7 +6,7 @@
 import os
 import sys
 from collections.abc import Callable, Sequence
-from typing import Any, cast
+from typing import Any
 
 import fiona
 import matplotlib.pyplot as plt
@@ -344,8 +344,7 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
         Raises:
             IndexError: if query is not found in the index
         """
-        hits = self.index.intersection(tuple(query), objects=True)
-        filepaths = cast(list[dict[str, str]], [hit.object for hit in hits])
+        filepaths = self.filespaths_intersecting_query(query)
 
         sample = {'image': [], 'mask': [], 'crs': self.crs, 'bounds': query}
 

torchgeo/datasets/geo.py

@@ -27,6 +27,8 @@
 from rasterio.io import DatasetReader
 from rasterio.vrt import WarpedVRT
 from rtree.index import Index, Property
+from rtree.index import Item as IndexItem
+from shapely import MultiPolygon, Polygon
 from torch import Tensor
 from torch.utils.data import Dataset
 from torchvision.datasets import ImageFolder
@@ -35,10 +37,13 @@
 from .errors import DatasetNotFoundError
 from .utils import (
     BoundingBox,
+    IndexData,
     Path,
     array_to_tensor,
+    calc_valid_data_footprint_from_datasource,
     concat_samples,
     disambiguate_timestamp,
+    get_valid_footprint_between_datasets,
     merge_samples,
     path_is_vsi,
 )
@@ -322,6 +327,31 @@ def files(self) -> list[Path]:
         # Sort the output to enforce deterministic behavior.
         return sorted(files)
 
+    def filepath_for_hit(self, hit: IndexItem) -> Path:
+        """Utility method for fetching filepath from rtee index item.
+
+        Alleviates the type casting from the user.
+
+        Args:
+            hit: Item of rtree Index
+
+        Returns:
+            filepath for the hit
+        """
+        return cast(IndexData, hit.object)['filepath']
+
+    def filespaths_intersecting_query(self, query: BoundingBox) -> list[Path]:
+        """Find all filepaths that intersects with query.
+
+        Args:
+            query: BoundingBox to intersect with
+
+        Returns:
+            list of all filepaths in rtree that intersects with query
+        """
+        hits = self.index.intersection(tuple(query), objects=True)
+        return [self.filepath_for_hit(hit) for hit in hits]
+
 
 class RasterDataset(GeoDataset):
     """Abstract base class for :class:`GeoDataset` stored as raster files."""
@@ -462,6 +492,15 @@ def __init__(
                         if crs is None:
                             crs = src.crs
 
+                        valid_footprint = calc_valid_data_footprint_from_datasource(
+                            masks=src.read_masks(),
+                            src_crs=src.crs,
+                            src_transform=src.transform,
+                            raster_width=src.width,
+                            raster_resolution_x=src.res[0],
+                            dst_crs=crs,
+                        )
+
                         with WarpedVRT(src, crs=crs) as vrt:
                             minx, miny, maxx, maxy = vrt.bounds
                             if res is None:
@@ -482,7 +521,13 @@ def __init__(
                         _, maxt = disambiguate_timestamp(stop, self.date_format)
 
                     coords = (minx, maxx, miny, maxy, mint, maxt)
-                    self.index.insert(i, coords, filepath)
+
+                    index_object: IndexData = {
+                        'filepath': filepath,
+                        'valid_footprint': valid_footprint,
+                    }
+
+                    self.index.insert(i, coords, index_object)
                     i += 1
 
         if i == 0:
@@ -517,8 +562,7 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
         Raises:
             IndexError: if query is not found in the index
         """
-        hits = self.index.intersection(tuple(query), objects=True)
-        filepaths = cast(list[Path], [hit.object for hit in hits])
+        filepaths = self.filespaths_intersecting_query(query)
 
         if not filepaths:
             raise IndexError(
@@ -609,11 +653,21 @@ def _load_warp_file(self, filepath: Path) -> DatasetReader:
         Returns:
             file handle of warped VRT
         """
+        # Todo: need to update meta kwarg `nodata` if it is not set.
+        #  create memory-file?
+        #  https://rasterio.groups.io/g/main/topic/change_the_nodata_value_in_a/28801885?p=
         src = rasterio.open(filepath)
 
         # Only warp if necessary
         if src.crs != self.crs:
-            vrt = WarpedVRT(src, crs=self.crs)
+            valid_nodatavals = [val for val in src.nodatavals if val is not None]
+            if not valid_nodatavals:  # Case for Sentinel2 L1C
+                nodata = 0.0  # Is it safe to assume? For Sentinel2 L1C it is 0.0
+            else:
+                # Get the first valid nodata value.
+                # Usualy the same value for all bands
+                nodata = valid_nodatavals[0]
+            vrt = WarpedVRT(src, nodata=nodata, crs=self.crs)
             src.close()
             return vrt
         else:
@@ -708,7 +762,7 @@ def __init__(
                         date = match.group('date')
                         mint, maxt = disambiguate_timestamp(date, self.date_format)
                     coords = (minx, maxx, miny, maxy, mint, maxt)
-                    self.index.insert(i, coords, filepath)
+                    self.index.insert(i, coords, {'filepath': filepath})
                     i += 1
 
         if i == 0:
@@ -729,8 +783,7 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
         Raises:
             IndexError: if query is not found in the index
         """
-        hits = self.index.intersection(tuple(query), objects=True)
-        filepaths = [hit.object for hit in hits]
+        filepaths = self.filespaths_intersecting_query(query)
 
         if not filepaths:
             raise IndexError(
@@ -990,14 +1043,22 @@ def _merge_dataset_indices(self) -> None:
         """Create a new R-tree out of the individual indices from two datasets."""
         i = 0
         ds1, ds2 = self.datasets
+
         for hit1 in ds1.index.intersection(ds1.index.bounds, objects=True):
             for hit2 in ds2.index.intersection(hit1.bounds, objects=True):
                 box1 = BoundingBox(*hit1.bounds)
                 box2 = BoundingBox(*hit2.bounds)
-                box3 = box1 & box2
+                box_intersection = box1 & box2
                 # Skip 0 area overlap (unless 0 area dataset)
-                if box3.area > 0 or box1.area == 0 or box2.area == 0:
-                    self.index.insert(i, tuple(box3))
+                if box_intersection.area > 0 or box1.area == 0 or box2.area == 0:
+                    valid_footprint: Polygon | MultiPolygon | None = (
+                        get_valid_footprint_between_datasets(
+                            ds1.index if isinstance(ds1, RasterDataset) else None,
+                            ds2.index if isinstance(ds2, RasterDataset) else None,
+                            box_intersection,
+                        )
+                    )
+                    self.index.insert(i, tuple(box_intersection), valid_footprint)
                     i += 1
 
         if i == 0:

torchgeo/datasets/globbiomass.py

@@ -7,7 +7,7 @@
 import os
 import pathlib
 from collections.abc import Callable, Iterable
-from typing import Any, cast
+from typing import Any
 
 import matplotlib.pyplot as plt
 import torch
@@ -192,8 +192,7 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
         Raises:
             IndexError: if query is not found in the index
         """
-        hits = self.index.intersection(tuple(query), objects=True)
-        filepaths = cast(list[Path], [hit.object for hit in hits])
+        filepaths = self.filespaths_intersecting_query(query)
 
         if not filepaths:
             raise IndexError(

torchgeo/datasets/l7irish.py

@@ -8,7 +8,7 @@
 import pathlib
 import re
 from collections.abc import Callable, Iterable, Sequence
-from typing import Any, cast
+from typing import Any
 
 import matplotlib.pyplot as plt
 import torch
@@ -94,7 +94,7 @@ def __init__(
         filename_regex = re.compile(L7IrishImage.filename_regex, re.VERBOSE)
         index = Index(interleaved=False, properties=Property(dimension=3))
         for hit in self.index.intersection(self.index.bounds, objects=True):
-            dirname = os.path.dirname(cast(Path, hit.object))
+            dirname = os.path.dirname(self.filepath_for_hit(hit))
             image = glob.glob(os.path.join(dirname, L7IrishImage.filename_glob))[0]
             minx, maxx, miny, maxy, mint, maxt = hit.bounds
             if match := re.match(filename_regex, os.path.basename(image)):

torchgeo/datasets/landcoverai.py

@@ -9,7 +9,7 @@
 import os
 from collections.abc import Callable
 from functools import lru_cache
-from typing import Any, cast
+from typing import Any
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -253,8 +253,7 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
         Raises:
             IndexError: if query is not found in the index
         """
-        hits = self.index.intersection(tuple(query), objects=True)
-        img_filepaths = cast(list[Path], [hit.object for hit in hits])
+        img_filepaths = self.filespaths_intersecting_query(query)
         mask_filepaths = [
             str(path).replace('images', 'masks') for path in img_filepaths
         ]

torchgeo/datasets/openbuildings.py

@@ -9,7 +9,7 @@
 import pathlib
 import sys
 from collections.abc import Callable, Iterable
-from typing import Any, cast
+from typing import Any
 
 import fiona
 import fiona.transform
@@ -304,8 +304,7 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
         Raises:
             IndexError: if query is not found in the index
         """
-        hits = self.index.intersection(tuple(query), objects=True)
-        filepaths = cast(list[Path], [hit.object for hit in hits])
+        filepaths = self.filespaths_intersecting_query(query)
 
         if not filepaths:
             raise IndexError(

torchgeo/datasets/south_africa_crop_type.py

@@ -7,7 +7,7 @@
 import pathlib
 import re
 from collections.abc import Callable, Iterable
-from typing import Any, cast
+from typing import Any
 
 import matplotlib.pyplot as plt
 import torch
@@ -164,8 +164,7 @@ def __getitem__(self, query: BoundingBox) -> dict[str, Any]:
         assert isinstance(self.paths, str | pathlib.Path)
 
         # Get all files matching the given query
-        hits = self.index.intersection(tuple(query), objects=True)
-        filepaths = cast(list[Path], [hit.object for hit in hits])
+        filepaths = self.filespaths_intersecting_query(query)
 
         if not filepaths:
             raise IndexError(