Sentinel-2 vessel attribute prediction

data/sentinel2_vessel_attribute/config.json

@@ -0,0 +1,93 @@
+{
+  "layers": {
+    "info": {
+      "format": {
+        "name": "geojson"
+      },
+      "type": "vector"
+    },
+    "output": {
+      "type": "vector"
+    },
+    "sentinel2": {
+      "band_sets": [
+        {
+          "bands": [
+            "B02",
+            "B03",
+            "B04",
+            "B08"
+          ],
+          "dtype": "uint16",
+          "format": {
+            "geotiff_options": {
+              "compress": "zstd",
+              "predictor": 2,
+              "zstd_level": 1
+            },
+            "name": "geotiff"
+          }
+        },
+        {
+          "bands": [
+            "B05",
+            "B06",
+            "B07",
+            "B8A",
+            "B11",
+            "B12"
+          ],
+          "dtype": "uint16",
+          "format": {
+            "geotiff_options": {
+              "compress": "zstd",
+              "predictor": 2,
+              "zstd_level": 1
+            },
+            "name": "geotiff"
+          },
+          "zoom_offset": -1
+        },
+        {
+          "bands": [
+            "B01",
+            "B09",
+            "B10"
+          ],
+          "dtype": "uint16",
+          "format": {
+            "geotiff_options": {
+              "compress": "zstd",
+              "predictor": 2,
+              "zstd_level": 1
+            },
+            "name": "geotiff"
+          },
+          "zoom_offset": -2
+        }
+      ],
+      "data_source": {
+        "harmonize": true,
+        "index_cache_dir": "cache/sentinel2/",
+        "max_time_delta": "0d",
+        "name": "rslearn.data_sources.gcp_public_data.Sentinel2",
+        "query_config": {
+          "max_matches": 1
+        },
+        "use_rtree_index": false
+      },
+      "type": "raster"
+    }
+  },
+  "tile_store": {
+    "class_path": "rslearn.tile_stores.default.DefaultTileStore",
+    "init_args": {
+      "geotiff_options": {
+        "compress": "zstd",
+        "predictor": 2,
+        "zstd_level": 1
+      },
+      "path_suffix": "gs://rslearn-eai/datasets/sentinel2_vessel_attribute/dataset_v1/20241212/tiles"
+    }
+  }
+}

data/sentinel2_vessel_attribute/config.yaml

@@ -0,0 +1,187 @@
+model:
+  class_path: rslp.sentinel2_vessel_attribute.train.VesselAttributeLightningModule
+  init_args:
+    model:
+      class_path: rslearn.models.multitask.MultiTaskModel
+      init_args:
+        encoder:
+          - class_path: rslearn.models.satlaspretrain.SatlasPretrain
+            init_args:
+              model_identifier: "Sentinel2_SwinB_SI_MS"
+              fpn: true
+        decoders:
+          length:
+            - class_path: rslearn.models.pooling_decoder.PoolingDecoder
+              init_args:
+                in_channels: 128
+                out_channels: 1
+                num_conv_layers: 1
+                num_fc_layers: 2
+            - class_path: rslearn.train.tasks.regression.RegressionHead
+          width:
+            - class_path: rslearn.models.pooling_decoder.PoolingDecoder
+              init_args:
+                in_channels: 128
+                out_channels: 1
+                num_conv_layers: 1
+                num_fc_layers: 2
+            - class_path: rslearn.train.tasks.regression.RegressionHead
+          speed:
+            - class_path: rslearn.models.pooling_decoder.PoolingDecoder
+              init_args:
+                in_channels: 128
+                out_channels: 1
+                num_conv_layers: 1
+                num_fc_layers: 2
+            - class_path: rslearn.train.tasks.regression.RegressionHead
+          heading_x:
+            - class_path: rslearn.models.pooling_decoder.PoolingDecoder
+              init_args:
+                in_channels: 128
+                out_channels: 1
+                num_conv_layers: 1
+                num_fc_layers: 2
+            - class_path: rslearn.train.tasks.regression.RegressionHead
+          heading_y:
+            - class_path: rslearn.models.pooling_decoder.PoolingDecoder
+              init_args:
+                in_channels: 128
+                out_channels: 1
+                num_conv_layers: 1
+                num_fc_layers: 2
+            - class_path: rslearn.train.tasks.regression.RegressionHead
+          ship_type:
+            - class_path: rslearn.models.pooling_decoder.PoolingDecoder
+              init_args:
+                in_channels: 128
+                out_channels: 9
+                num_conv_layers: 1
+                num_fc_layers: 2
+            - class_path: rslearn.train.tasks.classification.ClassificationHead
+    lr: 0.0001
+    plateau_factor: 0.2
+    plateau_patience: 2
+    plateau_min_lr: 0
+    plateau_cooldown: 10
+data:
+  class_path: rslearn.train.data_module.RslearnDataModule
+  init_args:
+    path: gs://rslearn-eai/datasets/sentinel2_vessel_attribute/dataset_v1/20250205/
+    inputs:
+      image:
+        data_type: "raster"
+        layers: ["sentinel2"]
+        bands: ["B04", "B03", "B02", "B05", "B06", "B07", "B08", "B11", "B12"]
+        passthrough: true
+        dtype: FLOAT32
+      info:
+        data_type: "vector"
+        layers: ["info"]
+        is_target: true
+    task:
+      class_path: rslp.sentinel2_vessel_attribute.train.VesselAttributeMultiTask
+      init_args:
+        length_buckets: [10, 20, 30, 50, 75, 100, 150, 200]
+        width_buckets: [5, 10, 20]
+        speed_buckets: [2, 4, 8]
+        tasks:
+          length:
+            class_path: rslearn.train.tasks.regression.RegressionTask
+            init_args:
+              property_name: "length"
+              allow_invalid: true
+              scale_factor: 0.01
+              metric_mode: l1
+          width:
+            class_path: rslearn.train.tasks.regression.RegressionTask
+            init_args:
+              property_name: "width"
+              allow_invalid: true
+              scale_factor: 0.01
+              metric_mode: l1
+          speed:
+            class_path: rslearn.train.tasks.regression.RegressionTask
+            init_args:
+              property_name: "sog"
+              allow_invalid: true
+              scale_factor: 0.01
+              metric_mode: l1
+          heading_x:
+            class_path: rslearn.train.tasks.regression.RegressionTask
+            init_args:
+              property_name: "cog_x"
+              allow_invalid: true
+              metric_mode: l1
+          heading_y:
+            class_path: rslearn.train.tasks.regression.RegressionTask
+            init_args:
+              property_name: "cog_y"
+              allow_invalid: true
+              metric_mode: l1
+          ship_type:
+            class_path: rslearn.train.tasks.classification.ClassificationTask
+            init_args:
+              property_name: "type"
+              allow_invalid: true
+              classes: ["cargo", "tanker", "passenger", "service", "tug", "pleasure", "fishing", "enforcement", "sar"]
+              metric_kwargs:
+                average: "micro"
+        input_mapping:
+          length:
+            info: "targets"
+          width:
+            info: "targets"
+          speed:
+            info: "targets"
+          heading_x:
+            info: "targets"
+          heading_y:
+            info: "targets"
+          ship_type:
+            info: "targets"
+    batch_size: 32
+    num_workers: 64
+    default_config:
+      transforms:
+        - class_path: rslearn.train.transforms.normalize.Normalize
+          init_args:
+            mean: 0
+            std: 10000
+            valid_range: [0, 1]
+    train_config:
+      transforms:
+        - class_path: rslearn.train.transforms.normalize.Normalize
+          init_args:
+            mean: 0
+            std: 10000
+            valid_range: [0, 1]
+        - class_path: rslp.sentinel2_vessel_attribute.train.VesselAttributeFlip
+      tags:
+        split: "train"
+    val_config:
+      tags:
+        split: "val"
+    test_config:
+      tags:
+        split: "val"
+    predict_config:
+      groups: ["attribute_predict"]
+      skip_targets: true
+trainer:
+  max_epochs: 100
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+      init_args:
+        logging_interval: "epoch"
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        save_top_k: 1
+        save_last: true
+        monitor: val_loss
+        mode: min
+    - class_path: rslearn.train.prediction_writer.RslearnWriter
+      init_args:
+        path: placeholder
+        output_layer: output
+rslp_project: sentinel2_vessel_attribute
+rslp_experiment: data_20250205_regress_00