rslearn 0.0.9__tar.gz → 0.0.12__tar.gz

rslearn-0.0.12/NOTICE

@@ -0,0 +1,115 @@
+rslearn is released under Apache License 2.0
+Copyright 2025 Allen Institute for AI
+
+The following third party code is included in this repository.
+
+====================
+
+rslearn.models.detr is adapted from https://github.com/facebookresearch/detr which is
+released under Apache License 2.0.
+
+Copyright 2020 - present, Facebook, Inc
+
+====================
+
+rslearn.models.use_croma is copied from https://github.com/antofuller/CROMA
+
+MIT License
+
+Copyright (c) 2023 Anthony Fuller
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+====================
+
+rslearn.models.galileo is adapted from https://github.com/nasaharvest/galileo
+
+MIT License
+
+Copyright (c) 2024 Presto Authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+====================
+
+rslearn.models.presto is adapted from https://github.com/nasaharvest/presto
+
+MIT License
+
+Copyright (c) 2024 Presto Authors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+====================
+
+rslearn.models.prithvi includes code adapted from https://github.com/NASA-IMPACT/Prithvi-WxC
+
+MIT License
+
+Copyright (c) 2024 Inter Agency Implementation and Advanced Concepts
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

{rslearn-0.0.9/rslearn.egg-info → rslearn-0.0.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.9
+Version: 0.0.12
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License
@@ -211,6 +211,7 @@ Project-URL: repository, https://github.com/allenai/rslearn
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
+License-File: NOTICE
 Requires-Dist: boto3>=1.39
 Requires-Dist: fiona>=1.10
 Requires-Dist: fsspec>=2025.9.0
@@ -243,6 +244,7 @@ Requires-Dist: planetary_computer>=1.0; extra == "extra"
 Requires-Dist: pycocotools>=2.0; extra == "extra"
 Requires-Dist: pystac_client>=0.9; extra == "extra"
 Requires-Dist: rtree>=1.4; extra == "extra"
+Requires-Dist: termcolor>=3.0; extra == "extra"
 Requires-Dist: satlaspretrain_models>=0.3; extra == "extra"
 Requires-Dist: scipy>=1.16; extra == "extra"
 Requires-Dist: terratorch>=1.0.2; extra == "extra"

{rslearn-0.0.9 → rslearn-0.0.12}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rslearn"
-version = "0.0.9"
+version = "0.0.12"
 description = "A library for developing remote sensing datasets and models"
 authors = [
     { name = "OlmoEarth Team" },
@@ -47,6 +47,8 @@ extra = [
     "pycocotools>=2.0",
     "pystac_client>=0.9",
     "rtree>=1.4",
+    # Needed by DINOv3.
+    "termcolor>=3.0",
     "satlaspretrain_models>=0.3",
     "scipy>=1.16",
     "terratorch>=1.0.2",

{rslearn-0.0.9 → rslearn-0.0.12}/rslearn/models/anysat.py

@@ -1,4 +1,8 @@
-"""AnySat model."""
+"""AnySat model.
+
+This code loads the AnySat model from torch hub. See
+https://github.com/gastruc/AnySat for applicable license and copyright information.
+"""
 
 from typing import Any
 

{rslearn-0.0.9 → rslearn-0.0.12}/rslearn/models/dinov3.py

@@ -1,4 +1,9 @@
-"""DinoV3 model."""
+"""DinoV3 model.
+
+This code loads the DINOv3 model. You must obtain the model separately from Meta to use
+it. See https://github.com/facebookresearch/dinov3 for applicable license and copyright
+information.
+"""
 
 from enum import StrEnum
 from pathlib import Path

rslearn-0.0.12/rslearn/models/feature_center_crop.py

@@ -0,0 +1,50 @@
+"""Apply center cropping on a feature map."""
+
+from typing import Any
+
+import torch
+
+
+class FeatureCenterCrop(torch.nn.Module):
+    """Apply center cropping on the input feature maps."""
+
+    def __init__(
+        self,
+        sizes: list[tuple[int, int]],
+    ) -> None:
+        """Create a new FeatureCenterCrop.
+
+        Only the center of each feature map will be retained and passed to the next
+        module.
+
+        Args:
+            sizes: a list of (height, width) tuples, with one tuple for each input
+                feature map.
+        """
+        super().__init__()
+        self.sizes = sizes
+
+    def forward(
+        self, features: list[torch.Tensor], inputs: list[dict[str, Any]]
+    ) -> list[torch.Tensor]:
+        """Apply center cropping on the feature maps.
+
+        Args:
+            features: list of feature maps at different resolutions.
+            inputs: original inputs (ignored).
+
+        Returns:
+            center cropped feature maps.
+        """
+        new_features = []
+        for i, feat in enumerate(features):
+            height, width = self.sizes[i]
+            if feat.shape[2] < height or feat.shape[3] < width:
+                raise ValueError(
+                    "feature map is smaller than the desired height and width"
+                )
+            start_h = feat.shape[2] // 2 - height // 2
+            start_w = feat.shape[3] // 2 - width // 2
+            feat = feat[:, :, start_h : start_h + height, start_w : start_w + width]
+            new_features.append(feat)
+        return new_features

rslearn-0.0.12/rslearn/models/olmoearth_pretrain/__init__.py

@@ -0,0 +1 @@
+"""OlmoEarth model architecture."""

rslearn-0.0.12/rslearn/models/olmoearth_pretrain/model.py

@@ -0,0 +1,263 @@
+"""OlmoEarth model wrapper for fine-tuning in rslearn."""
+
+import json
+from contextlib import nullcontext
+from typing import Any
+
+import torch
+from einops import rearrange
+from olmo_core.config import Config
+from olmo_core.distributed.checkpoint import load_model_and_optim_state
+from olmoearth_pretrain.data.constants import Modality
+from olmoearth_pretrain.model_loader import (
+    ModelID,
+    load_model_from_id,
+    load_model_from_path,
+)
+from olmoearth_pretrain.nn.flexihelios import Encoder, TokensAndMasks
+from olmoearth_pretrain.train.masking import MaskedOlmoEarthSample, MaskValue
+from upath import UPath
+
+from rslearn.log_utils import get_logger
+
+logger = get_logger(__name__)
+
+MODALITY_NAMES = [
+    "sentinel2_l2a",
+    "sentinel1",
+    "worldcover",
+    "openstreetmap_raster",
+    "landsat",
+]
+
+AUTOCAST_DTYPE_MAP = {
+    "bfloat16": torch.bfloat16,
+    "float16": torch.float16,
+    "float32": torch.float32,
+}
+
+EMBEDDING_SIZES = {
+    ModelID.OLMOEARTH_V1_NANO: 128,
+    ModelID.OLMOEARTH_V1_TINY: 192,
+    ModelID.OLMOEARTH_V1_BASE: 768,
+}
+
+
+class OlmoEarth(torch.nn.Module):
+    """A wrapper to support the OlmoEarth model."""
+
+    def __init__(
+        self,
+        patch_size: int,
+        model_id: ModelID | None = None,
+        model_path: str | None = None,
+        checkpoint_path: str | None = None,
+        selector: list[str | int] = ["encoder"],
+        forward_kwargs: dict[str, Any] = {},
+        random_initialization: bool = False,
+        embedding_size: int | None = None,
+        autocast_dtype: str | None = "bfloat16",
+    ):
+        """Create a new OlmoEarth model.
+
+        Args:
+            patch_size: token spatial patch size to use.
+            model_id: the model ID to load. One of model_id or model_path or checkpoint_path must be
+                set.
+            model_path: the path to load the model from. One of model_id or model_path or checkpoint_path must be
+                set. Same structure as the HF-hosted `model_id` models: bundle with a config.json and weights.pth.
+            checkpoint_path: the checkpoint directory to load from, if model_id or model_path is not
+                set. It should contain a distributed checkpoint with a config.json file as well as model_and_optim
+                folder.
+            selector: an optional sequence of attribute names or list indices to select
+                the sub-module that should be applied on the input images. Defaults to
+                ["encoder"] to select only the transformer encoder.
+            forward_kwargs: additional arguments to pass to forward pass besides the
+                 MaskedOlmoEarthSample.
+            random_initialization: whether to skip loading the checkpoint so the
+                weights are randomly initialized. In this case, the checkpoint is only
+                used to define the model architecture.
+            embedding_size: optional embedding size to report via
+                get_backbone_channels (if model_id is not set).
+            autocast_dtype: which dtype to use for autocasting, or set None to disable.
+        """
+        if (
+            sum(
+                [
+                    model_id is not None,
+                    model_path is not None,
+                    checkpoint_path is not None,
+                ]
+            )
+            != 1
+        ):
+            raise ValueError(
+                "exactly one of model_id, model_path, or checkpoint_path must be set"
+            )
+
+        super().__init__()
+        self.patch_size = patch_size
+        self.forward_kwargs = forward_kwargs
+        self.embedding_size = embedding_size
+
+        if autocast_dtype is not None:
+            self.autocast_dtype = AUTOCAST_DTYPE_MAP[autocast_dtype]
+        else:
+            self.autocast_dtype = None
+
+        if model_id is not None:
+            # Load from Hugging Face.
+            model = load_model_from_id(model_id, load_weights=not random_initialization)
+            if self.embedding_size is None and model_id in EMBEDDING_SIZES:
+                self.embedding_size = EMBEDDING_SIZES[model_id]
+
+        elif model_path is not None:
+            # Load from path.
+            model = load_model_from_path(
+                UPath(model_path), load_weights=not random_initialization
+            )
+
+        else:
+            # Load the distributed model checkpoint by path through Olmo Core
+            model = self._load_model_from_checkpoint(
+                UPath(checkpoint_path), random_initialization
+            )
+
+        # Select just the portion of the model that we actually want to use.
+        for part in selector:
+            if isinstance(part, str):
+                model = getattr(model, part)
+            else:
+                model = model[part]
+        self.model = model
+
+    def _load_model_from_checkpoint(
+        self, checkpoint_upath: UPath, random_initialization: bool
+    ) -> torch.nn.Module:
+        """Load the OlmoEarth pre-trained model from a distributed checkpoint folder.
+
+        The folder should contain config.json as well as the model_and_optim folder
+        that contains the distributed checkpoint. This is the format produced by
+        pre-training runs in olmoearth_pretrain.
+        """
+        # Load the model config and initialize it.
+        # We avoid loading the train module here because it depends on running within
+        # olmo_core.
+        with (checkpoint_upath / "config.json").open() as f:
+            config_dict = json.load(f)
+            model_config = Config.from_dict(config_dict["model"])
+
+        model = model_config.build()
+
+        # Load the checkpoint.
+        if not random_initialization:
+            train_module_dir = checkpoint_upath / "model_and_optim"
+            if train_module_dir.exists():
+                load_model_and_optim_state(str(train_module_dir), model)
+                logger.info(f"loaded OlmoEarth encoder from {train_module_dir}")
+            else:
+                logger.info(f"could not find OlmoEarth encoder at {train_module_dir}")
+
+        return model
+
+    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+        """Compute feature maps from the OlmoEarth backbone.
+
+        Inputs:
+            inputs: input dicts. It should include keys corresponding to the modalities
+                that should be passed to the OlmoEarth model.
+        """
+        kwargs = {}
+        present_modalities = []
+        device = None
+        # Handle the case where some modalities are multitemporal and some are not.
+        # We assume all multitemporal modalities have the same number of timesteps.
+        max_timesteps = 1
+        for modality in MODALITY_NAMES:
+            if modality not in inputs[0]:
+                continue
+            present_modalities.append(modality)
+            cur = torch.stack([inp[modality] for inp in inputs], dim=0)
+            device = cur.device
+            # Check if it's single or multitemporal, and reshape accordingly
+            num_bands = Modality.get(modality).num_bands
+            num_timesteps = cur.shape[1] // num_bands
+            max_timesteps = max(max_timesteps, num_timesteps)
+            cur = rearrange(cur, "b (t c) h w -> b h w t c", t=num_timesteps)
+            kwargs[modality] = cur
+            # Create mask array which is BHWTS (without channels but with band sets).
+            num_band_sets = len(Modality.get(modality).band_sets)
+            mask_shape = cur.shape[0:4] + (num_band_sets,)
+            mask = (
+                torch.ones(mask_shape, dtype=torch.int32, device=device)
+                * MaskValue.ONLINE_ENCODER.value
+            )
+            kwargs[f"{modality}_mask"] = mask
+
+        # Timestamps is required.
+        # Note that only months (0 to 11) are used in OlmoEarth position encoding.
+        # For now, we assign same timestamps to all inputs, but later we should handle varying timestamps per input.
+        timestamps = torch.zeros(
+            (len(inputs), max_timesteps, 3), dtype=torch.int32, device=device
+        )
+        timestamps[:, :, 0] = 1  # day
+        timestamps[:, :, 1] = torch.arange(max_timesteps, device=device)[
+            None, :
+        ]  # month
+        timestamps[:, :, 2] = 2024  # year
+        kwargs["timestamps"] = timestamps
+
+        sample = MaskedOlmoEarthSample(**kwargs)
+
+        # Decide context based on self.autocast_dtype.
+        if self.autocast_dtype is None:
+            context = nullcontext()
+        else:
+            assert device is not None
+            context = torch.amp.autocast(
+                device_type=device.type, dtype=self.autocast_dtype
+            )
+
+        with context:
+            # Currently we assume the provided model always returns a TokensAndMasks object.
+            tokens_and_masks: TokensAndMasks
+            if isinstance(self.model, Encoder):
+                # Encoder has a fast_pass argument to indicate mask is not needed.
+                tokens_and_masks = self.model(
+                    sample,
+                    fast_pass=True,
+                    patch_size=self.patch_size,
+                    **self.forward_kwargs,
+                )["tokens_and_masks"]
+            else:
+                # Other models like STEncoder do not have this option supported.
+                tokens_and_masks = self.model(
+                    sample, patch_size=self.patch_size, **self.forward_kwargs
+                )["tokens_and_masks"]
+
+        # Apply temporal/modality pooling so we just have one feature per patch.
+        features = []
+        for modality in present_modalities:
+            modality_features = getattr(tokens_and_masks, modality)
+            # Pool over band sets and timesteps (BHWTSC -> BHWC).
+            pooled = modality_features.mean(dim=[3, 4])
+            # We want BHWC -> BCHW.
+            pooled = rearrange(pooled, "b h w c -> b c h w")
+            features.append(pooled)
+        # Pool over the modalities, so we get one BCHW feature map.
+        pooled = torch.stack(features, dim=0).mean(dim=0)
+        return [pooled]
+
+    def get_backbone_channels(self) -> list:
+        """Returns the output channels of this model when used as a backbone.
+
+        The output channels is a list of (downsample_factor, depth) that corresponds
+        to the feature maps that the backbone returns. For example, an element [2, 32]
+        indicates that the corresponding feature map is 1/2 the input resolution and
+        has 32 channels.
+
+        Returns:
+            the output channels of the backbone as a list of (downsample_factor, depth)
+            tuples.
+        """
+        return [(self.patch_size, self.embedding_size)]

rslearn-0.0.12/rslearn/models/olmoearth_pretrain/norm.py

@@ -0,0 +1,84 @@
+"""Normalization transforms."""
+
+import json
+from typing import Any
+
+from olmoearth_pretrain.data.normalize import load_computed_config
+
+from rslearn.log_utils import get_logger
+from rslearn.train.transforms.transform import Transform
+
+logger = get_logger(__file__)
+
+
+class OlmoEarthNormalize(Transform):
+    """Normalize using OlmoEarth JSON config.
+
+    For Sentinel-1 data, the values should be converted to decibels before being passed
+    to this transform.
+    """
+
+    def __init__(
+        self,
+        band_names: dict[str, list[str]],
+        std_multiplier: float | None = 2,
+        config_fname: str | None = None,
+    ) -> None:
+        """Initialize a new OlmoEarthNormalize.
+
+        Args:
+            band_names: map from modality name to the list of bands in that modality in
+                the order they are being loaded. Note that this order must match the
+                expected order for the OlmoEarth model.
+            std_multiplier: the std multiplier matching the one used for the model
+                training in OlmoEarth.
+            config_fname: load the normalization configuration from this file, instead
+                of getting it from OlmoEarth.
+        """
+        super().__init__()
+        self.band_names = band_names
+        self.std_multiplier = std_multiplier
+
+        if config_fname is None:
+            self.norm_config = load_computed_config()
+        else:
+            logger.warning(
+                f"Loading normalization config from {config_fname}. This argument is deprecated and will be removed in a future version."
+            )
+            with open(config_fname) as f:
+                self.norm_config = json.load(f)
+
+    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
+        """
+        for modality_name, cur_band_names in self.band_names.items():
+            band_norms = self.norm_config[modality_name]
+            image = input_dict[modality_name]
+            # Keep a set of indices to make sure that we normalize all of them.
+            needed_band_indices = set(range(image.shape[0]))
+            num_timesteps = image.shape[0] // len(cur_band_names)
+
+            for band, norm_dict in band_norms.items():
+                # If multitemporal, normalize each timestep separately.
+                for t in range(num_timesteps):
+                    band_idx = cur_band_names.index(band) + t * len(cur_band_names)
+                    min_val = norm_dict["mean"] - self.std_multiplier * norm_dict["std"]
+                    max_val = norm_dict["mean"] + self.std_multiplier * norm_dict["std"]
+                    image[band_idx] = (image[band_idx] - min_val) / (max_val - min_val)
+                    needed_band_indices.remove(band_idx)
+
+            if len(needed_band_indices) > 0:
+                raise ValueError(
+                    f"for modality {modality_name}, bands {needed_band_indices} were unexpectedly not normalized"
+                )
+
+        return input_dict, target_dict

{rslearn-0.0.9 → rslearn-0.0.12}/rslearn/models/pooling_decoder.py

@@ -76,3 +76,46 @@ class PoolingDecoder(torch.nn.Module):
         features = torch.amax(features, dim=(2, 3))
         features = self.fc_layers(features)
         return self.output_layer(features)
+
+
+class SegmentationPoolingDecoder(PoolingDecoder):
+    """Like PoolingDecoder, but copy output to all pixels.
+
+    This allows for the model to produce a global output while still being compatible
+    with SegmentationTask. This only makes sense for very small windows, since the
+    output probabilities will be the same at all pixels. The main use case is to train
+    for a classification-like task on small windows, but still produce a raster during
+    inference on large windows.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        image_key: str = "image",
+        **kwargs: Any,
+    ):
+        """Create a new SegmentationPoolingDecoder.
+
+        Args:
+            in_channels: input channels (channels in the last feature map passed to
+                this module)
+            out_channels: channels for the output flat feature vector
+            image_key: the key in inputs for the image from which the expected width
+                and height is derived.
+            kwargs: other arguments to pass to PoolingDecoder.
+        """
+        super().__init__(in_channels=in_channels, out_channels=out_channels, **kwargs)
+        self.image_key = image_key
+
+    def forward(
+        self, features: list[torch.Tensor], inputs: list[dict[str, Any]]
+    ) -> torch.Tensor:
+        """Extend PoolingDecoder forward to upsample the output to a segmentation mask.
+
+        This only works when all of the pixels have the same segmentation target.
+        """
+        output_probs = super().forward(features, inputs)
+        # BC -> BCHW
+        h, w = inputs[0][self.image_key].shape[1:3]
+        return output_probs[:, :, None, None].repeat([1, 1, h, w])

{rslearn-0.0.9 → rslearn-0.0.12}/rslearn/models/prithvi.py

@@ -1,4 +1,12 @@
-"""Prithvi V2."""
+"""Prithvi V2.
+
+This code is adapted from https://github.com/NASA-IMPACT/Prithvi-WxC
+
+The code is released under:
+
+MIT License
+Copyright (c) 2024 Inter Agency Implementation and Advanced Concepts
+"""
 
 import json
 import logging

{rslearn-0.0.9 → rslearn-0.0.12}/rslearn/train/lightning_module.py

@@ -94,7 +94,6 @@ class RslearnLightningModule(L.LightningModule):
         restore_config: RestoreConfig | None = None,
         print_parameters: bool = False,
         print_model: bool = False,
-        strict_loading: bool = True,
         # Deprecated options.
         lr: float = 1e-3,
         plateau: bool = False,
@@ -118,7 +117,6 @@ class RslearnLightningModule(L.LightningModule):
             print_parameters: whether to print the list of model parameters after model
                 initialization
             print_model: whether to print the model after model initialization
-            strict_loading: whether to strictly load the model parameters.
             lr: deprecated.
             plateau: deprecated.
             plateau_factor: deprecated.
@@ -132,7 +130,6 @@ class RslearnLightningModule(L.LightningModule):
         self.visualize_dir = visualize_dir
         self.metrics_file = metrics_file
         self.restore_config = restore_config
-        self.strict_loading = strict_loading
 
         self.scheduler_factory: SchedulerFactory | None = None
         if scheduler:

{rslearn-0.0.9 → rslearn-0.0.12}/rslearn/train/tasks/classification.py

@@ -49,8 +49,8 @@ class ClassificationTask(BasicTask):
                 features with matching properties.
             read_class_id: whether to read an integer class ID instead of the class
                 name.
-            allow_invalid: instead of throwing error when no regression label is found
-                at a window, simply mark the example invalid for this task
+            allow_invalid: instead of throwing error when no classification label is
+                found at a window, simply mark the example invalid for this task
             skip_unknown_categories: whether to skip examples with categories that are
                 not passed via classes, instead of throwing error
             prob_property: when predicting, write probabilities in addition to class ID