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

rslearn/models/panopticon_data/sensors/modis_terra.yaml

@@ -0,0 +1,99 @@
+instrument: 'MSI'  # You may want to update this based on the actual instrument
+processing_level: 'L1C'  # You may want to update this based on the actual processing level
+
+bands:
+  B01:
+    name: 'Band 1'  # You may want to provide a descriptive name
+    gaussian:
+      mu: 411.633593
+      sigma: 14.652
+
+  B02:
+    name: 'Band 2' # coastal aerosol
+    gaussian:
+      mu: 442.155
+      sigma: 8.434933
+
+  B03:
+    name: 'Band 3'
+    gaussian:
+      mu: 466.122
+      sigma: 18.894
+
+  B04:
+    name: 'Band 4' # ~ blue
+    gaussian:
+      mu: 487.078
+      sigma: 10.633
+
+  B05:
+    name: 'Band 5'
+    gaussian:
+      mu: 529.783
+      sigma: 10.782657
+
+  B06:
+    name: 'Band 6'
+    gaussian:
+      mu: 546.981
+      sigma: 10.331
+
+  B07:
+    name: 'Band 7' # green
+    gaussian:
+      mu: 554.026
+      sigma: 17.766939
+
+  B08:
+    name: 'Band 8'
+    gaussian:
+      mu: 644.898
+      sigma: 34.650989
+
+  B09:
+    name: 'Band 9' # red
+    gaussian:
+      mu: 665.695
+      sigma: 10.117
+
+  B10:
+    name: 'Band 10'
+    gaussian:
+      mu: 677.068
+      sigma: 9.007702
+
+  B11:
+    name: 'Band 11'
+    gaussian:
+      mu: 746.736
+      sigma: 9.952
+
+  B12:
+    name: 'Band 12'
+    gaussian:
+      mu: 857.323
+      sigma: 34.696136
+
+  B13:
+    name: 'Band 13'
+    gaussian:
+      mu: 866.55
+      sigma: 13.421421
+
+  B14:
+    name: 'Band 14'
+    gaussian:
+      mu: 1241.597
+      sigma: 23.356
+
+  B15:
+    name: 'Band 15'
+    gaussian:
+      mu: 1627.972
+      sigma: 27.593
+
+  B16:
+    name: 'Band 16'
+    gaussian:
+      mu: 2113.124
+      sigma: 47.341375

rslearn/models/panopticon_data/sensors/qb2_ge1.yaml

@@ -0,0 +1,34 @@
+instrument: '4-band'
+processing_level: 'SR'
+
+# QuickBird-2 and GeoEye-1
+# In the corresponding dataset, either of qb2 or ge1 are used. However, since
+# their bands have the same properties, we use the same file for both
+bands:
+  B01:
+    name: 'blue'
+    gaussian:
+      mu: 482.417803
+      sigma: 45.029733
+    GSD: 1.24
+
+  B02:
+    name: 'green'
+    gaussian:
+      mu: 547.272289
+      sigma: 58.295452
+    GSD: 1.24
+
+  B03:
+    name: 'red'
+    gaussian:
+      mu: 665.031763
+      sigma: 35.775908
+    GSD: 1.24
+
+  B04:
+    name: 'nir'
+    gaussian:
+      mu: 840.773239
+      sigma: 83.700732
+    GSD: 1.24

rslearn/models/panopticon_data/sensors/sentinel1.yaml

@@ -0,0 +1,85 @@
+bands:
+# Modify the band numbering to be the Helios band strings in uppercase
+  VV:
+    name: VV
+    gaussian:
+      mu: -1
+      sigma: -1
+    orbit: BOTH
+
+  VH:
+    name: VH
+    gaussian:
+      mu: -2
+      sigma: -1
+    orbit: BOTH
+
+  HH:
+    name: HH
+    gaussian:
+      mu: -3
+      sigma: -1
+    orbit: BOTH
+
+  HV:
+    name: HV
+    gaussian:
+      mu: -4
+      sigma: -1
+    orbit: BOTH
+
+  VV_ASCENDING:
+    name: VV
+    gaussian:
+      mu: -5
+      sigma: -1
+    orbit: ASCENDING
+
+  VH_ASCENDING:
+    name: VH
+    gaussian:
+      mu: -6
+      sigma: -1
+    orbit: ASCENDING
+
+  HH_ASCENDING:
+    name: HH
+    gaussian:
+      mu: -7
+      sigma: -1
+    orbit: ASCENDING
+
+  HV_ASCENDING:
+    name: HV
+    gaussian:
+      mu: -8
+      sigma: -1
+    orbit: ASCENDING
+
+  VV_DESCENDING:
+    name: VV
+    gaussian:
+      mu: -9
+      sigma: -1
+    orbit: DESCENDING
+
+  VH_DESCENDING:
+    name: VH
+    gaussian:
+      mu: -10
+      sigma: -1
+    orbit: DESCENDING
+
+  HH_DESCENDING:
+    name: HH
+    gaussian:
+      mu: -11
+      sigma: -1
+    orbit: DESCENDING
+
+  HV_DESCENDING:
+    name: HV
+    gaussian:
+      mu: -12
+      sigma: -1
+    orbit: DESCENDING

rslearn/models/panopticon_data/sensors/sentinel2.yaml

@@ -0,0 +1,97 @@
+instrument: MSI
+level: L1C
+
+srf_filename: rfs_sentinel2_a_13b.npy
+
+bands:
+  B01:
+    name: '01 - Coastal aerosol'
+    gaussian:
+      mu: 442.922568734037
+      sigma: 7.248330717861807
+    GSD: 60
+
+  B02:
+    name: '02 - Blue'
+    gaussian:
+      mu: 492.9971095687347
+      sigma: 23.810316659477703
+    GSD: 10
+
+  B03:
+    name: '03 - Green'
+    gaussian:
+      mu: 559.5987534818435
+      sigma: 12.768882177939654
+    GSD: 10
+
+  B04:
+    name: '04 - Red'
+    gaussian:
+      mu: 664.6300422881802
+      sigma: 11.757355524910432
+    GSD: 10
+
+  B05:
+    name: '05 - Vegetation Red Edge'
+    gaussian:
+      mu: 704.0059319834206
+      sigma: 5.362493403740522
+    GSD: 20
+
+  B06:
+    name: '06 - Vegetation Red Edge'
+    gaussian:
+      mu: 740.5521320760564
+      sigma: 5.2330999827526155
+    GSD: 20
+
+  B07:
+    name: '07 - Vegetation Red Edge'
+    gaussian:
+      mu: 782.4190761493182
+      sigma: 7.212484180540051
+    GSD: 20
+
+  B08:
+    name: '08 - NIR'
+    gaussian:
+      mu: 827.5394062383036
+      sigma: 36.79409520400872
+    GSD: 10
+
+  B8A:
+    name: '08A - Vegetation Red Edge'
+    gaussian:
+      mu: 864.7801257644385
+      sigma: 8.07210759526792
+    GSD: 20
+
+  B09:
+    name: '09 - Water vapour'
+    gaussian:
+      mu: 945.0294901407692
+      sigma: 7.518965324285279
+    GSD: 60
+
+  B10:
+    name: '10 - SWIR - Cirrus'
+    gaussian:
+      mu: 1373.3636762095748
+      sigma: 11.163498916290587
+    GSD: 60
+    p: 2
+
+  B11:
+    name: '11 - SWIR'
+    gaussian:
+      mu: 1613.8624163477282
+      sigma: 34.4986558584479
+    GSD: 20
+
+  B12:
+    name: '12 - SWIR'
+    gaussian:
+      mu: 2203.6182057820033
+      sigma: 64.60648125885301
+    GSD: 20

rslearn/models/panopticon_data/sensors/superdove.yaml

@@ -0,0 +1,60 @@
+instrument: 'PSB.SD'
+processing_level: 'NA'
+
+bands:
+  B01:
+    name: 'Coastal Blue'
+    gaussian:
+      mu: 443.704
+      sigma: 7.9672
+    GSD: 3.7
+
+  B02:
+    name: 'Blue'
+    gaussian:
+      mu: 490.973
+      sigma: 20.5096
+    GSD: 3.7
+
+  B03:
+    name: 'Green I'
+    gaussian:
+      mu: 532.719
+      sigma: 14.4789
+    GSD: 3.7
+
+  B04:
+    name: 'Green'
+    gaussian:
+      mu: 565.811
+      sigma: 15.2825
+    GSD: 3.7
+
+  B05:
+    name: 'Yellow'
+    gaussian:
+      mu: 611.587
+      sigma: 9.33594
+    GSD: 3.7
+
+  B06:
+    name: 'Red'
+    gaussian:
+      mu: 665.751
+      sigma: 12.6253
+    GSD: 3.7
+
+
+  B07:
+    name: 'Red Edge'
+    gaussian:
+      mu: 706.918
+      sigma: 6.92817
+    GSD: 3.7
+
+  B08:
+    name: 'NIR'
+    gaussian:
+      mu: 864.831
+      sigma: 15.2059
+    GSD: 3.7

rslearn/models/panopticon_data/sensors/wv23.yaml

@@ -0,0 +1,63 @@
+processing_level: 'SR'
+
+srf_filename: rfs_wv23_recon.npy
+
+# Worldview-2 and Worldview-3_VNIR
+# In the corresponding dataset, either of the two above satellites are used. However, since
+# their bands have the same properties, we use the same file for both
+bands:
+  B01:
+    name: 'coastal'
+    gaussian:
+      mu: 427.911967712222
+      sigma: 17.620786889126904
+    GSD: 1.24
+
+  B02:
+    name: 'blue'
+    gaussian:
+      mu: 482.40648216687816
+      sigma: 22.189227543486883
+    GSD: 1.24
+
+  B03:
+    name: 'green'
+    gaussian:
+      mu: 545.1346759174888
+      sigma: 27.270655243664613
+    GSD: 1.24
+
+  B04:
+    name: 'yellow'
+    gaussian:
+      mu: 604.6891589644367
+      sigma: 15.166919163740687
+    GSD: 1.24
+
+  B05:
+    name: 'red'
+    gaussian:
+      mu: 660.5315665213377
+      sigma: 23.075009737550587
+    GSD: 1.24
+
+  B06:
+    name: 'red edge'
+    gaussian:
+      mu: 723.1823149413602
+      sigma: 15.151759763702627
+    GSD: 1.24
+
+  B07:
+    name: 'nir1'
+    gaussian:
+      mu: 823.9274208290032
+      sigma: 42.09302701870739
+    GSD: 1.24
+
+  B08:
+    name: 'nir2'
+    gaussian:
+      mu: 906.4611534199017
+      sigma: 36.61665833552878
+    GSD: 1.24

rslearn/models/pick_features.py

@@ -2,10 +2,15 @@
 
 from typing import Any
 
-import torch
+from rslearn.train.model_context import ModelContext
 
+from .component import (
+    FeatureMaps,
+    IntermediateComponent,
+)
 
-class PickFeatures(torch.nn.Module):
+
+class PickFeatures(IntermediateComponent):
     """Picks a subset of feature maps in a multi-scale feature map list."""
 
     def __init__(self, indexes: list[int]):
@@ -19,15 +24,17 @@ class PickFeatures(torch.nn.Module):
 
     def forward(
         self,
-        features: list[torch.Tensor],
-        inputs: list[dict[str, Any]] | None = None,
-        targets: list[dict[str, Any]] | None = None,
-    ) -> list[torch.Tensor]:
+        intermediates: Any,
+        context: ModelContext,
+    ) -> FeatureMaps:
         """Pick a subset of the features.
 
         Args:
-            features: input features
-            inputs: raw inputs, not used
-            targets: targets, not used
+            intermediates: the output from the previous component, which must be a FeatureMaps.
+            context: the model context.
         """
-        return [features[idx] for idx in self.indexes]
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to PickFeatures must be FeatureMaps")
+
+        new_features = [intermediates.feature_maps[idx] for idx in self.indexes]
+        return FeatureMaps(new_features)

rslearn/models/pooling_decoder.py

@@ -4,8 +4,16 @@ from typing import Any
 
 import torch
 
+from rslearn.train.model_context import ModelContext
 
-class PoolingDecoder(torch.nn.Module):
+from .component import (
+    FeatureMaps,
+    FeatureVector,
+    IntermediateComponent,
+)
+
+
+class PoolingDecoder(IntermediateComponent):
     """Decoder that computes flat vector from a 2D feature map.
 
     It inputs multi-scale features, but only uses the last feature map. Then applies a
@@ -21,7 +29,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,20 +65,65 @@ 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, intermediates: Any, context: ModelContext) -> Any:
         """Compute flat output vector from multi-scale feature map.
 
         Args:
-            features: list of feature maps at different resolutions.
-            inputs: original inputs (ignored).
+            intermediates: the output from the previous component, which must be a FeatureMaps.
+            context: the model context.
 
         Returns:
             flat feature vector
         """
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to PoolingDecoder must be a FeatureMaps")
+
         # Only use last feature map.
-        features = features[-1]
+        features = intermediates.feature_maps[-1]
 
         features = self.conv_layers(features)
         features = torch.amax(features, dim=(2, 3))
         features = self.fc_layers(features)
-        return self.output_layer(features)
+        return FeatureVector(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, intermediates: Any, context: ModelContext) -> Any:
+        """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(intermediates, context)
+        # BC -> BCHW
+        h, w = context.inputs[0][self.image_key].shape[1:3]
+        feat_map = output_probs.feature_vector[:, :, None, None].repeat([1, 1, h, w])
+        return FeatureMaps([feat_map])

rslearn/models/presto/__init__.py

@@ -0,0 +1,5 @@
+"""Presto."""
+
+from .presto import Presto
+
+__all__ = ["Presto"]