rslearn 0.0.24__py3-none-any.whl → 0.0.26__py3-none-any.whl

rslearn/train/dataset_index.py

@@ -0,0 +1,156 @@
+"""Dataset index for caching window lists to speed up ModelDataset initialization."""
+
+import hashlib
+import json
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+from upath import UPath
+
+from rslearn.dataset.window import Window
+from rslearn.log_utils import get_logger
+from rslearn.utils.fsspec import open_atomic
+
+if TYPE_CHECKING:
+    from rslearn.dataset.storage.storage import WindowStorage
+
+logger = get_logger(__name__)
+
+# Increment this when the index format changes to force rebuild
+INDEX_VERSION = 1
+
+# Directory name for storing index files
+INDEX_DIR_NAME = ".rslearn_dataset_index"
+
+
+class DatasetIndex:
+    """Manages indexed window lists for faster ModelDataset initialization.
+
+    Note: The index does NOT automatically detect when windows are added or removed
+    from the dataset. Use refresh=True after modifying dataset windows.
+    """
+
+    def __init__(
+        self,
+        storage: "WindowStorage",
+        dataset_path: UPath,
+        groups: list[str] | None,
+        names: list[str] | None,
+        tags: dict[str, Any] | None,
+        num_samples: int | None,
+        skip_targets: bool,
+        inputs: dict[str, Any],
+    ) -> None:
+        """Initialize DatasetIndex with specific configuration.
+
+        Args:
+            storage: WindowStorage for deserializing windows.
+            dataset_path: Path to the dataset directory.
+            groups: list of window groups to include.
+            names: list of window names to include.
+            tags: tags to filter windows by.
+            num_samples: limit on number of samples.
+            skip_targets: whether targets are skipped.
+            inputs: dict mapping input names to DataInput objects.
+        """
+        self.storage = storage
+        self.dataset_path = dataset_path
+        self.index_dir = dataset_path / INDEX_DIR_NAME
+
+        # Compute index key from configuration
+        inputs_data = {}
+        for name, inp in inputs.items():
+            inputs_data[name] = {
+                "layers": inp.layers,
+                "required": inp.required,
+                "load_all_layers": inp.load_all_layers,
+                "is_target": inp.is_target,
+            }
+
+        key_data = {
+            "groups": groups,
+            "names": names,
+            "tags": tags,
+            "num_samples": num_samples,
+            "skip_targets": skip_targets,
+            "inputs": inputs_data,
+        }
+        self.index_key = hashlib.sha256(
+            json.dumps(key_data, sort_keys=True).encode()
+        ).hexdigest()
+
+    def _get_config_hash(self) -> str:
+        """Get hash of config.json for quick validation.
+
+        Returns:
+            A 16-character hex string hash of the config, or empty string if no config.
+        """
+        config_path = self.dataset_path / "config.json"
+        if config_path.exists():
+            with config_path.open() as f:
+                return hashlib.sha256(f.read().encode()).hexdigest()[:16]
+        return ""
+
+    def load_windows(self, refresh: bool = False) -> list[Window] | None:
+        """Load indexed window list if valid, else return None.
+
+        Args:
+            refresh: If True, ignore existing index and return None.
+
+        Returns:
+            List of Window objects if index is valid, None otherwise.
+        """
+        if refresh:
+            logger.info("refresh=True, rebuilding index")
+            return None
+
+        index_file = self.index_dir / f"{self.index_key}.json"
+        if not index_file.exists():
+            logger.info(f"No index found at {index_file}, will build")
+            return None
+
+        try:
+            with index_file.open() as f:
+                index_data = json.load(f)
+        except (OSError, json.JSONDecodeError):
+            logger.warning(f"Corrupted index file at {index_file}, will rebuild")
+            return None
+
+        # Check index version
+        if index_data.get("version") != INDEX_VERSION:
+            logger.info(
+                f"Index version mismatch (got {index_data.get('version')}, "
+                f"expected {INDEX_VERSION}), will rebuild"
+            )
+            return None
+
+        # Quick validation: check config hash
+        if index_data.get("config_hash") != self._get_config_hash():
+            logger.info("Config hash mismatch, index invalidated")
+            return None
+
+        # Deserialize windows
+        return [Window.from_metadata(self.storage, w) for w in index_data["windows"]]
+
+    def save_windows(self, windows: list[Window]) -> None:
+        """Save processed windows to index with atomic write.
+
+        Args:
+            windows: List of Window objects to index.
+        """
+        self.index_dir.mkdir(parents=True, exist_ok=True)
+        index_file = self.index_dir / f"{self.index_key}.json"
+
+        # Serialize windows
+        serialized_windows = [w.get_metadata() for w in windows]
+
+        index_data = {
+            "version": INDEX_VERSION,
+            "config_hash": self._get_config_hash(),
+            "created_at": datetime.now().isoformat(),
+            "num_windows": len(windows),
+            "windows": serialized_windows,
+        }
+        with open_atomic(index_file, "w") as f:
+            json.dump(index_data, f)
+        logger.info(f"Saved {len(windows)} windows to index at {index_file}")

rslearn/train/model_context.py

@@ -43,6 +43,22 @@ class RasterImage:
             raise ValueError(f"Expected a single timestep, got {self.image.shape[1]}")
         return self.image[:, 0]
 
+    def get_hw_tensor(self) -> torch.Tensor:
+        """Get a 2D HW tensor from a single-channel, single-timestep RasterImage.
+
+        This function checks that C=1 and T=1, then returns the HW tensor.
+        Useful for per-pixel labels like segmentation masks.
+        """
+        if self.image.shape[0] != 1:
+            raise ValueError(
+                f"Expected single channel (C=1), got {self.image.shape[0]}"
+            )
+        if self.image.shape[1] != 1:
+            raise ValueError(
+                f"Expected single timestep (T=1), got {self.image.shape[1]}"
+            )
+        return self.image[0, 0]
+
 
 @dataclass
 class SampleMetadata:

rslearn/train/tasks/detection.py

@@ -14,27 +14,10 @@ from torchmetrics import Metric, MetricCollection
 
 from rslearn.train.model_context import RasterImage, SampleMetadata
 from rslearn.utils import Feature, STGeometry
+from rslearn.utils.colors import DEFAULT_COLORS
 
 from .task import BasicTask
 
-DEFAULT_COLORS = [
-    (255, 0, 0),
-    (0, 255, 0),
-    (0, 0, 255),
-    (255, 255, 0),
-    (0, 255, 255),
-    (255, 0, 255),
-    (0, 128, 0),
-    (255, 160, 122),
-    (139, 69, 19),
-    (128, 128, 128),
-    (255, 255, 255),
-    (143, 188, 143),
-    (95, 158, 160),
-    (255, 200, 0),
-    (128, 0, 0),
-]
-
 
 class DetectionTask(BasicTask):
     """A point or bounding box detection task."""

rslearn/train/tasks/per_pixel_regression.py

@@ -66,20 +66,18 @@ class PerPixelRegressionTask(BasicTask):
             return {}, {}
 
         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
+        labels = raw_inputs["targets"].get_hw_tensor().float() * self.scale_factor
 
         if self.nodata_value is not None:
-            valid = (
-                raw_inputs["targets"].image[0, 0, :, :] != self.nodata_value
-            ).float()
+            valid = (raw_inputs["targets"].get_hw_tensor() != self.nodata_value).float()
         else:
             valid = torch.ones(labels.shape, dtype=torch.float32)
 
+        # Wrap in RasterImage with CTHW format (C=1, T=1) so values and valid can be
+        # used in image transforms.
         return {}, {
-            "values": labels,
-            "valid": valid,
+            "values": RasterImage(labels[None, None, :, :], timestamps=None),
+            "valid": RasterImage(valid[None, None, :, :], timestamps=None),
         }
 
     def process_output(
@@ -121,7 +119,7 @@ class PerPixelRegressionTask(BasicTask):
         image = super().visualize(input_dict, target_dict, output)["image"]
         if target_dict is None:
             raise ValueError("target_dict is required for visualization")
-        gt_values = target_dict["classes"].cpu().numpy()
+        gt_values = target_dict["values"].get_hw_tensor().cpu().numpy()
         pred_values = output.cpu().numpy()[0, :, :]
         gt_vis = np.clip(gt_values * 255, 0, 255).astype(np.uint8)
         pred_vis = np.clip(pred_values * 255, 0, 255).astype(np.uint8)
@@ -210,8 +208,10 @@ class PerPixelRegressionHead(Predictor):
 
         losses = {}
         if targets:
-            labels = torch.stack([target["values"] for target in targets])
-            mask = torch.stack([target["valid"] for target in targets])
+            labels = torch.stack(
+                [target["values"].get_hw_tensor() for target in targets]
+            )
+            mask = torch.stack([target["valid"].get_hw_tensor() for target in targets])
 
             if self.loss_mode == "mse":
                 scores = torch.square(outputs - labels)
@@ -262,14 +262,14 @@ class PerPixelRegressionMetricWrapper(Metric):
         """
         if not isinstance(preds, torch.Tensor):
             preds = torch.stack(preds)
-        labels = torch.stack([target["values"] for target in targets])
+        labels = torch.stack([target["values"].get_hw_tensor() for target in targets])
 
         # Sub-select the valid labels.
         # We flatten the prediction and label images at valid pixels.
         if len(preds.shape) == 4:
             assert preds.shape[1] == 1
             preds = preds[:, 0, :, :]
-        mask = torch.stack([target["valid"] > 0 for target in targets])
+        mask = torch.stack([target["valid"].get_hw_tensor() > 0 for target in targets])
         preds = preds[mask]
         labels = labels[mask]
         if len(preds) == 0:

rslearn/train/tasks/segmentation.py

@@ -17,28 +17,10 @@ from rslearn.train.model_context import (
     SampleMetadata,
 )
 from rslearn.utils import Feature
+from rslearn.utils.colors import DEFAULT_COLORS
 
 from .task import BasicTask
 
-# TODO: This is duplicated code fix it
-DEFAULT_COLORS = [
-    (255, 0, 0),
-    (0, 255, 0),
-    (0, 0, 255),
-    (255, 255, 0),
-    (0, 255, 255),
-    (255, 0, 255),
-    (0, 128, 0),
-    (255, 160, 122),
-    (139, 69, 19),
-    (128, 128, 128),
-    (255, 255, 255),
-    (143, 188, 143),
-    (95, 158, 160),
-    (255, 200, 0),
-    (128, 0, 0),
-]
-
 
 class SegmentationTask(BasicTask):
     """A segmentation (per-pixel classification) task."""
@@ -146,9 +128,7 @@ class SegmentationTask(BasicTask):
             return {}, {}
 
         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()
+        labels = raw_inputs["targets"].get_hw_tensor().long()
 
         if self.class_id_mapping is not None:
             new_labels = labels.clone()
@@ -164,9 +144,11 @@ class SegmentationTask(BasicTask):
         else:
             valid = torch.ones(labels.shape, dtype=torch.float32)
 
+        # Wrap in RasterImage with CTHW format (C=1, T=1) so classes and valid can be
+        # used in image transforms.
         return {}, {
-            "classes": labels,
-            "valid": valid,
+            "classes": RasterImage(labels[None, None, :, :], timestamps=None),
+            "valid": RasterImage(valid[None, None, :, :], timestamps=None),
         }
 
     def process_output(
@@ -224,7 +206,7 @@ class SegmentationTask(BasicTask):
         image = super().visualize(input_dict, target_dict, output)["image"]
         if target_dict is None:
             raise ValueError("target_dict is required for visualization")
-        gt_classes = target_dict["classes"].cpu().numpy()
+        gt_classes = target_dict["classes"].get_hw_tensor().cpu().numpy()
         pred_classes = output.cpu().numpy().argmax(axis=0)
         gt_vis = np.zeros((gt_classes.shape[0], gt_classes.shape[1], 3), dtype=np.uint8)
         pred_vis = np.zeros(
@@ -309,12 +291,19 @@ class SegmentationTask(BasicTask):
 class SegmentationHead(Predictor):
     """Head for segmentation task."""
 
-    def __init__(self, weights: list[float] | None = None, dice_loss: bool = False):
+    def __init__(
+        self,
+        weights: list[float] | None = None,
+        dice_loss: bool = False,
+        temperature: float = 1.0,
+    ):
         """Initialize a new SegmentationTask.
 
         Args:
             weights: weights for cross entropy loss (Tensor of size C)
             dice_loss: weather to add dice loss to cross entropy
+            temperature: temperature scaling for softmax, does not affect the loss,
+                only the predictor outputs
         """
         super().__init__()
         if weights is not None:
@@ -322,6 +311,7 @@ class SegmentationHead(Predictor):
         else:
             self.weights = None
         self.dice_loss = dice_loss
+        self.temperature = temperature
 
     def forward(
         self,
@@ -350,12 +340,16 @@ class SegmentationHead(Predictor):
             )
 
         logits = intermediates.feature_maps[0]
-        outputs = torch.nn.functional.softmax(logits, dim=1)
+        outputs = torch.nn.functional.softmax(logits / self.temperature, dim=1)
 
         losses = {}
         if targets:
-            labels = torch.stack([target["classes"] for target in targets], dim=0)
-            mask = torch.stack([target["valid"] for target in targets], dim=0)
+            labels = torch.stack(
+                [target["classes"].get_hw_tensor() for target in targets], dim=0
+            )
+            mask = torch.stack(
+                [target["valid"].get_hw_tensor() for target in targets], dim=0
+            )
             per_pixel_loss = torch.nn.functional.cross_entropy(
                 logits, labels, weight=self.weights, reduction="none"
             )
@@ -368,7 +362,8 @@ class SegmentationHead(Predictor):
                 # the summed mask loss be zero.
                 losses["cls"] = torch.sum(per_pixel_loss * mask)
             if self.dice_loss:
-                dice_loss = DiceLoss()(outputs, labels, mask)
+                softmax_woT = torch.nn.functional.softmax(logits, dim=1)
+                dice_loss = DiceLoss()(softmax_woT, labels, mask)
                 losses["dice"] = dice_loss
 
         return ModelOutput(
@@ -419,12 +414,12 @@ class SegmentationMetric(Metric):
         """
         if not isinstance(preds, torch.Tensor):
             preds = torch.stack(preds)
-        labels = torch.stack([target["classes"] for target in targets])
+        labels = torch.stack([target["classes"].get_hw_tensor() for target in targets])
 
         # Sub-select the valid labels.
         # We flatten the prediction and label images at valid pixels.
         # Prediction is changed from BCHW to BHWC so we can select the valid BHW mask.
-        mask = torch.stack([target["valid"] > 0 for target in targets])
+        mask = torch.stack([target["valid"].get_hw_tensor() > 0 for target in targets])
         preds = preds.permute(0, 2, 3, 1)[mask]
         labels = labels[mask]
         if len(preds) == 0:

rslearn/train/transforms/concatenate.py

@@ -54,36 +54,26 @@ class Concatenate(Transform):
             target_dict: the target
 
         Returns:
-            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.
+            (input_dicts, target_dicts) where the entry corresponding to
+            output_selector contains the concatenated RasterImage.
         """
-        images = []
-        return_raster_image: bool = False
+        tensors: list[torch.Tensor] = []
         timestamps: list[tuple[datetime, datetime]] | None = None
+
         for selector, wanted_bands in self.selections.items():
             image = read_selector(input_dict, target_dict, selector)
-            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)
+            if wanted_bands:
+                tensors.append(image.image[wanted_bands, :, :])
+            else:
+                tensors.append(image.image)
+            if timestamps is None and image.timestamps is not None:
+                # assume all concatenated modalities have the same
+                # number of timestamps
+                timestamps = image.timestamps
+
+        result = RasterImage(
+            torch.concatenate(tensors, dim=self.concatenate_dim),
+            timestamps=timestamps,
+        )
         write_selector(input_dict, target_dict, self.output_selector, result)
         return input_dict, target_dict

rslearn/train/transforms/crop.py

@@ -71,9 +71,7 @@ class Crop(Transform):
             "remove_from_top": remove_from_top,
         }
 
-    def apply_image(
-        self, image: RasterImage | torch.Tensor, state: dict[str, Any]
-    ) -> RasterImage | torch.Tensor:
+    def apply_image(self, image: RasterImage, state: dict[str, Any]) -> RasterImage:
         """Apply the sampled state on the specified image.
 
         Args:
@@ -84,22 +82,13 @@ 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]
-        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,
-            )
+        image.image = torchvision.transforms.functional.crop(
+            image.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:

rslearn/train/transforms/flip.py

@@ -57,16 +57,10 @@ class Flip(Transform):
             image: the image to transform.
             state: the sampled state.
         """
-        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])
+        if state["horizontal"]:
+            image.image = torch.flip(image.image, dims=[-1])
+        if state["vertical"]:
+            image.image = torch.flip(image.image, dims=[-2])
         return image
 
     def apply_boxes(

rslearn/train/transforms/mask.py

@@ -1,7 +1,5 @@
 """Mask transform."""
 
-import torch
-
 from rslearn.train.model_context import RasterImage
 from rslearn.train.transforms.transform import Transform, read_selector
 
@@ -32,9 +30,7 @@ class Mask(Transform):
         self.mask_selector = mask_selector
         self.mask_value = mask_value
 
-    def apply_image(
-        self, image: torch.Tensor | RasterImage, mask: torch.Tensor | RasterImage
-    ) -> torch.Tensor | RasterImage:
+    def apply_image(self, image: RasterImage, mask: RasterImage) -> RasterImage:
         """Apply the mask on the image.
 
         Args:
@@ -44,21 +40,19 @@ class Mask(Transform):
         Returns:
             masked image
         """
-        # Tile the mask to have same number of bands as the image.
-        if isinstance(mask, RasterImage):
-            mask = mask.image
+        # Extract the mask tensor (CTHW format)
+        mask_tensor = mask.image
 
-        if image.shape[0] != mask.shape[0]:
-            if mask.shape[0] != 1:
+        # Tile the mask to have same number of bands (C dimension) as the image.
+        if image.shape[0] != mask_tensor.shape[0]:
+            if mask_tensor.shape[0] != 1:
                 raise ValueError(
                     "expected mask to either have same bands as image, or one band"
                 )
-            mask = mask.repeat(image.shape[0], 1, 1)
+            # Repeat along C dimension, keep T, H, W the same
+            mask_tensor = mask_tensor.repeat(image.shape[0], 1, 1, 1)
 
-        if isinstance(image, torch.Tensor):
-            image[mask == 0] = self.mask_value
-        else:
-            image.image[mask == 0] = self.mask_value
+        image.image[mask_tensor == 0] = self.mask_value
         return image
 
     def forward(self, input_dict: dict, target_dict: dict) -> tuple[dict, dict]: