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

rslearn/train/prediction_writer.py

@@ -12,10 +12,8 @@ from lightning.pytorch.callbacks import BasePredictionWriter
 from upath import UPath
 
 from rslearn.config import (
+    LayerConfig,
     LayerType,
-    RasterFormatConfig,
-    RasterLayerConfig,
-    VectorLayerConfig,
 )
 from rslearn.dataset import Dataset, Window
 from rslearn.log_utils import get_logger
@@ -25,9 +23,8 @@ from rslearn.utils.geometry import PixelBounds
 from rslearn.utils.raster_format import (
     RasterFormat,
     adjust_projection_and_bounds_for_array,
-    load_raster_format,
 )
-from rslearn.utils.vector_format import VectorFormat, load_vector_format
+from rslearn.utils.vector_format import VectorFormat
 
 from .lightning_module import RslearnLightningModule
 from .tasks.task import Task
@@ -150,7 +147,7 @@ class RslearnWriter(BasePredictionWriter):
         selector: list[str] | None = None,
         merger: PatchPredictionMerger | None = None,
         output_path: str | Path | None = None,
-        layer_config: RasterLayerConfig | VectorLayerConfig | None = None,
+        layer_config: LayerConfig | None = None,
     ):
         """Create a new RslearnWriter.
 
@@ -178,43 +175,31 @@ class RslearnWriter(BasePredictionWriter):
         )
 
         # Handle dataset and layer config
-        self.layer_config: RasterLayerConfig | VectorLayerConfig
+        self.layer_config: LayerConfig
         if layer_config:
             self.layer_config = layer_config
-            self.dataset = None if self.output_path else Dataset(self.path)
         else:
-            self.dataset = Dataset(self.path)
-            if self.output_layer not in self.dataset.layers:
+            dataset = Dataset(self.path)
+            if self.output_layer not in dataset.layers:
                 raise KeyError(
                     f"Output layer '{self.output_layer}' not found in dataset layers."
                 )
-            raw_layer_config = self.dataset.layers[self.output_layer]
-            # Type narrowing to ensure compatibility
-            if isinstance(raw_layer_config, (RasterLayerConfig | VectorLayerConfig)):
-                self.layer_config = raw_layer_config
-            else:
-                raise ValueError(
-                    f"Layer config must be RasterLayerConfig or VectorLayerConfig, got {type(raw_layer_config)}"
-                )
+            self.layer_config = dataset.layers[self.output_layer]
 
         self.format: RasterFormat | VectorFormat
-        if self.layer_config.layer_type == LayerType.RASTER:
-            assert isinstance(self.layer_config, RasterLayerConfig)
+        if self.layer_config.type == LayerType.RASTER:
             band_cfg = self.layer_config.band_sets[0]
-            self.format = load_raster_format(
-                RasterFormatConfig(band_cfg.format["name"], band_cfg.format)
-            )
-        elif self.layer_config.layer_type == LayerType.VECTOR:
-            assert isinstance(self.layer_config, VectorLayerConfig)
-            self.format = load_vector_format(self.layer_config.format)
+            self.format = band_cfg.instantiate_raster_format()
+        elif self.layer_config.type == LayerType.VECTOR:
+            self.format = self.layer_config.instantiate_vector_format()
         else:
-            raise ValueError(f"invalid layer type {self.layer_config.layer_type}")
+            raise ValueError(f"invalid layer type {self.layer_config.type}")
 
         if merger is not None:
             self.merger = merger
-        elif self.layer_config.layer_type == LayerType.RASTER:
+        elif self.layer_config.type == LayerType.RASTER:
             self.merger = RasterMerger()
-        elif self.layer_config.layer_type == LayerType.VECTOR:
+        elif self.layer_config.type == LayerType.VECTOR:
             self.merger = VectorMerger()
 
         # Map from window name to pending data to write.
@@ -337,8 +322,7 @@ class RslearnWriter(BasePredictionWriter):
         logger.debug(f"Merging and writing for window {window.name}")
         merged_output = self.merger.merge(window, pending_output)
 
-        if self.layer_config.layer_type == LayerType.RASTER:
-            assert isinstance(self.layer_config, RasterLayerConfig)
+        if self.layer_config.type == LayerType.RASTER:
             raster_dir = window.get_raster_dir(
                 self.output_layer, self.layer_config.band_sets[0].bands
             )
@@ -351,7 +335,7 @@ class RslearnWriter(BasePredictionWriter):
             )
             self.format.encode_raster(raster_dir, projection, bounds, merged_output)
 
-        elif self.layer_config.layer_type == LayerType.VECTOR:
+        elif self.layer_config.type == LayerType.VECTOR:
             layer_dir = window.get_layer_dir(self.output_layer)
             assert isinstance(self.format, VectorFormat)
             self.format.encode_vector(layer_dir, merged_output)

rslearn/train/tasks/classification.py

@@ -26,7 +26,7 @@ class ClassificationTask(BasicTask):
     def __init__(
         self,
         property_name: str,
-        classes: list,  # TODO: Should this be a list of str or int or can it be both?
+        classes: list[str],
         filters: list[tuple[str, str]] = [],
         read_class_id: bool = False,
         allow_invalid: bool = False,
@@ -176,6 +176,7 @@ class ClassificationTask(BasicTask):
             # For multiclass classification or when using the default threshold
             class_idx = probs.argmax().item()
 
+        value: str | int
         if not self.read_class_id:
             value = self.classes[class_idx]  # type: ignore
         else:

rslearn/utils/fsspec.py

@@ -156,3 +156,23 @@ def open_rasterio_upath_writer(
         with path.open("wb") as f:
             with rasterio.open(f, "w", **kwargs) as raster:
                 yield raster
+
+
+def get_relative_suffix(base_dir: UPath, fname: UPath) -> str:
+    """Get the suffix of fname relative to base_dir.
+
+    Args:
+        base_dir: the base directory.
+        fname: a filename within the base directory.
+
+    Returns:
+        the suffix on base_dir that would yield the given filename.
+    """
+    if not fname.path.startswith(base_dir.path):
+        raise ValueError(
+            f"filename {fname.path} must start with base directory {base_dir.path}"
+        )
+    suffix = fname.path[len(base_dir.path) :]
+    if suffix.startswith("/"):
+        suffix = suffix[1:]
+    return suffix

rslearn/utils/jsonargparse.py

@@ -1,7 +1,18 @@
 """Custom serialization for jsonargparse."""
 
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
 import jsonargparse
 from rasterio.crs import CRS
+from upath import UPath
+
+from rslearn.config.dataset import LayerConfig
+
+if TYPE_CHECKING:
+    from rslearn.data_sources.data_source import DataSourceContext
+
+INITIALIZED = False
 
 
 def crs_serializer(v: CRS) -> str:
@@ -28,6 +39,74 @@ def crs_deserializer(v: str) -> CRS:
     return CRS.from_string(v)
 
 
+def datetime_serializer(v: datetime) -> str:
+    """Serialize datetime for jsonargparse.
+
+    Args:
+        v: the datetime object.
+
+    Returns:
+        the datetime encoded to string
+    """
+    return v.isoformat()
+
+
+def datetime_deserializer(v: str) -> datetime:
+    """Deserialize datetime for jsonargparse.
+
+    Args:
+        v: the encoded datetime.
+
+    Returns:
+        the decoded datetime object
+    """
+    return datetime.fromisoformat(v)
+
+
+def data_source_context_serializer(v: "DataSourceContext") -> dict[str, Any]:
+    """Serialize DataSourceContext for jsonargparse."""
+    x = {
+        "ds_path": (str(v.ds_path) if v.ds_path is not None else None),
+        "layer_config": (
+            v.layer_config.model_dump(mode="json")
+            if v.layer_config is not None
+            else None
+        ),
+    }
+    return x
+
+
+def data_source_context_deserializer(v: dict[str, Any]) -> "DataSourceContext":
+    """Deserialize DataSourceContext for jsonargparse."""
+    # We lazily import these to avoid cyclic dependency.
+    from rslearn.data_sources.data_source import DataSourceContext
+
+    return DataSourceContext(
+        ds_path=(UPath(v["ds_path"]) if v["ds_path"] is not None else None),
+        layer_config=(
+            LayerConfig.model_validate(v["layer_config"])
+            if v["layer_config"] is not None
+            else None
+        ),
+    )
+
+
 def init_jsonargparse() -> None:
     """Initialize custom jsonargparse serializers."""
+    global INITIALIZED
+    if INITIALIZED:
+        return
     jsonargparse.typing.register_type(CRS, crs_serializer, crs_deserializer)
+    jsonargparse.typing.register_type(
+        datetime, datetime_serializer, datetime_deserializer
+    )
+
+    from rslearn.data_sources.data_source import DataSourceContext
+
+    jsonargparse.typing.register_type(
+        DataSourceContext,
+        data_source_context_serializer,
+        data_source_context_deserializer,
+    )
+
+    INITIALIZED = True

rslearn/utils/raster_format.py

@@ -2,8 +2,7 @@
 
 import hashlib
 import json
-from collections.abc import Callable
-from typing import Any, BinaryIO, TypeVar
+from typing import Any, BinaryIO
 
 import affine
 import numpy as np
@@ -14,34 +13,12 @@ from rasterio.crs import CRS
 from rasterio.enums import Resampling
 from upath import UPath
 
-from rslearn.config import RasterFormatConfig
 from rslearn.const import TILE_SIZE
 from rslearn.log_utils import get_logger
 from rslearn.utils.fsspec import open_rasterio_upath_reader, open_rasterio_upath_writer
 
 from .geometry import PixelBounds, Projection
 
-_RasterFormatT = TypeVar("_RasterFormatT", bound="RasterFormat")
-
-
-class _RasterFormatRegistry(dict[str, type["RasterFormat"]]):
-    """Registry for RasterFormat classes."""
-
-    def register(
-        self, name: str
-    ) -> Callable[[type[_RasterFormatT]], type[_RasterFormatT]]:
-        """Decorator to register a raster format class."""
-
-        def decorator(cls: type[_RasterFormatT]) -> type[_RasterFormatT]:
-            self[name] = cls
-            return cls
-
-        return decorator
-
-
-RasterFormats = _RasterFormatRegistry()
-
-
 logger = get_logger(__name__)
 
 
@@ -219,7 +196,6 @@ class RasterFormat:
         raise NotImplementedError
 
 
-@RasterFormats.register("image_tile")
 class ImageTileRasterFormat(RasterFormat):
     """A RasterFormat that stores data in image tiles corresponding to grid cells.
 
@@ -468,7 +444,6 @@ class ImageTileRasterFormat(RasterFormat):
         )
 
 
-@RasterFormats.register("geotiff")
 class GeotiffRasterFormat(RasterFormat):
     """A raster format that uses one big, tiled GeoTIFF with small block size."""
 
@@ -623,7 +598,6 @@ class GeotiffRasterFormat(RasterFormat):
         return GeotiffRasterFormat(**kwargs)
 
 
-@RasterFormats.register("single_image")
 class SingleImageRasterFormat(RasterFormat):
     """A raster format that produces a single image called image.png/jpg.
 
@@ -775,17 +749,3 @@ class SingleImageRasterFormat(RasterFormat):
         if "format" in config:
             kwargs["format"] = config["format"]
         return SingleImageRasterFormat(**kwargs)
-
-
-def load_raster_format(config: RasterFormatConfig) -> RasterFormat:
-    """Loads a RasterFormat from a RasterFormatConfig.
-
-    Args:
-        config: the RasterFormatConfig configuration object specifying the
-            RasterFormat.
-
-    Returns:
-        the loaded RasterFormat implementation
-    """
-    cls = RasterFormats[config.name]
-    return cls.from_config(config.name, config.config_dict)

rslearn/utils/vector_format.py

@@ -1,15 +1,13 @@
 """Classes for writing vector data to a UPath."""
 
 import json
-from collections.abc import Callable
 from enum import Enum
-from typing import Any, TypeVar
+from typing import Any
 
 import shapely
 from rasterio.crs import CRS
 from upath import UPath
 
-from rslearn.config import VectorFormatConfig
 from rslearn.const import WGS84_PROJECTION
 from rslearn.log_utils import get_logger
 from rslearn.utils.fsspec import open_atomic
@@ -18,25 +16,6 @@ from .feature import Feature
 from .geometry import PixelBounds, Projection, STGeometry, safely_reproject_and_clip
 
 logger = get_logger(__name__)
-_VectorFormatT = TypeVar("_VectorFormatT", bound="VectorFormat")
-
-
-class _VectorFormatRegistry(dict[str, type["VectorFormat"]]):
-    """Registry for VectorFormat classes."""
-
-    def register(
-        self, name: str
-    ) -> Callable[[type[_VectorFormatT]], type[_VectorFormatT]]:
-        """Decorator to register a vector format class."""
-
-        def decorator(cls: type[_VectorFormatT]) -> type[_VectorFormatT]:
-            self[name] = cls
-            return cls
-
-        return decorator
-
-
-VectorFormats = _VectorFormatRegistry()
 
 
 class VectorFormat:
@@ -85,7 +64,6 @@ class VectorFormat:
         raise NotImplementedError
 
 
-@VectorFormats.register("tile")
 class TileVectorFormat(VectorFormat):
     """TileVectorFormat stores data in GeoJSON files corresponding to grid cells.
 
@@ -275,7 +253,6 @@ class GeojsonCoordinateMode(Enum):
     WGS84 = "wgs84"
 
 
-@VectorFormats.register("geojson")
 class GeojsonVectorFormat(VectorFormat):
     """A vector format that uses one big GeoJSON."""
 
@@ -429,17 +406,3 @@ class GeojsonVectorFormat(VectorFormat):
         if "coordinate_mode" in config:
             kwargs["coordinate_mode"] = GeojsonCoordinateMode(config["coordinate_mode"])
         return GeojsonVectorFormat(**kwargs)
-
-
-def load_vector_format(config: VectorFormatConfig) -> VectorFormat:
-    """Loads a VectorFormat from a VectorFormatConfig.
-
-    Args:
-        config: the VectorFormatConfig configuration object specifying the
-            VectorFormat.
-
-    Returns:
-        the loaded VectorFormat implementation
-    """
-    cls = VectorFormats[config.name]
-    return cls.from_config(config.name, config.config_dict)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rslearn
-Version: 0.0.15
+Version: 0.0.17
 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
             }