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

rslearn/models/dinov3.py

@@ -0,0 +1,177 @@
+"""DinoV3 model.
+
+This code loads the DINOv3 model. You must obtain the model separately from Meta to use
+it. See https://github.com/facebookresearch/dinov3 for applicable license and copyright
+information.
+"""
+
+from enum import StrEnum
+from pathlib import Path
+from typing import Any
+
+import torch
+import torchvision
+from einops import rearrange
+
+from rslearn.train.model_context import ModelContext
+from rslearn.train.transforms.normalize import Normalize
+from rslearn.train.transforms.transform import Transform
+
+from .component import FeatureExtractor, FeatureMaps
+
+
+class DinoV3Models(StrEnum):
+    """Names for different DinoV3 images on torch hub."""
+
+    SMALL_WEB = "dinov3_vits16"
+    SMALL_PLUS_WEB = "dinov3_vits16plus"
+    BASE_WEB = "dinov3_vitb16"
+    LARGE_WEB = "dinov3_vitl16"
+    HUGE_PLUS_WEB = "dinov3_vith16plus"
+    FULL_7B_WEB = "dinov3_vit7b16"
+    LARGE_SATELLITE = "dinov3_vitl16_sat"
+    FULL_7B_SATELLITE = "dinov3_vit7b16_sat"
+
+
+DINOV3_PTHS: dict[str, str] = {
+    DinoV3Models.LARGE_SATELLITE: "dinov3_vitl16_pretrain_sat493m-eadcf0ff.pth",
+    DinoV3Models.FULL_7B_SATELLITE: "dinov3_vit7b16_pretrain_sat493m-a6675841.pth",
+    DinoV3Models.BASE_WEB: "dinov3_vitb16_pretrain_lvd1689m-73cec8be.pth",
+    DinoV3Models.LARGE_WEB: "dinov3_vitl16_pretrain_lvd1689m-8aa4cbdd.pth",
+    DinoV3Models.HUGE_PLUS_WEB: "dinov3_vith16plus_pretrain_lvd1689m-7c1da9a5.pth",
+    DinoV3Models.FULL_7B_WEB: "dinov3_vit7b16_pretrain_lvd1689m-a955f4.pth",
+}
+
+
+class DinoV3(FeatureExtractor):
+    """DinoV3 Backbones.
+
+    Must have the pretrained weights downloaded in checkpoint_dir for them to be loaded.
+    See https://github.com/facebookresearch/dinov3?tab=readme-ov-file#pretrained-models
+
+    Only takes RGB as input. Expects normalized data (use the below normalizer).
+
+    Uses patch size 16. The input is resized to 256x256; when applying DinoV3 on
+    segmentation or detection tasks with inputs larger than 256x256, it may be best to
+    train and predict on 256x256 crops (using SplitConfig.patch_size argument).
+    """
+
+    image_size: int = 256
+    patch_size: int = 16
+    output_dim: int = 1024
+
+    def _load_model(self, size: str, checkpoint_dir: str | None) -> torch.nn.Module:
+        model_name = size.replace("_sat", "")
+        if checkpoint_dir is not None:
+            weights = str(Path(checkpoint_dir) / DINOV3_PTHS[size])
+            return torch.hub.load(
+                "facebookresearch/dinov3",
+                model_name,
+                weights=weights,
+            )  # nosec
+        return torch.hub.load("facebookresearch/dinov3", model_name, pretrained=False)  # nosec
+
+    def __init__(
+        self,
+        checkpoint_dir: str | None,
+        size: str = DinoV3Models.LARGE_SATELLITE,
+        use_cls_token: bool = False,
+        do_resizing: bool = True,
+    ) -> None:
+        """Instantiate a new DinoV3 instance.
+
+        Args:
+            checkpoint_dir: the local path to the pretrained weight dir. If None, we load the architecture
+                only (randomly initialized).
+            size: the model size, see class for various models.
+            use_cls_token: use pooled class token (for classification), otherwise returns spatial feature map.
+            do_resizing: whether to resize inputs to 256x256. Default true.
+        """
+        super().__init__()
+        self.size = size
+        self.checkpoint_dir = checkpoint_dir
+        self.use_cls_token = use_cls_token
+        self.do_resizing = do_resizing
+        self.model = self._load_model(size, checkpoint_dir)
+
+    def forward(self, context: ModelContext) -> FeatureMaps:
+        """Forward pass for the dinov3 model.
+
+        Args:
+            context: the model context. Input dicts must include "image" key.
+
+        Returns:
+            a FeatureMaps with one feature map.
+        """
+        cur = torch.stack(
+            [inp["image"].single_ts_to_chw_tensor() for inp in context.inputs],
+            dim=0,
+        )  # (B, C, H, W)
+
+        if self.do_resizing and (
+            cur.shape[2] != self.image_size or cur.shape[3] != self.image_size
+        ):
+            cur = torchvision.transforms.functional.resize(
+                cur,
+                [self.image_size, self.image_size],
+            )
+
+        if self.use_cls_token:
+            features = self.model(cur)
+        else:
+            features = self.model.forward_features(cur)["x_norm_patchtokens"]
+            batch_size, num_patches, _ = features.shape
+            height, width = int(num_patches**0.5), int(num_patches**0.5)
+            features = rearrange(features, "b (h w) d -> b d h w", h=height, w=width)
+
+        return FeatureMaps([features])
+
+    def get_backbone_channels(self) -> list:
+        """Returns the output channels of this model when used as a backbone.
+
+        The output channels is a list of (downsample_factor, depth) that corresponds
+        to the feature maps that the backbone returns. For example, an element [2, 32]
+        indicates that the corresponding feature map is 1/2 the input resolution and
+        has 32 channels.
+        """
+        return [(self.patch_size, self.output_dim)]
+
+
+class DinoV3Normalize(Transform):
+    """Normalize inputs using DinoV3 normalization.
+
+    Normalize "image" key in input according to Dino statistics from pretraining. Satellite pretraining has slightly different normalizing than the base image model so set 'satellite' depending on what pretrained model you are using.
+
+    Input "image" should be RGB-like image between 0-255.
+    """
+
+    def __init__(self, satellite: bool = True):
+        """Initialize a new DinoV3Normalize."""
+        super().__init__()
+        self.satellite = satellite
+        if satellite:
+            mean = [0.430, 0.411, 0.296]
+            std = [0.213, 0.156, 0.143]
+        else:
+            mean = [0.485, 0.456, 0.406]
+            std = [0.229, 0.224, 0.225]
+
+        self.normalize = Normalize(
+            [value * 255 for value in mean],
+            [value * 255 for value in std],
+            num_bands=3,
+        )
+
+    def forward(
+        self, input_dict: dict[str, Any], target_dict: dict[str, Any]
+    ) -> tuple[dict[str, Any], dict[str, Any]]:
+        """Normalize the specified image with DinoV3 normalization.
+
+        Args:
+            input_dict: the input dictionary.
+            target_dict: the target dictionary.
+
+        Returns:
+            normalized (input_dicts, target_dicts) tuple
+        """
+        return self.normalize(input_dict, target_dict)

rslearn/models/faster_rcnn.py

@@ -6,14 +6,24 @@ from typing import Any
 import torch
 import torchvision
 
+from rslearn.train.model_context import ModelContext, ModelOutput
+
+from .component import FeatureMaps, Predictor
+
 
 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__()
 
+        # We initialize a GeneralizedRCNNTransform just to use its batch_images
+        # function, which concatenates the images (padding to the dimensions of the
+        # largest image as needed) to the form needed by the Faster R-CNN head.
+        # We pass an arbitrary min_size and max_size here, but these are ignored since
+        # we call GeneralizedRCNNTransform.batch_images directly rather than calling
+        # its forward function.
         self.transform = (
             torchvision.models.detection.transform.GeneralizedRCNNTransform(
                 min_size=800,
@@ -39,32 +49,17 @@ class NoopTransform(torch.nn.Module):
         Returns:
             wrapped images and unmodified targets
         """
+        # See comment above, this just pads/concatenates the images without resizing.
         images = self.transform.batch_images(images, size_divisible=32)
+        # Now convert to ImageList object needed by Faster R-CNN head.
         image_sizes = [(image.shape[1], image.shape[2]) for image in images]
         image_list = torchvision.models.detection.image_list.ImageList(
             images, image_sizes
         )
         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):
+class FasterRCNN(Predictor):
     """Faster R-CNN head for predicting bounding boxes.
 
     It inputs multi-scale features, using each feature map to predict ROIs and then
@@ -80,7 +75,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:
@@ -185,20 +180,23 @@ class FasterRCNN(torch.nn.Module):
 
     def forward(
         self,
-        features: list[torch.Tensor],
-        inputs: list[dict[str, Any]],
+        intermediates: Any,
+        context: ModelContext,
         targets: list[dict[str, Any]] | None = None,
-    ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+    ) -> ModelOutput:
         """Compute the detection outputs and loss from features.
 
         Args:
-            features: multi-scale feature maps.
-            inputs: original inputs, should cotnain image key for original image size.
+            intermediates: the output from the previous component, which must be a FeatureMaps.
+            context: the model context. Input dicts must contain image key for original image size.
             targets: should contain class key that stores the class label.
 
         Returns:
             tuple of outputs and loss dict
         """
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to FasterRCNN must be FeatureMaps")
+
         # Fix target labels to be 1 size in case it's empty.
         # For some reason this is needed.
         if targets:
@@ -212,11 +210,12 @@ class FasterRCNN(torch.nn.Module):
                     ),
                 )
 
-        image_list = [inp["image"] for inp in inputs]
+        # take the first (and assumed to be only) timestep
+        image_list = [inp["image"].image[:, 0] for inp in context.inputs]
         images, targets = self.noop_transform(image_list, targets)
 
         feature_dict = collections.OrderedDict()
-        for i, feat_map in enumerate(features):
+        for i, feat_map in enumerate(intermediates.feature_maps):
             feature_dict[f"feat{i}"] = feat_map
 
         proposals, proposal_losses = self.rpn(images, feature_dict, targets)
@@ -228,4 +227,7 @@ class FasterRCNN(torch.nn.Module):
         losses.update(proposal_losses)
         losses.update(detector_losses)
 
-        return detections, losses
+        return ModelOutput(
+            outputs=detections,
+            loss_dict=losses,
+        )

rslearn/models/feature_center_crop.py

@@ -0,0 +1,53 @@
+"""Apply center cropping on a feature map."""
+
+from typing import Any
+
+from rslearn.train.model_context import ModelContext
+
+from .component import FeatureMaps, IntermediateComponent
+
+
+class FeatureCenterCrop(IntermediateComponent):
+    """Apply center cropping on the input feature maps."""
+
+    def __init__(
+        self,
+        sizes: list[tuple[int, int]],
+    ) -> None:
+        """Create a new FeatureCenterCrop.
+
+        Only the center of each feature map will be retained and passed to the next
+        module.
+
+        Args:
+            sizes: a list of (height, width) tuples, with one tuple for each input
+                feature map.
+        """
+        super().__init__()
+        self.sizes = sizes
+
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
+        """Apply center cropping on the feature maps.
+
+        Args:
+            intermediates: output from the previous model component, which must be a FeatureMaps.
+            context: the model context.
+
+        Returns:
+            center cropped feature maps.
+        """
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to FeatureCenterCrop must be FeatureMaps")
+
+        new_features = []
+        for i, feat in enumerate(intermediates.feature_maps):
+            height, width = self.sizes[i]
+            if feat.shape[2] < height or feat.shape[3] < width:
+                raise ValueError(
+                    "feature map is smaller than the desired height and width"
+                )
+            start_h = feat.shape[2] // 2 - height // 2
+            start_w = feat.shape[3] // 2 - width // 2
+            feat = feat[:, :, start_h : start_h + height, start_w : start_w + width]
+            new_features.append(feat)
+        return FeatureMaps(new_features)

rslearn/models/fpn.py

@@ -1,12 +1,16 @@
 """Feature pyramid network."""
 
 import collections
+from typing import Any
 
-import torch
 import torchvision
 
+from rslearn.train.model_context import ModelContext
 
-class Fpn(torch.nn.Module):
+from .component import FeatureMaps, IntermediateComponent
+
+
+class Fpn(IntermediateComponent):
     """A feature pyramid network (FPN).
 
     The FPN inputs a multi-scale feature map. At each scale, it computes new features
@@ -32,20 +36,27 @@ class Fpn(torch.nn.Module):
             in_channels_list=in_channels, out_channels=out_channels
         )
 
-    def forward(self, x: list[torch.Tensor]):
+    def forward(self, intermediates: Any, context: ModelContext) -> FeatureMaps:
         """Compute outputs of the FPN.
 
         Args:
-            x: the multi-scale feature maps
+            intermediates: the output from the previous component, which must be a FeatureMaps.
+            context: the model context.
 
         Returns:
-            new multi-scale feature maps from the FPN
+            new multi-scale feature maps from the FPN.
         """
-        inp = collections.OrderedDict([(f"feat{i}", el) for i, el in enumerate(x)])
+        if not isinstance(intermediates, FeatureMaps):
+            raise ValueError("input to Fpn must be FeatureMaps")
+
+        feature_maps = intermediates.feature_maps
+        inp = collections.OrderedDict(
+            [(f"feat{i}", el) for i, el in enumerate(feature_maps)]
+        )
         output = self.fpn(inp)
         output = list(output.values())
 
         if self.prepend:
-            return output + x
+            return FeatureMaps(output + feature_maps)
         else:
-            return output
+            return FeatureMaps(output)

rslearn/models/galileo/__init__.py

@@ -0,0 +1,5 @@
+"""Galileo model."""
+
+from .galileo import GalileoModel, GalileoSize
+
+__all__ = ["GalileoModel", "GalileoSize"]