rslearn 0.0.14__py3-none-any.whl → 0.0.15__py3-none-any.whl

rslearn/dataset/manage.py

@@ -124,12 +124,24 @@ def prepare_dataset_windows(
             )
             continue
         data_source_cfg = layer_cfg.data_source
+        min_matches = data_source_cfg.query_config.min_matches
 
         # Get windows that need to be prepared for this layer.
+        # Also track which windows are skipped vs previously rejected.
         needed_windows = []
+        windows_skipped = 0
+        windows_rejected = 0
         for window in windows:
             layer_datas = window.load_layer_datas()
             if layer_name in layer_datas and not force:
+                # Window already has layer data - check if it was previously rejected
+                layer_data = layer_datas[layer_name]
+                if len(layer_data.serialized_item_groups) == 0 and min_matches > 0:
+                    # Previously rejected due to min_matches
+                    windows_rejected += 1
+                else:
+                    # Successfully prepared previously
+                    windows_skipped += 1
                 continue
             needed_windows.append(window)
         logger.info(f"Preparing {len(needed_windows)} windows for layer {layer_name}")
@@ -141,8 +153,8 @@ def prepare_dataset_windows(
                     data_source_name=data_source_cfg.name,
                     duration_seconds=time.monotonic() - layer_start_time,
                     windows_prepared=0,
-                    windows_skipped=len(windows),
-                    windows_rejected=0,
+                    windows_skipped=windows_skipped,
+                    windows_rejected=windows_rejected,
                     get_items_attempts=0,
                 )
             )
@@ -184,8 +196,6 @@ def prepare_dataset_windows(
         )
 
         windows_prepared = 0
-        windows_rejected = 0
-        min_matches = data_source_cfg.query_config.min_matches
         for window, result in zip(needed_windows, results):
             layer_datas = window.load_layer_datas()
             layer_datas[layer_name] = WindowLayerData(
@@ -202,8 +212,6 @@ def prepare_dataset_windows(
             else:
                 windows_prepared += 1
 
-        windows_skipped = len(windows) - len(needed_windows)
-
         layer_summaries.append(
             LayerPrepareSummary(
                 layer_name=layer_name,

rslearn/models/clay/clay.py

@@ -8,6 +8,7 @@ from importlib.resources import files
 from typing import Any
 
 import torch
+import torch.nn.functional as F
 import yaml
 from einops import rearrange
 from huggingface_hub import hf_hub_download
@@ -30,6 +31,7 @@ PATCH_SIZE = 8
 CLAY_MODALITIES = ["sentinel-2-l2a", "sentinel-1-rtc", "landsat-c2l1", "naip"]
 CONFIG_DIR = files("rslearn.models.clay.configs")
 CLAY_METADATA_PATH = str(CONFIG_DIR / "metadata.yaml")
+DEFAULT_IMAGE_RESOLUTION = 128  # image resolution during pretraining
 
 
 def get_clay_checkpoint_path(
@@ -49,6 +51,7 @@ class Clay(torch.nn.Module):
         modality: str = "sentinel-2-l2a",
         checkpoint_path: str | None = None,
         metadata_path: str = CLAY_METADATA_PATH,
+        do_resizing: bool = False,
     ) -> None:
         """Initialize the Clay model.
 
@@ -57,6 +60,7 @@ class Clay(torch.nn.Module):
             modality: The modality to use (subset of CLAY_MODALITIES).
             checkpoint_path: Path to clay-v1.5.ckpt, if None, fetch from HF Hub.
             metadata_path: Path to metadata.yaml.
+            do_resizing: Whether to resize the image to the input resolution.
         """
         super().__init__()
 
@@ -95,6 +99,14 @@ class Clay(torch.nn.Module):
 
         self.model_size = model_size
         self.modality = modality
+        self.do_resizing = do_resizing
+
+    def _resize_image(self, image: torch.Tensor, original_hw: int) -> torch.Tensor:
+        """Resize the image to the input resolution."""
+        new_hw = self.patch_size if original_hw == 1 else DEFAULT_IMAGE_RESOLUTION
+        return F.interpolate(
+            image, size=(new_hw, new_hw), mode="bilinear", align_corners=False
+        )
 
     def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
         """Forward pass for the Clay model.
@@ -114,7 +126,8 @@ class Clay(torch.nn.Module):
         chips = torch.stack(
             [inp[self.modality] for inp in inputs], dim=0
         )  # (B, C, H, W)
-
+        if self.do_resizing:
+            chips = self._resize_image(chips, chips.shape[2])
         order = self.metadata[self.modality]["band_order"]
         wavelengths = []
         for band in self.metadata[self.modality]["band_order"]:

rslearn/models/croma.py

@@ -7,6 +7,7 @@ from enum import Enum
 from typing import Any
 
 import torch
+import torch.nn.functional as F
 from einops import rearrange
 from upath import UPath
 
@@ -99,6 +100,7 @@ class Croma(torch.nn.Module):
         modality: CromaModality,
         pretrained_path: str | None = None,
         image_resolution: int = DEFAULT_IMAGE_RESOLUTION,
+        do_resizing: bool = False,
     ) -> None:
         """Instantiate a new Croma instance.
 
@@ -107,12 +109,21 @@ class Croma(torch.nn.Module):
             modality: the modalities to configure the model to accept.
             pretrained_path: the local path to the pretrained weights. Otherwise it is
                 downloaded and cached in temp directory.
-            image_resolution: the width and height of the input images.
+            image_resolution: the width and height of the input images passed to the model. if do_resizing is True, the image will be resized to this resolution.
+            do_resizing: Whether to resize the image to the input resolution.
         """
         super().__init__()
         self.size = size
         self.modality = modality
-        self.image_resolution = image_resolution
+        self.do_resizing = do_resizing
+        if not do_resizing:
+            self.image_resolution = image_resolution
+        else:
+            # With single pixel input, we always resample to the patch size.
+            if image_resolution == 1:
+                self.image_resolution = PATCH_SIZE
+            else:
+                self.image_resolution = DEFAULT_IMAGE_RESOLUTION
 
         # Cache the CROMA weights to a deterministic path in temporary directory if the
         # path is not provided by the user.
@@ -137,7 +148,16 @@ class Croma(torch.nn.Module):
             pretrained_path=pretrained_path,
             size=size.value,
             modality=modality.value,
-            image_resolution=image_resolution,
+            image_resolution=self.image_resolution,
+        )
+
+    def _resize_image(self, image: torch.Tensor) -> torch.Tensor:
+        """Resize the image to the input resolution."""
+        return F.interpolate(
+            image,
+            size=(self.image_resolution, self.image_resolution),
+            mode="bilinear",
+            align_corners=False,
         )
 
     def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
@@ -151,8 +171,11 @@ class Croma(torch.nn.Module):
         sentinel2: torch.Tensor | None = None
         if self.modality in [CromaModality.BOTH, CromaModality.SENTINEL1]:
             sentinel1 = torch.stack([inp["sentinel1"] for inp in inputs], dim=0)
+            sentinel1 = self._resize_image(sentinel1) if self.do_resizing else sentinel1
         if self.modality in [CromaModality.BOTH, CromaModality.SENTINEL2]:
             sentinel2 = torch.stack([inp["sentinel2"] for inp in inputs], dim=0)
+            sentinel2 = self._resize_image(sentinel2) if self.do_resizing else sentinel2
+
         outputs = self.model(
             SAR_images=sentinel1,
             optical_images=sentinel2,

rslearn/models/satlaspretrain.py

@@ -4,15 +4,14 @@ from typing import Any
 
 import satlaspretrain_models
 import torch
+import torch.nn.functional as F
 
 
 class SatlasPretrain(torch.nn.Module):
     """SatlasPretrain backbones."""
 
     def __init__(
-        self,
-        model_identifier: str,
-        fpn: bool = False,
+        self, model_identifier: str, fpn: bool = False, resize_to_pretrain: bool = False
     ) -> None:
         """Instantiate a new SatlasPretrain instance.
 
@@ -21,6 +20,8 @@ class SatlasPretrain(torch.nn.Module):
                 https://github.com/allenai/satlaspretrain_models
             fpn: whether to include the feature pyramid network, otherwise only the
                 Swin-v2-Transformer is used.
+            resize_to_pretrain: whether to resize inputs to the pretraining input
+                size (512 x 512)
         """
         super().__init__()
         weights_manager = satlaspretrain_models.Weights()
@@ -49,6 +50,19 @@ class SatlasPretrain(torch.nn.Module):
                 [16, 1024],
                 [32, 2048],
             ]
+        self.resize_to_pretrain = resize_to_pretrain
+
+    def maybe_resize(self, data: torch.Tensor) -> list[torch.Tensor]:
+        """Resize to pretraining sizes if resize_to_pretrain == True."""
+        if self.resize_to_pretrain:
+            return F.interpolate(
+                data,
+                size=(512, 512),
+                mode="bilinear",
+                align_corners=False,
+            )
+        else:
+            return data
 
     def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
         """Compute feature maps from the SatlasPretrain backbone.
@@ -58,7 +72,7 @@ class SatlasPretrain(torch.nn.Module):
                 process.
         """
         images = torch.stack([inp["image"] for inp in inputs], dim=0)
-        return self.model(images)
+        return self.model(self.maybe_resize(images))
 
     def get_backbone_channels(self) -> list:
         """Returns the output channels of this model when used as a backbone.

rslearn/models/terramind.py

@@ -4,6 +4,7 @@ from enum import Enum
 from typing import Any
 
 import torch
+import torch.nn.functional as F
 from einops import rearrange
 from terratorch.registry import BACKBONE_REGISTRY
 
@@ -18,6 +19,8 @@ class TerramindSize(str, Enum):
     LARGE = "large"
 
 
+# Pretraining image size for Terramind
+IMAGE_SIZE = 224
 # Default patch size for Terramind
 PATCH_SIZE = 16
 
@@ -89,12 +92,14 @@ class Terramind(torch.nn.Module):
         self,
         model_size: TerramindSize,
         modalities: list[str] = ["S2L2A"],
+        do_resizing: bool = False,
     ) -> None:
         """Initialize the Terramind model.
 
         Args:
             model_size: The size of the Terramind model.
             modalities: The modalities to use.
+            do_resizing: Whether to resize the input images to the pretraining resolution.
         """
         super().__init__()
 
@@ -116,6 +121,7 @@ class Terramind(torch.nn.Module):
 
         self.model_size = model_size
         self.modalities = modalities
+        self.do_resizing = do_resizing
 
     def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
         """Forward pass for the Terramind model.
@@ -132,6 +138,19 @@ class Terramind(torch.nn.Module):
             if modality not in inputs[0]:
                 continue
             cur = torch.stack([inp[modality] for inp in inputs], dim=0)  # (B, C, H, W)
+            if self.do_resizing and (
+                cur.shape[2] != IMAGE_SIZE or cur.shape[3] != IMAGE_SIZE
+            ):
+                if cur.shape[2] == 1 and cur.shape[3] == 1:
+                    new_height, new_width = PATCH_SIZE, PATCH_SIZE
+                else:
+                    new_height, new_width = IMAGE_SIZE, IMAGE_SIZE
+                cur = F.interpolate(
+                    cur,
+                    size=(new_height, new_width),
+                    mode="bilinear",
+                    align_corners=False,
+                )
             model_inputs[modality] = cur
 
         # By default, the patch embeddings are averaged over all modalities to reduce output tokens

{rslearn-0.0.14.dist-info → rslearn-0.0.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.14
+Version: 0.0.15
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License

{rslearn-0.0.14.dist-info → rslearn-0.0.15.dist-info}/RECORD

@@ -40,7 +40,7 @@ rslearn/dataset/add_windows.py,sha256=pwCEvwLE1jQCoqQxw6CJ-sP46ayWppFa2hGYIB6VVk
 rslearn/dataset/dataset.py,sha256=bjf9nI55j-MF0bIQWSNPjNbpfqnLK4jy-96TAcwO0MM,5214
 rslearn/dataset/handler_summaries.py,sha256=wI99RDk5erCWkzl1A7Uc4chatQ9KWIr4F_0Hxr9Co6s,2607
 rslearn/dataset/index.py,sha256=Wni5m6h4gisRB54fPLnCfUrRTEsJ5EvwS0fs9sYc2wg,6025
-rslearn/dataset/manage.py,sha256=owelBiBqvoIQYLhFMDK4ULzcoGBNE27JV8kl68jf3wg,18563
+rslearn/dataset/manage.py,sha256=IURlbCtm9a5f4d52AXfte1yyodlf6MgjfYn3__GdIL4,19062
 rslearn/dataset/materialize.py,sha256=-z47svc_JqGhzkp8kq5Hd9fykWNqFEUCQezo887TWBw,22056
 rslearn/dataset/remap.py,sha256=6MaImsY02GNACpvRM81RvWmjZWRfAHxo_R3Ox6XLF6A,2723
 rslearn/dataset/window.py,sha256=I5RqZ12jlIXhohw4qews1x_I4tSDpml709DZRtLiN24,12546
@@ -48,7 +48,7 @@ rslearn/models/__init__.py,sha256=_vWoF9d2Slah8-6XhYhdU4SRsy_CNxXjCGQTD2yvu3Q,22
 rslearn/models/anysat.py,sha256=3Oh2gWxicVdUzOjevBEZf0PuolmCy0KC5Ad7JY-0Plc,7949
 rslearn/models/clip.py,sha256=u5aqYnVB4Jag7o1h8EzPDAc1t2BAHeALA9FcUwP5tfo,2238
 rslearn/models/conv.py,sha256=fWyByeswIOKKzyPmP3erYUlZaKEV0huWHA4CyKTBbfY,1703
-rslearn/models/croma.py,sha256=cOazTp3l2PNJltKrmPqD5Gy4pi3CI03-X9G4T10cX2k,9529
+rslearn/models/croma.py,sha256=n7yunpT7lo8vWWaOpx4yt8jZSXjgWqfgZcZKFW5zuEQ,10591
 rslearn/models/dinov3.py,sha256=9k9kNlXCorQQwKjLGptooANd48TUBsITQ1e4fUomlM4,6337
 rslearn/models/faster_rcnn.py,sha256=uaxX6-E1f0BibaA9sorEg3be83C7kTdTc39pC5jRqwE,8286
 rslearn/models/feature_center_crop.py,sha256=24eOrvLEGGVWPw7kPHyUes5HtYNAX7GZ_NpqDGMILEY,1553
@@ -63,18 +63,18 @@ rslearn/models/prithvi.py,sha256=AIzcO5xk1ggR0MjbfhIzqPVgUKFN7odxygmgyAelfW8,401
 rslearn/models/registry.py,sha256=yCcrOvLkbn07Xtln1j7hAB_kmGw0MGsiR2TloJq9Bmk,504
 rslearn/models/resize_features.py,sha256=asKXWrLHIBrU6GaAV0Ory9YuK7IK104XjhkB4ljzI3A,1289
 rslearn/models/sam2_enc.py,sha256=gNlPokr7eNxO2KvnzDMXNxYM2WRO0YkQPjR4110n6cw,3508
-rslearn/models/satlaspretrain.py,sha256=YpjXl-uClhTZMDmyhN64Fg3AszzT-ymZgJB0fO9RyoY,2419
+rslearn/models/satlaspretrain.py,sha256=b6FR_il6MnWU4UpB9OxInZSK9n0IS0PcQuLrWH4YD8g,3046
 rslearn/models/simple_time_series.py,sha256=oTg_akabYFBExJu7JCpbuM211-ZgQS4WerG2nEYrIZY,12774
 rslearn/models/singletask.py,sha256=z4vN9Yvzz0I-U4KJdVZxLJK2ZV-MIv9tzwCGcOWoUPY,1604
 rslearn/models/ssl4eo_s12.py,sha256=sOGEHcDo-rNdmEyoLu2AVEqfxRM_cv6zpfAmyn5c6tw,3553
 rslearn/models/swin.py,sha256=bMlGePXMFou4A_YSUZzjHgN9NniGXaCWdGQ31xHDKis,5511
 rslearn/models/task_embedding.py,sha256=Z6sf61BLCtvdrdnvjh8500b-KiFp3GeWbT4mOqpaCKk,9100
-rslearn/models/terramind.py,sha256=kipar8sMaHJJ3b8vIgL0-s4qhHcA0Vb854vmlZ9cWh4,7524
+rslearn/models/terramind.py,sha256=5POVk_y29LlbVswa6ojd9gdB70iO41yB9Y2aqVY4WdQ,8327
 rslearn/models/trunk.py,sha256=H1QPQGAKsmocq3OiF66GW8MQI4LffupTDrgZR4Ta7QM,4708
 rslearn/models/unet.py,sha256=WUgLgvvlgV8l_6MIDBl6aX1HNOkb24DfnVRIyYXHCjo,6865
 rslearn/models/upsample.py,sha256=3kWbyWZIk56JJxj8en9pieitbrk3XnbIsTKlEkiDQQY,938
 rslearn/models/use_croma.py,sha256=OSBqMuLp-pDtqPNWAVBfmX4wckmyYCKtUDdGCjJk_K8,17966
-rslearn/models/clay/clay.py,sha256=5RO5H8EM0tKjCwWMQ4xDkKkUCwKpm2K_Yw1alnhvVhU,7773
+rslearn/models/clay/clay.py,sha256=29CGCOysx9duEX4Y6LUNHXck_sHjCFrlV4w8CP_hKmI,8460
 rslearn/models/clay/configs/metadata.yaml,sha256=rZTFh4Yb9htEfbQNOPl4HTbFogEhzwIRqFzG-1uT01Y,4652
 rslearn/models/detr/__init__.py,sha256=GGAnTIhyuvl34IRrJ_4gXjm_01OlM5rbQQ3c3TGfbK8,84
 rslearn/models/detr/box_ops.py,sha256=ORCF6EwMpMBB_VgQT05SjR47dCR2rN2gPhL_gsuUWJs,3236
@@ -154,10 +154,10 @@ rslearn/utils/spatial_index.py,sha256=eomJAUgzmjir8j9HZnSgQoJHwN9H0wGTjmJkMkLLfs
 rslearn/utils/sqlite_index.py,sha256=YGOJi66544e6JNtfSft6YIlHklFdSJO2duxQ4TJ2iu4,2920
 rslearn/utils/time.py,sha256=2ilSLG94_sxLP3y5RSV5L5CG8CoND_dbdzYEHVtN-I8,387
 rslearn/utils/vector_format.py,sha256=EIChYCL6GLOILS2TO2JBkca1TuaWsSubWv6iRS3P2ds,16139
-rslearn-0.0.14.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
-rslearn-0.0.14.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
-rslearn-0.0.14.dist-info/METADATA,sha256=Jbm6ySbM4gkT_5o-RWbRr8APS8TYXq3Q-bWyeda-Uc8,36319
-rslearn-0.0.14.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-rslearn-0.0.14.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
-rslearn-0.0.14.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
-rslearn-0.0.14.dist-info/RECORD,,
+rslearn-0.0.15.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
+rslearn-0.0.15.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
+rslearn-0.0.15.dist-info/METADATA,sha256=HRkJjQfvxCEosmCBvLcLd9nZnXjbmfAgPIknMy_ORBo,36319
+rslearn-0.0.15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+rslearn-0.0.15.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
+rslearn-0.0.15.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
+rslearn-0.0.15.dist-info/RECORD,,