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

rslearn/models/molmo.py

@@ -0,0 +1,65 @@
+"""Molmo model."""
+
+from typing import Any
+
+import torch
+from transformers import AutoModelForCausalLM, AutoProcessor
+
+
+class Molmo(torch.nn.Module):
+    """Molmo image encoder."""
+
+    def __init__(
+        self,
+        model_name: str,
+    ):
+        """Instantiate a new Molmo instance.
+
+        Args:
+            model_name: the model name like "allenai/Molmo-7B-D-0924".
+        """
+        super().__init__()
+
+        self.processor = AutoProcessor.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            torch_dtype="auto",
+            device_map="cpu",
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            trust_remote_code=True,
+            torch_dtype="auto",
+            device_map="cpu",
+        )
+        self.encoder = model.model.vision_backbone
+
+    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+        """Compute outputs from the backbone.
+
+        Inputs:
+            inputs: input dicts that must include "image" key containing the image to
+                process. The images should have values 0-255.
+
+        Returns:
+            list of feature maps. Molmo produces features at one scale, so the list
+                contains a single Bx24x24x2048 tensor.
+        """
+        device = inputs[0]["image"].device
+        molmo_inputs_list = []
+        # Process each one so we can isolate just the full image without any crops.
+        for inp in inputs:
+            image = inp["image"].cpu().numpy().transpose(1, 2, 0)
+            processed = self.processor.process(
+                images=[image],
+                text="",
+            )
+            molmo_inputs_list.append(processed["images"][0])
+        molmo_inputs: torch.Tensor = torch.stack(molmo_inputs_list, dim=0).unsqueeze(1)
+
+        image_features, _ = self.encoder.encode_image(molmo_inputs.to(device))
+
+        # 576x2048 -> 24x24x2048
+        return [
+            image_features[:, 0, :, :].reshape(-1, 24, 24, 2048).permute(0, 3, 1, 2)
+        ]

rslearn/models/multitask.py

@@ -44,7 +44,7 @@ class MultiTaskModel(torch.nn.Module):
             tuple (outputs, loss_dict) from the last module.
         """
         features = self.encoder(inputs)
-        outputs = [{} for _ in inputs]
+        outputs: list[dict[str, Any]] = [{} for _ in inputs]
         losses = {}
         for name, decoder in self.decoders.items():
             cur = features

rslearn/models/pooling_decoder.py

@@ -21,7 +21,7 @@ class PoolingDecoder(torch.nn.Module):
         num_fc_layers: int = 0,
         conv_channels: int = 128,
         fc_channels: int = 512,
-    ):
+    ) -> None:
         """Initialize a PoolingDecoder.
 
         Args:
@@ -57,7 +57,9 @@ class PoolingDecoder(torch.nn.Module):
 
         self.output_layer = torch.nn.Linear(prev_channels, out_channels)
 
-    def forward(self, features: list[torch.Tensor], inputs: list[dict[str, Any]]):
+    def forward(
+        self, features: list[torch.Tensor], inputs: list[dict[str, Any]]
+    ) -> torch.Tensor:
         """Compute flat output vector from multi-scale feature map.
 
         Args:

rslearn/models/satlaspretrain.py

@@ -13,7 +13,7 @@ class SatlasPretrain(torch.nn.Module):
         self,
         model_identifier: str,
         fpn: bool = False,
-    ):
+    ) -> None:
         """Instantiate a new SatlasPretrain instance.
 
         Args:
@@ -25,7 +25,7 @@ class SatlasPretrain(torch.nn.Module):
         super().__init__()
         weights_manager = satlaspretrain_models.Weights()
         self.model = weights_manager.get_pretrained_model(
-            model_identifier=model_identifier, fpn=fpn
+            model_identifier=model_identifier, fpn=fpn, device="cpu"
         )
 
         if "SwinB" in model_identifier:
@@ -50,20 +50,17 @@ class SatlasPretrain(torch.nn.Module):
                 [32, 2048],
             ]
 
-    def forward(
-        self, inputs: list[dict[str, Any]], targets: list[dict[str, Any]] = None
-    ):
+    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
         """Compute feature maps from the SatlasPretrain backbone.
 
         Inputs:
             inputs: input dicts that must include "image" key containing the image to
                 process.
-            targets: target dicts that are ignored
         """
         images = torch.stack([inp["image"] for inp in inputs], dim=0)
         return self.model(images)
 
-    def get_backbone_channels(self):
+    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

rslearn/models/simple_time_series.py

@@ -24,7 +24,7 @@ class SimpleTimeSeries(torch.nn.Module):
         op: str = "max",
         groups: list[list[int]] | None = None,
         num_layers: int | None = None,
-    ):
+    ) -> None:
         """Create a new SimpleTimeSeries.
 
         Args:
@@ -55,63 +55,69 @@ class SimpleTimeSeries(torch.nn.Module):
         else:
             self.num_groups = 1
 
-        if self.op == "convrnn":
-            rnn_kernel_size = 3
-            self.rnn_layers = []
-            for _, count in out_channels:
-                cur_layer = [
-                    torch.nn.Sequential(
-                        torch.nn.Conv2d(
-                            2 * count, count, rnn_kernel_size, padding="same"
-                        ),
-                        torch.nn.ReLU(inplace=True),
-                    )
-                ]
-                for _ in range(num_layers - 1):
-                    cur_layer.append(
+        if self.op in ["convrnn", "conv3d", "conv1d"]:
+            if num_layers is None:
+                raise ValueError(f"num_layers must be specified for {self.op} op")
+
+            if self.op == "convrnn":
+                rnn_kernel_size = 3
+                rnn_layers = []
+                for _, count in out_channels:
+                    cur_layer = [
                         torch.nn.Sequential(
                             torch.nn.Conv2d(
-                                count, count, rnn_kernel_size, padding="same"
+                                2 * count, count, rnn_kernel_size, padding="same"
                             ),
                             torch.nn.ReLU(inplace=True),
                         )
-                    )
-                cur_layer = torch.nn.Sequential(*cur_layer)
-                self.rnn_layers.append(cur_layer)
-            self.rnn_layers = torch.nn.ModuleList(self.rnn_layers)
-
-        elif self.op == "conv3d":
-            self.conv3d_layers = []
-            for _, count in out_channels:
-                cur_layer = [
-                    torch.nn.Sequential(
-                        torch.nn.Conv3d(count, count, 3, padding=1, stride=(2, 1, 1)),
-                        torch.nn.ReLU(inplace=True),
-                    )
-                    for _ in range(num_layers)
-                ]
-                cur_layer = torch.nn.Sequential(*cur_layer)
-                self.conv3d_layers.append(cur_layer)
-            self.conv3d_layers = torch.nn.ModuleList(self.conv3d_layers)
-
-        elif self.op == "conv1d":
-            self.conv1d_layers = []
-            for _, count in out_channels:
-                cur_layer = [
-                    torch.nn.Sequential(
-                        torch.nn.Conv1d(count, count, 3, padding=1, stride=2),
-                        torch.nn.ReLU(inplace=True),
-                    )
-                    for _ in range(num_layers)
-                ]
-                cur_layer = torch.nn.Sequential(*cur_layer)
-                self.conv1d_layers.append(cur_layer)
-            self.conv1d_layers = torch.nn.ModuleList(self.conv1d_layers)
+                    ]
+                    for _ in range(num_layers - 1):
+                        cur_layer.append(
+                            torch.nn.Sequential(
+                                torch.nn.Conv2d(
+                                    count, count, rnn_kernel_size, padding="same"
+                                ),
+                                torch.nn.ReLU(inplace=True),
+                            )
+                        )
+                    cur_layer = torch.nn.Sequential(*cur_layer)
+                    rnn_layers.append(cur_layer)
+                self.rnn_layers = torch.nn.ModuleList(rnn_layers)
+
+            elif self.op == "conv3d":
+                conv3d_layers = []
+                for _, count in out_channels:
+                    cur_layer = [
+                        torch.nn.Sequential(
+                            torch.nn.Conv3d(
+                                count, count, 3, padding=1, stride=(2, 1, 1)
+                            ),
+                            torch.nn.ReLU(inplace=True),
+                        )
+                        for _ in range(num_layers)
+                    ]
+                    cur_layer = torch.nn.Sequential(*cur_layer)
+                    conv3d_layers.append(cur_layer)
+                self.conv3d_layers = torch.nn.ModuleList(conv3d_layers)
+
+            elif self.op == "conv1d":
+                conv1d_layers = []
+                for _, count in out_channels:
+                    cur_layer = [
+                        torch.nn.Sequential(
+                            torch.nn.Conv1d(count, count, 3, padding=1, stride=2),
+                            torch.nn.ReLU(inplace=True),
+                        )
+                        for _ in range(num_layers)
+                    ]
+                    cur_layer = torch.nn.Sequential(*cur_layer)
+                    conv1d_layers.append(cur_layer)
+                self.conv1d_layers = torch.nn.ModuleList(conv1d_layers)
 
         else:
             assert self.op in ["max", "mean"]
 
-    def get_backbone_channels(self):
+    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
@@ -129,14 +135,14 @@ class SimpleTimeSeries(torch.nn.Module):
         return out_channels
 
     def forward(
-        self, inputs: list[dict[str, Any]], targets: list[dict[str, Any]] = None
-    ):
+        self,
+        inputs: list[dict[str, Any]],
+    ) -> list[torch.Tensor]:
         """Compute outputs from the backbone.
 
         Inputs:
             inputs: input dicts that must include "image" key containing the image time
                 series to process (with images concatenated on the channel dimension).
-            targets: target dicts that are ignored unless
         """
         # First get features of each image.
         # To do so, we need to split up each grouped image into its component images (which have had their channels stacked).
@@ -171,13 +177,13 @@ class SimpleTimeSeries(torch.nn.Module):
         for feature_idx in range(len(all_features)):
             aggregated_features = []
             for group in groups:
-                group_features = []
+                group_features_list = []
                 for image_idx in group:
-                    group_features.append(
+                    group_features_list.append(
                         all_features[feature_idx][:, image_idx, :, :, :]
                     )
                 # Resulting group features are (depth, batch, C, height, width).
-                group_features = torch.stack(group_features, dim=0)
+                group_features = torch.stack(group_features_list, dim=0)
 
                 if self.op == "max":
                     group_features = torch.amax(group_features, dim=0)

rslearn/models/ssl4eo_s12.py

@@ -14,7 +14,7 @@ class Ssl4eoS12(torch.nn.Module):
         backbone_ckpt_path: str,
         arch: str = "resnet50",
         output_layers: list[int] = [0, 1, 2, 3],
-    ):
+    ) -> None:
         """Instantiate a new Swin instance.
 
         Args:
@@ -51,7 +51,7 @@ class Ssl4eoS12(torch.nn.Module):
                 f"warning: got missing_keys={missing_keys}, unexpected_keys={unexpected_keys} when loading SSL4EO-S12 state dict"
             )
 
-    def get_backbone_channels(self):
+    def get_backbone_channels(self) -> list[tuple[int, int]]:
         """Returns the output channels of this model when used as a backbone.
 
         The output channels is a list of (downsample_factor, depth) that corresponds
@@ -65,16 +65,17 @@ class Ssl4eoS12(torch.nn.Module):
         """
         if self.arch == "resnet50":
             all_out_channels = [
-                [4, 256],
-                [8, 512],
-                [16, 1024],
-                [32, 2048],
+                (4, 256),
+                (8, 512),
+                (16, 1024),
+                (32, 2048),
             ]
         return [all_out_channels[idx] for idx in self.output_layers]
 
     def forward(
-        self, inputs: list[dict[str, Any]], targets: list[dict[str, Any]] = None
-    ):
+        self,
+        inputs: list[dict[str, Any]],
+    ) -> list[torch.Tensor]:
         """Compute outputs from the backbone.
 
         If output_layers is set, then the outputs are multi-scale feature maps;
@@ -84,7 +85,6 @@ class Ssl4eoS12(torch.nn.Module):
         Inputs:
             inputs: input dicts that must include "image" key containing the image to
                 process.
-            targets: target dicts that are ignored unless
         """
         x = torch.stack([inp["image"] for inp in inputs], dim=0)
         x = self.model.conv1(x)

rslearn/models/swin.py

@@ -28,7 +28,7 @@ class Swin(torch.nn.Module):
         input_channels: int = 3,
         output_layers: list[int] | None = None,
         num_outputs: int = 1000,
-    ):
+    ) -> None:
         """Instantiate a new Swin instance.
 
         Args:
@@ -89,7 +89,7 @@ class Swin(torch.nn.Module):
         if num_outputs != self.model.head.out_features:
             self.model.head = torch.nn.Linear(self.model.head.in_features, num_outputs)
 
-    def get_backbone_channels(self):
+    def get_backbone_channels(self) -> list[tuple[int, int]]:
         """Returns the output channels of this model when used as a backbone.
 
         The output channels is a list of (downsample_factor, depth) that corresponds
@@ -105,31 +105,33 @@ class Swin(torch.nn.Module):
 
         if self.arch in ["swin_b", "swin_v2_b"]:
             all_out_channels = [
-                [4, 128],
-                [4, 128],
-                [8, 256],
-                [8, 256],
-                [16, 512],
-                [16, 512],
-                [32, 1024],
-                [32, 1024],
+                (4, 128),
+                (4, 128),
+                (4, 128),
+                (8, 256),
+                (8, 256),
+                (16, 512),
+                (16, 512),
+                (32, 1024),
+                (32, 1024),
             ]
         elif self.arch in ["swin_s", "swin_v2_s", "swin_t", "swin_v2_t"]:
             all_out_channels = [
-                [4, 96],
-                [4, 96],
-                [8, 192],
-                [8, 192],
-                [16, 384],
-                [16, 384],
-                [32, 768],
-                [32, 768],
+                (4, 96),
+                (4, 96),
+                (8, 192),
+                (8, 192),
+                (16, 384),
+                (16, 384),
+                (32, 768),
+                (32, 768),
             ]
         return [all_out_channels[idx] for idx in self.output_layers]
 
     def forward(
-        self, inputs: list[dict[str, Any]], targets: list[dict[str, Any]] = None
-    ):
+        self,
+        inputs: list[dict[str, Any]],
+    ) -> list[torch.Tensor]:
         """Compute outputs from the backbone.
 
         If output_layers is set, then the outputs are multi-scale feature maps;
@@ -139,7 +141,6 @@ class Swin(torch.nn.Module):
         Inputs:
             inputs: input dicts that must include "image" key containing the image to
                 process.
-            targets: target dicts that are ignored unless
         """
         images = torch.stack([inp["image"] for inp in inputs], dim=0)
 

rslearn/models/unet.py

@@ -19,7 +19,7 @@ class UNetDecoder(torch.nn.Module):
         out_channels: int,
         conv_layers_per_resolution: int = 1,
         kernel_size: int = 3,
-    ):
+    ) -> None:
         """Initialize a UNetDecoder.
 
         Args:
@@ -110,7 +110,9 @@ class UNetDecoder(torch.nn.Module):
         layers.append(torch.nn.Sequential(*cur_layers))
         self.layers = torch.nn.ModuleList(layers)
 
-    def forward(self, in_features: list[torch.Tensor], inputs: list[dict[str, Any]]):
+    def forward(
+        self, in_features: list[torch.Tensor], inputs: list[dict[str, Any]]
+    ) -> torch.Tensor:
         """Compute output from multi-scale feature map.
 
         Args:

rslearn/models/upsample.py

@@ -0,0 +1,35 @@
+"""An upsampling layer."""
+
+import torch
+
+
+class Upsample(torch.nn.Module):
+    """Upsamples each input feature map by the same factor."""
+
+    def __init__(
+        self,
+        scale_factor: int,
+        mode: str = "bilinear",
+    ):
+        """Initialize an Upsample.
+
+        Args:
+            scale_factor: the upsampling factor, e.g. 2 to double the size.
+            mode: "nearest" or "bilinear".
+        """
+        super().__init__()
+        self.layer = torch.nn.Upsample(scale_factor=scale_factor, mode=mode)
+
+    def forward(
+        self, features: list[torch.Tensor], inputs: list[torch.Tensor]
+    ) -> list[torch.Tensor]:
+        """Compute flat output vector from multi-scale feature map.
+
+        Args:
+            features: list of feature maps at different resolutions.
+            inputs: original inputs (ignored).
+
+        Returns:
+            flat feature vector
+        """
+        return [self.layer(feat_map) for feat_map in features]

rslearn/tile_stores/file.py

@@ -11,7 +11,6 @@ from rslearn.utils import Feature, PixelBounds, Projection
 from rslearn.utils.fsspec import open_atomic
 from rslearn.utils.raster_format import (
     GeotiffRasterFormat,
-    RasterFormat,
     load_raster_format,
 )
 from rslearn.utils.vector_format import (
@@ -30,7 +29,7 @@ class FileTileStoreLayer(TileStoreLayer):
         self,
         path: UPath,
         projection: Projection | None = None,
-        raster_format: RasterFormat = GeotiffRasterFormat(),
+        raster_format: GeotiffRasterFormat = GeotiffRasterFormat(),
         vector_format: VectorFormat = GeojsonVectorFormat(),
     ):
         """Creates a new FileTileStoreLayer.
@@ -70,6 +69,8 @@ class FileTileStoreLayer(TileStoreLayer):
             bounds: the bounds of the raster
             array: the raster data
         """
+        if self.projection is None:
+            raise ValueError("need a projection to encode and write raster data")
         self.raster_format.encode_raster(self.path, self.projection, bounds, array)
 
     def get_raster_bounds(self) -> PixelBounds:
@@ -93,6 +94,8 @@ class FileTileStoreLayer(TileStoreLayer):
         Args:
             data: the vector data
         """
+        if self.projection is None:
+            raise ValueError("need a projection to encode and write vector data")
         self.vector_format.encode_vector(self.path, self.projection, data)
 
     def get_metadata(self) -> LayerMetadata:
@@ -125,7 +128,7 @@ class FileTileStore(TileStore):
     def __init__(
         self,
         path: UPath,
-        raster_format: RasterFormat = GeotiffRasterFormat(),
+        raster_format: GeotiffRasterFormat = GeotiffRasterFormat(),
         vector_format: VectorFormat = GeojsonVectorFormat(),
     ):
         """Initialize a new FileTileStore.

rslearn/tile_stores/tile_store.py

@@ -1,5 +1,6 @@
 """Base class for tile stores."""
 
+from abc import ABC, abstractmethod
 from datetime import datetime
 from typing import Any
 
@@ -52,12 +53,13 @@ class LayerMetadata:
         )
 
 
-class TileStoreLayer:
+class TileStoreLayer(ABC):
     """An abstract class for a layer in a tile store.
 
     The layer can store one or more raster and vector datas.
     """
 
+    @abstractmethod
     def read_raster(self, bounds: PixelBounds) -> npt.NDArray[Any] | None:
         """Read raster data from the store.
 
@@ -67,8 +69,9 @@ class TileStoreLayer:
         Returns:
             the raster data
         """
-        raise NotImplementedError
+        pass
 
+    @abstractmethod
     def write_raster(self, bounds: PixelBounds, array: npt.NDArray[Any]) -> None:
         """Write raster data to the store.
 
@@ -76,8 +79,9 @@ class TileStoreLayer:
             bounds: the bounds of the raster
             array: the raster data
         """
-        raise NotImplementedError
+        pass
 
+    @abstractmethod
     def read_vector(self, bounds: PixelBounds) -> list[Feature]:
         """Read vector data from the store.
 
@@ -87,20 +91,23 @@ class TileStoreLayer:
         Returns:
             the vector data
         """
-        raise NotImplementedError
+        pass
 
+    @abstractmethod
     def write_vector(self, data: list[Feature]) -> None:
         """Save vector tiles to the store.
 
         Args:
             data: the vector data
         """
-        raise NotImplementedError
+        pass
 
+    @abstractmethod
     def get_metadata(self) -> LayerMetadata:
         """Get the LayerMetadata associated with this layer."""
-        raise NotImplementedError
+        pass
 
+    @abstractmethod
     def set_property(self, key: str, value: Any) -> None:
         """Set a property in the metadata for this layer.
 
@@ -108,7 +115,12 @@ class TileStoreLayer:
             key: the property key
             value: the property value
         """
-        raise NotImplementedError
+        pass
+
+    @abstractmethod
+    def get_raster_bounds(self) -> PixelBounds:
+        """Get the bounds of the raster data in the store."""
+        pass
 
 
 class TileStore:

rslearn/train/callbacks/freeze_unfreeze.py

@@ -14,7 +14,7 @@ class FreezeUnfreeze(BaseFinetuning):
         module_selector: list[str | int],
         unfreeze_at_epoch: int | None = None,
         unfreeze_lr_factor: float = 1,
-    ):
+    ) -> None:
         """Creates a new FreezeUnfreeze.
 
         Args:
@@ -40,7 +40,7 @@ class FreezeUnfreeze(BaseFinetuning):
                 target_module = getattr(target_module, k)
         return target_module
 
-    def freeze_before_training(self, pl_module: LightningModule):
+    def freeze_before_training(self, pl_module: LightningModule) -> None:
         """Freeze the model at the beginning of training.
 
         Args:
@@ -51,7 +51,7 @@ class FreezeUnfreeze(BaseFinetuning):
 
     def finetune_function(
         self, pl_module: LightningModule, current_epoch: int, optimizer: Optimizer
-    ):
+    ) -> None:
         """Check whether we should unfreeze the model on each epoch.
 
         Args: