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

rslearn/main.py

@@ -1,39 +1,43 @@
 """Entrypoint for the rslearn command-line interface."""
 
 import argparse
-import logging
 import multiprocessing
 import random
 import sys
 from collections.abc import Callable
 from datetime import datetime, timezone
 from pathlib import Path
+from typing import Any, TypeVar
 
 import tqdm
 import wandb
-from lightning.pytorch.cli import LightningCLI
+from lightning.pytorch.cli import LightningArgumentParser, LightningCLI
 from rasterio.crs import CRS
 from upath import UPath
 
 from rslearn.config import LayerConfig
 from rslearn.const import WGS84_EPSG
 from rslearn.data_sources import Item, data_source_from_config
-from rslearn.dataset import Dataset, Window
+from rslearn.dataset import Dataset, Window, WindowLayerData
 from rslearn.dataset.add_windows import add_windows_from_box, add_windows_from_file
 from rslearn.dataset.manage import materialize_dataset_windows, prepare_dataset_windows
+from rslearn.log_utils import get_logger
 from rslearn.tile_stores import get_tile_store_for_layer
 from rslearn.train.data_module import RslearnDataModule
 from rslearn.train.lightning_module import RslearnLightningModule
-from rslearn.utils import Projection, STGeometry
+from rslearn.utils import Projection, STGeometry, parse_disabled_layers
+
+logger = get_logger(__name__)
 
-logging.basicConfig()
 handler_registry = {}
 
+ItemType = TypeVar("ItemType", bound="Item")
+
 
-def register_handler(category, command):
+def register_handler(category: Any, command: str) -> Callable:
     """Register a new handler for a command."""
 
-    def decorator(f):
+    def decorator(f: Callable) -> Callable:
         handler_registry[(category, command)] = f
         return f
 
@@ -61,7 +65,7 @@ def parse_time_range(
 
 
 @register_handler("dataset", "add_windows")
-def add_windows():
+def add_windows() -> None:
     """Handler for the rslearn dataset add_windows command."""
     parser = argparse.ArgumentParser(
         prog="rslearn dataset add_windows",
@@ -156,7 +160,13 @@ def add_windows():
     )
     args = parser.parse_args(args=sys.argv[3:])
 
-    def parse_projection(crs_str, resolution, x_res, y_res, default_crs=None):
+    def parse_projection(
+        crs_str: str | None,
+        resolution: float | None,
+        x_res: float,
+        y_res: float,
+        default_crs: CRS | None = None,
+    ) -> Projection | None:
         if not crs_str:
             if default_crs:
                 crs = default_crs
@@ -197,7 +207,8 @@ def add_windows():
         box = [float(value) for value in args.box.split(",")]
 
         windows = add_windows_from_box(
-            box=box,
+            # TODO: we should have an object for box
+            box=box,  # type: ignore
             src_projection=parse_projection(
                 args.src_crs, args.src_resolution, args.src_x_res, args.src_y_res
             ),
@@ -210,10 +221,10 @@ def add_windows():
     else:
         raise Exception("one of box or fname must be specified")
 
-    print(f"created {len(windows)} windows")
+    logger.info(f"created {len(windows)} windows")
 
 
-def add_apply_on_windows_args(parser: argparse.ArgumentParser):
+def add_apply_on_windows_args(parser: argparse.ArgumentParser) -> None:
     """Add arguments for handlers that use the apply_on_windows helper.
 
     Args:
@@ -263,7 +274,7 @@ def apply_on_windows(
     batch_size: int = 1,
     jobs_per_process: int | None = None,
     use_initial_job: bool = True,
-):
+) -> None:
     """A helper to apply a function on windows in a dataset.
 
     Args:
@@ -293,11 +304,11 @@ def apply_on_windows(
     windows = dataset.load_windows(
         groups=groups, names=names, workers=workers, show_progress=True
     )
-    print(f"found {len(windows)} windows")
+    logger.info(f"found {len(windows)} windows")
 
     if hasattr(f, "get_jobs"):
         jobs = f.get_jobs(windows, workers)
-        print(f"got {len(jobs)} jobs")
+        logger.info(f"got {len(jobs)} jobs")
     else:
         jobs = windows
 
@@ -323,9 +334,9 @@ def apply_on_windows(
     p.close()
 
 
-def apply_on_windows_args(f: Callable[[list[Window]], None], args: argparse.Namespace):
+def apply_on_windows_args(f: Callable[..., None], args: argparse.Namespace) -> None:
     """Call apply_on_windows with arguments passed via command-line interface."""
-    dataset = Dataset(UPath(args.root))
+    dataset = Dataset(UPath(args.root), args.disabled_layers)
     apply_on_windows(
         f,
         dataset,
@@ -341,16 +352,16 @@ def apply_on_windows_args(f: Callable[[list[Window]], None], args: argparse.Name
 class PrepareHandler:
     """apply_on_windows handler for the rslearn dataset prepare command."""
 
-    def __init__(self, force: bool):
+    def __init__(self, force: bool) -> None:
         """Initialize a new PrepareHandler.
 
         Args:
             force: force prepare
         """
         self.force = force
-        self.dataset = None
+        self.dataset: Dataset | None = None
 
-    def set_dataset(self, dataset: Dataset):
+    def set_dataset(self, dataset: Dataset) -> None:
         """Captures the dataset from apply_on_windows_args.
 
         Args:
@@ -358,13 +369,16 @@ class PrepareHandler:
         """
         self.dataset = dataset
 
-    def __call__(self, windows: list[Window]):
+    def __call__(self, windows: list[Window]) -> None:
         """Prepares the windows from apply_on_windows."""
+        logger.info(f"Running prepare on {len(windows)} windows")
+        if self.dataset is None:
+            raise ValueError("dataset not set")
         prepare_dataset_windows(self.dataset, windows, self.force)
 
 
 @register_handler("dataset", "prepare")
-def dataset_prepare():
+def dataset_prepare() -> None:
     """Handler for the rslearn dataset prepare command."""
     parser = argparse.ArgumentParser(
         prog="rslearn dataset prepare",
@@ -377,6 +391,12 @@ def dataset_prepare():
         action=argparse.BooleanOptionalAction,
         help="Prepare windows even if they were previously prepared",
     )
+    parser.add_argument(
+        "--disabled-layers",
+        type=parse_disabled_layers,
+        default="",
+        help="List of layers to disable e.g 'layer1,layer2'",
+    )
     add_apply_on_windows_args(parser)
     args = parser.parse_args(args=sys.argv[3:])
 
@@ -384,7 +404,9 @@ def dataset_prepare():
     apply_on_windows_args(fn, args)
 
 
-def _load_window_layer_datas(window: Window):
+def _load_window_layer_datas(
+    window: Window,
+) -> tuple[Window, dict[str, WindowLayerData]]:
     # Helper for IngestHandler to use with multiprocessing.
     return window, window.load_layer_datas()
 
@@ -392,11 +414,12 @@ def _load_window_layer_datas(window: Window):
 class IngestHandler:
     """apply_on_windows handler for the rslearn dataset ingest command."""
 
-    def __init__(self):
+    def __init__(self, ignore_errors: bool = False) -> None:
         """Initialize a new IngestHandler."""
-        self.dataset = None
+        self.dataset: Dataset | None = None
+        self.ignore_errors = ignore_errors
 
-    def set_dataset(self, dataset: Dataset):
+    def set_dataset(self, dataset: Dataset) -> None:
         """Captures the dataset from apply_on_windows_args.
 
         Args:
@@ -404,7 +427,9 @@ class IngestHandler:
         """
         self.dataset = dataset
 
-    def __call__(self, jobs: list[tuple[str, LayerConfig, Item, list[STGeometry]]]):
+    def __call__(
+        self, jobs: list[tuple[str, LayerConfig, Item, list[STGeometry]]]
+    ) -> None:
         """Ingest the specified items.
 
         The items are computed from list of windows via IngestHandler.get_jobs.
@@ -412,13 +437,16 @@ class IngestHandler:
         Args:
             jobs: list of (layer_name, item, geometries) tuples to ingest.
         """
+        logger.info(f"Running ingest for {len(jobs)} jobs")
         import gc
 
+        if self.dataset is None:
+            raise ValueError("dataset not set")
         tile_store = self.dataset.get_tile_store()
 
         # Group jobs by layer name.
-        jobs_by_layer = {}
-        configs_by_layer = {}
+        jobs_by_layer: dict = {}
+        configs_by_layer: dict = {}
         for layer_name, layer_cfg, item, geometries in jobs:
             if layer_name not in jobs_by_layer:
                 jobs_by_layer[layer_name] = []
@@ -437,13 +465,31 @@ class IngestHandler:
                     geometries=[geometries for _, geometries in items_and_geometries],
                 )
             except Exception as e:
-                print(
+                if not self.ignore_errors:
+                    raise
+
+                logger.error(
                     "warning: got error while ingesting "
                     + f"{len(items_and_geometries)} items: {e}"
                 )
 
         gc.collect()
 
+    def _load_layer_data_for_windows(
+        self, windows: list[Window], workers: int
+    ) -> list[tuple[Window, dict[str, WindowLayerData]]]:
+        if workers == 0:
+            return [(_load_window_layer_datas(window)) for window in windows]
+        p = multiprocessing.Pool(workers)
+        outputs = p.imap_unordered(_load_window_layer_datas, windows)
+        windows_and_layer_datas = []
+        for window, layer_datas in tqdm.tqdm(
+            outputs, total=len(windows), desc="Loading window layer datas"
+        ):
+            windows_and_layer_datas.append((window, layer_datas))
+        p.close()
+        return windows_and_layer_datas
+
     def get_jobs(
         self, windows: list[Window], workers: int
     ) -> list[tuple[str, LayerConfig, Item, list[STGeometry]]]:
@@ -455,17 +501,12 @@ class IngestHandler:
         This makes sure that jobs are grouped by item rather than by window, which
         makes sense because there's no reason to ingest the same item twice.
         """
+        if self.dataset is None:
+            raise ValueError("dataset not set")
         # TODO: avoid duplicating ingest_dataset_windows...
 
         # Load layer datas of each window.
-        p = multiprocessing.Pool(workers)
-        outputs = p.imap_unordered(_load_window_layer_datas, windows)
-        windows_and_layer_datas = []
-        for window, layer_datas in tqdm.tqdm(
-            outputs, total=len(windows), desc="Loading window layer datas"
-        ):
-            windows_and_layer_datas.append((window, layer_datas))
-        p.close()
+        windows_and_layer_datas = self._load_layer_data_for_windows(windows, workers)
 
         jobs: list[tuple[str, LayerConfig, Item, list[STGeometry]]] = []
         for layer_name, layer_cfg in self.dataset.layers.items():
@@ -476,7 +517,7 @@ class IngestHandler:
 
             data_source = data_source_from_config(layer_cfg, self.dataset.path)
 
-            geometries_by_item = {}
+            geometries_by_item: dict = {}
             for window, layer_datas in windows_and_layer_datas:
                 if layer_name not in layer_datas:
                     continue
@@ -484,7 +525,9 @@ class IngestHandler:
                 layer_data = layer_datas[layer_name]
                 for group in layer_data.serialized_item_groups:
                     for serialized_item in group:
-                        item = data_source.deserialize_item(serialized_item)
+                        item = data_source.deserialize_item(  # type: ignore
+                            serialized_item
+                        )
                         if item not in geometries_by_item:
                             geometries_by_item[item] = []
                         geometries_by_item[item].append(geometry)
@@ -492,32 +535,45 @@ class IngestHandler:
             for item, geometries in geometries_by_item.items():
                 jobs.append((layer_name, layer_cfg, item, geometries))
 
-        print(f"computed {len(jobs)} ingest jobs from {len(windows)} windows")
+        logger.info(f"computed {len(jobs)} ingest jobs from {len(windows)} windows")
         return jobs
 
 
 @register_handler("dataset", "ingest")
-def dataset_ingest():
+def dataset_ingest() -> None:
     """Handler for the rslearn dataset ingest command."""
     parser = argparse.ArgumentParser(
         prog="rslearn dataset ingest",
         description="rslearn dataset ingest: ingest items in retrieved data sources",
     )
+    parser.add_argument(
+        "--disabled-layers",
+        type=parse_disabled_layers,
+        default="",
+        help="List of layers to disable e.g 'layer1,layer2'",
+    )
+    parser.add_argument(
+        "--ignore-errors",
+        type=bool,
+        default=False,
+        help="Ignore ingestion errors in individual jobs",
+        action=argparse.BooleanOptionalAction,
+    )
     add_apply_on_windows_args(parser)
     args = parser.parse_args(args=sys.argv[3:])
 
-    fn = IngestHandler()
+    fn = IngestHandler(ignore_errors=args.ignore_errors)
     apply_on_windows_args(fn, args)
 
 
 class MaterializeHandler:
     """apply_on_windows handler for the rslearn dataset materialize command."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize a MaterializeHandler."""
-        self.dataset = None
+        self.dataset: Dataset | None = None
 
-    def set_dataset(self, dataset: Dataset):
+    def set_dataset(self, dataset: Dataset) -> None:
         """Captures the dataset from apply_on_windows_args.
 
         Args:
@@ -525,13 +581,16 @@ class MaterializeHandler:
         """
         self.dataset = dataset
 
-    def __call__(self, windows: list[Window]):
+    def __call__(self, windows: list[Window]) -> None:
         """Materializes the windows from apply_on_windows."""
+        logger.info(f"Running Materialize with {len(windows)} windows")
+        if self.dataset is None:
+            raise ValueError("dataset not set")
         materialize_dataset_windows(self.dataset, windows)
 
 
 @register_handler("dataset", "materialize")
-def dataset_materialize():
+def dataset_materialize() -> None:
     """Handler for the rslearn dataset materialize command."""
     parser = argparse.ArgumentParser(
         prog="rslearn dataset materialize",
@@ -540,9 +599,14 @@ def dataset_materialize():
             + "materialize data from retrieved data sources"
         ),
     )
+    parser.add_argument(
+        "--disabled-layers",
+        type=parse_disabled_layers,
+        default="",
+        help="List of layers to disable e.g 'layer1,layer2'",
+    )
     add_apply_on_windows_args(parser)
     args = parser.parse_args(args=sys.argv[3:])
-
     fn = MaterializeHandler()
     apply_on_windows_args(fn, args)
 
@@ -550,7 +614,7 @@ def dataset_materialize():
 class RslearnLightningCLI(LightningCLI):
     """LightningCLI that links data.tasks to model.tasks."""
 
-    def add_arguments_to_parser(self, parser) -> None:
+    def add_arguments_to_parser(self, parser: LightningArgumentParser) -> None:
         """Link data.tasks to model.tasks.
 
         Args:
@@ -572,7 +636,7 @@ class RslearnLightningCLI(LightningCLI):
             help="Whether to resume from specified wandb_run_id",
         )
 
-    def before_instantiate_classes(self):
+    def before_instantiate_classes(self) -> None:
         """Called before Lightning class initialization.
 
         Sets up wandb_run_id / wandb_resume arguments.
@@ -585,7 +649,7 @@ class RslearnLightningCLI(LightningCLI):
             artifact_id = (
                 f"{c.trainer.logger.init_args.project}/model-{c.wandb_run_id}:latest"
             )
-            print(f"restoring from artifact {artifact_id} on wandb")
+            logger.info(f"restoring from artifact {artifact_id} on wandb")
             artifact = api.artifact(artifact_id, type="model")
             artifact_dir = artifact.download()
             c.ckpt_path = str(Path(artifact_dir) / "model.ckpt")
@@ -606,7 +670,7 @@ class RslearnLightningCLI(LightningCLI):
             prediction_writer_callback.init_args.path = c.data.init_args.path
 
 
-def model_handler():
+def model_handler() -> None:
     """Handler for any rslearn model X commands."""
     RslearnLightningCLI(
         model_class=RslearnLightningModule,
@@ -619,30 +683,30 @@ def model_handler():
 
 
 @register_handler("model", "fit")
-def model_fit():
+def model_fit() -> None:
     """Handler for rslearn model fit."""
     model_handler()
 
 
 @register_handler("model", "validate")
-def model_validate():
+def model_validate() -> None:
     """Handler for rslearn model validate."""
     model_handler()
 
 
 @register_handler("model", "test")
-def model_test():
+def model_test() -> None:
     """Handler for rslearn model test."""
     model_handler()
 
 
 @register_handler("model", "predict")
-def model_predict():
+def model_predict() -> None:
     """Handler for rslearn model predict."""
     model_handler()
 
 
-def main():
+def main() -> None:
     """CLI entrypoint."""
     parser = argparse.ArgumentParser(description="rslearn")
     parser.add_argument(
@@ -653,7 +717,7 @@ def main():
 
     handler = handler_registry.get((args.category, args.command))
     if handler is None:
-        print(f"Unknown command: {args.category} {args.command}", file=sys.stderr)
+        logger.error(f"Unknown command: {args.category} {args.command}")
         sys.exit(1)
 
     handler()

rslearn/models/clip.py

@@ -0,0 +1,62 @@
+"""OpenAI CLIP models."""
+
+from typing import Any
+
+import torch
+from transformers import AutoModelForZeroShotImageClassification, AutoProcessor
+
+
+class CLIP(torch.nn.Module):
+    """CLIP image encoder."""
+
+    def __init__(
+        self,
+        model_name: str,
+    ):
+        """Instantiate a new CLIP instance.
+
+        Args:
+            model_name: the model name like "openai/clip-vit-large-patch14-336".
+        """
+        super().__init__()
+
+        self.processor = AutoProcessor.from_pretrained(model_name)
+        model = AutoModelForZeroShotImageClassification.from_pretrained(model_name)
+        self.encoder = model.vision_model
+
+        # Get number of features and token map size from encoder attributes.
+        self.num_features = self.encoder.post_layernorm.normalized_shape[0]
+        crop_size = self.processor.image_processor.crop_size
+        stride = self.encoder.embeddings.patch_embedding.stride
+        self.height = crop_size["height"] // stride[0]
+        self.width = crop_size["width"] // stride[1]
+
+    def forward(self, inputs: list[dict[str, Any]]) -> list[torch.Tensor]:
+        """Compute outputs from the backbone.
+
+        Inputs:
+            inputs: input dicts that must include "image" key containing the image to
+                process. The images should have values 0-255.
+
+        Returns:
+            list of feature maps. The ViT produces features at one scale, so the list
+                contains a single Bx24x24x1024 feature map.
+        """
+        device = inputs[0]["image"].device
+        clip_inputs = self.processor(
+            images=[inp["image"].cpu().numpy().transpose(1, 2, 0) for inp in inputs],
+            return_tensors="pt",
+            padding=True,
+        )
+        pixel_values = clip_inputs["pixel_values"].to(device)
+        output = self.encoder(pixel_values=pixel_values)
+        # Ignore class token output which is before the patch tokens.
+        image_features = output.last_hidden_state[:, 1:, :]
+        batch_size = image_features.shape[0]
+
+        # 576x1024 -> HxWxC
+        return [
+            image_features.reshape(
+                batch_size, self.height, self.width, self.num_features
+            ).permute(0, 3, 1, 2)
+        ]

rslearn/models/conv.py

@@ -0,0 +1,56 @@
+"""A single convolutional layer."""
+
+import torch
+
+
+class Conv(torch.nn.Module):
+    """A single convolutional layer.
+
+    It inputs a set of feature maps; the conv layer is applied to each feature map
+    independently, and list of outputs is returned.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        padding: str = "same",
+        stride: int = 1,
+        activation: torch.nn.Module = torch.nn.ReLU(inplace=True),
+    ):
+        """Initialize a Conv.
+
+        Args:
+            in_channels: number of input channels.
+            out_channels: number of output channels.
+            kernel_size: kernel size
+            padding: either "same" or "valid" to control padding
+            stride: stride to apply.
+            activation: activation to apply after convolution
+        """
+        super().__init__()
+
+        self.layer = torch.nn.Conv2d(
+            in_channels, out_channels, kernel_size, padding=padding, stride=stride
+        )
+        self.activation = activation
+
+    def forward(
+        self, features: list[torch.Tensor], inputs: list[torch.Tensor]
+    ) -> list[torch.Tensor]:
+        """Compute flat output vector from multi-scale feature map.
+
+        Args:
+            features: list of feature maps at different resolutions.
+            inputs: original inputs (ignored).
+
+        Returns:
+            flat feature vector
+        """
+        new_features = []
+        for feat_map in features:
+            feat_map = self.layer(feat_map)
+            feat_map = self.activation(feat_map)
+            new_features.append(feat_map)
+        return new_features

rslearn/models/faster_rcnn.py

@@ -10,7 +10,7 @@ import torchvision
 class NoopTransform(torch.nn.Module):
     """A placeholder transform used with torchvision detection model."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Create a new NoopTransform."""
         super().__init__()
 
@@ -46,23 +46,6 @@ class NoopTransform(torch.nn.Module):
         )
         return image_list, targets
 
-    def postprocess(
-        self, detections: dict[str, torch.Tensor], image_sizes, orig_sizes
-    ) -> dict[str, torch.Tensor]:
-        """Post-process the detections to reflect original image size.
-
-        Since we didn't transform the images, we don't need to do anything here.
-
-        Args:
-            detections: the raw detections
-            image_sizes: the transformed image sizes
-            orig_sizes: the original image sizes
-
-        Returns:
-            the post-processed detections (unmodified from the provided detections)
-        """
-        return detections
-
 
 class FasterRCNN(torch.nn.Module):
     """Faster R-CNN head for predicting bounding boxes.
@@ -80,7 +63,7 @@ class FasterRCNN(torch.nn.Module):
         anchor_sizes: list[list[int]],
         instance_segmentation: bool = False,
         box_score_thresh: float = 0.05,
-    ):
+    ) -> None:
         """Create a new FasterRCNN.
 
         Args:

rslearn/models/fpn.py

@@ -32,7 +32,7 @@ class Fpn(torch.nn.Module):
             in_channels_list=in_channels, out_channels=out_channels
         )
 
-    def forward(self, x: list[torch.Tensor]):
+    def forward(self, x: list[torch.Tensor]) -> list[torch.Tensor]:
         """Compute outputs of the FPN.
 
         Args:

rslearn/models/module_wrapper.py

@@ -0,0 +1,43 @@
+"""Module wrappers."""
+
+import torch
+
+
+class DecoderModuleWrapper(torch.nn.Module):
+    """Wrapper for a module that processes features to work in decoder.
+
+    The module should input feature map and produce a new feature map.
+
+    We wrap it to process each feature map in multi-scale features which is what's used
+    for most decoders.
+    """
+
+    def __init__(
+        self,
+        module: torch.nn.Module,
+    ):
+        """Initialize a DecoderModuleWrapper.
+
+        Args:
+            module: the module to wrap
+        """
+        super().__init__()
+        self.module = module
+
+    def forward(
+        self, features: list[torch.Tensor], inputs: list[torch.Tensor]
+    ) -> list[torch.Tensor]:
+        """Apply the wrapped module on each feature map.
+
+        Args:
+            features: list of feature maps at different resolutions.
+            inputs: original inputs (ignored).
+
+        Returns:
+            new features
+        """
+        new_features = []
+        for feat_map in features:
+            feat_map = self.module(feat_map)
+            new_features.append(feat_map)
+        return new_features