rslearn 0.0.16__py3-none-any.whl → 0.0.18__py3-none-any.whl

{rslearn-0.0.16.dist-info → rslearn-0.0.18.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.16
+Version: 0.0.18
 Summary: A library for developing remote sensing datasets and models
 Author: OlmoEarth Team
 License:                                  Apache License
@@ -343,10 +343,12 @@ directory `/path/to/dataset` and corresponding configuration file at
                 "bands": ["R", "G", "B"]
             }],
             "data_source": {
-                "name": "rslearn.data_sources.gcp_public_data.Sentinel2",
-                "index_cache_dir": "cache/sentinel2/",
-                "sort_by": "cloud_cover",
-                "use_rtree_index": false
+                "class_path": "rslearn.data_sources.gcp_public_data.Sentinel2",
+                "init_args": {
+                    "index_cache_dir": "cache/sentinel2/",
+                    "sort_by": "cloud_cover",
+                    "use_rtree_index": false
+                }
             }
         }
     }
@@ -453,8 +455,10 @@ automate this process. Update the dataset `config.json` with a new layer:
         }],
         "resampling_method": "nearest",
         "data_source": {
-            "name": "rslearn.data_sources.local_files.LocalFiles",
-            "src_dir": "file:///path/to/world_cover_tifs/"
+            "class_path": "rslearn.data_sources.local_files.LocalFiles",
+            "init_args": {
+                "src_dir": "file:///path/to/world_cover_tifs/"
+            }
         }
     }
 },
@@ -516,8 +520,7 @@ model:
 data:
   class_path: rslearn.train.data_module.RslearnDataModule
   init_args:
-    # Replace this with the dataset path.
-    path: /path/to/dataset/
+    path: ${DATASET_PATH}
     # This defines the layers that should be read for each window.
     # The key ("image" / "targets") is what the data will be called in the model,
     # while the layers option specifies which layers will be read.
@@ -615,7 +618,9 @@ trainer:
       ...
     - class_path: rslearn.train.prediction_writer.RslearnWriter
       init_args:
-        path: /path/to/dataset/
+        # We need to include this argument, but it will be overridden with the dataset
+        # path from data.init_args.path.
+        path: placeholder
         output_layer: output
 ```
 
@@ -768,24 +773,43 @@ This will produce PNGs in the vis directory. The visualizations are produced by
 SegmentationTask and overriding the visualize function.
 
 
-### Logging to Weights & Biases
+### Checkpoint and Logging Management
+
+Above, we needed to configure the checkpoint directory in the model config (the
+`dirpath` option under `lightning.pytorch.callbacks.ModelCheckpoint`), and explicitly
+specify the checkpoint path when applying the model. Additionally, metrics are logged
+to the local filesystem and not well organized.
 
-We can log to W&B by setting the logger under trainer in the model configuration file:
+We can instead let rslearn automatically manage checkpoints, along with logging to
+Weights & Biases. To do so, we add project_name, run_name, and management_dir options
+to the model config. The project_name corresponds to the W&B project, and the run name
+corresponds to the W&B name. The management_dir is a directory to store project data;
+rslearn determines a per-project directory at `{management_dir}/{project_name}/{run_name}/`
+and uses it to store checkpoints.
 
 ```yaml
+model:
+  # ...
+data:
+  # ...
 trainer:
   # ...
-  logger:
-    class_path: lightning.pytorch.loggers.WandbLogger
-    init_args:
-      project: land_cover_model
-      name: version_00
+project_name: land_cover_model
+run_name: version_00
+# This sets the option via the MANAGEMENT_DIR environment variable.
+management_dir: ${MANAGEMENT_DIR}
 ```
 
-Now, runs with this model configuration should show on W&B. For `model fit` runs,
-the training and validation loss and accuracy metric will be logged. The accuracy
-metric is provided by SegmentationTask, and additional metrics can be enabled by
-passing the relevant init_args to the task, e.g. mean IoU and F1:
+Now, set the `MANAGEMENT_DIR` environment variable and run `model fit`:
+
+```
+export MANAGEMENT_DIR=./project_data
+rslearn model fit --config land_cover_model.yaml
+```
+
+The training and validation loss and accuracy metric should now be logged to W&B. The
+accuracy metric is provided by SegmentationTask, and additional metrics can be enabled
+by passing the relevant init_args to the task, e.g. mean IoU and F1:
 
 ```yaml
       class_path: rslearn.train.tasks.segmentation.SegmentationTask
@@ -796,6 +820,13 @@ passing the relevant init_args to the task, e.g. mean IoU and F1:
         enable_f1_metric: true
 ```
 
+When calling `model test` and `model predict` with management_dir set, rslearn will
+automatically load the best checkpoint from the project directory, or raise an error if
+no existing checkpoint exists. This behavior can be overridden with the
+`--load_checkpoint_mode` and `--load_checkpoint_required` options (see `--help` for
+details). Logging will be enabled during fit but not test/predict, and this can also
+be overridden, using `--log_mode`.
+
 
 ### Inputting Multiple Sentinel-2 Images
 
@@ -818,10 +849,12 @@ query_config section. This can replace the sentinel2 layer:
             "bands": ["R", "G", "B"]
         }],
         "data_source": {
-            "name": "rslearn.data_sources.gcp_public_data.Sentinel2",
-            "index_cache_dir": "cache/sentinel2/",
-            "sort_by": "cloud_cover",
-            "use_rtree_index": false,
+            "class_path": "rslearn.data_sources.gcp_public_data.Sentinel2",
+            "init_args": {
+              "index_cache_dir": "cache/sentinel2/",
+              "sort_by": "cloud_cover",
+              "use_rtree_index": false
+            },
             "query_config": {
                 "max_matches": 3
             }

{rslearn-0.0.16.dist-info → rslearn-0.0.18.dist-info}/RECORD

@@ -1,13 +1,13 @@
 rslearn/__init__.py,sha256=fFmAen3vxZyosEfPbG0W46IttujYGVxzrGkJ0YutmmY,73
 rslearn/arg_parser.py,sha256=GNlJncO6Ck_dCNrcg7z_SSG61I-2gKn3Ix2tAxIk9CI,1428
 rslearn/const.py,sha256=FUCfsvFAs-QarEDJ0grdy0C1HjUjLpNFYGo5I2Vpc5Y,449
-rslearn/lightning_cli.py,sha256=x8i2QJvEBaYdqh2_f0-ety7_sNEH9UCKRZUPkqWYZdU,17169
+rslearn/lightning_cli.py,sha256=Cihdf3dOQ17b_n4432Y6LmCQ5XFDghW4rGb4fqw-b6g,17525
 rslearn/log_utils.py,sha256=unD9gShiuO7cx5Nnq8qqVQ4qrbOOwFVgcHxN5bXuiAo,941
-rslearn/main.py,sha256=0g1SRO975eC9DTzKqJnwlWHgVo2Pvotyr72KoJBgjew,29060
+rslearn/main.py,sha256=CLVLpJJXc8f8BCG9SBdOuw_taADu3majD1TNSOVC6Ws,28565
 rslearn/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 rslearn/template_params.py,sha256=Vop0Ha-S44ctCa9lvSZRjrMETznJZlR5y_gJrVIwrPg,791
-rslearn/config/__init__.py,sha256=a8xTvYSnpfIzniHgcnSeob5jo5PVBfacpakA_150MME,434
-rslearn/config/dataset.py,sha256=iUFuwzlM9z6n1pGCd40SmAVj3fG6zTXWrlH6eenfon8,21143
+rslearn/config/__init__.py,sha256=n1qpZ0ImshTtLYl5mC73BORYyUcjPJyHiyZkqUY1hiY,474
+rslearn/config/dataset.py,sha256=fp7HT8FeKyOmkR77pR5tVZYAAnuetCLY6zGFVcdPHyY,23087
 rslearn/data_sources/__init__.py,sha256=zzuZUxrlEIw84YpD2I0HJvCoLDB29LbmnKTXiJykzGU,660
 rslearn/data_sources/aws_landsat.py,sha256=0ZQtmd2NCnvLy4vFSB1AlmoguJbiQB_e_T4eS1tnW9Q,20443
 rslearn/data_sources/aws_open_data.py,sha256=lrHnMJTH3NAaRdNjxwCIxSq8rq90IvV4ho-qAG6Hdgc,29348
@@ -34,59 +34,61 @@ rslearn/data_sources/worldcover.py,sha256=n7gi-JRytxkvkUhKT--dVziMcWSSyMbZA7ZCzL
 rslearn/data_sources/worldpop.py,sha256=S3RSc5kTwSs2bcREjNarBsqf3MBX5CN0eHj7Qkx4K74,5625
 rslearn/data_sources/xyz_tiles.py,sha256=P601CvUmoVDC_ZRVhPKaIoPYYCq5UhZ-v7DaGlN5y_0,13797
 rslearn/dataset/__init__.py,sha256=bHtBlEEBCekO-gaJqiww0-VjvZTE5ahx0llleo8bfP8,289
-rslearn/dataset/add_windows.py,sha256=pwCEvwLE1jQCoqQxw6CJ-sP46ayWppFa2hGYIB6VVkc,8494
-rslearn/dataset/dataset.py,sha256=qmZmFfQOHoKVx6_sBYtBR5H4GTNMgETvq0S4XrqafQU,5165
+rslearn/dataset/add_windows.py,sha256=NwIvku6zxCJ9kgVFa5phJc0Gj1Y1bCzh6TLb9nEGl0s,8462
+rslearn/dataset/dataset.py,sha256=A-iXXdGvTCus1kCKC58dk1mEidMe29X3Lox9JEto_TQ,3112
 rslearn/dataset/handler_summaries.py,sha256=wI99RDk5erCWkzl1A7Uc4chatQ9KWIr4F_0Hxr9Co6s,2607
-rslearn/dataset/index.py,sha256=Wni5m6h4gisRB54fPLnCfUrRTEsJ5EvwS0fs9sYc2wg,6025
 rslearn/dataset/manage.py,sha256=-lGSIgk6Z7-verF_POwe4n5w9eSpgyt4nEOcOj382rc,18971
-rslearn/dataset/materialize.py,sha256=x7FewmdqFUviLtPZGZIfAcw1rd0wfKZhk_N-uN-tQms,20922
+rslearn/dataset/materialize.py,sha256=o05OeLk_wWEOsw15oc5yjpD4J-twGCTfXAtxyAQsQ9I,20974
 rslearn/dataset/remap.py,sha256=6MaImsY02GNACpvRM81RvWmjZWRfAHxo_R3Ox6XLF6A,2723
-rslearn/dataset/window.py,sha256=I5RqZ12jlIXhohw4qews1x_I4tSDpml709DZRtLiN24,12546
+rslearn/dataset/window.py,sha256=X4q8YzcSOTtwKxCPf71QLMoyKUtYMSnZu0kPnmVSUx4,10644
+rslearn/dataset/storage/__init__.py,sha256=R50AVV5LH2g7ol0-jyvGcB390VsclXGbJXz4fmkn9as,52
+rslearn/dataset/storage/file.py,sha256=g9HZ3CD4QcgyVNsBaXhjIKQgDOAeZ4R08sJ7ntx4wo8,6815
+rslearn/dataset/storage/storage.py,sha256=DxZ7iwV938PiLwdQzb5EXSb4Mj8bRGmOTmA9fzq_Ge8,4840
 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/concatenate_features.py,sha256=qQPwKF-wz18hC1EA1OS81C6RPjYggPrZ5grCuHIM4aY,3434
-rslearn/models/conv.py,sha256=fWyByeswIOKKzyPmP3erYUlZaKEV0huWHA4CyKTBbfY,1703
-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
-rslearn/models/fpn.py,sha256=s3cz29I14FaSuvBvLOcwCrqVsaRBxG5GjLlqap4WgPc,1603
-rslearn/models/module_wrapper.py,sha256=H2zb-8Au4t31kawW_4JEKHsaXFjpYDawb31ZEauKcxU,2728
-rslearn/models/molmo.py,sha256=mVrARBhZciMzOgOOjGB5AHlPIf2iO9IBSJmdyKSl1L8,2061
-rslearn/models/multitask.py,sha256=j2Kiwj_dUiUp_CIUr25bS8HiyeoFlr1PGqjTfpgIGLc,14672
-rslearn/models/panopticon.py,sha256=woNEs53wVc5D-NxbSDEPRZ_mYe8vllnuldmADjvhfDQ,5806
-rslearn/models/pick_features.py,sha256=y8e4tJFhyG7ZuVSElWhQ5-Aer4ZKJCEH9wLGJU7WqGI,1551
-rslearn/models/pooling_decoder.py,sha256=unr2fSE_QmJHPi3dKtopqMtb1Kn-2h94LgwwAVP9vZg,4437
-rslearn/models/prithvi.py,sha256=AIzcO5xk1ggR0MjbfhIzqPVgUKFN7odxygmgyAelfW8,40143
-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=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/anysat.py,sha256=-1uE2kSfR34lPld5AXEuU8KCfVIa3YAPqs-filxNiWY,8026
+rslearn/models/clip.py,sha256=TgCPA7IEsnONFbWQxQlLguvAXzwp_Y90-Vk9MZnwXak,2337
+rslearn/models/component.py,sha256=1vRK9K7hEgm0eyRhyPyLX4SBWZxiCXzO0fb_bwPQjQg,3015
+rslearn/models/concatenate_features.py,sha256=Attemr5KurxlOojpclD0Pd5Cu2KHpNdpXe8jCSjpJ9U,3818
+rslearn/models/conv.py,sha256=dEAAfhPo4bFlZPSAQjzqZTpP-hdJ394TytYssVK-fDA,2001
+rslearn/models/croma.py,sha256=x8bTFOJB-yR9PydSEjCeV9WnzpalT9QLtozlICNDhyE,10820
+rslearn/models/dinov3.py,sha256=zlICRIeOaxREpR75bzF4nS8P-ZUvLdrJLBU3cBTYgec,6458
+rslearn/models/faster_rcnn.py,sha256=yYRk3attz_GyhJA6jE1ss4ybT_knbLNT1lMRrkz22PI,8614
+rslearn/models/feature_center_crop.py,sha256=_Mu3E4iJLBug9I4ZIBIpB_VJo-xGterHmhtIFGaHR34,1808
+rslearn/models/fpn.py,sha256=qm7nKMgsZrCoAdz8ASmNKU2nvZ6USm5CedMfy_w_gwE,2079
+rslearn/models/module_wrapper.py,sha256=XjRgmhss9_beBEE77t9iySsIvYpV7R-DeyBlOBG000I,2004
+rslearn/models/molmo.py,sha256=pPtCy7eg-xN-iRKCtNL_hpxagfJRyYc9o3NVb-ifonI,2182
+rslearn/models/multitask.py,sha256=bpFxvtFowRyT-tvRSdY7AKbEx_i1y7sToEzZgTMcF4s,16264
+rslearn/models/panopticon.py,sha256=4-KHbgAHUk-ZX04PFRIJF_BRO_euOZO0NfzjB3CEP7Y,5891
+rslearn/models/pick_features.py,sha256=fI9SYubqpCWOAHYGVUSg5sgD31dsnAR9mNuLmqfIeL8,1110
+rslearn/models/pooling_decoder.py,sha256=zrMH6wUExCa-XD1q9CIFD2ScgiasapyJs9plhcUxhIs,4767
+rslearn/models/prithvi.py,sha256=SzdWm0CjYCZIviES3yBXec6TmvrWGQ6pf37w4BaODlY,40326
+rslearn/models/resize_features.py,sha256=U7ZIVwwToJJnwchFG59wLWWP9eikHDB_1c4OtpubxHU,1693
+rslearn/models/sam2_enc.py,sha256=iq2JBqloJ7xkg7Qz5M1HymRfEkLGGn2msF00R_nGBgM,3593
+rslearn/models/satlaspretrain.py,sha256=ytlke5IlxlUEhUoIRcm0pz7umppppZBskCux9YE8mVg,3281
+rslearn/models/simple_time_series.py,sha256=bmy3SsFocpbQwTk4s9gBMAYvz89IlUL_2pAFMZLew_8,13402
+rslearn/models/singletask.py,sha256=9DM9a9-Mv3vVQqRhPOIXG2HHuVqVa_zuvgafeeYh4r0,1903
+rslearn/models/ssl4eo_s12.py,sha256=FTDSkHCaQxw0B-R-8NIe50lZRBYFdI-v68lk9NrnwII,3737
+rslearn/models/swin.py,sha256=QKiV00bu6dw9p0a9PbqDOe5zsnaFHuu6dOfEn41ms7I,6017
 rslearn/models/task_embedding.py,sha256=Z6sf61BLCtvdrdnvjh8500b-KiFp3GeWbT4mOqpaCKk,9100
-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/terramind.py,sha256=R_daJOIzCnaC_Z3s5n9YtBa9Gf_HgQatgCOwwy4MtSA,8732
+rslearn/models/trunk.py,sha256=1GCH9iyLIytoHVntLSMwfH9duQpe1W4DPmOClLpPKjc,4778
+rslearn/models/unet.py,sha256=HuuINvkB1-5w9ZOTXZCWkVxJShruPKCol8pKeA3zw_4,7251
+rslearn/models/upsample.py,sha256=JvfnktT6Dgcql3cSoySWXZ7dmkDkfpRo6vDkpz8KFAQ,1326
 rslearn/models/use_croma.py,sha256=OSBqMuLp-pDtqPNWAVBfmX4wckmyYCKtUDdGCjJk_K8,17966
-rslearn/models/clay/clay.py,sha256=29CGCOysx9duEX4Y6LUNHXck_sHjCFrlV4w8CP_hKmI,8460
+rslearn/models/clay/clay.py,sha256=q2vcqyRFhCJvzLHwHLxHfWGyvVt819Xmk7TzuhXdPFI,8478
 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
-rslearn/models/detr/detr.py,sha256=otLmmyUm05e4MUyvQBoqo-RKnx3hbodTXvfPQWvuTEI,18737
+rslearn/models/detr/detr.py,sha256=ZiYNwJ3zWqZdvvzc0CXDzdvcwSCO_wSOCWlChemTsX8,19178
 rslearn/models/detr/matcher.py,sha256=4h_xFlgTMEJvJ6aLZUamrKZ72L5hDk9wPglNZ81JBg8,4533
 rslearn/models/detr/position_encoding.py,sha256=8FFoBT-Jtgqk7D4qDBTbVLOeAdmjdjtJTC608TaX6yY,3869
 rslearn/models/detr/transformer.py,sha256=aK4HO7AkCZn7xGHP3Iq91w2iFPVshugOILYAjVjroCw,13971
 rslearn/models/detr/util.py,sha256=NMHhHbkIo7PoBUVbDqa2ZknJBTswmaxFCGHrPtFKnGg,676
 rslearn/models/galileo/__init__.py,sha256=QQa0C29nuPRva0KtGiMHQ2ZB02n9SSwj_wqTKPz18NM,112
-rslearn/models/galileo/galileo.py,sha256=jUHA64YvVC3Fz5fevc_9dFJfZaINODRDrhSGLIiOZcw,21115
+rslearn/models/galileo/galileo.py,sha256=PR90sPKxBfrpYolAfDf-AyBkNVARr6UbQmYc2XzcNPc,21365
 rslearn/models/galileo/single_file_galileo.py,sha256=l5tlmmdr2eieHNH-M7rVIvcptkv0Fuk3vKXFW691ezA,56143
 rslearn/models/olmoearth_pretrain/__init__.py,sha256=AjRvbjBdadCdPh-EdvySH76sVAQ8NGQaJt11Tsn1D5I,36
-rslearn/models/olmoearth_pretrain/model.py,sha256=ZJgoyy7vwB0PUMJtHF-sdJ-uSBqnUXMDBco0Dx4cAes,10670
+rslearn/models/olmoearth_pretrain/model.py,sha256=mGGoe45WEhFUJMQBYgKyJ_peDtsXKxNrgiYLz9u8I9o,10882
 rslearn/models/olmoearth_pretrain/norm.py,sha256=rHjFyWkpNLYMx9Ow7TsU-jGm9Sjx7FVf0p4R__ohx2c,3266
 rslearn/models/panopticon_data/sensors/drone.yaml,sha256=xqWS-_QMtJyRoWXJm-igoSur9hAmCFdqkPin8DT5qpw,431
 rslearn/models/panopticon_data/sensors/enmap.yaml,sha256=b2j6bSgYR2yKR9DRm3SPIzSVYlHf51ny_p-1B4B9sB4,13431
@@ -101,18 +103,19 @@ rslearn/models/panopticon_data/sensors/sentinel2.yaml,sha256=qYJ92x-GHO0ZdCrTtCj
 rslearn/models/panopticon_data/sensors/superdove.yaml,sha256=QpIRyopdV4hAez_EIsDwhGFT4VtTk7UgzQveyc8t8fc,795
 rslearn/models/panopticon_data/sensors/wv23.yaml,sha256=SWYSlkka6UViKAz6YI8aqwQ-Ayo-S5kmNa9rO3iGW6o,1172
 rslearn/models/presto/__init__.py,sha256=eZrB-XKi_vYqZhpyAOwppJi4dRuMtYVAdbq7KRygze0,64
-rslearn/models/presto/presto.py,sha256=8mZnc0jk_r_JikybHQNyyHg6t7JNPmoPmgoivyNf-U8,9177
-rslearn/models/presto/single_file_presto.py,sha256=Kbwp8V7pO8HHM2vlCPpjekQiFiDryW8zQkWmt1g05BY,30381
+rslearn/models/presto/presto.py,sha256=SQCxY-jibe_PQs5yvoEOzjUhv8jaCoEkDE0pvSVhEUY,9542
+rslearn/models/presto/single_file_presto.py,sha256=-P00xjhj9dx3O6HqWpQmG9dPk_i6bT_t8vhX4uQm5tA,30242
 rslearn/tile_stores/__init__.py,sha256=-cW1J7So60SEP5ZLHCPdaFBV5CxvV3QlOhaFnUkhTJ0,1675
 rslearn/tile_stores/default.py,sha256=PYaDNvBxhJTDKJGw0EjDTSE1OKajR7_iJpMbOjj-mE8,15054
 rslearn/tile_stores/tile_store.py,sha256=9AeYduDYPp_Ia2NMlq6osptpz_AFGIOQcLJrqZ_m-z0,10469
 rslearn/train/__init__.py,sha256=fnJyY4aHs5zQqbDKSfXsJZXY_M9fbTsf7dRYaPwZr2M,30
-rslearn/train/all_patches_dataset.py,sha256=xFJ96HU3CodrUBzXTsgrmEShosKH79T2SxI0xDVSH3Q,18217
+rslearn/train/all_patches_dataset.py,sha256=qgx1tHZOGOFxUB-HQ7jDxJXE_cvySQfDE6cMr60VS7s,18206
 rslearn/train/data_module.py,sha256=pgut8rEWHIieZ7RR8dUvhtlNqk0egEdznYF3tCvqdHg,23552
-rslearn/train/dataset.py,sha256=OLPBtVf7tmHodoMnB_gI-jLQq2xQ9aXz38Hq8kBgbp0,33944
-rslearn/train/lightning_module.py,sha256=ZLBiId3secUlVs2yzkN-mwVv4rMdh5TkdZYl4vv_Cw0,14466
+rslearn/train/dataset.py,sha256=bsgpUoUKv9AeXzCzBqumiq8b-NZ5PjkFTvLNG82Vx2Q,34179
+rslearn/train/lightning_module.py,sha256=HA9e-74oUZR5s7piQP9Mwxwz0vw0-p4HdmijfgBSpU0,14776
+rslearn/train/model_context.py,sha256=TJHCw8xXLZrRHqGzJr30QARATgG_pJLbMBmyiIcZXRM,1449
 rslearn/train/optimizer.py,sha256=EKSqkmERalDA0bF32Gey7n6z69KLyaUWKlRsGJfKBmE,927
-rslearn/train/prediction_writer.py,sha256=D2CCLlwlElMoMxnPiI6B9Q9HafGspuwoqYD8TKq98pk,13173
+rslearn/train/prediction_writer.py,sha256=rW0BUaYT_F1QqmpnQlbrLiLya1iBfC5Pb78G_NlF-vA,15956
 rslearn/train/scheduler.py,sha256=wFbmycMHgL6nRYeYalDjb0G8YVo8VD3T3sABS61jJ7c,2318
 rslearn/train/callbacks/__init__.py,sha256=VNV0ArZyYMvl3dGK2wl6F046khYJ1dEBlJS6G_SYNm0,47
 rslearn/train/callbacks/adapters.py,sha256=yfv8nyCj3jmo2_dNkFrjukKxh0MHsf2xKqWwMF0QUtY,1869
@@ -120,14 +123,14 @@ rslearn/train/callbacks/freeze_unfreeze.py,sha256=8fIzBMhCKKjpTffIeAdhdSjsBd8NjT
 rslearn/train/callbacks/gradients.py,sha256=4YqCf0tBb6E5FnyFYbveXfQFlgNPyxIXb2FCWX4-6qs,5075
 rslearn/train/callbacks/peft.py,sha256=wEOKsS3RhsRaZTXn_Kz2wdsZdIiIaZPdCJWtdJBurT8,4156
 rslearn/train/tasks/__init__.py,sha256=dag1u72x1-me6y0YcOubUo5MYZ0Tjf6-dOir9UeFNMs,75
-rslearn/train/tasks/classification.py,sha256=8nSv0caf2PzV3Pmme_iN4WQIac4ry3hdW6FRHbh4L1M,13152
-rslearn/train/tasks/detection.py,sha256=9j9webusrjGexvUmZ7gl3NTBS63Qq511VFlB2WbLi5Y,22302
-rslearn/train/tasks/embedding.py,sha256=DK3l1aQ3d5gQUT1h3cD6vcUaNKvSsH26RHx2Bbzutbg,3667
-rslearn/train/tasks/multi_task.py,sha256=dBWsnbvQ0CReNsbDHmZ_-sXjUE0H4S2OPcbJwMquG9g,6016
-rslearn/train/tasks/per_pixel_regression.py,sha256=W8dbLyIiPgFI3gA_aZQX0pSFRWLP2v6tthsFbKhcDVg,8783
-rslearn/train/tasks/regression.py,sha256=zZhrrZ1qxjrdLjKWC9McRivDXCcKiYfdLC-kaMeVkDc,11547
-rslearn/train/tasks/segmentation.py,sha256=xEni3CLDyetviv84XrpJg5xeJU87WHGFKTVfIeemGIY,21868
-rslearn/train/tasks/task.py,sha256=4w2xKL_U5JAtdj2dYoVv82h6xTtgUsA3IvIOcXyZecs,3887
+rslearn/train/tasks/classification.py,sha256=2_Iz3g9ifdtMo6a2sRtnAoYGP2Om8JPP7rm2AfwoDLc,14190
+rslearn/train/tasks/detection.py,sha256=DrPLF_63WU99Qh1yILxSJWrYrl_2mCGxlTX2SznCei0,21938
+rslearn/train/tasks/embedding.py,sha256=98ykdmfaxQjsH0UrUdTGmz1f0hCMPcNYTzt1YFqNQwQ,3869
+rslearn/train/tasks/multi_task.py,sha256=piauvjg4j6eBEZnIFrKKxyYPWKdAJ4yUCFxt_ngsxDY,6125
+rslearn/train/tasks/per_pixel_regression.py,sha256=jME-AFC74dfNAF5cBBmdmIgOV-KfkHt_z1GIOifuPJw,9975
+rslearn/train/tasks/regression.py,sha256=8U0bcSofefmP9Drpbo7PO9xAkVptAuso2rOQyPOUhzo,12624
+rslearn/train/tasks/segmentation.py,sha256=D9VaxsuiF0R8nvOKL821wPHey2GTMCgIqm8rZtnAXdk,22778
+rslearn/train/tasks/task.py,sha256=CwGEFXquSjWXoJqbpBE1mQFCJl7NmMvgreAydphvu6U,3942
 rslearn/train/transforms/__init__.py,sha256=BkCAzm4f-8TEhPIuyvCj7eJGh36aMkZFYlq-H_jkSvY,778
 rslearn/train/transforms/concatenate.py,sha256=sdVLJIyr9Nj2tzXEzvWFQnjJjyRSuhR_Faf6UlMIvbg,1568
 rslearn/train/transforms/crop.py,sha256=4jA3JJsC0ghicPHbfsNJ0d3WpChyvftY73ONiwQaif0,4214
@@ -153,10 +156,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=4ZDYpfBLLxguJkiIaavTagiQK2Sv4Rz9NumbHlq-3Lw,15041
-rslearn-0.0.16.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
-rslearn-0.0.16.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
-rslearn-0.0.16.dist-info/METADATA,sha256=h0p9V4jlSLDsrC2_owCn0xEKL7Kka74mEsE_pj-tJf0,36319
-rslearn-0.0.16.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-rslearn-0.0.16.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
-rslearn-0.0.16.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
-rslearn-0.0.16.dist-info/RECORD,,
+rslearn-0.0.18.dist-info/licenses/LICENSE,sha256=_99ZWPoLdlUbqZoSC5DF4ihiNwl5rTEmBaq2fACecdg,11352
+rslearn-0.0.18.dist-info/licenses/NOTICE,sha256=wLPr6rwV_jCg-xEknNGwhnkfRfuoOE9MZ-lru2yZyLI,5070
+rslearn-0.0.18.dist-info/METADATA,sha256=n-P9W0g_cQJ1BFYHxpIvC6EcJjbFOCoNlmHdSiN7sNo,37853
+rslearn-0.0.18.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+rslearn-0.0.18.dist-info/entry_points.txt,sha256=doTBQ57NT7nq-dgYGgTTw6mafcGWb_4PWYtYR4rGm50,46
+rslearn-0.0.18.dist-info/top_level.txt,sha256=XDKo90WBH8P9RQumHxo0giLJsoufT4r9odv-WE6Ahk4,8
+rslearn-0.0.18.dist-info/RECORD,,

rslearn/dataset/index.py

@@ -1,173 +0,0 @@
-"""Index about windows in the dataset."""
-
-import json
-import multiprocessing
-from typing import TYPE_CHECKING
-
-import tqdm
-from upath import UPath
-
-from .window import (
-    Window,
-    WindowLayerData,
-)
-
-if TYPE_CHECKING:
-    from .dataset import Dataset
-
-
-def get_window_layer_datas(window: Window) -> list[WindowLayerData]:
-    """Helper function for multiprocessing to load window layer datas."""
-    return list(window.load_layer_datas().values())
-
-
-def get_window_completed_layers(window: Window) -> list[tuple[str, int]]:
-    """Helper function for multiprocessing to load window completed layers."""
-    return window.list_completed_layers()
-
-
-class DatasetIndex:
-    """Manage an index about windows in the dataset.
-
-    The index is just a single file containing information about all windows in the
-    dataset, so that this information does not need to be loaded from per-window files.
-
-    The information includes the window metadata, the window layer datas (matching data
-    source items), and the completed layers.
-
-    Currently the index must be manually maintained. It can be created for relatively
-    static datasets, and updated each time the dataset is modified.
-    """
-
-    FNAME = "dataset_index.json"
-
-    def __init__(
-        self,
-        windows: list[Window],
-        layer_datas: dict[str, list[WindowLayerData]],
-        completed_layers: dict[str, list[tuple[str, int]]],
-    ):
-        """Create a new DatasetIndex.
-
-        Args:
-            windows: the windows in the dataset.
-            layer_datas: map from window name to the layer datas for that window.
-            completed_layers: map from window name to its list of completed layers.
-                Each element is (layer_name, group_idx).
-        """
-        self.windows = windows
-        self.layer_datas = layer_datas
-        self.completed_layers = completed_layers
-
-        for window in self.windows:
-            window.index = self
-
-    def get_windows(
-        self,
-        groups: list[str] | None = None,
-        names: list[str] | None = None,
-    ) -> list[Window]:
-        """Get the windows in the dataset.
-
-        Args:
-            groups: an optional list of groups to filter by
-            names: an optional list of window names to filter by
-        """
-        windows = self.windows
-        if groups is not None:
-            group_set = set(groups)
-            windows = [window for window in windows if window.group in group_set]
-        if names is not None:
-            name_set = set(names)
-            windows = [window for window in windows if window.name in name_set]
-        return windows
-
-    def save_index(self, ds_path: UPath) -> None:
-        """Save the index to the specified file."""
-        encoded_windows = [window.get_metadata() for window in self.windows]
-
-        encoded_layer_datas = {}
-        for window_name, layer_data_list in self.layer_datas.items():
-            encoded_layer_datas[window_name] = [
-                layer_data.serialize() for layer_data in layer_data_list
-            ]
-
-        encoded_index = {
-            "windows": encoded_windows,
-            "layer_datas": encoded_layer_datas,
-            "completed_layers": self.completed_layers,
-        }
-        with (ds_path / self.FNAME).open("w") as f:
-            json.dump(encoded_index, f)
-
-    @staticmethod
-    def load_index(ds_path: UPath) -> "DatasetIndex":
-        """Load the DatasetIndex for the specified dataset."""
-        with (ds_path / DatasetIndex.FNAME).open() as f:
-            encoded_index = json.load(f)
-
-        windows = []
-        for encoded_window in encoded_index["windows"]:
-            window = Window.from_metadata(
-                path=Window.get_window_root(
-                    ds_path, encoded_window["group"], encoded_window["name"]
-                ),
-                metadata=encoded_window,
-            )
-            windows.append(window)
-
-        layer_datas = {}
-        for window_name, encoded_layer_data_list in encoded_index[
-            "layer_datas"
-        ].items():
-            layer_datas[window_name] = [
-                WindowLayerData.deserialize(encoded_layer_data)
-                for encoded_layer_data in encoded_layer_data_list
-            ]
-
-        completed_layers = {}
-        for window_name, encoded_layer_list in encoded_index[
-            "completed_layers"
-        ].items():
-            completed_layers[window_name] = [
-                (layer_name, group_idx)
-                for (layer_name, group_idx) in encoded_layer_list
-            ]
-
-        return DatasetIndex(
-            windows=windows,
-            layer_datas=layer_datas,
-            completed_layers=completed_layers,
-        )
-
-    @staticmethod
-    def build_index(dataset: "Dataset", workers: int) -> "DatasetIndex":
-        """Build a new DatasetIndex for the specified dataset."""
-        # Load windows.
-        windows = dataset.load_windows(workers=workers, no_index=True)
-
-        # Load layer datas.
-        p = multiprocessing.Pool(workers)
-        layer_data_outputs = p.imap(get_window_layer_datas, windows)
-        layer_data_outputs = tqdm.tqdm(
-            layer_data_outputs, total=len(windows), desc="Loading window layer datas"
-        )
-        layer_datas = {}
-        for window, cur_layer_datas in zip(windows, layer_data_outputs):
-            layer_datas[window.name] = cur_layer_datas
-
-        # Load completed layers.
-        completed_layer_outputs = p.imap(get_window_completed_layers, windows)
-        completed_layer_outputs = tqdm.tqdm(
-            completed_layer_outputs, total=len(windows), desc="Loading completed layers"
-        )
-        completed_layers = {}  # window name -> list of (layer name, group idx)
-        for window, cur_completed_layers in zip(windows, completed_layer_outputs):
-            completed_layers[window.name] = cur_completed_layers
-        p.close()
-
-        return DatasetIndex(
-            windows=windows,
-            layer_datas=layer_datas,
-            completed_layers=completed_layers,
-        )

rslearn/models/registry.py

@@ -1,22 +0,0 @@
-"""Model registry."""
-
-from collections.abc import Callable
-from typing import Any, TypeVar
-
-_ModelT = TypeVar("_ModelT")
-
-
-class _ModelRegistry(dict[str, type[Any]]):
-    """Registry for Model classes."""
-
-    def register(self, name: str) -> Callable[[type[_ModelT]], type[_ModelT]]:
-        """Decorator to register a model class."""
-
-        def decorator(cls: type[_ModelT]) -> type[_ModelT]:
-            self[name] = cls
-            return cls
-
-        return decorator
-
-
-Models = _ModelRegistry()