mapdet3d/model/mapdet3d.py

"""Map-Det3D."""

from __future__ import annotations

import time
from typing import NamedTuple

import torch
import torch.nn.functional as F
from huggingface_hub import PyTorchModelHubMixin
from torch import Tensor, nn
from uniception.models.encoders import ViTEncoderInput
from uniception.models.info_sharing.base import MultiViewTransformerInput
from uniception.models.prediction_heads.base import (
    AdaptorInput,
    PredictionHeadTokenInput,
)

from mapdet3d.op.mapanything.geometry import (
    convert_ray_dirs_depth_along_ray_pose_trans_quats_to_pointmap,
)
from mapdet3d.op.mapanything.image import preprocess_inputs
from mapdet3d.op.mapanything.inference import (
    postprocess_model_outputs_for_inference,
    preprocess_input_views_for_inference,
    validate_input_views_for_inference,
)
from mapdet3d.op.mapdet3d.head import Box3DHead, RoI2Det
from mapdet3d.state.track3d import MapDet3DTrackGraph

from .mapanything import MapAnything


class MapDet3DOut(NamedTuple):
    """Output of the 4D detection model."""

    boxes2d: list[Tensor]  # (B, num_queries, 4)
    boxes3d: list[Tensor]  # (B, num_queries, 10)
    class_ids: list[Tensor]  # (B, num_queries)
    track_ids: list[Tensor]  # (B, num_queries)
    scores: list[Tensor]  # (B, num_queries)


class MapDet3DTrainOut(NamedTuple):
    """Output of the Map-Det3D for training."""

    all_layers_cls_scores: list[Tensor]
    all_layers_bbox_preds: list[Tensor]
    all_layers_outputs_3d: list[Tensor] | None
    enc_outputs_class: Tensor
    enc_outputs_coord: Tensor
    enc_outputs_3d: Tensor | None
    scales: Tensor  # (B, 1)
    dn_meta: dict | None  # Denoising outputs


class MapDet3D(
    nn.Module,
    PyTorchModelHubMixin,
    library_name="mapdet3d",
    repo_url="https://github.com/cvg/Map-Det3D",
    paper_url="https://arxiv.org/abs/2608.12179",
    license="apache-2.0",
    pipeline_tag="object-detection",
    tags=["object-detection", "arxiv:2608.12179"],
):
    """Map-Det3D model."""

    def __init__(
        self,
        window_size: int = 5,
        use_intrinsics: bool = True,
        use_extrinsics: bool = True,
        roi2det: RoI2Det | None = None,
        box3d_head: Box3DHead | None = None,
        use_checkpoint: bool = True,
        track_whole_scene: bool = False,
        track_graph: MapDet3DTrackGraph | None = None,
        compute_fps: bool = False,
    ) -> None:
        """Init."""
        super().__init__()
        self.mapa = MapAnything.from_pretrained("facebook/map-anything")

        # Multi-view settings
        self.views = []
        self.cam_poses = []

        # Compute FPS
        self.compute_fps = compute_fps
        self.dinov2_features = []
        self.additonal_tokens = []

        self.window_size = window_size
        self.use_intrinsics = use_intrinsics
        self.use_extrinsics = use_extrinsics

        # Gradient checkpointing for memory efficiency
        self.use_checkpoint = use_checkpoint

        # 3D Detection
        self.box_head = Box3DHead() if box3d_head is None else box3d_head

        self.roi2det = RoI2Det() if roi2det is None else roi2det

        self._freeze_mapa()

        # Tracking
        self.track_whole_scene = track_whole_scene

        self.track_graph = (
            MapDet3DTrackGraph() if track_graph is None else track_graph
        )

    def _freeze_mapa(self) -> None:
        """Freeze the MapAnything model."""
        mapa = [
            # DINOv2
            self.mapa.encoder,
            # DPT
            self.mapa.dpt_feature_head,
            self.mapa.dpt_regressor_head,
            # Depth
            self.mapa.dense_head,
            self.mapa.dense_adaptor,
            # Pose
            self.mapa.pose_head,
            self.mapa.pose_adaptor,
        ]

        # Freeze multi-modal fusion
        mapa += [
            self.mapa.ray_dirs_encoder,
            self.mapa.depth_encoder,
            self.mapa.depth_scale_encoder,
            self.mapa.cam_rot_encoder,
            self.mapa.cam_trans_encoder,
            self.mapa.cam_trans_scale_encoder,
            self.mapa.fusion_norm_layer,
        ]

        # Multi-view information sharing
        if self.use_checkpoint:
            for i, block in enumerate(
                self.mapa.info_sharing.self_attention_blocks
            ):
                self.mapa.info_sharing.self_attention_blocks[i] = (
                    self.mapa.info_sharing.wrap_module_with_gradient_checkpointing(
                        block
                    )
                )

        for model in mapa:
            model.eval()
            for param in model.parameters():
                param.requires_grad = False

    def _forward_train(
        self,
        input_views,
        boxes2d: list[list[Tensor]] | None = None,
        class_ids: list[list[Tensor]] | None = None,
        input_hw: list[list[tuple[int, int]]] | None = None,
    ) -> MapDet3DTrainOut:
        """MapAnthing forward."""
        batch_size = len(input_views)
        seq_len = len(input_views[0])

        processed_views = []
        for i in range(seq_len):
            single_view = {
                "img": torch.cat([b[i]["img"] for b in input_views], 0),
                "data_norm_type": ["dinov2"] * batch_size,
            }

            if self.use_intrinsics:
                single_view["intrinsics"] = torch.cat(
                    [b[i]["intrinsics"] for b in input_views], 0
                )

            if self.use_extrinsics:
                single_view["camera_poses"] = torch.cat(
                    [b[i]["camera_poses"] for b in input_views], 0
                )

            if self.use_extrinsics:
                single_view["is_metric_scale"] = torch.cat(
                    [b[i]["is_metric_scale"] for b in input_views], 0
                )

            processed_views.append(single_view)

        validated_views = validate_input_views_for_inference(processed_views)

        # Pre-process the input views
        views = preprocess_input_views_for_inference(validated_views)

        num_views = len(views)

        batch_size_per_view, _, height, width = views[0]["img"].shape
        img_shape = (int(height), int(width))

        # Run the image encoder on all the input views
        with torch.no_grad():
            (
                all_encoder_features_across_views,
                all_encoder_registers_across_views,
            ) = self.mapa._encode_n_views(views)

        # Encode the optional geometric inputs and fuse with the encoded
        # features from the N input views
        # Use high precision to prevent NaN values after layer norm in dense
        # representation encoder (due to high variance in last dim of features)
        with torch.autocast("cuda", enabled=False):
            with torch.no_grad():
                all_encoder_features_across_views = (
                    self.mapa._encode_and_fuse_optional_geometric_inputs(
                        views, all_encoder_features_across_views
                    )
                )

        # Expand the scale token to match the batch size
        input_scale_token = (
            self.mapa.scale_token.unsqueeze(0)
            .unsqueeze(-1)
            .repeat(batch_size_per_view, 1, 1)
        )  # (B, C, 1)

        # Combine all images into view-centric representation
        # Output is a list containing the encoded features for all N views
        # after information sharing.
        info_sharing_input = MultiViewTransformerInput(
            features=all_encoder_features_across_views,
            additional_input_tokens_per_view=all_encoder_registers_across_views,
            additional_input_tokens=input_scale_token,
        )

        # With intermediate_features
        (
            final_info_sharing_multi_view_feat,
            intermediate_info_sharing_multi_view_feat,
        ) = self.mapa.info_sharing(info_sharing_input)

        # VGGT pred head: dpt+pose + use_encoder_features_for_dpt
        dense_head_inputs_list = []

        # Stack all the image encoder features for all views
        stacked_encoder_features = torch.cat(
            all_encoder_features_across_views, dim=0
        )
        dense_head_inputs_list.append(stacked_encoder_features)
        # Stack the first intermediate features for all views
        stacked_intermediate_features_1 = torch.cat(
            intermediate_info_sharing_multi_view_feat[0].features,
            dim=0,
        )
        dense_head_inputs_list.append(stacked_intermediate_features_1)
        # Stack the second intermediate features for all views
        stacked_intermediate_features_2 = torch.cat(
            intermediate_info_sharing_multi_view_feat[1].features,
            dim=0,
        )
        dense_head_inputs_list.append(stacked_intermediate_features_2)
        # Stack the last layer features for all views
        stacked_final_features = torch.cat(
            final_info_sharing_multi_view_feat.features, dim=0
        )
        dense_head_inputs_list.append(stacked_final_features)

        with torch.autocast("cuda", enabled=False):
            # Prepare inputs for the downstream heads
            dense_head_inputs = dense_head_inputs_list

            scale_head_inputs = (
                final_info_sharing_multi_view_feat.additional_token_features
            )

            # Scale prediction
            scale_head_output = self.mapa.scale_head(
                PredictionHeadTokenInput(last_feature=scale_head_inputs)
            )
            scale_final_output = self.mapa.scale_adaptor(
                AdaptorInput(
                    adaptor_feature=scale_head_output.decoded_channels,
                    output_shape_hw=img_shape,
                )
            )

            scales = scale_final_output.value.squeeze(
                -1
            )  # (B, 1, 1) -> (B, 1)

            (
                enc_outputs_class,
                enc_outputs_coord,
                enc_ouptputs_3d,
                all_layers_cls_scores,
                all_layers_bbox_preds,
                all_layers_outputs_3d,
                dn_meta,
            ) = self.box_head(
                dense_head_inputs,
                batch_size_per_view,
                num_views=num_views,
                boxes=boxes2d,
                class_ids=class_ids,
                input_hw=input_hw,
            )

        return MapDet3DTrainOut(
            all_layers_cls_scores=all_layers_cls_scores,
            all_layers_bbox_preds=all_layers_bbox_preds,
            all_layers_outputs_3d=all_layers_outputs_3d,
            enc_outputs_class=enc_outputs_class,
            enc_outputs_coord=enc_outputs_coord,
            enc_outputs_3d=enc_ouptputs_3d,
            scales=torch.tile(scales.squeeze(-1), (num_views, 1)),
            dn_meta=dn_meta,
        )

    def _encode_n_views(self, views) -> list[Tensor]:
        """Encode N views with cached DINOv2 features."""
        num_views = len(views)
        data_norm_type = views[0]["data_norm_type"][0]
        imgs_list = [view["img"] for view in views]
        all_imgs_across_views = torch.cat(imgs_list, dim=0)
        encoder_input = ViTEncoderInput(
            image=all_imgs_across_views[-1].unsqueeze(0),
            data_norm_type=data_norm_type,
        )
        encoder_output = self.mapa.encoder(encoder_input)

        all_encoder_features_across_views = encoder_output.features.chunk(
            num_views, dim=0
        )

        self.dinov2_features.append(all_encoder_features_across_views[0])

        if len(self.dinov2_features) > len(imgs_list):
            self.dinov2_features.pop(0)

        all_encoder_features_across_views = self.dinov2_features

        all_encoder_registers_across_views = None
        if (
            self.mapa.use_register_tokens_from_encoder
            and encoder_output.registers is not None
        ):
            all_encoder_registers_across_views = (
                encoder_output.registers.chunk(num_views, dim=0)
            )
            self.additonal_tokens.append(all_encoder_registers_across_views[0])

            if len(self.additonal_tokens) > len(imgs_list):
                self.additonal_tokens.pop(0)

            all_encoder_registers_across_views = self.additonal_tokens

        return (
            all_encoder_features_across_views,
            all_encoder_registers_across_views,
        )

    def _forward_test(
        self,
        images: Tensor,
        frame_ids: list[int],
        intrinsics: Tensor | None = None,
        extrinsics: Tensor | None = None,
    ) -> MapDet3DOut:
        """Forward for testing."""
        assert len(images) == 1, "Only support batch size 1 now."
        images = images[0]

        if frame_ids[0] == 0:
            self.views.clear()
            self.cam_poses.clear()
            self.dinov2_features.clear()
            self.additonal_tokens.clear()

        input_view = {
            "img": images[0].permute(1, 2, 0),
            "intrinsics": intrinsics[0],
            "data_norm_type": ["dinov2"] * images.shape[0],
        }

        if self.use_extrinsics:
            input_view["camera_poses"] = extrinsics[0]

        if self.use_extrinsics:
            input_view["is_metric_scale"] = torch.tensor([True])

        self.views.append(input_view)
        self.cam_poses.append(extrinsics[0])

        if len(self.views) > self.window_size:
            self.views.pop(0)
            self.cam_poses.pop(0)

        processed_views = preprocess_inputs(self.views, padding_mode=True)

        if self.use_intrinsics:
            intrinsics_list = [
                v["intrinsics"].to(images.device, non_blocking=True)
                for v in processed_views
            ]
        else:
            intrinsics_list = [
                v.pop("intrinsics").to(images.device, non_blocking=True)
                for v in processed_views
            ]

        pad_info_list = [v.pop("pad_info", None) for v in processed_views]
        pad_info = pad_info_list[-1]

        # Validate the input views
        validated_views = validate_input_views_for_inference(processed_views)

        # Transfer the views to the same device as the model
        ignore_keys = set(
            [
                "instance",
                "idx",
                "true_shape",
                "data_norm_type",
            ]
        )
        for view in validated_views:
            for name in view.keys():
                if name in ignore_keys:
                    continue
                view[name] = view[name].to(images.device, non_blocking=True)

        # Pre-process the input views
        views = preprocess_input_views_for_inference(validated_views)
        num_views = len(views)

        batch_size_per_view, _, height, width = views[0]["img"].shape
        img_shape = (int(height), int(width))

        # Run the image encoder on all the input views
        if self.compute_fps:
            start_time = time.time()
            (
                all_encoder_features_across_views,
                all_encoder_registers_across_views,
            ) = self._encode_n_views(views)
        else:
            (
                all_encoder_features_across_views,
                all_encoder_registers_across_views,
            ) = self.mapa._encode_n_views(views)

        # Encode the optional geometric inputs and fuse with the encoded
        # features from the N input views
        # Use high precision to prevent NaN values after layer norm in dense
        # representation encoder (due to high variance in last dim of features)
        with torch.autocast("cuda", enabled=False):
            all_encoder_features_across_views = (
                self.mapa._encode_and_fuse_optional_geometric_inputs(
                    views, all_encoder_features_across_views
                )
            )

        # Expand the scale token to match the batch size
        input_scale_token = (
            self.mapa.scale_token.unsqueeze(0)
            .unsqueeze(-1)
            .repeat(batch_size_per_view, 1, 1)
        )  # (B, C, 1)

        # Combine all images into view-centric representation
        # Output is a list containing the encoded features for all N views
        # after information sharing.
        info_sharing_input = MultiViewTransformerInput(
            features=all_encoder_features_across_views,
            additional_input_tokens_per_view=all_encoder_registers_across_views,
            additional_input_tokens=input_scale_token,
        )

        # With intermediate_features
        (
            final_info_sharing_multi_view_feat,
            intermediate_info_sharing_multi_view_feat,
        ) = self.mapa.info_sharing(info_sharing_input)

        # VGGT pred head: dpt+pose + use_encoder_features_for_dpt
        dense_head_inputs_list = []

        # Stack all the image encoder features for all views
        stacked_encoder_features = torch.cat(
            all_encoder_features_across_views, dim=0
        )
        dense_head_inputs_list.append(stacked_encoder_features)
        # Stack the first intermediate features for all views
        stacked_intermediate_features_1 = torch.cat(
            intermediate_info_sharing_multi_view_feat[0].features,
            dim=0,
        )
        dense_head_inputs_list.append(stacked_intermediate_features_1)
        # Stack the second intermediate features for all views
        stacked_intermediate_features_2 = torch.cat(
            intermediate_info_sharing_multi_view_feat[1].features,
            dim=0,
        )
        dense_head_inputs_list.append(stacked_intermediate_features_2)
        # Stack the last layer features for all views
        stacked_final_features = torch.cat(
            final_info_sharing_multi_view_feat.features, dim=0
        )
        dense_head_inputs_list.append(stacked_final_features)

        with torch.autocast("cuda", enabled=False):
            # Prepare inputs for the downstream heads
            dense_head_inputs = dense_head_inputs_list

            scale_head_inputs = (
                final_info_sharing_multi_view_feat.additional_token_features
            )

            # Scale prediction is lightweight, so we can run it in one go
            scale_head_output = self.mapa.scale_head(
                PredictionHeadTokenInput(last_feature=scale_head_inputs)
            )
            scale_final_output = self.mapa.scale_adaptor(
                AdaptorInput(
                    adaptor_feature=scale_head_output.decoded_channels,
                    output_shape_hw=img_shape,
                )
            )
            scale_final_output = scale_final_output.value.squeeze(
                -1
            )  # (B, 1, 1) -> (B, 1)

            # 3D Head
            (
                _,
                _,
                _,
                all_layers_cls_scores,
                all_layers_bbox_preds,
                all_layers_outputs_3d,
                _,
            ) = self.box_head(
                dense_head_inputs, batch_size_per_view, num_views=num_views
            )

            # [B, 1] -> [1]
            scale = scale_final_output[0]

            # Get the last frame's detections
            cls_scores = all_layers_cls_scores[-1][-1]
            bbox_preds = all_layers_bbox_preds[-1][-1]
            outputs_3d = all_layers_outputs_3d[-1][-1]

            det_bboxes, det_bboxes3d, scores, det_labels = self.roi2det(
                cls_scores,
                bbox_preds,
                outputs_3d,
                scale=scale,
                intrinsics=intrinsics_list[-1][0],
                img_shape=img_shape,
                ori_shape=(images.shape[2], images.shape[3]),
                pad_info=pad_info,
            )

            if self.track_whole_scene:
                tracks = self.track_graph(
                    det_bboxes3d, scores, extrinsics[0], frame_ids[0]
                )

                det_bboxes3d = tracks.boxes_3d_world
                scores = tracks.scores
                track_ids = tracks.track_ids
                det_labels = torch.zeros_like(track_ids)

                # A scene-level track has no single 2D box.
                det_bboxes = det_bboxes3d.new_empty((0, 4))
            else:
                track_ids = torch.arange(len(det_labels))

        if self.compute_fps:
            print(f"Inference time: {time.time() - start_time:.2f} seconds")

            print(
                f"Peak GPU memory usage: {torch.cuda.max_memory_allocated() / 1024**3:.2f} GB"
            )

        return MapDet3DOut(
            boxes2d=[det_bboxes],
            boxes3d=[det_bboxes3d],
            class_ids=[det_labels],
            track_ids=[track_ids],
            scores=[scores],
        )

    def forward(
        self,
        images: list[list[Tensor]] | None = None,
        views=None,
        boxes2d: list[list[Tensor]] | None = None,
        class_ids: list[list[Tensor]] | None = None,
        input_hw: list[list[tuple[int, int]]] | None = None,
        intrinsics: list[list[Tensor]] | None = None,
        extrinsics: list[list[Tensor]] | None = None,
        frame_ids: list[list[int]] | None = None,
    ) -> MapDet3DOut | MapDet3DTrainOut:
        """Forward."""
        if self.training:
            return self._forward_train(views, boxes2d, class_ids, input_hw)
        else:
            return self._forward_test(
                images, frame_ids, intrinsics, extrinsics
            )