rslearn 0.0.1__py3-none-any.whl → 0.0.21__py3-none-any.whl

rslearn/train/transforms/select_bands.py

@@ -0,0 +1,76 @@
+"""The SelectBands transform."""
+
+from typing import Any
+
+from rslearn.train.model_context import RasterImage
+
+from .transform import Transform, read_selector, write_selector
+
+
+class SelectBands(Transform):
+    """Select a subset of bands from an image."""
+
+    def __init__(
+        self,
+        band_indices: list[int],
+        input_selector: str = "image",
+        output_selector: str = "image",
+        num_bands_per_timestep: int | None = None,
+    ):
+        """Initialize a new Concatenate.
+
+        Args:
+            band_indices: the bands to select.
+            input_selector: the selector to read the input image.
+            output_selector: the output selector under which to save the output image.
+            num_bands_per_timestep: the number of bands per image, to distinguish
+                between stacked images in an image time series. If set, then the
+                band_indices are selected for each image in the time series.
+        """
+        super().__init__()
+        self.input_selector = input_selector
+        self.output_selector = output_selector
+        self.band_indices = band_indices
+        self.num_bands_per_timestep = num_bands_per_timestep
+
+    def forward(
+        self, input_dict: dict[str, Any], target_dict: dict[str, Any]
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Apply concatenation over the inputs and targets.
+
+        Args:
+            input_dict: the input
+            target_dict: the target
+
+        Returns:
+            normalized (input_dicts, target_dicts) tuple
+        """
+        image = read_selector(input_dict, target_dict, self.input_selector)
+        num_bands_per_timestep = (
+            self.num_bands_per_timestep
+            if self.num_bands_per_timestep is not None
+            else image.shape[0]
+        )
+        if isinstance(image, RasterImage):
+            assert num_bands_per_timestep == image.shape[0], (
+                "Expect a seperate dimension for timesteps in RasterImages."
+            )
+
+        if image.shape[0] % num_bands_per_timestep != 0:
+            raise ValueError(
+                f"channel dimension {image.shape[0]} is not multiple of bands per timestep {num_bands_per_timestep}"
+            )
+
+        # Copy the band indices for each timestep in the input.
+        wanted_bands: list[int] = []
+        for start_channel_idx in range(0, image.shape[0], num_bands_per_timestep):
+            wanted_bands.extend(
+                [(start_channel_idx + band_idx) for band_idx in self.band_indices]
+            )
+
+        if isinstance(image, RasterImage):
+            image.image = image.image[wanted_bands]
+        else:
+            image = image[wanted_bands]
+        write_selector(input_dict, target_dict, self.output_selector, image)
+        return input_dict, target_dict

rslearn/train/transforms/sentinel1.py

@@ -0,0 +1,75 @@
+"""Transforms related to Sentinel-1 data."""
+
+from typing import Any
+
+import torch
+
+from rslearn.train.model_context import RasterImage
+
+from .transform import Transform
+
+
+class Sentinel1ToDecibels(Transform):
+    """Convert Sentinel-1 data from raw intensity to or from decibels."""
+
+    def __init__(
+        self,
+        selectors: list[str] = ["image"],
+        from_decibels: bool = False,
+        epsilon: float = 1e-6,
+    ):
+        """Initialize a new Sentinel1ToDecibels.
+
+        Args:
+            selectors: the input selectors to apply the transform on.
+            from_decibels: convert from decibels to intensities instead of intensity to
+                decibels.
+            epsilon: when converting to decibels, clip the intensities to this minimum
+                value to avoid log issues. This is mostly to avoid pixels that have no
+                data with no data value being 0.
+        """
+        super().__init__()
+        self.selectors = selectors
+        self.from_decibels = from_decibels
+        self.epsilon = epsilon
+
+    def apply_image(
+        self, image: torch.Tensor | RasterImage
+    ) -> torch.Tensor | RasterImage:
+        """Normalize the specified image.
+
+        Args:
+            image: the image to transform.
+        """
+        if isinstance(image, torch.Tensor):
+            image_to_process = image
+        else:
+            image_to_process = image.image
+        if self.from_decibels:
+            # Decibels to linear scale.
+            image_to_process = torch.pow(10.0, image_to_process / 10.0)
+        else:
+            # Linear scale to decibels.
+            image_to_process = 10 * torch.log10(
+                torch.clamp(image_to_process, min=self.epsilon)
+            )
+        if isinstance(image, torch.Tensor):
+            return image_to_process
+        else:
+            image.image = image_to_process
+            return image
+
+    def forward(
+        self, input_dict: dict[str, Any], target_dict: dict[str, Any]
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Apply normalization over the inputs and targets.
+
+        Args:
+            input_dict: the input
+            target_dict: the target
+
+        Returns:
+            normalized (input_dicts, target_dicts) tuple
+        """
+        self.apply_fn(self.apply_image, input_dict, target_dict, self.selectors)
+        return input_dict, target_dict

rslearn/train/transforms/transform.py

@@ -6,96 +6,115 @@ from typing import Any
 import torch
 
 
-class Transform(torch.nn.Module):
-    """An rslearn transform.
-
-    Provides helper functions for subclasses to select input and target keys and to
-    transform them.
-    """
+def get_dict_and_subselector(
+    input_dict: dict[str, Any], target_dict: dict[str, Any], selector: str
+) -> tuple[dict[str, Any], str]:
+    """Determine whether to use input or target dict, and the sub-selector.
 
-    def get_dict_and_subselector(
-        self, input_dict: dict[str, Any], target_dict: dict[str, Any], selector: str
-    ) -> tuple[dict[str, Any], str]:
-        """Determine whether to use input or target dict, and the sub-selector.
+    For example, if the selector is "input/x", then we use input dict and the
+    sub-selector is "x".
 
-        For example, if the selector is "input/x", then we use input dict and the
-        sub-selector is "x".
+    If neither input/ nor target/ prefixes are present, then we assume it is for
+    input dict.
 
-        If neither input/ nor target/ prefixes are present, then we assume it is for
-        input dict.
+    Args:
+        input_dict: the input dict
+        target_dict: the target dict
+        selector: the full selector configured by the user
 
-        Args:
-            input_dict: the input dict
-            target_dict: the target dict
-            selector: the full selector configured by the user
+    Returns:
+        a tuple (referenced dict, sub-selector string)
+    """
+    input_prefix = "input/"
+    target_prefix = "target/"
 
-        Returns:
-            a tuple (referenced dict, sub-selector string)
-        """
-        input_prefix = "input/"
-        target_prefix = "target/"
+    if selector.startswith(input_prefix):
+        d = input_dict
+        selector = selector[len(input_prefix) :]
+    elif selector.startswith(target_prefix):
+        d = target_dict
+        selector = selector[len(target_prefix) :]
+    else:
+        d = input_dict
 
-        if selector.startswith(input_prefix):
-            d = input_dict
-            selector = selector[len(input_prefix) :]
-        elif selector.startswith(target_prefix):
-            d = target_dict
-            selector = selector[len(target_prefix) :]
-        else:
-            d = input_dict
+    return d, selector
 
-        return d, selector
 
-    def read_selector(
-        self, input_dict: dict[str, Any], target_dict: dict[str, Any], selector: str
-    ) -> Any:
-        """Read the item referenced by the selector.
+def read_selector(
+    input_dict: dict[str, Any], target_dict: dict[str, Any], selector: str
+) -> Any:
+    """Read the item referenced by the selector.
 
-        Args:
-            input_dict: the input dict
-            target_dict: the target dict
-            selector: the selector specifying the item to read
+    Args:
+        input_dict: the input dict
+        target_dict: the target dict
+        selector: the selector specifying the item to read
 
-        Returns:
-            the item specified by the selector
-        """
-        d, selector = self.get_dict_and_subselector(input_dict, target_dict, selector)
-        parts = selector.split("/")
-        cur = d
-        for part in parts:
-            cur = cur[part]
-        return cur
-
-    def write_selector(
-        self,
-        input_dict: dict[str, Any],
-        target_dict: dict[str, Any],
-        selector: str,
-        v: Any,
-    ):
-        """Write the item to the specified selector.
-
-        Args:
-            input_dict: the input dict
-            target_dict: the target dict
-            selector: the selector specifying the item to write
-            v: the value to write
-        """
-        d, selector = self.get_dict_and_subselector(input_dict, target_dict, selector)
+    Returns:
+        the item specified by the selector
+    """
+    d, selector = get_dict_and_subselector(input_dict, target_dict, selector)
+    parts = selector.split("/") if selector else []
+    cur = d
+    for part in parts:
+        cur = cur[part]
+    return cur
+
+
+def write_selector(
+    input_dict: dict[str, Any],
+    target_dict: dict[str, Any],
+    selector: str,
+    v: Any,
+) -> None:
+    """Write the item to the specified selector.
+
+    Args:
+        input_dict: the input dict
+        target_dict: the target dict
+        selector: the selector specifying the item to write
+        v: the value to write
+    """
+    d, selector = get_dict_and_subselector(input_dict, target_dict, selector)
+    if selector:
         parts = selector.split("/")
         cur = d
         for part in parts[:-1]:
             cur = cur[part]
         cur[parts[-1]] = v
+    else:
+        # If the selector references the input or target dictionary directly, then we
+        # have a special case where instead of overwriting with v, we replace the keys
+        # with those in v. v must be a dictionary here, not a tensor, since otherwise
+        # it wouldn't match the type of the input or target dictionary.
+        if not isinstance(v, dict):
+            raise ValueError(
+                "when directly specifying the input or target dict, expected the value to be a dict"
+            )
+        if d == v:
+            # This may happen if the writer did not make a copy of the dictionary. In
+            # this case the code below would not update d correctly since it would also
+            # clear v.
+            return
+        d.clear()
+        d.update(v)
+
+
+class Transform(torch.nn.Module):
+    """An rslearn transform.
+
+    Provides helper functions for subclasses to select input and target keys and to
+    transform them.
+    """
 
     def apply_fn(
         self,
-        fn: Callable[[Any], Any],
+        fn: Callable,
         input_dict: dict[str, Any],
         target_dict: dict[str, Any],
         selectors: list[str],
         **kwargs: dict[str, Any],
-    ):
+    ) -> None:
         """Apply the specified function on the selectors in input/target dicts.
 
         Args:
@@ -106,9 +125,9 @@ class Transform(torch.nn.Module):
             kwargs: additional arguments to pass to the function
         """
         for selector in selectors:
-            v = self.read_selector(input_dict, target_dict, selector)
+            v = read_selector(input_dict, target_dict, selector)
             v = fn(v, **kwargs)
-            self.write_selector(input_dict, target_dict, selector, v)
+            write_selector(input_dict, target_dict, selector, v)
 
 
 class Identity(Transform):

rslearn/utils/__init__.py

@@ -1,7 +1,6 @@
 """rslearn utilities."""
 
-import logging
-import os
+from rslearn.log_utils import get_logger
 
 from .feature import Feature
 from .geometry import (
@@ -14,10 +13,8 @@ from .geometry import (
 from .get_utm_ups_crs import get_utm_ups_crs
 from .grid_index import GridIndex
 from .time import daterange
-from .utils import open_atomic
 
-logger = logging.getLogger(__name__)
-logger.setLevel(os.environ.get("RSLEARN_LOGLEVEL", "INFO").upper())
+logger = get_logger(__name__)
 
 __all__ = (
     "Feature",
@@ -29,6 +26,5 @@ __all__ = (
     "get_utm_ups_crs",
     "is_same_resolution",
     "logger",
-    "open_atomic",
     "shp_intersects",
 )

rslearn/utils/array.py

@@ -1,17 +1,19 @@
 """Array util functions."""
 
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy.typing as npt
-import torch
+
+if TYPE_CHECKING:
+    import torch
 
 
 def copy_spatial_array(
-    src: torch.Tensor | npt.NDArray[Any],
-    dst: torch.Tensor | npt.NDArray[Any],
+    src: "torch.Tensor | npt.NDArray[Any]",
+    dst: "torch.Tensor | npt.NDArray[Any]",
     src_offset: tuple[int, int],
     dst_offset: tuple[int, int],
-):
+) -> None:
     """Copy image content from a source array onto a destination array.
 
     The source and destination might be in the same coordinate system. Only the portion
@@ -59,4 +61,4 @@ def copy_spatial_array(
             src_col_offset : src_col_offset + col_overlap,
         ]
     else:
-        assert False
+        raise ValueError(f"Unsupported src shape: {src.shape}")

rslearn/utils/feature.py

@@ -11,7 +11,7 @@ from .geometry import Projection, STGeometry
 class Feature:
     """A GeoJSON-like feature that contains one vector geometry."""
 
-    def __init__(self, geometry: STGeometry, properties: dict[str, Any] | None = {}):
+    def __init__(self, geometry: STGeometry, properties: dict[str, Any] = {}):
         """Initialize a new Feature.
 
         Args:
@@ -41,7 +41,7 @@ class Feature:
         return Feature(self.geometry.to_projection(projection), self.properties)
 
     @staticmethod
-    def from_geojson(projection: Projection, d: dict[str, Any]):
+    def from_geojson(projection: Projection, d: dict[str, Any]) -> "Feature":
         """Construct a Feature from a GeoJSON encoding.
 
         Args:

rslearn/utils/fsspec.py

@@ -6,9 +6,15 @@ from collections.abc import Generator
 from contextlib import contextmanager
 from typing import Any
 
+import rasterio
+import rasterio.io
 from fsspec.implementations.local import LocalFileSystem
 from upath import UPath
 
+from rslearn.log_utils import get_logger
+
+logger = get_logger(__name__)
+
 
 @contextmanager
 def get_upath_local(
@@ -65,7 +71,7 @@ def join_upath(path: UPath, suffix: str) -> UPath:
 
 
 @contextmanager
-def open_atomic(path: UPath, *args: list[Any], **kwargs: dict[str, Any]):
+def open_atomic(path: UPath, *args: Any, **kwargs: Any) -> Generator[Any, None, None]:
     """Open a path for atomic writing.
 
     If it is local filesystem, we will write to a temporary file, and rename it to the
@@ -79,11 +85,94 @@ def open_atomic(path: UPath, *args: list[Any], **kwargs: dict[str, Any]):
         **kwargs: any valid keyword arguments for :code:`open`
     """
     if isinstance(path.fs, LocalFileSystem):
+        logger.debug("open_atomic: writing atomically to local file at %s", path)
         tmppath = path.path + ".tmp." + str(os.getpid())
         with open(tmppath, *args, **kwargs) as file:
             yield file
         os.rename(tmppath, path.path)
 
     else:
+        logger.debug("open_atomic: writing to remote file at %s", path)
         with path.open(*args, **kwargs) as file:
             yield file
+
+
+@contextmanager
+def open_rasterio_upath_reader(
+    path: UPath, **kwargs: Any
+) -> Generator[rasterio.io.DatasetReader, None, None]:
+    """Open a raster for reading.
+
+    If the UPath is local, then we open with rasterio directly, since this is much
+    faster. Otherwise, we open the file stream first and then use rasterio with file
+    stream.
+
+    Args:
+        path: the path to read.
+        **kwargs: additional keyword arguments for :code:`rasterio.open`
+    """
+    if isinstance(path.fs, LocalFileSystem):
+        logger.debug("reading from local rasterio dataset at %s", path)
+        with rasterio.open(path.path, **kwargs) as raster:
+            yield raster
+
+    else:
+        logger.debug("reading from remote rasterio dataset at %s", path)
+        with path.open("rb") as f:
+            with rasterio.open(f, **kwargs) as raster:
+                yield raster
+
+
+@contextmanager
+def open_rasterio_upath_writer(
+    path: UPath, **kwargs: Any
+) -> Generator[rasterio.io.DatasetWriter, None, None]:
+    """Open a raster for writing.
+
+    If the UPath is local, then we open with rasterio directly, since this is much
+    faster. We also write atomically by writing to temporary file and then renaming it,
+    to avoid concurrency issues. Otherwise, we open the file stream first and then use
+    rasterio with file stream (and assume that it is object storage so the write will
+    be atomic).
+
+    Args:
+        path: the path to write.
+        **kwargs: additional keyword arguments for :code:`rasterio.open`
+    """
+    if isinstance(path.fs, LocalFileSystem):
+        logger.debug(
+            "open_rasterio_upath_writer: writing atomically to local rasterio dataset at %s",
+            path,
+        )
+        tmppath = path.path + ".tmp." + str(os.getpid())
+        with rasterio.open(tmppath, "w", **kwargs) as raster:
+            yield raster
+        os.rename(tmppath, path.path)
+
+    else:
+        logger.debug(
+            "open_rasterio_upath_writer: writing to remote rasterio dataset at %s", path
+        )
+        with path.open("wb") as f:
+            with rasterio.open(f, "w", **kwargs) as raster:
+                yield raster
+
+
+def get_relative_suffix(base_dir: UPath, fname: UPath) -> str:
+    """Get the suffix of fname relative to base_dir.
+
+    Args:
+        base_dir: the base directory.
+        fname: a filename within the base directory.
+
+    Returns:
+        the suffix on base_dir that would yield the given filename.
+    """
+    if not fname.path.startswith(base_dir.path):
+        raise ValueError(
+            f"filename {fname.path} must start with base directory {base_dir.path}"
+        )
+    suffix = fname.path[len(base_dir.path) :]
+    if suffix.startswith("/"):
+        suffix = suffix[1:]
+    return suffix