add vessel attribute code

allenai · Jan 29, 2025 · 782c60f · 782c60f
1 parent 8a39971
commit 782c60f
Show file tree

Hide file tree

Showing 20 changed files with 3,589 additions and 0 deletions.
diff --git a/data/sentinel2_vessel_attribute/config.json b/data/sentinel2_vessel_attribute/config.json
@@ -0,0 +1,89 @@
+{
+  "layers": {
+    "info": {
+      "format": {
+        "name": "geojson"
+      },
+      "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
+      }
+    }
+  }
+}
diff --git a/data/sentinel2_vessel_attribute/config.yaml b/data/sentinel2_vessel_attribute/config.yaml
@@ -0,0 +1,187 @@
+model:
+  class_path: rslearn.train.lightning_module.RslearnLightningModule
+  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/20241212/
+    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:
+        tasks:
+          length:
+            class_path: rslearn.train.tasks.regression.RegressionTask
+            init_args:
+              property_name: "length"
+              allow_invalid: true
+              scale_factor: 0.01
+              metric_mode: l1
+              use_within_factor_metric: true
+              within_factor: 0.25
+          width:
+            class_path: rslearn.train.tasks.regression.RegressionTask
+            init_args:
+              property_name: "width"
+              allow_invalid: true
+              scale_factor: 0.01
+              metric_mode: l1
+              use_within_factor_metric: true
+              within_factor: 0.25
+          speed:
+            class_path: rslearn.train.tasks.regression.RegressionTask
+            init_args:
+              property_name: "sog"
+              allow_invalid: true
+              scale_factor: 0.01
+              metric_mode: l1
+              use_within_factor_metric: true
+              within_factor: 0.25
+          heading_x:
+            class_path: rslearn.train.tasks.regression.RegressionTask
+            init_args:
+              property_name: "cog_x"
+              allow_invalid: true
+              metric_mode: l1
+              use_within_factor_metric: true
+              within_factor: 0.25
+          heading_y:
+            class_path: rslearn.train.tasks.regression.RegressionTask
+            init_args:
+              property_name: "cog_y"
+              allow_invalid: true
+              metric_mode: l1
+              use_within_factor_metric: true
+              within_factor: 0.25
+          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"
+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
+rslp_project: sentinel2_vessel_attribute
+rslp_experiment: data_20241212_model_20241212_ms_flip_satlaspretrain_00