Source code for torchgeo.datasets.geo
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
"""Base classes for all :mod:`torchgeo` datasets."""
import abc
import functools
import glob
import os
import re
import sys
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, cast
import fiona
import fiona.transform
import matplotlib.pyplot as plt
import numpy as np
import pyproj
import rasterio
import rasterio.merge
import shapely
import torch
from rasterio.crs import CRS
from rasterio.io import DatasetReader
from rasterio.vrt import WarpedVRT
from rasterio.windows import from_bounds
from rtree.index import Index, Property
from torch import Tensor
from torch.utils.data import Dataset
from torchvision.datasets import ImageFolder
from torchvision.datasets.folder import default_loader as pil_loader
from .utils import BoundingBox, concat_samples, disambiguate_timestamp, merge_samples
# https://github.com/pytorch/pytorch/issues/60979
# https://github.com/pytorch/pytorch/pull/61045
Dataset.__module__ = "torch.utils.data"
ImageFolder.__module__ = "torchvision.datasets"
class GeoDataset(Dataset[Dict[str, Any]], abc.ABC):
"""Abstract base class for datasets containing geospatial information.
Geospatial information includes things like:
* coordinates (latitude, longitude)
* :term:`coordinate reference system (CRS)`
* resolution
:class:`GeoDataset` is a special class of datasets. Unlike :class:`VisionDataset`,
the presence of geospatial information allows two or more datasets to be combined
based on latitude/longitude. This allows users to do things like:
* Combine image and target labels and sample from both simultaneously
(e.g. Landsat and CDL)
* Combine datasets for multiple image sources for multimodal learning or data fusion
(e.g. Landsat and Sentinel)
These combinations require that all queries are present in *both* datasets,
and can be combined using an :class:`IntersectionDataset`:
.. code-block:: python
dataset = landsat & cdl
Users may also want to:
* Combine datasets for multiple image sources and treat them as equivalent
(e.g. Landsat 7 and Landsat 8)
* Combine datasets for disparate geospatial locations
(e.g. Chesapeake NY and PA)
These combinations require that all queries are present in *at least one* dataset,
and can be combined using a :class:`UnionDataset`:
.. code-block:: python
dataset = landsat7 | landsat8
"""
#: Resolution of the dataset in units of CRS.
res: float
_crs: CRS
# NOTE: according to the Python docs:
#
# * https://docs.python.org/3/library/exceptions.html#NotImplementedError
#
# the correct way to handle __add__ not being supported is to set it to None,
# not to return NotImplemented or raise NotImplementedError. The downside of
# this is that we have no way to explain to a user why they get an error and
# what they should do instead (use __and__ or __or__).
#: :class:`GeoDataset` addition can be ambiguous and is no longer supported.
#: Users should instead use the intersection or union operator.
__add__ = None # type: ignore[assignment]
[docs] def __init__(
self, transforms: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None
) -> None:
"""Initialize a new Dataset instance.
Args:
transforms: a function/transform that takes an input sample
and returns a transformed version
"""
self.transforms = transforms
# Create an R-tree to index the dataset
self.index = Index(interleaved=False, properties=Property(dimension=3))
[docs] @abc.abstractmethod
def __getitem__(self, query: BoundingBox) -> Dict[str, Any]:
"""Retrieve image/mask and metadata indexed by query.
Args:
query: (minx, maxx, miny, maxy, mint, maxt) coordinates to index
Returns:
sample of image/mask and metadata at that index
Raises:
IndexError: if query is not found in the index
"""
[docs] def __and__(self, other: "GeoDataset") -> "IntersectionDataset":
"""Take the intersection of two :class:`GeoDataset`.
Args:
other: another dataset
Returns:
a single dataset
Raises:
ValueError: if other is not a :class:`GeoDataset`
.. versionadded:: 0.2
"""
return IntersectionDataset(self, other)
[docs] def __or__(self, other: "GeoDataset") -> "UnionDataset":
"""Take the union of two GeoDatasets.
Args:
other: another dataset
Returns:
a single dataset
Raises:
ValueError: if other is not a :class:`GeoDataset`
.. versionadded:: 0.2
"""
return UnionDataset(self, other)
[docs] def __len__(self) -> int:
"""Return the number of files in the dataset.
Returns:
length of the dataset
"""
count: int = self.index.get_size()
return count
[docs] def __str__(self) -> str:
"""Return the informal string representation of the object.
Returns:
informal string representation
"""
return f"""\
{self.__class__.__name__} Dataset
type: GeoDataset
bbox: {self.bounds}
size: {len(self)}"""
# NOTE: This hack should be removed once the following issue is fixed:
# https://github.com/Toblerity/rtree/issues/87
[docs] def __getstate__(
self,
) -> Tuple[
Dict[Any, Any],
List[Tuple[int, Tuple[float, float, float, float, float, float], str]],
]:
"""Define how instances are pickled.
Returns:
the state necessary to unpickle the instance
"""
objects = self.index.intersection(self.index.bounds, objects=True)
tuples = [(item.id, item.bounds, item.object) for item in objects]
return self.__dict__, tuples
[docs] def __setstate__(
self,
state: Tuple[
Dict[Any, Any],
List[Tuple[int, Tuple[float, float, float, float, float, float], str]],
],
) -> None:
"""Define how to unpickle an instance.
Args:
state: the state of the instance when it was pickled
"""
attrs, tuples = state
self.__dict__.update(attrs)
for item in tuples:
self.index.insert(*item)
@property
def bounds(self) -> BoundingBox:
"""Bounds of the index.
Returns:
(minx, maxx, miny, maxy, mint, maxt) of the dataset
"""
return BoundingBox(*self.index.bounds)
@property
def crs(self) -> CRS:
""":term:`coordinate reference system (CRS)` for the dataset.
Returns:
the :term:`coordinate reference system (CRS)`
.. versionadded:: 0.2
"""
return self._crs
@crs.setter
def crs(self, new_crs: CRS) -> None:
"""Change the :term:`coordinate reference system (CRS)` of a GeoDataset.
If ``new_crs == self.crs``, does nothing, otherwise updates the R-tree index.
Args:
new_crs: new :term:`coordinate reference system (CRS)`
.. versionadded:: 0.2
"""
if new_crs == self._crs:
return
new_index = Index(interleaved=False, properties=Property(dimension=3))
project = pyproj.Transformer.from_crs(
pyproj.CRS(str(self._crs)), pyproj.CRS(str(new_crs)), always_xy=True
).transform
for hit in self.index.intersection(self.index.bounds, objects=True):
old_minx, old_maxx, old_miny, old_maxy, mint, maxt = hit.bounds
old_box = shapely.geometry.box(old_minx, old_miny, old_maxx, old_maxy)
new_box = shapely.ops.transform(project, old_box)
new_minx, new_miny, new_maxx, new_maxy = new_box.bounds
new_bounds = (new_minx, new_maxx, new_miny, new_maxy, mint, maxt)
new_index.insert(hit.id, new_bounds, hit.object)
self._crs = new_crs
self.index = new_index
class RasterDataset(GeoDataset):
"""Abstract base class for :class:`GeoDataset` stored as raster files."""
#: Glob expression used to search for files.
#:
#: This expression should be specific enough that it will not pick up files from
#: other datasets. It should not include a file extension, as the dataset may be in
#: a different file format than what it was originally downloaded as.
filename_glob = "*"
#: Regular expression used to extract date from filename.
#:
#: The expression should use named groups. The expression may contain any number of
#: groups. The following groups are specifically searched for by the base class:
#:
#: * ``date``: used to calculate ``mint`` and ``maxt`` for ``index`` insertion
#:
#: When :attr:`separate_files`` is True, the following additional groups are
#: searched for to find other files:
#:
#: * ``band``: replaced with requested band name
#: * ``resolution``: replaced with a glob character
filename_regex = ".*"
#: Date format string used to parse date from filename.
#:
#: Not used if :attr:`filename_regex` does not contain a ``date`` group.
date_format = "%Y%m%d"
#: True if dataset contains imagery, False if dataset contains mask
is_image = True
#: True if data is stored in a separate file for each band, else False.
separate_files = False
#: Names of all available bands in the dataset
all_bands: List[str] = []
#: Names of RGB bands in the dataset, used for plotting
rgb_bands: List[str] = []
#: If True, stretch the image from the 2nd percentile to the 98th percentile,
#: used for plotting
stretch = False
#: Color map for the dataset, used for plotting
cmap: Dict[int, Tuple[int, int, int, int]] = {}
[docs] def __init__(
self,
root: str,
crs: Optional[CRS] = None,
res: Optional[float] = None,
transforms: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
cache: bool = True,
) -> None:
"""Initialize a new Dataset instance.
Args:
root: root directory where dataset can be found
crs: :term:`coordinate reference system (CRS)` to warp to
(defaults to the CRS of the first file found)
res: resolution of the dataset in units of CRS
(defaults to the resolution of the first file found)
transforms: a function/transform that takes an input sample
and returns a transformed version
cache: if True, cache file handle to speed up repeated sampling
Raises:
FileNotFoundError: if no files are found in ``root``
"""
super().__init__(transforms)
self.root = root
self.cache = cache
# Populate the dataset index
i = 0
pathname = os.path.join(root, "**", self.filename_glob)
filename_regex = re.compile(self.filename_regex, re.VERBOSE)
for filepath in glob.iglob(pathname, recursive=True):
match = re.match(filename_regex, os.path.basename(filepath))
if match is not None:
try:
with rasterio.open(filepath) as src:
# See if file has a color map
if len(self.cmap) == 0:
try:
self.cmap = src.colormap(1)
except ValueError:
pass
if crs is None:
crs = src.crs
if res is None:
res = src.res[0]
with WarpedVRT(src, crs=crs) as vrt:
minx, miny, maxx, maxy = vrt.bounds
except rasterio.errors.RasterioIOError:
# Skip files that rasterio is unable to read
continue
else:
mint: float = 0
maxt: float = sys.maxsize
if "date" in match.groupdict():
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)
i += 1
if i == 0:
raise FileNotFoundError(
f"No {self.__class__.__name__} data was found in '{root}'"
)
self._crs = cast(CRS, crs)
self.res = cast(float, res)
[docs] def __getitem__(self, query: BoundingBox) -> Dict[str, Any]:
"""Retrieve image/mask and metadata indexed by query.
Args:
query: (minx, maxx, miny, maxy, mint, maxt) coordinates to index
Returns:
sample of image/mask and metadata at that index
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]
if not filepaths:
raise IndexError(
f"query: {query} not found in index with bounds: {self.bounds}"
)
if self.separate_files:
data_list: List[Tensor] = []
filename_regex = re.compile(self.filename_regex, re.VERBOSE)
for band in getattr(self, "bands", self.all_bands):
band_filepaths = []
for filepath in filepaths:
filename = os.path.basename(filepath)
directory = os.path.dirname(filepath)
match = re.match(filename_regex, filename)
if match:
if "date" in match.groupdict():
start = match.start("band")
end = match.end("band")
filename = filename[:start] + band + filename[end:]
if "resolution" in match.groupdict():
start = match.start("resolution")
end = match.end("resolution")
filename = filename[:start] + "*" + filename[end:]
filepath = glob.glob(os.path.join(directory, filename))[0]
band_filepaths.append(filepath)
data_list.append(self._merge_files(band_filepaths, query))
data = torch.cat(data_list)
else:
data = self._merge_files(filepaths, query)
key = "image" if self.is_image else "mask"
sample = {key: data, "crs": self.crs, "bbox": query}
if self.transforms is not None:
sample = self.transforms(sample)
return sample
def _merge_files(self, filepaths: Sequence[str], query: BoundingBox) -> Tensor:
"""Load and merge one or more files.
Args:
filepaths: one or more files to load and merge
query: (minx, maxx, miny, maxy, mint, maxt) coordinates to index
Returns:
image/mask at that index
"""
if self.cache:
vrt_fhs = [self._cached_load_warp_file(fp) for fp in filepaths]
else:
vrt_fhs = [self._load_warp_file(fp) for fp in filepaths]
bounds = (query.minx, query.miny, query.maxx, query.maxy)
if len(vrt_fhs) == 1:
src = vrt_fhs[0]
out_width = int(round((query.maxx - query.minx) / self.res))
out_height = int(round((query.maxy - query.miny) / self.res))
out_shape = (src.count, out_height, out_width)
dest = src.read(
out_shape=out_shape, window=from_bounds(*bounds, src.transform)
)
else:
dest, _ = rasterio.merge.merge(vrt_fhs, bounds, self.res)
# fix numpy dtypes which are not supported by pytorch tensors
if dest.dtype == np.uint16:
dest = dest.astype(np.int32)
elif dest.dtype == np.uint32:
dest = dest.astype(np.int64)
tensor = torch.tensor(dest)
return tensor
@functools.lru_cache(maxsize=128)
def _cached_load_warp_file(self, filepath: str) -> DatasetReader:
"""Cached version of :meth:`_load_warp_file`.
Args:
filepath: file to load and warp
Returns:
file handle of warped VRT
"""
return self._load_warp_file(filepath)
def _load_warp_file(self, filepath: str) -> DatasetReader:
"""Load and warp a file to the correct CRS and resolution.
Args:
filepath: file to load and warp
Returns:
file handle of warped VRT
"""
src = rasterio.open(filepath)
# Only warp if necessary
if src.crs != self.crs:
vrt = WarpedVRT(src, crs=self.crs)
src.close()
return vrt
else:
return src
[docs] def plot(self, data: Tensor) -> None:
"""Plot a data sample.
Args:
data: the data to plot
Raises:
AssertionError: if ``is_image`` is True and ``data`` has a different number
of channels than expected
"""
array = data.squeeze().numpy()
if self.is_image:
bands = getattr(self, "bands", self.all_bands)
assert array.shape[0] == len(bands)
# Only plot RGB bands
if bands and self.rgb_bands:
indices: "np.typing.NDArray[np.int_]" = np.array(
[bands.index(band) for band in self.rgb_bands]
)
array = array[indices]
# Convert from CxHxW to HxWxC
array = np.rollaxis(array, 0, 3)
if self.cmap:
# Convert from class labels to RGBA values
cmap: "np.typing.NDArray[np.int_]" = np.array(
[self.cmap[i] for i in range(len(self.cmap))]
)
array = cmap[array]
if self.stretch:
# Stretch to the range of 2nd to 98th percentile
per02 = np.percentile(array, 2)
per98 = np.percentile(array, 98)
array = (array - per02) / (per98 - per02)
array = np.clip(array, 0, 1)
# Plot the data
ax = plt.axes()
ax.imshow(array)
ax.axis("off")
plt.show()
plt.close()
class VectorDataset(GeoDataset):
"""Abstract base class for :class:`GeoDataset` stored as vector files."""
#: Glob expression used to search for files.
#:
#: This expression should be specific enough that it will not pick up files from
#: other datasets. It should not include a file extension, as the dataset may be in
#: a different file format than what it was originally downloaded as.
filename_glob = "*"
[docs] def __init__(
self,
root: str = "data",
crs: Optional[CRS] = None,
res: float = 0.0001,
transforms: Optional[Callable[[Dict[str, Any]], Dict[str, Any]]] = None,
) -> None:
"""Initialize a new Dataset instance.
Args:
root: root directory where dataset can be found
crs: :term:`coordinate reference system (CRS)` to warp to
(defaults to the CRS of the first file found)
res: resolution of the dataset in units of CRS
transforms: a function/transform that takes input sample and its target as
entry and returns a transformed version
Raises:
FileNotFoundError: if no files are found in ``root``
"""
super().__init__(transforms)
self.root = root
self.res = res
# Populate the dataset index
i = 0
pathname = os.path.join(root, "**", self.filename_glob)
for filepath in glob.iglob(pathname, recursive=True):
try:
with fiona.open(filepath) as src:
if crs is None:
crs = CRS.from_dict(src.crs)
minx, miny, maxx, maxy = src.bounds
(minx, maxx), (miny, maxy) = fiona.transform.transform(
src.crs, crs.to_dict(), [minx, maxx], [miny, maxy]
)
except fiona.errors.FionaValueError:
# Skip files that fiona is unable to read
continue
else:
mint = 0
maxt = sys.maxsize
coords = (minx, maxx, miny, maxy, mint, maxt)
self.index.insert(i, coords, filepath)
i += 1
if i == 0:
raise FileNotFoundError(
f"No {self.__class__.__name__} data was found in '{root}'"
)
self._crs = crs
[docs] def __getitem__(self, query: BoundingBox) -> Dict[str, Any]:
"""Retrieve image/mask and metadata indexed by query.
Args:
query: (minx, maxx, miny, maxy, mint, maxt) coordinates to index
Returns:
sample of image/mask and metadata at that index
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]
if not filepaths:
raise IndexError(
f"query: {query} not found in index with bounds: {self.bounds}"
)
shapes = []
for filepath in filepaths:
with fiona.open(filepath) as src:
# We need to know the bounding box of the query in the source CRS
(minx, maxx), (miny, maxy) = fiona.transform.transform(
self.crs.to_dict(),
src.crs,
[query.minx, query.maxx],
[query.miny, query.maxy],
)
# Filter geometries to those that intersect with the bounding box
for feature in src.filter(bbox=(minx, miny, maxx, maxy)):
# Warp geometries to requested CRS
shape = fiona.transform.transform_geom(
src.crs, self.crs.to_dict(), feature["geometry"]
)
shapes.append(shape)
# Rasterize geometries
width = (query.maxx - query.minx) / self.res
height = (query.maxy - query.miny) / self.res
transform = rasterio.transform.from_bounds(
query.minx, query.miny, query.maxx, query.maxy, width, height
)
if shapes:
masks = rasterio.features.rasterize(
shapes, out_shape=(int(height), int(width)), transform=transform
)
else:
# If no features are found in this query, return an empty mask
# with the default fill value and dtype used by rasterize
masks = np.zeros((int(height), int(width)), dtype=np.uint8)
sample = {"mask": torch.tensor(masks), "crs": self.crs, "bbox": query}
if self.transforms is not None:
sample = self.transforms(sample)
return sample
[docs] def plot(self, data: Tensor) -> None:
"""Plot a data sample.
Args:
data: the data to plot
"""
array = data.squeeze().numpy()
# Plot the image
ax = plt.axes()
ax.imshow(array)
ax.axis("off")
plt.show()
plt.close()
class VisionDataset(Dataset[Dict[str, Any]], abc.ABC):
"""Abstract base class for datasets lacking geospatial information.
This base class is designed for datasets with pre-defined image chips.
"""
[docs] @abc.abstractmethod
def __getitem__(self, index: int) -> Dict[str, Any]:
"""Return an index within the dataset.
Args:
index: index to return
Returns:
data and labels at that index
Raises:
IndexError: if index is out of range of the dataset
"""
[docs] @abc.abstractmethod
def __len__(self) -> int:
"""Return the length of the dataset.
Returns:
length of the dataset
"""
[docs] def __str__(self) -> str:
"""Return the informal string representation of the object.
Returns:
informal string representation
"""
return f"""\
{self.__class__.__name__} Dataset
type: VisionDataset
size: {len(self)}"""
class VisionClassificationDataset(VisionDataset, ImageFolder): # type: ignore[misc]
"""Abstract base class for classification datasets lacking geospatial information.
This base class is designed for datasets with pre-defined image chips which
are separated into separate folders per class.
"""
[docs] def __init__(
self,
root: str,
transforms: Optional[Callable[[Dict[str, Tensor]], Dict[str, Tensor]]] = None,
loader: Optional[Callable[[str], Any]] = pil_loader,
is_valid_file: Optional[Callable[[str], bool]] = None,
) -> None:
"""Initialize a new VisionClassificationDataset instance.
Args:
root: root directory where dataset can be found
transforms: a function/transform that takes input sample and its target as
entry and returns a transformed version
loader: a callable function which takes as input a path to an image and
returns a PIL Image or numpy array
is_valid_file: A function that takes the path of an Image file and checks if
the file is a valid file
"""
# When transform & target_transform are None, ImageFolder.__getitem__(index)
# returns a PIL.Image and int for image and label, respectively
super().__init__(
root=root,
transform=None,
target_transform=None,
loader=loader,
is_valid_file=is_valid_file,
)
# Must be set after calling super().__init__()
self.transforms = transforms
[docs] def __getitem__(self, index: int) -> Dict[str, Tensor]:
"""Return an index within the dataset.
Args:
index: index to return
Returns:
data and label at that index
"""
image, label = self._load_image(index)
sample = {"image": image, "label": label}
if self.transforms is not None:
sample = self.transforms(sample)
return sample
[docs] def __len__(self) -> int:
"""Return the number of data points in the dataset.
Returns:
length of the dataset
"""
return len(self.imgs)
def _load_image(self, index: int) -> Tuple[Tensor, Tensor]:
"""Load a single image and it's class label.
Args:
index: index to return
Returns:
the image
the image class label
"""
img, label = ImageFolder.__getitem__(self, index)
array: "np.typing.NDArray[np.int_]" = np.array(img)
tensor = torch.from_numpy(array)
# Convert from HxWxC to CxHxW
tensor = tensor.permute((2, 0, 1))
label = torch.tensor(label)
return tensor, label
class IntersectionDataset(GeoDataset):
"""Dataset representing the intersection of two GeoDatasets.
This allows users to do things like:
* Combine image and target labels and sample from both simultaneously
(e.g. Landsat and CDL)
* Combine datasets for multiple image sources for multimodal learning or data fusion
(e.g. Landsat and Sentinel)
These combinations require that all queries are present in *both* datasets,
and can be combined using an :class:`IntersectionDataset`:
.. code-block:: python
dataset = landsat & cdl
.. versionadded:: 0.2
"""
[docs] def __init__(
self,
dataset1: GeoDataset,
dataset2: GeoDataset,
collate_fn: Callable[
[Sequence[Dict[str, Any]]], Dict[str, Any]
] = concat_samples,
) -> None:
"""Initialize a new Dataset instance.
Args:
dataset1: the first dataset
dataset2: the second dataset
collate_fn: function used to collate samples
Raises:
ValueError: if either dataset is not a :class:`GeoDataset`
"""
super().__init__()
self.datasets = [dataset1, dataset2]
self.collate_fn = collate_fn
for ds in self.datasets:
if not isinstance(ds, GeoDataset):
raise ValueError("IntersectionDataset only supports GeoDatasets")
self._crs = dataset1.crs
self.res = dataset1.res
# Force dataset2 to have the same CRS/res as dataset1
if dataset1.crs != dataset2.crs:
print(
f"Converting {dataset2.__class__.__name__} CRS from "
f"{dataset2.crs} to {dataset1.crs}"
)
dataset2.crs = dataset1.crs
if dataset1.res != dataset2.res:
print(
f"Converting {dataset2.__class__.__name__} resolution from "
f"{dataset2.res} to {dataset1.res}"
)
dataset2.res = dataset1.res
# Merge dataset indices into a single index
self._merge_dataset_indices()
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)
self.index.insert(i, tuple(box1 & box2))
i += 1
[docs] def __getitem__(self, query: BoundingBox) -> Dict[str, Any]:
"""Retrieve image and metadata indexed by query.
Args:
query: (minx, maxx, miny, maxy, mint, maxt) coordinates to index
Returns:
sample of data/labels and metadata at that index
Raises:
IndexError: if query is not within bounds of the index
"""
if not query.intersects(self.bounds):
raise IndexError(
f"query: {query} not found in index with bounds: {self.bounds}"
)
# All datasets are guaranteed to have a valid query
samples = [ds[query] for ds in self.datasets]
return self.collate_fn(samples)
[docs] def __str__(self) -> str:
"""Return the informal string representation of the object.
Returns:
informal string representation
"""
return f"""\
{self.__class__.__name__} Dataset
type: IntersectionDataset
bbox: {self.bounds}
size: {len(self)}"""
class UnionDataset(GeoDataset):
"""Dataset representing the union of two GeoDatasets.
This allows users to do things like:
* Combine datasets for multiple image sources and treat them as equivalent
(e.g. Landsat 7 and Landsat 8)
* Combine datasets for disparate geospatial locations
(e.g. Chesapeake NY and PA)
These combinations require that all queries are present in *at least one* dataset,
and can be combined using a :class:`UnionDataset`:
.. code-block:: python
dataset = landsat7 | landsat8
.. versionadded:: 0.2
"""
[docs] def __init__(
self,
dataset1: GeoDataset,
dataset2: GeoDataset,
collate_fn: Callable[
[Sequence[Dict[str, Any]]], Dict[str, Any]
] = merge_samples,
) -> None:
"""Initialize a new Dataset instance.
Args:
dataset1: the first dataset
dataset2: the second dataset
collate_fn: function used to collate samples
Raises:
ValueError: if either dataset is not a :class:`GeoDataset`
"""
super().__init__()
self.datasets = [dataset1, dataset2]
self.collate_fn = collate_fn
for ds in self.datasets:
if not isinstance(ds, GeoDataset):
raise ValueError("UnionDataset only supports GeoDatasets")
self._crs = dataset1.crs
self.res = dataset1.res
# Force dataset2 to have the same CRS/res as dataset1
if dataset1.crs != dataset2.crs:
print(
f"Converting {dataset2.__class__.__name__} CRS from "
f"{dataset2.crs} to {dataset1.crs}"
)
dataset2.crs = dataset1.crs
if dataset1.res != dataset2.res:
print(
f"Converting {dataset2.__class__.__name__} resolution from "
f"{dataset2.res} to {dataset1.res}"
)
dataset2.res = dataset1.res
# Merge dataset indices into a single index
self._merge_dataset_indices()
def _merge_dataset_indices(self) -> None:
"""Create a new R-tree out of the individual indices from two datasets."""
i = 0
for ds in self.datasets:
hits = ds.index.intersection(ds.index.bounds, objects=True)
for hit in hits:
self.index.insert(i, hit.bounds)
i += 1
[docs] def __getitem__(self, query: BoundingBox) -> Dict[str, Any]:
"""Retrieve image and metadata indexed by query.
Args:
query: (minx, maxx, miny, maxy, mint, maxt) coordinates to index
Returns:
sample of data/labels and metadata at that index
Raises:
IndexError: if query is not within bounds of the index
"""
if not query.intersects(self.bounds):
raise IndexError(
f"query: {query} not found in index with bounds: {self.bounds}"
)
# Not all datasets are guaranteed to have a valid query
samples = []
for ds in self.datasets:
if ds.index.intersection(tuple(query)):
samples.append(ds[query])
return self.collate_fn(samples)
[docs] def __str__(self) -> str:
"""Return the informal string representation of the object.
Returns:
informal string representation
"""
return f"""\
{self.__class__.__name__} Dataset
type: UnionDataset
bbox: {self.bounds}
size: {len(self)}"""
