rslearn 0.0.19__py3-none-any.whl → 0.0.20__py3-none-any.whl

rslearn/train/tasks/classification.py

@@ -16,7 +16,12 @@ from torchmetrics.classification import (
 )
 
 from rslearn.models.component import FeatureVector, Predictor
-from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
+from rslearn.train.model_context import (
+    ModelContext,
+    ModelOutput,
+    RasterImage,
+    SampleMetadata,
+)
 from rslearn.utils import Feature, STGeometry
 
 from .task import BasicTask
@@ -99,7 +104,7 @@ class ClassificationTask(BasicTask):
 
     def process_inputs(
         self,
-        raw_inputs: dict[str, torch.Tensor | list[Feature]],
+        raw_inputs: dict[str, RasterImage | list[Feature]],
         metadata: SampleMetadata,
         load_targets: bool = True,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
@@ -118,6 +123,7 @@ class ClassificationTask(BasicTask):
             return {}, {}
 
         data = raw_inputs["targets"]
+        assert isinstance(data, list)
         for feat in data:
             if feat.properties is None:
                 continue

rslearn/train/tasks/detection.py

@@ -12,7 +12,7 @@ import torchmetrics.classification
 import torchvision
 from torchmetrics import Metric, MetricCollection
 
-from rslearn.train.model_context import SampleMetadata
+from rslearn.train.model_context import RasterImage, SampleMetadata
 from rslearn.utils import Feature, STGeometry
 
 from .task import BasicTask
@@ -127,7 +127,7 @@ class DetectionTask(BasicTask):
 
     def process_inputs(
         self,
-        raw_inputs: dict[str, torch.Tensor | list[Feature]],
+        raw_inputs: dict[str, RasterImage | list[Feature]],
         metadata: SampleMetadata,
         load_targets: bool = True,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
@@ -152,6 +152,7 @@ class DetectionTask(BasicTask):
         valid = 1
 
         data = raw_inputs["targets"]
+        assert isinstance(data, list)
         for feat in data:
             if feat.properties is None:
                 continue

rslearn/train/tasks/multi_task.py

@@ -3,10 +3,9 @@
 from typing import Any
 
 import numpy.typing as npt
-import torch
 from torchmetrics import Metric, MetricCollection
 
-from rslearn.train.model_context import SampleMetadata
+from rslearn.train.model_context import RasterImage, SampleMetadata
 from rslearn.utils import Feature
 
 from .task import Task
@@ -30,7 +29,7 @@ class MultiTask(Task):
 
     def process_inputs(
         self,
-        raw_inputs: dict[str, torch.Tensor | list[Feature]],
+        raw_inputs: dict[str, RasterImage | list[Feature]],
         metadata: SampleMetadata,
         load_targets: bool = True,
     ) -> tuple[dict[str, Any], dict[str, Any]]:

rslearn/train/tasks/per_pixel_regression.py

@@ -9,7 +9,12 @@ import torchmetrics
 from torchmetrics import Metric, MetricCollection
 
 from rslearn.models.component import FeatureMaps, Predictor
-from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
+from rslearn.train.model_context import (
+    ModelContext,
+    ModelOutput,
+    RasterImage,
+    SampleMetadata,
+)
 from rslearn.utils.feature import Feature
 
 from .task import BasicTask
@@ -42,7 +47,7 @@ class PerPixelRegressionTask(BasicTask):
 
     def process_inputs(
         self,
-        raw_inputs: dict[str, torch.Tensor],
+        raw_inputs: dict[str, RasterImage | list[Feature]],
         metadata: SampleMetadata,
         load_targets: bool = True,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
@@ -60,11 +65,15 @@ class PerPixelRegressionTask(BasicTask):
         if not load_targets:
             return {}, {}
 
-        assert raw_inputs["targets"].shape[0] == 1
-        labels = raw_inputs["targets"][0, :, :].float() * self.scale_factor
+        assert isinstance(raw_inputs["targets"], RasterImage)
+        assert raw_inputs["targets"].image.shape[0] == 1
+        assert raw_inputs["targets"].image.shape[1] == 1
+        labels = raw_inputs["targets"].image[0, 0, :, :].float() * self.scale_factor
 
         if self.nodata_value is not None:
-            valid = (raw_inputs["targets"][0, :, :] != self.nodata_value).float()
+            valid = (
+                raw_inputs["targets"].image[0, 0, :, :] != self.nodata_value
+            ).float()
         else:
             valid = torch.ones(labels.shape, dtype=torch.float32)
 

rslearn/train/tasks/regression.py

@@ -11,7 +11,12 @@ from PIL import Image, ImageDraw
 from torchmetrics import Metric, MetricCollection
 
 from rslearn.models.component import FeatureVector, Predictor
-from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
+from rslearn.train.model_context import (
+    ModelContext,
+    ModelOutput,
+    RasterImage,
+    SampleMetadata,
+)
 from rslearn.utils.feature import Feature
 from rslearn.utils.geometry import STGeometry
 
@@ -63,7 +68,7 @@ class RegressionTask(BasicTask):
 
     def process_inputs(
         self,
-        raw_inputs: dict[str, torch.Tensor | list[Feature]],
+        raw_inputs: dict[str, RasterImage | list[Feature]],
         metadata: SampleMetadata,
         load_targets: bool = True,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
@@ -82,6 +87,7 @@ class RegressionTask(BasicTask):
             return {}, {}
 
         data = raw_inputs["targets"]
+        assert isinstance(data, list)
         for feat in data:
             if feat.properties is None or self.filters is None:
                 continue

rslearn/train/tasks/segmentation.py

@@ -1,5 +1,6 @@
 """Segmentation task."""
 
+from collections.abc import Mapping
 from typing import Any
 
 import numpy as np
@@ -9,7 +10,13 @@ import torchmetrics.classification
 from torchmetrics import Metric, MetricCollection
 
 from rslearn.models.component import FeatureMaps, Predictor
-from rslearn.train.model_context import ModelContext, ModelOutput, SampleMetadata
+from rslearn.train.model_context import (
+    ModelContext,
+    ModelOutput,
+    RasterImage,
+    SampleMetadata,
+)
+from rslearn.utils import Feature
 
 from .task import BasicTask
 
@@ -108,7 +115,7 @@ class SegmentationTask(BasicTask):
 
     def process_inputs(
         self,
-        raw_inputs: dict[str, torch.Tensor],
+        raw_inputs: Mapping[str, RasterImage | list[Feature]],
         metadata: SampleMetadata,
         load_targets: bool = True,
     ) -> tuple[dict[str, Any], dict[str, Any]]:
@@ -126,8 +133,10 @@ class SegmentationTask(BasicTask):
         if not load_targets:
             return {}, {}
 
-        assert raw_inputs["targets"].shape[0] == 1
-        labels = raw_inputs["targets"][0, :, :].long()
+        assert isinstance(raw_inputs["targets"], RasterImage)
+        assert raw_inputs["targets"].image.shape[0] == 1
+        assert raw_inputs["targets"].image.shape[1] == 1
+        labels = raw_inputs["targets"].image[0, 0, :, :].long()
 
         if self.class_id_mapping is not None:
             new_labels = labels.clone()

rslearn/train/tasks/task.py

@@ -7,7 +7,7 @@ import numpy.typing as npt
 import torch
 from torchmetrics import MetricCollection
 
-from rslearn.train.model_context import SampleMetadata
+from rslearn.train.model_context import RasterImage, SampleMetadata
 from rslearn.utils import Feature
 
 
@@ -21,7 +21,7 @@ class Task:
 
     def process_inputs(
         self,
-        raw_inputs: dict[str, torch.Tensor | list[Feature]],
+        raw_inputs: dict[str, RasterImage | list[Feature]],
         metadata: SampleMetadata,
         load_targets: bool = True,
     ) -> tuple[dict[str, Any], dict[str, Any]]:

rslearn/train/transforms/concatenate.py

@@ -1,12 +1,23 @@
 """Concatenate bands across multiple image inputs."""
 
+from datetime import datetime
+from enum import Enum
 from typing import Any
 
 import torch
 
+from rslearn.train.model_context import RasterImage
+
 from .transform import Transform, read_selector, write_selector
 
 
+class ConcatenateDim(Enum):
+    """Enum for concatenation dimensions."""
+
+    CHANNEL = 0
+    TIME = 1
+
+
 class Concatenate(Transform):
     """Concatenate bands across multiple image inputs."""
 
@@ -14,6 +25,7 @@ class Concatenate(Transform):
         self,
         selections: dict[str, list[int]],
         output_selector: str,
+        concatenate_dim: ConcatenateDim | int = ConcatenateDim.TIME,
     ):
         """Initialize a new Concatenate.
 
@@ -21,10 +33,16 @@ class Concatenate(Transform):
             selections: map from selector to list of band indices in that input to
                 retain, or empty list to use all bands.
             output_selector: the output selector under which to save the concatenate image.
+            concatenate_dim: the dimension against which to concatenate the inputs
         """
         super().__init__()
         self.selections = selections
         self.output_selector = output_selector
+        self.concatenate_dim = (
+            concatenate_dim.value
+            if isinstance(concatenate_dim, ConcatenateDim)
+            else concatenate_dim
+        )
 
     def forward(
         self, input_dict: dict[str, Any], target_dict: dict[str, Any]
@@ -36,14 +54,36 @@ class Concatenate(Transform):
             target_dict: the target
 
         Returns:
-            normalized (input_dicts, target_dicts) tuple
+            concatenated (input_dicts, target_dicts) tuple. If one of the
+            specified inputs is a RasterImage, a RasterImage will be returned.
+            Otherwise it will be a torch.Tensor.
         """
         images = []
+        return_raster_image: bool = False
+        timestamps: list[tuple[datetime, datetime]] | None = None
         for selector, wanted_bands in self.selections.items():
             image = read_selector(input_dict, target_dict, selector)
-            if wanted_bands:
-                image = image[wanted_bands, :, :]
-            images.append(image)
-        result = torch.concatenate(images, dim=0)
+            if isinstance(image, torch.Tensor):
+                if wanted_bands:
+                    image = image[wanted_bands, :, :]
+                images.append(image)
+            elif isinstance(image, RasterImage):
+                return_raster_image = True
+                if wanted_bands:
+                    images.append(image.image[wanted_bands, :, :])
+                else:
+                    images.append(image.image)
+                if timestamps is None:
+                    if image.timestamps is not None:
+                        # assume all concatenated modalities have the same
+                        # number of timestamps
+                        timestamps = image.timestamps
+        if return_raster_image:
+            result = RasterImage(
+                torch.concatenate(images, dim=self.concatenate_dim),
+                timestamps=timestamps,
+            )
+        else:
+            result = torch.concatenate(images, dim=self.concatenate_dim)
         write_selector(input_dict, target_dict, self.output_selector, result)
         return input_dict, target_dict

rslearn/train/transforms/crop.py

@@ -5,6 +5,8 @@ from typing import Any
 import torch
 import torchvision
 
+from rslearn.train.model_context import RasterImage
+
 from .transform import Transform, read_selector
 
 
@@ -69,7 +71,9 @@ class Crop(Transform):
             "remove_from_top": remove_from_top,
         }
 
-    def apply_image(self, image: torch.Tensor, state: dict[str, Any]) -> torch.Tensor:
+    def apply_image(
+        self, image: RasterImage | torch.Tensor, state: dict[str, Any]
+    ) -> RasterImage | torch.Tensor:
         """Apply the sampled state on the specified image.
 
         Args:
@@ -80,13 +84,23 @@ class Crop(Transform):
         crop_size = state["crop_size"] * image.shape[-1] // image_shape[1]
         remove_from_left = state["remove_from_left"] * image.shape[-1] // image_shape[1]
         remove_from_top = state["remove_from_top"] * image.shape[-2] // image_shape[0]
-        return torchvision.transforms.functional.crop(
-            image,
-            top=remove_from_top,
-            left=remove_from_left,
-            height=crop_size,
-            width=crop_size,
-        )
+        if isinstance(image, RasterImage):
+            image.image = torchvision.transforms.functional.crop(
+                image.image,
+                top=remove_from_top,
+                left=remove_from_left,
+                height=crop_size,
+                width=crop_size,
+            )
+        else:
+            image = torchvision.transforms.functional.crop(
+                image,
+                top=remove_from_top,
+                left=remove_from_left,
+                height=crop_size,
+                width=crop_size,
+            )
+        return image
 
     def apply_boxes(self, boxes: Any, state: dict[str, bool]) -> torch.Tensor:
         """Apply the sampled state on the specified image.

rslearn/train/transforms/flip.py

@@ -4,6 +4,8 @@ from typing import Any
 
 import torch
 
+from rslearn.train.model_context import RasterImage
+
 from .transform import Transform
 
 
@@ -48,17 +50,23 @@ class Flip(Transform):
             "vertical": vertical,
         }
 
-    def apply_image(self, image: torch.Tensor, state: dict[str, bool]) -> torch.Tensor:
+    def apply_image(self, image: RasterImage, state: dict[str, bool]) -> RasterImage:
         """Apply the sampled state on the specified image.
 
         Args:
             image: the image to transform.
             state: the sampled state.
         """
-        if state["horizontal"]:
-            image = torch.flip(image, dims=[-1])
-        if state["vertical"]:
-            image = torch.flip(image, dims=[-2])
+        if isinstance(image, RasterImage):
+            if state["horizontal"]:
+                image.image = torch.flip(image.image, dims=[-1])
+            if state["vertical"]:
+                image.image = torch.flip(image.image, dims=[-2])
+        elif isinstance(image, torch.Tensor):
+            if state["horizontal"]:
+                image = torch.flip(image, dims=[-1])
+            if state["vertical"]:
+                image = torch.flip(image, dims=[-2])
         return image
 
     def apply_boxes(

rslearn/train/transforms/mask.py

@@ -2,6 +2,7 @@
 
 import torch
 
+from rslearn.train.model_context import RasterImage
 from rslearn.train.transforms.transform import Transform, read_selector
 
 
@@ -31,7 +32,9 @@ class Mask(Transform):
         self.mask_selector = mask_selector
         self.mask_value = mask_value
 
-    def apply_image(self, image: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    def apply_image(
+        self, image: torch.Tensor | RasterImage, mask: torch.Tensor | RasterImage
+    ) -> torch.Tensor | RasterImage:
         """Apply the mask on the image.
 
         Args:
@@ -42,6 +45,9 @@ class Mask(Transform):
             masked image
         """
         # Tile the mask to have same number of bands as the image.
+        if isinstance(mask, RasterImage):
+            mask = mask.image
+
         if image.shape[0] != mask.shape[0]:
             if mask.shape[0] != 1:
                 raise ValueError(
@@ -49,7 +55,10 @@ class Mask(Transform):
                 )
             mask = mask.repeat(image.shape[0], 1, 1)
 
-        image[mask == 0] = self.mask_value
+        if isinstance(image, torch.Tensor):
+            image[mask == 0] = self.mask_value
+        else:
+            image.image[mask == 0] = self.mask_value
         return image
 
     def forward(self, input_dict: dict, target_dict: dict) -> tuple[dict, dict]:

rslearn/train/transforms/normalize.py

@@ -4,6 +4,8 @@ from typing import Any
 
 import torch
 
+from rslearn.train.model_context import RasterImage
+
 from .transform import Transform
 
 
@@ -55,7 +57,9 @@ class Normalize(Transform):
         self.bands = torch.tensor(bands) if bands is not None else None
         self.num_bands = num_bands
 
-    def apply_image(self, image: torch.Tensor) -> torch.Tensor:
+    def apply_image(
+        self, image: torch.Tensor | RasterImage
+    ) -> torch.Tensor | RasterImage:
         """Normalize the specified image.
 
         Args:
@@ -63,7 +67,7 @@ class Normalize(Transform):
         """
 
         def _repeat_mean_and_std(
-            image_channels: int, num_bands: int | None
+            image_channels: int, num_bands: int | None, is_raster_image: bool
         ) -> tuple[torch.Tensor, torch.Tensor]:
             """Get mean and std tensor that are suitable for applying on the image."""
             # We only need to repeat the tensor if both of these are true:
@@ -74,9 +78,16 @@ class Normalize(Transform):
             if num_bands is None:
                 return self.mean, self.std
             num_images = image_channels // num_bands
-            return self.mean.repeat(num_images)[:, None, None], self.std.repeat(
-                num_images
-            )[:, None, None]
+            if is_raster_image:
+                # add an extra T dimension, CTHW
+                return self.mean.repeat(num_images)[
+                    :, None, None, None
+                ], self.std.repeat(num_images)[:, None, None, None]
+            else:
+                # add an extra T dimension, CTHW
+                return self.mean.repeat(num_images)[:, None, None], self.std.repeat(
+                    num_images
+                )[:, None, None]
 
         if self.bands is not None:
             # User has provided band indices to normalize.
@@ -96,20 +107,40 @@ class Normalize(Transform):
             # We use len(self.bands) here because that is how many bands per timestep
             # we are actually processing with the mean/std.
             mean, std = _repeat_mean_and_std(
-                image_channels=len(band_indices), num_bands=len(self.bands)
+                image_channels=len(band_indices),
+                num_bands=len(self.bands),
+                is_raster_image=isinstance(image, RasterImage),
             )
-            image[band_indices] = (image[band_indices] - mean) / std
-            if self.valid_min is not None:
-                image[band_indices] = torch.clamp(
-                    image[band_indices], min=self.valid_min, max=self.valid_max
-                )
+            if isinstance(image, torch.Tensor):
+                image[band_indices] = (image[band_indices] - mean) / std
+                if self.valid_min is not None:
+                    image[band_indices] = torch.clamp(
+                        image[band_indices], min=self.valid_min, max=self.valid_max
+                    )
+            else:
+                image.image[band_indices] = (image.image[band_indices] - mean) / std
+                if self.valid_min is not None:
+                    image.image[band_indices] = torch.clamp(
+                        image.image[band_indices],
+                        min=self.valid_min,
+                        max=self.valid_max,
+                    )
         else:
             mean, std = _repeat_mean_and_std(
-                image_channels=image.shape[0], num_bands=self.num_bands
+                image_channels=image.shape[0],
+                num_bands=self.num_bands,
+                is_raster_image=isinstance(image, RasterImage),
             )
-            image = (image - mean) / std
-            if self.valid_min is not None:
-                image = torch.clamp(image, min=self.valid_min, max=self.valid_max)
+            if isinstance(image, torch.Tensor):
+                image = (image - mean) / std
+                if self.valid_min is not None:
+                    image = torch.clamp(image, min=self.valid_min, max=self.valid_max)
+            else:
+                image.image = (image.image - mean) / std
+                if self.valid_min is not None:
+                    image.image = torch.clamp(
+                        image.image, min=self.valid_min, max=self.valid_max
+                    )
         return image
 
     def forward(

rslearn/train/transforms/pad.py

@@ -5,6 +5,8 @@ from typing import Any
 import torch
 import torchvision
 
+from rslearn.train.model_context import RasterImage
+
 from .transform import Transform
 
 
@@ -48,7 +50,9 @@ class Pad(Transform):
         """
         return {"size": torch.randint(low=self.size[0], high=self.size[1], size=())}
 
-    def apply_image(self, image: torch.Tensor, state: dict[str, bool]) -> torch.Tensor:
+    def apply_image(
+        self, image: RasterImage | torch.Tensor, state: dict[str, bool]
+    ) -> RasterImage | torch.Tensor:
         """Apply the sampled state on the specified image.
 
         Args:
@@ -101,8 +105,16 @@ class Pad(Transform):
             horizontal_pad = (horizontal_half, horizontal_extra - horizontal_half)
             vertical_pad = (vertical_half, vertical_extra - vertical_half)
 
-        image = apply_padding(image, True, horizontal_pad[0], horizontal_pad[1])
-        image = apply_padding(image, False, vertical_pad[0], vertical_pad[1])
+        if isinstance(image, RasterImage):
+            image.image = apply_padding(
+                image.image, True, horizontal_pad[0], horizontal_pad[1]
+            )
+            image.image = apply_padding(
+                image.image, False, vertical_pad[0], vertical_pad[1]
+            )
+        else:
+            image = apply_padding(image, True, horizontal_pad[0], horizontal_pad[1])
+            image = apply_padding(image, False, vertical_pad[0], vertical_pad[1])
         return image
 
     def apply_boxes(self, boxes: Any, state: dict[str, bool]) -> torch.Tensor:

rslearn/train/transforms/resize.py

@@ -6,6 +6,8 @@ import torch
 import torchvision
 from torchvision.transforms import InterpolationMode
 
+from rslearn.train.model_context import RasterImage
+
 from .transform import Transform
 
 INTERPOLATION_MODES = {
@@ -38,7 +40,9 @@ class Resize(Transform):
         self.selectors = selectors
         self.interpolation = INTERPOLATION_MODES[interpolation]
 
-    def apply_resize(self, image: torch.Tensor) -> torch.Tensor:
+    def apply_resize(
+        self, image: torch.Tensor | RasterImage
+    ) -> torch.Tensor | RasterImage:
         """Apply resizing on the specified image.
 
         If the image is 2D, it is unsqueezed to 3D and then squeezed
@@ -47,16 +51,21 @@ class Resize(Transform):
         Args:
             image: the image to transform.
         """
-        if image.dim() == 2:
-            image = image.unsqueeze(0)  # (H, W) -> (1, H, W)
-            result = torchvision.transforms.functional.resize(
+        if isinstance(image, torch.Tensor):
+            if image.dim() == 2:
+                image = image.unsqueeze(0)  # (H, W) -> (1, H, W)
+                result = torchvision.transforms.functional.resize(
+                    image, self.target_size, self.interpolation
+                )
+                return result.squeeze(0)  # (1, H, W) -> (H, W)
+            return torchvision.transforms.functional.resize(
                 image, self.target_size, self.interpolation
             )
-            return result.squeeze(0)  # (1, H, W) -> (H, W)
-
-        return torchvision.transforms.functional.resize(
-            image, self.target_size, self.interpolation
-        )
+        else:
+            image.image = torchvision.transforms.functional.resize(
+                image.image, self.target_size, self.interpolation
+            )
+            return image
 
     def forward(
         self, input_dict: dict[str, Any], target_dict: dict[str, Any]

rslearn/train/transforms/select_bands.py

@@ -2,6 +2,8 @@
 
 from typing import Any
 
+from rslearn.train.model_context import RasterImage
+
 from .transform import Transform, read_selector, write_selector
 
 
@@ -49,6 +51,10 @@ class SelectBands(Transform):
             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(
@@ -62,6 +68,9 @@ class SelectBands(Transform):
                 [(start_channel_idx + band_idx) for band_idx in self.band_indices]
             )
 
-        result = image[wanted_bands]
-        write_selector(input_dict, target_dict, self.output_selector, result)
+        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