feat: Try to build everything locally.

.gitignore

@@ -0,0 +1,81 @@
+# no IntelliJ files
+.idea
+
+# don't upload macOS folder info
+*.DS_Store
+
+# don't upload node_modules from npm test
+node_modules/*
+flow-typed/*
+
+# potential files generated by golang
+bin/
+
+# don't upload webpack bundle file
+app/dist/
+
+# potential integration testing data directory
+# test_data/
+/data
+
+#python
+*.pyc
+__pycache__/
+
+# pytype
+.pytype
+
+# vscode sftp settings
+.vscode/sftp.json
+
+# vscode launch settings
+.vscode/launch.json
+
+# redis
+*.rdb
+
+# mypy
+.mypy_cache
+
+# jest coverage cache
+coverage/
+
+# downloaded repos and models
+scalabel/bot/experimental/*
+
+# python virtual environment
+env/
+
+# vscode workspace configuration
+*.code-workspace
+
+# sphinx build folder
+_build/
+
+# media files are not in this repo
+doc/media
+
+# ignore rope db cache
+.vscode/.ropeproject
+
+# python build
+build/
+dist/
+
+# coverage
+.coverage*
+
+# package default workspace
+vis4d-workspace
+
+*.tmp
+
+# local test logs and scripts
+log/
+/*.sh
+docs/source/api/*
+docs/source/tutorials/.ipynb_checkpoints/*
+wandb/
+
+# No lightning logs
+lightning_logs/

README.md

@@ -6,6 +6,7 @@ colorTo: yellow
 sdk: gradio
 sdk_version: 5.44.0
 app_file: app.py
+app_build_command: cd vis4d_cuda_ops && pip install -v .
 pinned: false
 ---
 

opendet3d/__init__.py

@@ -0,0 +1 @@
+"""3D-MOOD."""

opendet3d/common/parallel.py

@@ -0,0 +1,76 @@
+"""Utilities for cpu parallelization."""
+
+from __future__ import annotations
+
+import os
+from multiprocessing import Process, Queue
+
+# Disabling unused import becase we need Tuple in typing
+from typing import Callable, Iterable, List, Optional, Tuple, TypeVar
+
+from tqdm import tqdm
+
+Inputs = TypeVar("Inputs")
+Return = TypeVar("Return")
+
+cpu_num = os.cpu_count()
+NPROC: int = min(4, cpu_num if cpu_num else 1)
+
+
+def run(
+    func: Callable[[Inputs], Return],
+    q_in: "Queue[Tuple[int, Optional[Tuple[Inputs]]]]",
+    q_out: "Queue[Tuple[int, Return]]",
+) -> None:
+    """Run function on the inputs from the queue."""
+    while True:
+        i, x = q_in.get()
+        if i < 0 or x is None:
+            break
+        q_out.put((i, func(x[0])))
+
+
+def pmap(
+    func: Callable[[Inputs], Return],
+    inputs: Iterable[Inputs],
+    max_len: int,
+    nprocs: int = NPROC,
+) -> List[Return]:
+    """Parrell mapping func to arguments.
+
+    Different from the python pool map, this function will not hang if any of
+    the processes throws an exception.
+    """
+    q_in: "Queue[Tuple[int, Optional[Tuple[Inputs]]]]" = Queue(1)
+    q_out: "Queue[Tuple[int, Return]]" = Queue()
+
+    proc = [
+        Process(target=run, args=(func, q_in, q_out)) for _ in range(nprocs)
+    ]
+    for p in proc:
+        p.daemon = True
+        p.start()
+
+    count = 0
+    with tqdm(total=max_len) as pbar:
+        for i, x in enumerate(inputs):
+            q_in.put((i, (x,)))
+            count += 1
+
+            if count % nprocs == 0:
+                pbar.update()
+
+            pbar.refresh()
+
+        pbar.update()
+        pbar.refresh()
+
+    for _ in range(nprocs):
+        q_in.put((-1, None))
+
+    res = [q_out.get() for _ in range(count)]
+
+    for p in proc:
+        p.join()
+
+    return [x for _, x in sorted(res)]

opendet3d/data/datasets/argoverse.py

@@ -0,0 +1,94 @@
+"""Argoverse V2 Sensor dataset."""
+
+from __future__ import annotations
+
+from vis4d.common.typing import ArgsType, DictStrAny
+
+from .coco3d import COCO3DDataset
+
+TRAIN_SAMPLE_RATE = 10
+VAL_SAMPLE_RATE = 5
+ACC_FRAMES = 5
+
+
+av2_class_map = {
+    "regular vehicle": 0,
+    "pedestrian": 1,
+    "bicyclist": 2,
+    "motorcyclist": 3,
+    "wheeled rider": 4,
+    "bollard": 5,
+    "construction cone": 6,
+    "sign": 7,
+    "construction barrel": 8,
+    "stop sign": 9,
+    "mobile pedestrian crossing sign": 10,
+    "large vehicle": 11,
+    "bus": 12,
+    "box truck": 13,
+    "truck": 14,
+    "vehicular trailer": 15,
+    "truck cab": 16,
+    "school bus": 17,
+    "articulated bus": 18,
+    "message board trailer": 19,
+    "bicycle": 20,
+    "motorcycle": 21,
+    "wheeled device": 22,
+    "wheelchair": 23,
+    "stroller": 24,
+    "dog": 25,
+}
+
+av2_det_map = {
+    "regular vehicle": 0,
+    "pedestrian": 1,
+    "bicyclist": 2,
+    "motorcyclist": 3,
+    "wheeled rider": 4,
+    "bollard": 5,
+    "construction cone": 6,
+    "sign": 7,
+    "construction barrel": 8,
+    "stop sign": 9,
+    "mobile pedestrian crossing sign": 10,
+    "large vehicle": 11,
+    "bus": 12,
+    "box truck": 13,
+    "truck": 14,
+    "vehicular trailer": 15,
+    "truck cab": 16,
+    "school bus": 17,
+    "articulated bus": 18,
+    "bicycle": 19,
+    "motorcycle": 20,
+    "wheeled device": 21,
+    "stroller": 22,
+}
+
+
+class AV2SensorDataset(COCO3DDataset):
+    """Argoverse V2 Sensor dataset."""
+
+    def __init__(
+        self,
+        class_map: dict[str, int] = av2_class_map,
+        max_depth: float = 80.0,
+        depth_scale: float = 256.0,
+        **kwargs: ArgsType,
+    ) -> None:
+        """Creates an instance of the class."""
+        super().__init__(
+            class_map=class_map,
+            max_depth=max_depth,
+            depth_scale=depth_scale,
+            **kwargs,
+        )
+
+    def get_depth_filenames(self, img: DictStrAny) -> str | None:
+        """Get the depth filenames."""
+        return (
+            img["file_path"]
+            .replace("images", "depth")
+            .replace(".jpg", "_depth.png")
+        )

opendet3d/data/datasets/coco3d.py

@@ -0,0 +1,518 @@
+"""COCO 3D API."""
+
+from __future__ import annotations
+
+import contextlib
+import io
+import json
+import os
+import time
+from collections import defaultdict
+from collections.abc import Sequence
+
+import numpy as np
+from pycocotools.coco import COCO
+from pyquaternion import Quaternion
+from scipy.spatial.transform import Rotation as R
+from vis4d.common.logging import rank_zero_info, rank_zero_warn
+from vis4d.common.typing import ArgsType, DictStrAny
+from vis4d.data.const import AxisMode
+from vis4d.data.const import CommonKeys as K
+from vis4d.data.datasets.base import Dataset
+from vis4d.data.datasets.util import (
+    CacheMappingMixin,
+    im_decode,
+    print_class_histogram,
+)
+from vis4d.data.typing import DictData
+
+
+class COCO3DDataset(CacheMappingMixin, Dataset):
+    """3D Object Detection Dataset using coco annotation files."""
+
+    def __init__(
+        self,
+        data_root: str,
+        dataset_name: str,
+        class_map: dict[str, int],
+        det_map: dict[str, int],
+        keys_to_load: Sequence[str] = (K.images, K.boxes2d, K.boxes3d),
+        with_depth: bool = False,
+        max_depth: float = 80.0,
+        depth_scale: float = 256.0,
+        remove_empty: bool = False,
+        data_prefix: str | None = None,
+        text_prompt_mapping: dict[str, dict[str, str]] | None = None,
+        cache_as_binary: bool = False,
+        cached_file_path: str | None = None,
+        **kwargs: ArgsType,
+    ) -> None:
+        """Creates an instance of the class."""
+        super().__init__(**kwargs)
+        self.data_root = data_root
+        self.dataset_name = dataset_name
+        self.annotation_file = f"{dataset_name}.json"
+
+        self.keys_to_load = list(keys_to_load)
+        self.remove_empty = remove_empty
+
+        self.class_map = class_map  # Class mapping in the annotation file
+        self.det_map = det_map  # Class mapping for detection
+        self.categories = sorted(self.det_map, key=self.det_map.get)
+
+        self.data_prefix = data_prefix
+        self.text_prompt_mapping = text_prompt_mapping
+
+        # Metric Depth
+        if with_depth and not K.depth_maps in keys_to_load:
+            self.keys_to_load.append(K.depth_maps)
+
+        self.max_depth = max_depth
+        self.depth_scale = depth_scale
+
+        # Load annotations
+        self.samples, _ = self._load_mapping(
+            self._generate_data_mapping,
+            self._filter_data,
+            cache_as_binary=cache_as_binary,
+            cached_file_path=cached_file_path,
+        )
+
+    def __repr__(self) -> str:
+        """Concise representation of the dataset."""
+        return self.dataset_name
+
+    def _filter_data(self, data: list[DictStrAny]) -> list[DictStrAny]:
+        """Remove empty samples."""
+        samples = []
+
+        frequencies = {cat: 0 for cat in sorted(self.det_map)}
+
+        empty_samples = 0
+        no_depth_samples = 0
+        for sample in data:
+            if self.remove_empty and len(sample["anns"]) == 0:
+                empty_samples += 1
+                continue
+
+            if (
+                K.depth_maps in self.keys_to_load
+                and "depth_filename" not in sample
+            ):
+                empty_samples += 1
+                no_depth_samples += 1
+                continue
+
+            for ann in sample["anns"]:
+                frequencies[ann["category_name"]] += 1
+
+            samples.append(sample)
+
+        rank_zero_info(
+            f"Propocessing {self.dataset_name} with {len(samples)} samples."
+        )
+        rank_zero_info(f"No depth samples: {no_depth_samples}")
+        rank_zero_info(f"Filtered {empty_samples} empty samples")
+        print_class_histogram(frequencies)
+
+        return samples
+
+    def _get_cat_id(
+        self, img: DictStrAny, ann: DictStrAny, cat_name: str
+    ) -> None:
+        """Get the category id from the category name."""
+        ann["category_id"] = self.det_map[cat_name]
+
+    def _generate_data_mapping(self) -> list[DictStrAny]:
+        """Generates the data mapping."""
+        # Load annotations
+        with contextlib.redirect_stdout(io.StringIO()):
+            coco_api = COCO3D(
+                os.path.join(
+                    self.data_root, "annotations", self.annotation_file
+                ),
+                self.categories,
+            )
+
+        cats_map = {v: k for k, v in self.class_map.items()}
+
+        img_ids = sorted(coco_api.getImgIds())
+        imgs = coco_api.loadImgs(img_ids)
+
+        samples = []
+        for img_id, img in zip(img_ids, imgs):
+            # Fix file path for Omni3D
+            if self.data_prefix is not None:
+                img["file_path"] = os.path.join(
+                    self.data_prefix, img["file_path"]
+                )
+
+            valid_anns = []
+            anns = coco_api.imgToAnns[img_id]
+
+            boxes = []
+            boxes3d = np.empty((0, 10), dtype=np.float32)[1:]
+            class_ids = np.empty((0,), dtype=np.int64)[1:]
+            for ann in anns:
+                cat_name = cats_map[ann["category_id"]]
+                assert cat_name == ann["category_name"]
+
+                if cat_name in {"dontcare", "ignore", "void"}:
+                    continue
+
+                if ann["ignore"]:
+                    continue
+
+                self._get_cat_id(img, ann, cat_name)
+
+                # Box 2D
+                x1, y1, width, height = ann["bbox"]
+                x2, y2 = x1 + width, y1 + height
+                boxes.append((x1, y1, x2, y2))
+
+                # Class
+                class_ids = np.concatenate(
+                    [
+                        class_ids,
+                        np.array([ann["category_id"]], dtype=np.int64),
+                    ]
+                )
+
+                # Box 3D
+                center = ann["center_cam"]
+                width, height, length = ann["dimensions"]
+
+                # Check if the rotation matrix is valid
+                try:
+                    x, y, z, w = R.from_matrix(
+                        np.array(ann["R_cam"])
+                    ).as_quat()
+                except Exception as e:
+                    rank_zero_warn(
+                        f"Error processing rotation matrix for annotation {ann['id']}: {e}"
+                    )
+                    continue
+
+                orientation = Quaternion([w, x, y, z])
+
+                boxes3d = np.concatenate(
+                    [
+                        boxes3d,
+                        np.array(
+                            [
+                                [
+                                    *center,
+                                    width,
+                                    length,
+                                    height,
+                                    *orientation.elements,
+                                ]
+                            ],
+                            dtype=np.float32,
+                        ),
+                    ]
+                )
+
+                valid_anns.append(ann)
+
+            boxes2d = (
+                np.empty((0, 4), dtype=np.float32)
+                if not boxes
+                else np.array(boxes, dtype=np.float32)
+            )
+
+            depth_filename = self.get_depth_filenames(img)
+
+            sample = {
+                "img_id": img_id,
+                "img": img,
+                "anns": valid_anns,
+                "boxes2d": boxes2d,
+                "boxes3d": boxes3d,
+                "class_ids": class_ids,
+            }
+
+            if depth_filename is not None and self.data_backend.exists(
+                depth_filename
+            ):
+                sample["depth_filename"] = depth_filename
+
+            samples.append(sample)
+
+        return samples
+
+    def get_depth_filenames(self, img: DictStrAny) -> str | None:
+        """Get the depth filenames.
+
+        Since not every data has depth.
+        """
+        return None
+
+    def get_cat_ids(self, idx: int) -> list[int]:
+        """Return the samples."""
+        return self.samples[idx]["class_ids"].tolist()
+
+    def __len__(self) -> int:
+        """Total number of samples of data."""
+        return len(self.samples)
+
+    def get_depth_map(self, sample: DictStrAny) -> np.ndarray:
+        """Get the depth map."""
+        depth_bytes = self.data_backend.get(sample["depth_filename"])
+        depth_array = im_decode(depth_bytes)
+
+        depth = np.ascontiguousarray(depth_array, dtype=np.float32)
+
+        depth = depth / self.depth_scale
+
+        return depth
+
+    def __getitem__(self, idx: int) -> DictData:
+        """Get single sample.
+
+        Args:
+            idx (int): Index of sample.
+
+        Returns:
+            DictData: sample at index in Vis4D input format.
+        """
+        sample = self.samples[idx]
+        data_dict: DictData = {}
+
+        # Get image info
+        data_dict[K.sample_names] = sample["img_id"]
+
+        data_dict["dataset_name"] = self.dataset_name
+        data_dict[K.boxes2d_names] = self.categories
+        data_dict["text_prompt_mapping"] = self.text_prompt_mapping
+
+        if K.images in self.keys_to_load:
+            im_bytes = self.data_backend.get(sample["img"]["file_path"])
+            image = np.ascontiguousarray(
+                im_decode(im_bytes, mode=self.image_channel_mode),
+                dtype=np.float32,
+            )[None]
+
+            data_dict[K.images] = image
+            data_dict[K.input_hw] = (image.shape[1], image.shape[2])
+
+            data_dict[K.original_images] = image
+            data_dict[K.original_hw] = (image.shape[1], image.shape[2])
+
+            # Get camera info
+            intrinsics = np.array(sample["img"]["K"], dtype=np.float32)
+            data_dict[K.intrinsics] = intrinsics
+            data_dict["original_intrinsics"] = intrinsics
+
+        data_dict[K.boxes2d] = sample["boxes2d"]
+        data_dict[K.boxes2d_classes] = sample["class_ids"]
+        data_dict[K.boxes3d] = sample["boxes3d"]
+        data_dict[K.boxes3d_classes] = sample["class_ids"]
+        data_dict[K.axis_mode] = AxisMode.OPENCV
+
+        if K.depth_maps in self.keys_to_load:
+            depth = self.get_depth_map(sample)
+
+            depth[depth > self.max_depth] = 0
+
+            data_dict[K.depth_maps] = depth
+
+        data_dict["tokens_positive"] = None
+
+        self.data_backend.close()
+
+        return data_dict
+
+
+class COCO3D(COCO):
+    """COCO API with 3D annotations."""
+
+    def __init__(
+        self,
+        annotation_files: Sequence[str] | str,
+        category_names: Sequence[str] | None = None,
+        ignore_names: Sequence[str] = ("dontcare", "ignore", "void"),
+        truncation_thres: float = 0.33333333,
+        visibility_thres: float = 0.33333333,
+        min_height_thres: float = 0.0625,
+        max_height_thres: float = 1.50,
+        modal_2D_boxes: bool = False,
+        trunc_2D_boxes: bool = True,
+        max_depth: int = 1e8,
+    ) -> None:
+        """Creates an instance of the class."""
+        self.dataset, self.anns, self.cats, self.imgs = {}, {}, {}, {}
+        self.imgToAnns, self.catToImgs = defaultdict(list), defaultdict(list)
+
+        self.truncation_thres = truncation_thres
+        self.visibility_thres = visibility_thres
+        self.min_height_thres = min_height_thres
+        self.max_height_thres = max_height_thres
+        self.max_depth = max_depth
+
+        if isinstance(annotation_files, str):
+            annotation_files = [annotation_files]
+
+        cats_ids_master = []
+        cats_master = []
+
+        for annotation_file in annotation_files:
+            _, tail = os.path.split(annotation_file)
+            name, _ = os.path.splitext(tail)
+
+            print(f"loading {name} annotations into memory...")
+            tic = time.time()
+
+            with open(annotation_file, "r") as f:
+                dataset = json.load(f)
+
+            assert (
+                type(dataset) == dict
+            ), f"annotation file format {type(dataset)} not supported"
+            print(f"Done (t={time.time() - tic:.2f}s)")
+
+            if type(dataset["info"]) == list:
+                dataset["info"] = dataset["info"][0]
+
+            dataset["info"]["known_category_ids"] = [
+                cat["id"] for cat in dataset["categories"]
+            ]
+
+            # first dataset
+            if len(self.dataset) == 0:
+                self.dataset = dataset
+            # concatenate datasets
+            else:
+                if type(self.dataset["info"]) == dict:
+                    self.dataset["info"] = [self.dataset["info"]]
+
+                self.dataset["info"] += [dataset["info"]]
+                self.dataset["annotations"] += dataset["annotations"]
+                self.dataset["images"] += dataset["images"]
+
+            # sort through categories
+            for cat in dataset["categories"]:
+                if not cat["id"] in cats_ids_master:
+                    cats_ids_master.append(cat["id"])
+                    cats_master.append(cat)
+
+        # category names are provided to us
+        if category_names is not None:
+            self.dataset["categories"] = [
+                cats_master[i]
+                for i in np.argsort(cats_ids_master)
+                if cats_master[i]["name"] in category_names
+            ]
+        # no categories are provided, so assume use ALL available.
+        else:
+            self.dataset["categories"] = [
+                cats_master[i] for i in np.argsort(cats_ids_master)
+            ]
+
+            category_names = [
+                cat["name"] for cat in self.dataset["categories"]
+            ]
+
+        # determine which categories we may actually use for filtering.
+        trainable_cats = set(ignore_names) | set(category_names)
+
+        valid_anns = []
+        im_height_map = {}
+
+        for im_obj in self.dataset["images"]:
+            im_height_map[im_obj["id"]] = im_obj["height"]
+
+        # Filter out annotations
+        for anno_idx, anno in enumerate(self.dataset["annotations"]):
+
+            im_height = im_height_map[anno["image_id"]]
+
+            # tightly annotated 2D boxes are not always available.
+            if (
+                modal_2D_boxes
+                and "bbox2D_tight" in anno
+                and anno["bbox2D_tight"][0] != -1
+            ):
+                bbox2D = anno["bbox2D_tight"]
+            elif (
+                trunc_2D_boxes
+                and "bbox2D_trunc" in anno
+                and not np.all([val == -1 for val in anno["bbox2D_trunc"]])
+            ):
+                bbox2D = anno["bbox2D_trunc"]
+            elif anno["bbox2D_proj"][0] != -1:
+                bbox2D = anno["bbox2D_proj"]
+            elif anno["bbox2D_tight"][0] != -1:
+                bbox2D = anno["bbox2D_tight"]
+            else:
+                continue
+
+            # convert to xywh
+            bbox2D[2] = bbox2D[2] - bbox2D[0]
+            bbox2D[3] = bbox2D[3] - bbox2D[1]
+
+            ignore = self.is_ignore(anno, bbox2D, ignore_names, im_height)
+
+            width = bbox2D[2]
+            height = bbox2D[3]
+
+            self.dataset["annotations"][anno_idx]["area"] = width * height
+            self.dataset["annotations"][anno_idx]["iscrowd"] = False
+            self.dataset["annotations"][anno_idx]["ignore"] = ignore
+            self.dataset["annotations"][anno_idx]["ignore2D"] = ignore
+            self.dataset["annotations"][anno_idx]["ignore3D"] = ignore
+
+            self.dataset["annotations"][anno_idx]["bbox"] = bbox2D
+            self.dataset["annotations"][anno_idx]["bbox3D"] = anno[
+                "bbox3D_cam"
+            ]
+            self.dataset["annotations"][anno_idx]["depth"] = anno[
+                "center_cam"
+            ][2]
+
+            category_name = anno["category_name"]
+
+            if category_name in trainable_cats:
+                valid_anns.append(self.dataset["annotations"][anno_idx])
+
+        self.dataset["annotations"] = valid_anns
+
+        self.createIndex()
+
+    def is_ignore(
+        self,
+        anno,
+        bbox2D: list[float, float, float, float],
+        ignore_names: Sequence[str] | None,
+        image_height: int,
+    ) -> bool:
+        ignore = anno["behind_camera"]
+        ignore |= not bool(anno["valid3D"])
+
+        if ignore:
+            return ignore
+
+        ignore |= anno["dimensions"][0] <= 0
+        ignore |= anno["dimensions"][1] <= 0
+        ignore |= anno["dimensions"][2] <= 0
+        ignore |= anno["center_cam"][2] > self.max_depth
+        ignore |= anno["lidar_pts"] == 0
+        ignore |= anno["segmentation_pts"] == 0
+        ignore |= anno["depth_error"] > 0.5
+
+        ignore |= bbox2D[3] <= self.min_height_thres * image_height
+        ignore |= bbox2D[3] >= self.max_height_thres * image_height
+
+        ignore |= (
+            anno["truncation"] >= 0
+            and anno["truncation"] >= self.truncation_thres
+        )
+        ignore |= (
+            anno["visibility"] >= 0
+            and anno["visibility"] <= self.visibility_thres
+        )
+
+        if ignore_names is not None:
+            ignore |= anno["category_name"] in ignore_names
+
+        return ignore

opendet3d/data/datasets/odvg.py

@@ -0,0 +1,280 @@
+"""Object detection and visual grounding dataset."""
+
+from __future__ import annotations
+
+import json
+import os.path as osp
+
+import numpy as np
+from tqdm import tqdm
+from vis4d.common.logging import rank_zero_info
+from vis4d.common.typing import ArgsType, DictStrAny
+from vis4d.data.const import CommonKeys as K
+from vis4d.data.datasets.base import Dataset
+from vis4d.data.datasets.util import (
+    CacheMappingMixin,
+    im_decode,
+    print_class_histogram,
+)
+from vis4d.data.typing import DictData
+
+
+class ODVGDataset(CacheMappingMixin, Dataset):
+    """Object detection and visual grounding dataset."""
+
+    def __init__(
+        self,
+        data_root: str,
+        ann_file: str,
+        label_map_file: str | None = None,
+        dataset_type: str = "VG",
+        dataset_prefix: str | None = None,
+        remove_empty: bool = False,
+        cache_as_binary: bool = False,
+        cached_file_path: str | None = None,
+        **kwargs: ArgsType,
+    ) -> None:
+        """Create an object detection and visual grounding dataset."""
+        super().__init__(**kwargs)
+
+        self.data_root = data_root
+        self.ann_file = ann_file
+        self.dataset_type = dataset_type
+        self.dataset_prefix = dataset_prefix
+        self.remove_empty = remove_empty
+
+        if label_map_file is not None:
+            label_map_file = osp.join(self.data_root, label_map_file)
+
+            with open(label_map_file, "r") as file:
+                # dict[class_id (str): class_name (str)]
+                self.label_map = json.load(file)
+
+            self.dataset_type = "OD"
+
+            self.det_map = {v: int(k) for k, v in self.label_map.items()}
+            self.categories = sorted(self.det_map, key=self.det_map.get)
+        else:
+            self.label_map = None
+            self.dataset_type = "VG"
+
+        # Load annotations
+        self.samples, _ = self._load_mapping(
+            self._generate_data_mapping,
+            self._filter_data,
+            cache_as_binary=cache_as_binary,
+            cached_file_path=cached_file_path,
+        )
+
+    def __repr__(self) -> str:
+        """Concise representation of the dataset."""
+        return f"ODVGDataset({self.ann_file})"
+
+    def _filter_data(self, data: list[DictStrAny]) -> list[DictStrAny]:
+        """Remove empty samples."""
+        samples = []
+
+        if self.dataset_type == "OD":
+            frequencies = {cat: 0 for _, cat in self.label_map.items()}
+
+        empty_samples = 0
+        for sample in data:
+            if self.remove_empty and len(sample["anns"]) == 0:
+                empty_samples += 1
+                continue
+
+            if self.dataset_type == "OD":
+                for ann in sample["anns"]:
+                    frequencies[ann["category"]] += 1
+
+            samples.append(sample)
+
+        rank_zero_info(f"Propocessing {self} with {len(samples)} samples.")
+        rank_zero_info(f"Filtered {empty_samples} empty samples")
+
+        if self.dataset_type == "OD":
+            frequencies = dict(sorted(frequencies.items()))
+
+            print_class_histogram(frequencies)
+
+        return samples
+
+    def _generate_data_mapping(self) -> list[DictStrAny]:
+        """Generates the data mapping."""
+        with open(osp.join(self.data_root, self.ann_file), "r") as f:
+            data_list = [json.loads(line) for line in f]
+
+        if self.with_camera:
+            with open(osp.join(self.data_root, "cam_info.json"), "r") as f:
+                cameras = json.load(f)
+
+        samples = []
+        for data in tqdm(data_list):
+            data_info = {}
+
+            if self.dataset_prefix is not None:
+                img_path = osp.join(
+                    self.data_root, self.dataset_prefix, data["filename"]
+                )
+            else:
+                img_path = osp.join(self.data_root, data["filename"])
+
+            data_info["img_path"] = img_path
+
+            # Pseudo K
+            if self.with_camera:
+                data_info["K"] = cameras[img_path][0]
+
+            # Pseudo Depth Path
+            if self.dataset_prefix is not None:
+                depth_path = osp.join(
+                    self.data_root,
+                    f"{self.dataset_prefix}_depth",
+                    data["filename"].replace(".jpg", "_depth.png"),
+                )
+            else:
+                depth_path = osp.join(
+                    self.data_root,
+                    data["filename"].replace(".jpg", "_depth.png"),
+                )
+            data_info["depth_path"] = depth_path
+
+            data_info["height"] = data["height"]
+            data_info["width"] = data["width"]
+
+            valid_anns = []
+            boxes = []
+            class_ids = np.empty((0,), dtype=np.int64)[1:]
+            if self.dataset_type == "OD":
+                instances = data.get("detection", {}).get("instances", [])
+
+                for ann in instances:
+                    bbox = ann["bbox"]
+
+                    # Box 2D
+                    x1, y1, x2, y2 = bbox
+                    inter_w = max(0, min(x2, data["width"]) - max(x1, 0))
+                    inter_h = max(0, min(y2, data["height"]) - max(y1, 0))
+
+                    if inter_w * inter_h == 0:
+                        continue
+                    if (x2 - x1) < 1 or (y2 - y1) < 1:
+                        continue
+
+                    boxes.append(bbox)
+
+                    # Class
+                    class_ids = np.concatenate(
+                        [class_ids, np.array([ann["label"]], dtype=np.int64)]
+                    )
+
+                    valid_anns.append(ann)
+            else:
+                anno = data["grounding"]
+
+                caption = anno["caption"].lower().strip()
+                if not caption.endswith("."):
+                    caption = caption + ". "
+
+                data_info["caption"] = caption
+
+                regions = anno["regions"]
+                phrases = []
+                positive_positions = []
+                for i, region in enumerate(regions):
+                    bboxes = region["bbox"]
+
+                    if not isinstance(bboxes[0], list):
+                        bboxes = [bboxes]
+
+                    for bbox in bboxes:
+                        x1, y1, x2, y2 = bbox
+                        inter_w = max(0, min(x2, data["width"]) - max(x1, 0))
+                        inter_h = max(0, min(y2, data["height"]) - max(y1, 0))
+
+                        if inter_w * inter_h == 0:
+                            continue
+                        if (x2 - x1) < 1 or (y2 - y1) < 1:
+                            continue
+
+                        boxes.append(bbox)
+                        phrases.append(region["phrase"])
+                        positive_positions.append(region["tokens_positive"])
+                        valid_anns.append(region)
+
+                        class_ids = np.concatenate(
+                            [class_ids, np.array([i], dtype=np.int64)]
+                        )
+
+                data_info["phrases"] = phrases
+                data_info["positive_positions"] = positive_positions
+
+            boxes2d = (
+                np.empty((0, 4), dtype=np.float32)
+                if not boxes
+                else np.array(boxes, dtype=np.float32)
+            )
+
+            data_info["boxes2d"] = boxes2d
+            data_info["class_ids"] = class_ids
+            data_info["anns"] = valid_anns
+
+            samples.append(data_info)
+
+        del data_list
+        return samples
+
+    def get_cat_ids(self, idx: int) -> list[int]:
+        """Return the samples."""
+        return self.samples[idx]["class_ids"].tolist()
+
+    def __len__(self) -> int:
+        """Total number of samples of data."""
+        return len(self.samples)
+
+    def __getitem__(self, idx: int) -> DictData:
+        """Get single sample.
+
+        Args:
+            idx (int): Index of sample.
+
+        Returns:
+            DictData: sample at index in Vis4D input format.
+        """
+        sample = self.samples[idx]
+        data_dict: DictData = {}
+
+        # Get image info
+        sample_name = sample["img_path"].split("/")[-1]
+        data_dict[K.sample_names] = sample_name
+
+        im_bytes = self.data_backend.get(sample["img_path"])
+        image = np.ascontiguousarray(
+            im_decode(im_bytes, mode=self.image_channel_mode),
+            dtype=np.float32,
+        )[None]
+
+        data_dict[K.images] = image
+        data_dict[K.input_hw] = (image.shape[1], image.shape[2])
+
+        data_dict[K.original_images] = image
+        data_dict[K.original_hw] = (image.shape[1], image.shape[2])
+
+        data_dict[K.boxes2d] = sample["boxes2d"]
+        data_dict[K.boxes2d_classes] = sample["class_ids"]
+
+        if self.dataset_type == "OD":
+            data_dict[K.boxes2d_names] = self.categories
+            data_dict["phrases"] = None
+            data_dict["positive_positions"] = None
+        else:
+            data_dict[K.boxes2d_names] = sample["caption"]
+            data_dict["phrases"] = sample["phrases"]
+            data_dict["positive_positions"] = sample["positive_positions"]
+
+        data_dict["dataset_type"] = self.dataset_type
+        data_dict["label_map"] = self.label_map
+
+        self.data_backend.close()
+
+        return data_dict

opendet3d/data/datasets/omni3d/__init__.py

@@ -0,0 +1 @@
+"""Omni3D Dataset."""

opendet3d/data/datasets/omni3d/arkitscenes.py

@@ -0,0 +1,81 @@
+"""ARKitScenes from Omni3D."""
+
+from __future__ import annotations
+
+import os
+
+from vis4d.common.typing import ArgsType, DictStrAny
+
+from opendet3d.data.datasets.coco3d import COCO3DDataset
+
+from .omni3d_classes import omni3d_class_map
+
+arkitscenes_det_map = {
+    "bathtub": 0,
+    "bed": 1,
+    "cabinet": 2,
+    "chair": 3,
+    "fireplace": 4,
+    "machine": 5,
+    "oven": 6,
+    "refrigerator": 7,
+    "shelves": 8,
+    "sink": 9,
+    "sofa": 10,
+    "stove": 11,
+    "table": 12,
+    "television": 13,
+    "toilet": 14,
+}
+
+omni3d_arkitscenes_det_map = {
+    "table": 0,
+    "bed": 1,
+    "sofa": 2,
+    "television": 3,
+    "refrigerator": 4,
+    "chair": 5,
+    "oven": 6,
+    "machine": 7,
+    "stove": 8,
+    "shelves": 9,
+    "sink": 10,
+    "cabinet": 11,
+    "bathtub": 12,
+    "toilet": 13,
+}
+
+
+class ARKitScenes(COCO3DDataset):
+    """ARKitScenes Dataset."""
+
+    def __init__(
+        self,
+        class_map: dict[str, int] = omni3d_class_map,
+        max_depth: float = 10.0,
+        depth_scale: float = 1000.0,
+        **kwargs: ArgsType,
+    ) -> None:
+        """Creates an instance of the class."""
+        super().__init__(
+            class_map=class_map,
+            max_depth=max_depth,
+            depth_scale=depth_scale,
+            **kwargs,
+        )
+
+    def get_depth_filenames(self, img: DictStrAny) -> str | None:
+        """Get the depth filenames.
+
+        Since not every data has depth.
+        """
+        _, _, split, video_id, image_name = img["file_path"].split("/")
+
+        depth_filename = os.path.join(
+            "data/ARKitScenes_depth",
+            split,
+            video_id,
+            image_name.replace("jpg", "png"),
+        )
+
+        return depth_filename

opendet3d/data/datasets/omni3d/hypersim.py

@@ -0,0 +1,190 @@
+"""Hypersim from Omni3D."""
+
+from __future__ import annotations
+
+import os
+
+from vis4d.common.typing import ArgsType, DictStrAny
+
+from opendet3d.data.datasets.coco3d import COCO3DDataset
+
+from .omni3d_classes import omni3d_class_map
+
+hypersim_train_det_map = {
+    "bathtub": 0,
+    "bed": 1,
+    "blinds": 2,
+    "bookcase": 3,
+    "books": 4,
+    "box": 5,
+    "cabinet": 6,
+    "chair": 7,
+    "clothes": 8,
+    "counter": 9,
+    "curtain": 10,
+    "desk": 11,
+    "door": 12,
+    "dresser": 13,
+    "floor mat": 14,
+    "lamp": 15,
+    "mirror": 16,
+    "night stand": 17,
+    "person": 18,
+    "picture": 19,
+    "pillow": 20,
+    "refrigerator": 21,
+    "shelves": 22,
+    "sink": 23,
+    "sofa": 24,
+    "stationery": 25,
+    "table": 26,
+    "television": 27,
+    "toilet": 28,
+    "towel": 29,
+    "window": 30,
+}
+
+hypersim_val_det_map = {
+    "bathtub": 0,
+    "bed": 1,
+    "blinds": 2,
+    "bookcase": 3,
+    "books": 4,
+    "box": 5,
+    "cabinet": 6,
+    "chair": 7,
+    "clothes": 8,
+    "counter": 9,
+    "curtain": 10,
+    "desk": 11,
+    "door": 12,
+    "dresser": 13,
+    "floor mat": 14,
+    "lamp": 15,
+    "mirror": 16,
+    "night stand": 17,
+    "picture": 18,
+    "pillow": 19,
+    "refrigerator": 20,
+    "shelves": 21,
+    "sink": 22,
+    "sofa": 23,
+    "stationery": 24,
+    "table": 25,
+    "television": 26,
+    "toilet": 27,
+    "towel": 28,
+    "window": 29,
+}
+
+hypersim_test_det_map = {
+    "bathtub": 0,
+    "bed": 1,
+    "blinds": 2,
+    "board": 3,
+    "bookcase": 4,
+    "books": 5,
+    "box": 6,
+    "cabinet": 7,
+    "chair": 8,
+    "clothes": 9,
+    "counter": 10,
+    "curtain": 11,
+    "desk": 12,
+    "door": 13,
+    "floor mat": 14,
+    "lamp": 15,
+    "mirror": 16,
+    "night stand": 17,
+    "picture": 18,
+    "pillow": 19,
+    "refrigerator": 20,
+    "shelves": 21,
+    "sink": 22,
+    "sofa": 23,
+    "stationery": 24,
+    "table": 25,
+    "television": 26,
+    "towel": 27,
+    "window": 28,
+}
+
+
+omni3d_hypersim_det_map = {
+    "books": 0,
+    "chair": 1,
+    "towel": 2,
+    "blinds": 3,
+    "window": 4,
+    "lamp": 5,
+    "shelves": 6,
+    "mirror": 7,
+    "sink": 8,
+    "cabinet": 9,
+    "bathtub": 10,
+    "door": 11,
+    "desk": 12,
+    "box": 13,
+    "bookcase": 14,
+    "picture": 15,
+    "table": 16,
+    "counter": 17,
+    "bed": 18,
+    "night stand": 19,
+    "pillow": 20,
+    "sofa": 21,
+    "television": 22,
+    "floor mat": 23,
+    "curtain": 24,
+    "clothes": 25,
+    "stationery": 26,
+    "refrigerator": 27,
+}
+
+
+def get_hypersim_det_map(split: str) -> dict[str, int]:
+    """Get Hypersim detection map."""
+    assert split in {"train", "val", "test"}, f"Invalid split: {split}"
+
+    if split == "train":
+        return hypersim_train_det_map
+    elif split == "val":
+        return hypersim_val_det_map
+    elif split == "test":
+        return hypersim_test_det_map
+
+
+class Hypersim(COCO3DDataset):
+    """Hypersim Dataset."""
+
+    def __init__(
+        self,
+        class_map: dict[str, int] = omni3d_class_map,
+        max_depth: float = 50.0,
+        depth_scale: float = 1000.0,
+        **kwargs: ArgsType,
+    ) -> None:
+        """Creates an instance of the class."""
+        super().__init__(
+            class_map=class_map,
+            max_depth=max_depth,
+            depth_scale=depth_scale,
+            **kwargs,
+        )
+
+    def get_depth_filenames(self, img: DictStrAny) -> str | None:
+        """Get the depth filenames.
+
+        Since not every data has depth.
+        """
+        _, _, scene, _, img_dir, img_name = img["file_path"].split("/")
+
+        depth_filename = os.path.join(
+            "data/hypersim_depth",
+            scene,
+            "images",
+            img_dir,
+            img_name.replace("jpg", "png"),
+        )
+
+        return depth_filename

opendet3d/data/datasets/omni3d/kitti_object.py

@@ -0,0 +1,105 @@
+"""KITTI Object from Omni3D.
+
+KITTI Object Labels:
+Categories, -, -, alpha, x1, y1, x2, y2, h, w, l, x, botom_y, z, ry
+
+KITTI Object Categories:
+{
+    "Pedestrian": "pedestrian",
+    "Cyclist": "cyclist",
+    "Car": "car",
+    "Van": "car",
+    "Truck": "truck",
+    "Tram": "tram",
+    "Person": "pedestrian",
+    "Person_sitting": "pedestrian",
+    "Misc": "misc",
+    "DontCare": "dontcare",
+}
+"""
+
+from __future__ import annotations
+
+import os
+
+from vis4d.common.typing import ArgsType, DictStrAny
+
+from opendet3d.data.datasets.coco3d import COCO3DDataset
+
+from .omni3d_classes import omni3d_class_map
+
+kitti_train_det_map = kitti_test_det_map = {
+    "car": 0,
+    "cyclist": 1,
+    "pedestrian": 2,
+    "person": 3,
+    "tram": 4,
+    "truck": 5,
+    "van": 6,
+}
+
+kitti_val_det_map = {
+    "car": 0,
+    "cyclist": 1,
+    "pedestrian": 2,
+    "tram": 3,
+    "truck": 4,
+}
+
+# KITTI-Omni3D Mapping
+omni3d_kitti_det_map = {
+    "pedestrian": 0,
+    "car": 1,
+    "cyclist": 2,
+    "van": 3,
+    "truck": 4,
+}
+
+
+def get_kitti_det_map(split: str) -> dict[str, int]:
+    """Get the KITTI detection map."""
+    assert split in {"train", "val", "test"}, f"Invalid split: {split}"
+
+    if split == "val":
+        return kitti_val_det_map
+
+    # Train and Test are the same
+    return kitti_train_det_map
+
+
+class KITTIObject(COCO3DDataset):
+    """KITTI Object Dataset."""
+
+    def __init__(
+        self,
+        class_map: dict[str, int] = omni3d_class_map,
+        max_depth: float = 80.0,
+        depth_scale: float = 256.0,
+        depth_data_root: str = "data/KITTI_object_depth",
+        **kwargs: ArgsType,
+    ) -> None:
+        """Creates an instance of the class."""
+        self.depth_data_root = depth_data_root
+
+        super().__init__(
+            class_map=class_map,
+            max_depth=max_depth,
+            depth_scale=depth_scale,
+            **kwargs,
+        )
+
+    def get_depth_filenames(self, img: DictStrAny) -> str | None:
+        """Get the depth filenames.
+
+        Since not every data has depth.
+        """
+        _, _, split, image_id, img_filename = img["file_path"].split("/")
+
+        depth_filename = os.path.join(
+            self.depth_data_root,
+            split,
+            image_id,
+            img_filename.replace(".jpg", ".png"),
+        )
+
+        return depth_filename

opendet3d/data/datasets/omni3d/nuscenes.py

@@ -0,0 +1,62 @@
+"""nuScenes from Omni3D."""
+
+from __future__ import annotations
+
+from vis4d.common.typing import ArgsType, DictStrAny
+
+from opendet3d.data.datasets.coco3d import COCO3DDataset
+
+from .omni3d_classes import omni3d_class_map
+
+nusc_det_map = {
+    "bicycle": 0,
+    "motorcycle": 1,
+    "pedestrian": 2,
+    "bus": 3,
+    "car": 4,
+    "trailer": 5,
+    "truck": 6,
+    "traffic cone": 7,
+    "barrier": 8,
+}
+
+
+class nuScenes(COCO3DDataset):
+    """nuScenes dataset."""
+
+    def __init__(
+        self,
+        class_map: dict[str, int] = omni3d_class_map,
+        max_depth: float = 80.0,
+        depth_scale: float = 256.0,
+        **kwargs: ArgsType,
+    ) -> None:
+        """Creates an instance of the class."""
+        super().__init__(
+            class_map=class_map,
+            max_depth=max_depth,
+            depth_scale=depth_scale,
+            **kwargs,
+        )
+
+    def get_depth_filenames(self, img: DictStrAny) -> str | None:
+        """Get the depth filenames.
+
+        Since not every data has depth.
+        """
+        img["file_path"] = img["file_path"].replace("nuScenes", "nuscenes")
+
+        depth_filename = (
+            img["file_path"]
+            .replace("nuscenes", "nuscenes_depth")
+            .replace("jpg", "png")
+        )
+        return depth_filename
+
+    def get_cat_ids(self, idx: int) -> list[int]:
+        """Return the samples."""
+        return self.samples[idx]["class_ids"].tolist()
+
+    def __len__(self) -> int:
+        """Total number of samples of data."""
+        return len(self.samples)

opendet3d/data/datasets/omni3d/objectron.py

@@ -0,0 +1,56 @@
+"""Objectron from Omni3D."""
+
+from __future__ import annotations
+
+import os
+
+from vis4d.common.typing import ArgsType, DictStrAny
+
+from opendet3d.data.datasets.coco3d import COCO3DDataset
+
+from .omni3d_classes import omni3d_class_map
+
+objectron_det_map = {
+    "bicycle": 0,
+    "books": 1,
+    "bottle": 2,
+    "camera": 3,
+    "cereal box": 4,
+    "chair": 5,
+    "cup": 6,
+    "laptop": 7,
+    "shoes": 8,
+}
+
+
+class Objectron(COCO3DDataset):
+    """Objectron dataset."""
+
+    def __init__(
+        self,
+        class_map: dict[str, int] = omni3d_class_map,
+        max_depth: float = 12.0,
+        depth_scale: float = 1000.0,
+        **kwargs: ArgsType,
+    ) -> None:
+        """Creates an instance of the class."""
+        super().__init__(
+            class_map=class_map,
+            max_depth=max_depth,
+            depth_scale=depth_scale,
+            **kwargs,
+        )
+
+    def get_depth_filenames(self, img: DictStrAny) -> str | None:
+        """Get the depth filenames.
+
+        Since not every data has depth.
+        """
+        _, _, split, img_name = img["file_path"].split("/")
+
+        depth_filename = os.path.join(
+            "data/objectron_depth",
+            split,
+            img_name.replace(".jpg", "_depth.png"),
+        )
+        return depth_filename

opendet3d/data/datasets/omni3d/omni3d_classes.py

@@ -0,0 +1,156 @@
+"""Omni3D classes."""
+
+omni3d_class_map = {
+    "pedestrian": 0,
+    "car": 1,
+    "dontcare": 2,
+    "cyclist": 3,
+    "van": 4,
+    "truck": 5,
+    "tram": 6,
+    "person": 7,
+    "traffic cone": 8,
+    "barrier": 9,
+    "motorcycle": 10,
+    "bicycle": 11,
+    "bus": 12,
+    "trailer": 13,
+    "books": 14,
+    "bottle": 15,
+    "camera": 16,
+    "cereal box": 17,
+    "chair": 18,
+    "cup": 19,
+    "laptop": 20,
+    "shoes": 21,
+    "towel": 22,
+    "blinds": 23,
+    "window": 24,
+    "lamp": 25,
+    "shelves": 26,
+    "mirror": 27,
+    "sink": 28,
+    "cabinet": 29,
+    "bathtub": 30,
+    "door": 31,
+    "toilet": 32,
+    "desk": 33,
+    "box": 34,
+    "bookcase": 35,
+    "picture": 36,
+    "table": 37,
+    "counter": 38,
+    "bed": 39,
+    "night stand": 40,
+    "dresser": 41,
+    "pillow": 42,
+    "sofa": 43,
+    "television": 44,
+    "floor mat": 45,
+    "curtain": 46,
+    "clothes": 47,
+    "stationery": 48,
+    "refrigerator": 49,
+    "board": 50,
+    "kitchen pan": 51,
+    "bin": 52,
+    "stove": 53,
+    "microwave": 54,
+    "plates": 55,
+    "bowl": 56,
+    "oven": 57,
+    "vase": 58,
+    "faucet": 59,
+    "tissues": 60,
+    "machine": 61,
+    "printer": 62,
+    "monitor": 63,
+    "podium": 64,
+    "cart": 65,
+    "projector": 66,
+    "electronics": 67,
+    "computer": 68,
+    "air conditioner": 69,
+    "drawers": 70,
+    "coffee maker": 71,
+    "toaster": 72,
+    "potted plant": 73,
+    "painting": 74,
+    "bag": 75,
+    "tray": 76,
+    "keyboard": 77,
+    "blanket": 78,
+    "rack": 79,
+    "phone": 80,
+    "mouse": 81,
+    "fire extinguisher": 82,
+    "toys": 83,
+    "ladder": 84,
+    "fan": 85,
+    "glass": 86,
+    "clock": 87,
+    "toilet paper": 88,
+    "closet": 89,
+    "fume hood": 90,
+    "utensils": 91,
+    "soundsystem": 92,
+    "fire place": 93,
+    "shower curtain": 94,
+    "remote": 95,
+    "pen": 96,
+    "fireplace": 97,
+}
+
+# Used for Cube R-CNN and Omni3D benchmark
+omni3d_det_map = {
+    "pedestrian": 0,
+    "car": 1,
+    "cyclist": 2,
+    "van": 3,
+    "truck": 4,
+    "traffic cone": 5,
+    "barrier": 6,
+    "motorcycle": 7,
+    "bicycle": 8,
+    "bus": 9,
+    "trailer": 10,
+    "books": 11,
+    "bottle": 12,
+    "camera": 13,
+    "cereal box": 14,
+    "chair": 15,
+    "cup": 16,
+    "laptop": 17,
+    "shoes": 18,
+    "towel": 19,
+    "blinds": 20,
+    "window": 21,
+    "lamp": 22,
+    "shelves": 23,
+    "mirror": 24,
+    "sink": 25,
+    "cabinet": 26,
+    "bathtub": 27,
+    "door": 28,
+    "toilet": 29,
+    "desk": 30,
+    "box": 31,
+    "bookcase": 32,
+    "picture": 33,
+    "table": 34,
+    "counter": 35,
+    "bed": 36,
+    "night stand": 37,
+    "pillow": 38,
+    "sofa": 39,
+    "television": 40,
+    "floor mat": 41,
+    "curtain": 42,
+    "clothes": 43,
+    "stationery": 44,
+    "refrigerator": 45,
+    "bin": 46,
+    "stove": 47,
+    "oven": 48,
+    "machine": 49,
+}

opendet3d/data/datasets/omni3d/sunrgbd.py

@@ -0,0 +1,278 @@
+"""SUN RGB-D from Omni3D."""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+from vis4d.common.typing import ArgsType, DictStrAny
+from vis4d.data.datasets.util import im_decode
+
+from opendet3d.data.datasets.coco3d import COCO3DDataset
+
+from .omni3d_classes import omni3d_class_map
+
+# Train and Test are sharing the classes
+sun_rgbd_train_det_map = sun_rgbd_test_det_map = {
+    "air conditioner": 0,
+    "bag": 1,
+    "bathtub": 2,
+    "bed": 3,
+    "bicycle": 4,
+    "bin": 5,
+    "blanket": 6,
+    "blinds": 7,
+    "board": 8,
+    "bookcase": 9,
+    "books": 10,
+    "bottle": 11,
+    "bowl": 12,
+    "box": 13,
+    "cabinet": 14,
+    "cart": 15,
+    "chair": 16,
+    "clock": 17,
+    "closet": 18,
+    "clothes": 19,
+    "coffee maker": 20,
+    "computer": 21,
+    "counter": 22,
+    "cup": 23,
+    "curtain": 24,
+    "desk": 25,
+    "door": 26,
+    "drawers": 27,
+    "dresser": 28,
+    "electronics": 29,
+    "fan": 30,
+    "faucet": 31,
+    "fire extinguisher": 32,
+    "fire place": 33,
+    "floor mat": 34,
+    "fume hood": 35,
+    "glass": 36,
+    "keyboard": 37,
+    "kitchen pan": 38,
+    "ladder": 39,
+    "lamp": 40,
+    "laptop": 41,
+    "machine": 42,
+    "microwave": 43,
+    "mirror": 44,
+    "monitor": 45,
+    "mouse": 46,
+    "night stand": 47,
+    "oven": 48,
+    "painting": 49,
+    "pen": 50,
+    "person": 51,
+    "phone": 52,
+    "picture": 53,
+    "pillow": 54,
+    "plates": 55,
+    "podium": 56,
+    "potted plant": 57,
+    "printer": 58,
+    "projector": 59,
+    "rack": 60,
+    "refrigerator": 61,
+    "remote": 62,
+    "shelves": 63,
+    "shoes": 64,
+    "shower curtain": 65,
+    "sink": 66,
+    "sofa": 67,
+    "soundsystem": 68,
+    "stationery": 69,
+    "stove": 70,
+    "table": 71,
+    "television": 72,
+    "tissues": 73,
+    "toaster": 74,
+    "toilet": 75,
+    "toilet paper": 76,
+    "towel": 77,
+    "toys": 78,
+    "tray": 79,
+    "utensils": 80,
+    "vase": 81,
+    "window": 82,
+}
+
+sun_rgbd_val_det_map = {
+    "air conditioner": 0,
+    "bag": 1,
+    "bathtub": 2,
+    "bed": 3,
+    "bin": 4,
+    "blanket": 5,
+    "blinds": 6,
+    "board": 7,
+    "bookcase": 8,
+    "books": 9,
+    "bottle": 10,
+    "bowl": 11,
+    "box": 12,
+    "cabinet": 13,
+    "cart": 14,
+    "chair": 15,
+    "closet": 16,
+    "clothes": 17,
+    "coffee maker": 18,
+    "computer": 19,
+    "counter": 20,
+    "cup": 21,
+    "curtain": 22,
+    "desk": 23,
+    "door": 24,
+    "drawers": 25,
+    "dresser": 26,
+    "electronics": 27,
+    "fan": 28,
+    "faucet": 29,
+    "fire extinguisher": 30,
+    "fire place": 31,
+    "fume hood": 32,
+    "keyboard": 33,
+    "kitchen pan": 34,
+    "lamp": 35,
+    "laptop": 36,
+    "machine": 37,
+    "microwave": 38,
+    "mirror": 39,
+    "monitor": 40,
+    "night stand": 41,
+    "oven": 42,
+    "painting": 43,
+    "pen": 44,
+    "person": 45,
+    "phone": 46,
+    "picture": 47,
+    "pillow": 48,
+    "plates": 49,
+    "potted plant": 50,
+    "printer": 51,
+    "projector": 52,
+    "rack": 53,
+    "refrigerator": 54,
+    "shelves": 55,
+    "sink": 56,
+    "sofa": 57,
+    "soundsystem": 58,
+    "stationery": 59,
+    "stove": 60,
+    "table": 61,
+    "television": 62,
+    "tissues": 63,
+    "toaster": 64,
+    "toilet": 65,
+    "towel": 66,
+    "toys": 67,
+    "tray": 68,
+    "utensils": 69,
+    "vase": 70,
+    "window": 71,
+}
+
+omni3d_sun_rgbd_det_map = {
+    "bicycle": 0,
+    "books": 1,
+    "bottle": 2,
+    "chair": 3,
+    "cup": 4,
+    "laptop": 5,
+    "shoes": 6,
+    "towel": 7,
+    "blinds": 8,
+    "window": 9,
+    "lamp": 10,
+    "shelves": 11,
+    "mirror": 12,
+    "sink": 13,
+    "cabinet": 14,
+    "bathtub": 15,
+    "door": 16,
+    "toilet": 17,
+    "desk": 18,
+    "box": 19,
+    "bookcase": 20,
+    "picture": 21,
+    "table": 22,
+    "counter": 23,
+    "bed": 24,
+    "night stand": 25,
+    "pillow": 26,
+    "sofa": 27,
+    "television": 28,
+    "floor mat": 29,
+    "curtain": 30,
+    "clothes": 31,
+    "stationery": 32,
+    "refrigerator": 33,
+    "bin": 34,
+    "stove": 35,
+    "oven": 36,
+    "machine": 37,
+}
+
+
+def get_sunrgbd_det_map(split: str) -> dict[str, int]:
+    """Get the SUN RGB-D detection map."""
+    assert split in {"train", "val", "test"}, f"Invalid split: {split}"
+
+    if split == "train":
+        return sun_rgbd_train_det_map
+    elif split == "val":
+        return sun_rgbd_val_det_map
+    else:
+        return sun_rgbd_test_det_map
+
+
+class SUNRGBD(COCO3DDataset):
+    """SUN RGB-D Dataset."""
+
+    def __init__(
+        self,
+        class_map: dict[str, int] = omni3d_class_map,
+        max_depth: float = 8.0,
+        depth_scale: float = 1000.0,
+        **kwargs: ArgsType,
+    ) -> None:
+        """Initialize SUN RGB-D dataset."""
+        super().__init__(
+            class_map=class_map,
+            max_depth=max_depth,
+            depth_scale=depth_scale,
+            **kwargs,
+        )
+
+    def get_depth_filenames(self, img: DictStrAny) -> str | None:
+        """Get the depth filenames.
+
+        Since not every data has depth.
+        """
+        img["file_path"] = img["file_path"].replace("//", "/")
+
+        data_dir = img["file_path"].split("/image")[0]
+
+        depth_files = self.data_backend.listdir(
+            os.path.join(data_dir, "depth")
+        )
+        assert len(depth_files) == 1
+
+        depth_filename = os.path.join(data_dir, "depth", depth_files[0])
+
+        return depth_filename
+
+    def get_depth_map(self, sample: DictStrAny) -> np.ndarray:
+        """Get the depth map."""
+        depth_bytes = self.data_backend.get(sample["depth_filename"])
+        depth_array = im_decode(depth_bytes)
+
+        depth_array = depth_array >> 3 | depth_array << (16 - 3)
+
+        depth = np.ascontiguousarray(depth_array, dtype=np.float32)
+
+        depth = depth / self.depth_scale
+
+        return depth

opendet3d/data/datasets/omni3d/util.py

@@ -0,0 +1,74 @@
+"""Omni3D data util."""
+
+from __future__ import annotations
+
+from .arkitscenes import arkitscenes_det_map, omni3d_arkitscenes_det_map
+from .hypersim import get_hypersim_det_map, omni3d_hypersim_det_map
+from .kitti_object import get_kitti_det_map, omni3d_kitti_det_map
+from .nuscenes import nusc_det_map
+from .objectron import objectron_det_map
+from .sunrgbd import get_sunrgbd_det_map, omni3d_sun_rgbd_det_map
+
+DATASET_ID_MAP = {
+    0: "KITTI_train",
+    1: "KITTI_val",
+    2: "KITTI_test",
+    3: "nuScenes_train",
+    4: "nuScenes_val",
+    5: "nuScenes_test",
+    6: "Objectron_train",
+    7: "Objectron_val",
+    8: "Objectron_test",
+    9: "Hypersim_train",
+    10: "Hypersim_val",
+    11: "Hypersim_test",
+    12: "SUNRGBD_train",
+    13: "SUNRGBD_val",
+    14: "SUNRGBD_test",
+    15: "ARKitScenes_train",
+    16: "ARKitScenes_val",
+    17: "ARKitScenes_test",
+}
+
+
+def get_dataset_det_map(
+    dataset_name: str,
+    omni3d50: bool = True,
+) -> tuple[str, dict[str, int]]:
+    """Get the detection map."""
+    if "train" in dataset_name:
+        split = "train"
+    elif "val" in dataset_name:
+        split = "val"
+    elif "test" in dataset_name:
+        split = "test"
+    else:
+        raise ValueError(f"Unknown dataset_name: {dataset_name}")
+
+    if "nuScenes" in dataset_name:
+        det_map = nusc_det_map
+    elif "KITTI" in dataset_name:
+        if omni3d50:
+            det_map = omni3d_kitti_det_map
+        else:
+            det_map = get_kitti_det_map(split)
+    elif "Objectron" in dataset_name:
+        det_map = objectron_det_map
+    elif "SUNRGBD" in dataset_name:
+        if omni3d50:
+            det_map = omni3d_sun_rgbd_det_map
+        else:
+            det_map = get_sunrgbd_det_map(split)
+    elif "Hypersim" in dataset_name:
+        if omni3d50:
+            det_map = omni3d_hypersim_det_map
+        else:
+            det_map = get_hypersim_det_map(split)
+    elif "ARKitScenes" in dataset_name:
+        det_map = (
+            omni3d_arkitscenes_det_map if omni3d50 else arkitscenes_det_map
+        )
+    else:
+        raise ValueError(f"Unknown dataset_name: {dataset_name}")
+
+    return det_map

opendet3d/data/datasets/scannet.py

@@ -0,0 +1,449 @@
+"""ScanNet dataset."""
+
+from __future__ import annotations
+
+from vis4d.common.typing import ArgsType, DictStrAny
+
+from .coco3d import COCO3DDataset
+
+scannet_class_map = {
+    "cabinet": 3,
+    "bed": 4,
+    "chair": 5,
+    "sofa": 6,
+    "table": 7,
+    "door": 8,
+    "window": 9,
+    "bookshelf": 10,
+    "picture": 11,
+    "counter": 12,
+    "desk": 14,
+    "curtain": 16,
+    "refrigerator": 24,
+    "shower curtain": 28,
+    "toilet": 33,
+    "sink": 34,
+    "bathtub": 36,
+    "other furniture": 39,
+}
+
+scannet_det_map = {
+    "cabinet": 0,
+    "bed": 1,
+    "chair": 2,
+    "sofa": 3,
+    "table": 4,
+    "door": 5,
+    "window": 6,
+    "bookshelf": 7,
+    "picture": 8,
+    "counter": 9,
+    "desk": 10,
+    "curtain": 11,
+    "refrigerator": 12,
+    "shower curtain": 13,
+    "toilet": 14,
+    "sink": 15,
+    "bathtub": 16,
+    "other furniture": 17,
+}
+
+scannet200_class_map = {
+    "chair": 2,
+    "book": 22,
+    "door": 5,
+    "object": 1163,
+    "window": 16,
+    "table": 4,
+    "trash can": 56,
+    "pillow": 13,
+    "picture": 15,
+    "box": 26,
+    "doorframe": 161,
+    "monitor": 19,
+    "cabinet": 7,
+    "desk": 9,
+    "shelf": 8,
+    "office chair": 10,
+    "towel": 31,
+    "couch": 6,
+    "sink": 14,
+    "backpack": 48,
+    "lamp": 28,
+    "bed": 11,
+    "bookshelf": 18,
+    "mirror": 71,
+    "curtain": 21,
+    "plant": 40,
+    "whiteboard": 52,
+    "radiator": 96,
+    "kitchen cabinet": 29,
+    "toilet paper": 49,
+    "armchair": 23,
+    "shoe": 63,
+    "coffee table": 24,
+    "toilet": 17,
+    "bag": 47,
+    "clothes": 32,
+    "keyboard": 46,
+    "bottle": 65,
+    "recycling bin": 97,
+    "nightstand": 34,
+    "stool": 38,
+    "tv": 33,
+    "file cabinet": 75,
+    "dresser": 36,
+    "computer tower": 64,
+    "telephone": 101,
+    "cup": 130,
+    "refrigerator": 27,
+    "end table": 44,
+    "jacket": 131,
+    "shower curtain": 55,
+    "bathtub": 42,
+    "microwave": 59,
+    "kitchen counter": 159,
+    "sofa chair": 74,
+    "paper towel dispenser": 82,
+    "bathroom vanity": 1164,
+    "suitcase": 93,
+    "laptop": 77,
+    "ottoman": 67,
+    "shower wall": 128,
+    "printer": 50,
+    "counter": 35,
+    "board": 69,
+    "soap dispenser": 100,
+    "stove": 62,
+    "light": 105,
+    "closet wall": 1165,
+    "mini fridge": 165,
+    "fan": 76,
+    "tissue box": 230,
+    "blanket": 54,
+    "bathroom stall": 125,
+    "copier": 72,
+    "bench": 68,
+    "bar": 145,
+    "soap dish": 157,
+    "laundry hamper": 1166,
+    "storage bin": 132,
+    "bathroom stall door": 1167,
+    "light switch": 232,
+    "coffee maker": 134,
+    "tv stand": 51,
+    "decoration": 250,
+    "ceiling light": 1168,
+    "range hood": 342,
+    "blackboard": 89,
+    "clock": 103,
+    "wardrobe": 99,
+    "rail": 95,
+    "bulletin board": 154,
+    "mat": 140,
+    "trash bin": 1169,
+    "ledge": 193,
+    "seat": 116,
+    "mouse": 202,
+    "basket": 73,
+    "shower": 78,
+    "dumbbell": 1170,
+    "paper": 79,
+    "person": 80,
+    "windowsill": 141,
+    "closet": 57,
+    "bucket": 102,
+    "sign": 261,
+    "speaker": 118,
+    "dishwasher": 136,
+    "container": 98,
+    "stair rail": 1171,
+    "shower curtain rod": 170,
+    "tube": 1172,
+    "bathroom cabinet": 1173,
+    "storage container": 221,
+    "paper bag": 570,
+    "paper towel roll": 138,
+    "ball": 168,
+    "closet door": 276,
+    "laundry basket": 106,
+    "cart": 214,
+    "dish rack": 323,
+    "stairs": 58,
+    "blinds": 86,
+    "purse": 399,
+    "bicycle": 121,
+    "tray": 185,
+    "plunger": 300,
+    "paper cutter": 180,
+    "toilet paper dispenser": 163,
+    "bin": 66,
+    "toilet seat cover dispenser": 208,
+    "guitar": 112,
+    "mailbox": 540,
+    "handicap bar": 395,
+    "fire extinguisher": 166,
+    "ladder": 122,
+    "column": 120,
+    "pipe": 107,
+    "vacuum cleaner": 283,
+    "plate": 88,
+    "piano": 90,
+    "water cooler": 177,
+    "cd case": 1174,
+    "bowl": 562,
+    "closet rod": 1175,
+    "bathroom counter": 1156,
+    "oven": 84,
+    "stand": 104,
+    "scale": 229,
+    "washing machine": 70,
+    "broom": 325,
+    "hat": 169,
+    "guitar case": 331,
+    "rack": 87,
+    "water pitcher": 488,
+    "laundry detergent": 776,
+    "hair dryer": 370,
+    "pillar": 191,
+    "divider": 748,
+    "power outlet": 242,
+    "dining table": 45,
+    "shower floor": 417,
+    "shower door": 188,
+    "coffee kettle": 1176,
+    "structure": 1178,
+    "clothes dryer": 110,
+    "toaster": 148,
+    "ironing board": 155,
+    "alarm clock": 572,
+    "shower head": 1179,
+    "water bottle": 392,
+    "keyboard piano": 1180,
+    "projector screen": 609,
+    "case of water bottles": 1181,
+    "toaster oven": 195,
+    "music stand": 581,
+    "coat rack": 1182,
+    "storage organizer": 1183,
+    "machine": 139,
+    "folded chair": 1184,
+    "fire alarm": 1185,
+    "fireplace": 156,
+    "vent": 408,
+    "furniture": 213,
+    "power strip": 1186,
+    "calendar": 1187,
+    "poster": 1188,
+    "toilet paper holder": 115,
+    "potted plant": 1189,
+    "stuffed animal": 304,
+    "luggage": 1190,
+    "headphones": 312,
+    "crate": 233,
+    "candle": 286,
+    "projector": 264,
+    "mattress": 1191,
+    "dustpan": 356,
+    "cushion": 39,
+    "stick": 1163,
+}
+
+scannet200_det_map = {
+    "chair": 0,
+    "table": 1,
+    "door": 2,
+    "couch": 3,
+    "cabinet": 4,
+    "shelf": 5,
+    "desk": 6,
+    "office chair": 7,
+    "bed": 8,
+    "pillow": 9,
+    "sink": 10,
+    "picture": 11,
+    "window": 12,
+    "toilet": 13,
+    "bookshelf": 14,
+    "monitor": 15,
+    "curtain": 16,
+    "book": 17,
+    "armchair": 18,
+    "coffee table": 19,
+    "box": 20,
+    "refrigerator": 21,
+    "lamp": 22,
+    "kitchen cabinet": 23,
+    "towel": 24,
+    "clothes": 25,
+    "tv": 26,
+    "nightstand": 27,
+    "counter": 28,
+    "dresser": 29,
+    "stool": 30,
+    "plant": 31,
+    "bathtub": 32,
+    "end table": 33,
+    "dining table": 34,
+    "keyboard": 35,
+    "bag": 36,
+    "backpack": 37,
+    "toilet paper": 38,
+    "printer": 39,
+    "tv stand": 40,
+    "whiteboard": 41,
+    "blanket": 42,
+    "shower curtain": 43,
+    "trash can": 44,
+    "closet": 45,
+    "stairs": 46,
+    "microwave": 47,
+    "stove": 48,
+    "shoe": 49,
+    "computer tower": 50,
+    "bottle": 51,
+    "bin": 52,
+    "ottoman": 53,
+    "bench": 54,
+    "board": 55,
+    "washing machine": 56,
+    "mirror": 57,
+    "copier": 58,
+    "basket": 59,
+    "sofa chair": 60,
+    "file cabinet": 61,
+    "fan": 62,
+    "laptop": 63,
+    "shower": 64,
+    "paper": 65,
+    "person": 66,
+    "paper towel dispenser": 67,
+    "oven": 68,
+    "blinds": 69,
+    "rack": 70,
+    "plate": 71,
+    "blackboard": 72,
+    "piano": 73,
+    "suitcase": 74,
+    "rail": 75,
+    "radiator": 76,
+    "recycling bin": 77,
+    "container": 78,
+    "wardrobe": 79,
+    "soap dispenser": 80,
+    "telephone": 81,
+    "bucket": 82,
+    "clock": 83,
+    "stand": 84,
+    "light": 85,
+    "laundry basket": 86,
+    "pipe": 87,
+    "clothes dryer": 88,
+    "guitar": 89,
+    "toilet paper holder": 90,
+    "seat": 91,
+    "speaker": 92,
+    "column": 93,
+    "ladder": 94,
+    "cup": 95,
+    "jacket": 96,
+    "storage bin": 97,
+    "coffee maker": 98,
+    "dishwasher": 99,
+    "paper towel roll": 100,
+    "machine": 101,
+    "mat": 102,
+    "windowsill": 103,
+    "bar": 104,
+    "bulletin board": 105,
+    "ironing board": 106,
+    "fireplace": 107,
+    "soap dish": 108,
+    "kitchen counter": 109,
+    "doorframe": 110,
+    "toilet paper dispenser": 111,
+    "mini fridge": 112,
+    "fire extinguisher": 113,
+    "ball": 114,
+    "hat": 115,
+    "shower curtain rod": 116,
+    "water cooler": 117,
+    "paper cutter": 118,
+    "tray": 119,
+    "pillar": 120,
+    "ledge": 121,
+    "toaster oven": 122,
+    "mouse": 123,
+    "toilet seat cover dispenser": 124,
+    "cart": 125,
+    "scale": 126,
+    "tissue box": 127,
+    "light switch": 128,
+    "crate": 129,
+    "power outlet": 130,
+    "decoration": 131,
+    "sign": 132,
+    "projector": 133,
+    "closet door": 134,
+    "vacuum cleaner": 135,
+    "headphones": 136,
+    "dish rack": 137,
+    "broom": 138,
+    "range hood": 139,
+    "hair dryer": 140,
+    "water bottle": 141,
+    "vent": 142,
+    "mailbox": 143,
+    "bowl": 144,
+    "paper bag": 145,
+    "projector screen": 146,
+    "divider": 147,
+    "laundry detergent": 148,
+    "bathroom counter": 149,
+    "stick": 150,
+    "bathroom vanity": 151,
+    "closet wall": 152,
+    "laundry hamper": 153,
+    "bathroom stall door": 154,
+    "ceiling light": 155,
+    "trash bin": 156,
+    "dumbbell": 157,
+    "stair rail": 158,
+    "tube": 159,
+    "bathroom cabinet": 160,
+    "coffee kettle": 161,
+    "shower head": 162,
+    "case of water bottles": 163,
+    "power strip": 164,
+    "calendar": 165,
+    "poster": 166,
+    "mattress": 167,
+}
+
+
+class ScanNetDataset(COCO3DDataset):
+    """ScanNetV2 dataset."""
+
+    def __init__(
+        self,
+        class_map: dict[str, int] = scannet_class_map,
+        max_depth: float = 12.0,
+        depth_scale: float = 1000.0,
+        **kwargs: ArgsType,
+    ) -> None:
+        """Creates an instance of the class."""
+        super().__init__(
+            class_map=class_map,
+            max_depth=max_depth,
+            depth_scale=depth_scale,
+            **kwargs,
+        )
+
+    def get_depth_filenames(self, img: DictStrAny) -> str | None:
+        """Get the depth filenames.
+
+        Since not every data has depth.
+        """
+        return (
+            img["file_path"].replace("image", "depth").replace(".jpg", ".png")
+        )

opendet3d/data/transforms/crop.py

@@ -0,0 +1,43 @@
+"""Crop transforms."""
+
+from __future__ import annotations
+
+from vis4d.common.typing import (
+    NDArrayBool,
+    NDArrayF32,
+    NDArrayI64,
+)
+from vis4d.data.const import CommonKeys as K
+from vis4d.data.transforms.base import Transform
+
+
+@Transform(
+    in_keys=[
+        K.boxes3d,
+        K.boxes3d_classes,
+        K.boxes3d_track_ids,
+        "transforms.crop.keep_mask",
+    ],
+    out_keys=[K.boxes3d, K.boxes3d_classes, K.boxes3d_track_ids],
+)
+class CropBoxes3D:
+    """Crop 3D bounding boxes."""
+
+    def __call__(
+        self,
+        boxes_list: list[NDArrayF32],
+        classes_list: list[NDArrayI64],
+        track_ids_list: list[NDArrayI64] | None,
+        keep_mask_list: list[NDArrayBool],
+    ) -> tuple[list[NDArrayF32], list[NDArrayI64], list[NDArrayI64] | None]:
+        """Crop 3D bounding boxes."""
+        for i, (boxes, classes, keep_mask) in enumerate(
+            zip(boxes_list, classes_list, keep_mask_list)
+        ):
+            boxes_list[i] = boxes[keep_mask]
+            classes_list[i] = classes[keep_mask]
+
+            if track_ids_list is not None:
+                track_ids_list[i] = track_ids_list[i][keep_mask]
+
+        return boxes_list, classes_list, track_ids_list

opendet3d/data/transforms/language.py

@@ -0,0 +1,267 @@
+"""Language related transforms."""
+
+from __future__ import annotations
+
+import random
+import re
+
+import numpy as np
+from transformers import AutoTokenizer
+from vis4d.common.logging import rank_zero_warn
+from vis4d.common.typing import NDArrayF32, NDArrayI64
+from vis4d.data.const import CommonKeys as K
+from vis4d.data.transforms.base import Transform
+
+
+def clean_name(name: str) -> str:
+    """Clean the name."""
+    name = re.sub(r"\(.*\)", "", name)
+    name = re.sub(r"_", " ", name)
+    name = re.sub(r"  ", " ", name)
+    name = name.lower()
+    return name
+
+
+def generate_senetence_given_labels(
+    positive_label_list: list[int],
+    negative_label_list: list[str],
+    label_map: dict[str, str],
+) -> tuple[dict[int, list[list[int]]], str, dict[int, int]]:
+    """Generate a sentence given positive and negative labels."""
+    label_to_positions = {}
+
+    label_list = negative_label_list + positive_label_list
+
+    random.shuffle(label_list)
+
+    pheso_caption = ""
+
+    label_remap_dict = {}
+    for index, label in enumerate(label_list):
+        start_index = len(pheso_caption)
+
+        pheso_caption += clean_name(label_map[str(label)])
+
+        end_index = len(pheso_caption)
+
+        if label in positive_label_list:
+            label_to_positions[index] = [[start_index, end_index]]
+            label_remap_dict[int(label)] = index
+
+        pheso_caption += ". "
+
+    return label_to_positions, pheso_caption, label_remap_dict
+
+
+@Transform(
+    [
+        "dataset_type",
+        K.boxes2d,
+        K.boxes2d_classes,
+        K.boxes2d_names,
+        "label_map",
+        "positive_positions",
+    ],
+    [K.boxes2d, K.boxes2d_classes, K.boxes2d_names, "tokens_positive"],
+)
+class RandomSamplingNegPos:
+    """Randomly sample negative and positive labels for object detection."""
+
+    def __init__(
+        self,
+        tokenizer_name: str = "bert-base-uncased",
+        num_sample_negative: int = 85,
+        max_tokens: int = 256,
+        full_sampling_prob: float = 0.5,
+    ) -> None:
+        """Creates an instance of RandomSamplingNegPos."""
+        if AutoTokenizer is None:
+            raise RuntimeError(
+                "transformers is not installed, please install it by: "
+                "pip install transformers."
+            )
+
+        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
+        self.num_sample_negative = num_sample_negative
+        self.full_sampling_prob = full_sampling_prob
+        self.max_tokens = max_tokens
+
+    def __call__(
+        self,
+        dataset_type_list: list[str],
+        boxes_list: list[NDArrayF32],
+        class_ids_list: list[NDArrayI64],
+        texts_list: list[str] | None = None,
+        label_map_list: dict | None = None,
+        positive_positions_list: list[dict] | None = None,
+    ) -> tuple[
+        list[NDArrayF32],
+        list[NDArrayI64],
+        list[str],
+        list[dict[int, list[list[int]]]],
+    ]:
+        """Randomly sample negative and positive labels."""
+        new_texts_list = []
+        tokens_positive_list = []
+        for i, (boxes, class_ids) in enumerate(
+            zip(boxes_list, class_ids_list)
+        ):
+            if dataset_type_list[i] == "OD":
+                assert (
+                    label_map_list[i] is not None
+                ), "label_map should not be None"
+                boxes_list[i], class_ids_list[i], text, tokens_positive = (
+                    self.od_aug(boxes, class_ids, label_map_list[i])
+                )
+                new_texts_list.append(text)
+                tokens_positive_list.append(tokens_positive)
+            else:
+                assert (
+                    positive_positions_list[i] is not None
+                ), "positive_positions should not be None"
+                tokens_positive = self.vg_aug(
+                    class_ids, positive_positions_list[i]
+                )
+                new_texts_list.append(texts_list[i])
+                tokens_positive_list.append(tokens_positive)
+
+        return boxes_list, class_ids_list, new_texts_list, tokens_positive_list
+
+    def vg_aug(self, class_ids: NDArrayI64, positive_positions):
+        """Visual Genome data augmentation."""
+        positive_label_list = np.unique(class_ids).tolist()
+
+        label_to_positions = {}
+        for label in positive_label_list:
+            label_to_positions[label] = positive_positions[label]
+
+        return label_to_positions
+
+    def od_aug(
+        self,
+        boxes: NDArrayF32,
+        class_ids: NDArrayI64,
+        label_map: dict,
+    ) -> tuple[NDArrayF32, NDArrayI64, str, dict[int, list[list[int]]]]:
+        """Object detection data augmentation."""
+        original_box_num = len(class_ids)
+
+        # If the category name is in the format of 'a/b' (in object365),
+        # we randomly select one of them.
+        for key, value in label_map.items():
+            if "/" in value:
+                label_map[key] = random.choice(value.split("/")).strip()
+
+        keep_box_index, class_ids, positive_caption_length = (
+            self.check_for_positive_overflow(class_ids, label_map)
+        )
+
+        boxes = boxes[keep_box_index]
+
+        if len(boxes) < original_box_num:
+            rank_zero_warn(
+                f"Remove {original_box_num - len(boxes)} boxes due to "
+                "positive caption overflow."
+            )
+
+        valid_negative_indexes = list(label_map.keys())
+
+        positive_label_list = np.unique(class_ids).tolist()
+
+        full_negative = self.num_sample_negative
+        if full_negative > len(valid_negative_indexes):
+            full_negative = len(valid_negative_indexes)
+
+        outer_prob = random.random()
+
+        if outer_prob < self.full_sampling_prob:
+            # c. probability_full: add both all positive and all negatives
+            num_negatives = full_negative
+        else:
+            if random.random() < 1.0:
+                num_negatives = np.random.choice(max(1, full_negative)) + 1
+            else:
+                num_negatives = full_negative
+
+        # Keep some negatives
+        negative_label_list = set()
+        if num_negatives != -1:
+            if num_negatives > len(valid_negative_indexes):
+                num_negatives = len(valid_negative_indexes)
+
+            for i in np.random.choice(
+                valid_negative_indexes, size=num_negatives, replace=False
+            ):
+                if int(i) not in positive_label_list:
+                    negative_label_list.add(i)
+
+        random.shuffle(positive_label_list)
+
+        negative_label_list = list(negative_label_list)
+        random.shuffle(negative_label_list)
+
+        negative_max_length = self.max_tokens - positive_caption_length
+        screened_negative_label_list = []
+
+        for negative_label in negative_label_list:
+            label_text = clean_name(label_map[str(negative_label)]) + ". "
+
+            tokenized = self.tokenizer.tokenize(label_text)
+
+            negative_max_length -= len(tokenized)
+
+            if negative_max_length > 0:
+                screened_negative_label_list.append(negative_label)
+            else:
+                break
+
+        negative_label_list = screened_negative_label_list
+        label_to_positions, pheso_caption, label_remap_dict = (
+            generate_senetence_given_labels(
+                positive_label_list, negative_label_list, label_map
+            )
+        )
+
+        # label remap
+        if len(class_ids) > 0:
+            class_ids = np.vectorize(lambda x: label_remap_dict[x])(class_ids)
+
+        return boxes, class_ids, pheso_caption, label_to_positions
+
+    def check_for_positive_overflow(
+        self, class_ids: NDArrayI64, label_map: dict[str, str]
+    ) -> tuple[list[int], NDArrayI64, int]:
+        """Check if having too many positive labels."""
+        # generate a caption by appending the positive labels
+        positive_label_list = np.unique(class_ids).tolist()
+
+        # random shuffule so we can sample different annotations
+        # at different epochs
+        random.shuffle(positive_label_list)
+
+        kept_lables = []
+        length = 0
+        for _, label in enumerate(positive_label_list):
+            label_text = clean_name(label_map[str(label)]) + ". "
+
+            tokenized = self.tokenizer.tokenize(label_text)
+
+            length += len(tokenized)
+
+            if length > self.max_tokens:
+                break
+            else:
+                kept_lables.append(label)
+
+        keep_box_index = []
+        keep_gt_labels = []
+        for i, class_id in enumerate(class_ids):
+            if class_id in kept_lables:
+                keep_box_index.append(i)
+                keep_gt_labels.append(class_id)
+
+        return (
+            keep_box_index,
+            np.array(keep_gt_labels, dtype=np.int64),
+            length,
+        )

opendet3d/data/transforms/pad.py

@@ -0,0 +1,176 @@
+"""Pad transformation."""
+
+from __future__ import annotations
+
+from typing import TypedDict
+
+import torch
+import torch.nn.functional as F
+from vis4d.common.typing import NDArrayF32
+from vis4d.data.const import CommonKeys as K
+from vis4d.data.transforms.base import Transform
+from vis4d.data.transforms.pad import _get_max_shape
+
+
+class PadParam(TypedDict):
+    """Parameters for Reshape."""
+
+    pad_top: int
+    pad_bottom: int
+    pad_left: int
+    pad_right: int
+
+
+@Transform(
+    [K.images, K.input_hw],
+    [K.images, "transforms.pad", K.input_hw, "padding"],
+)
+class CenterPadImages:
+    """Pad batch of images at the bottom right."""
+
+    def __init__(
+        self,
+        stride: int = 32,
+        mode: str = "constant",
+        value: float = 0.0,
+        update_input_hw: bool = False,
+        shape: tuple[int, int] | None = None,
+        pad2square: bool = False,
+    ) -> None:
+        """Creates an instance of PadImage.
+
+        Args:
+            stride (int, optional): Chooses padding size so that the input will
+                be divisible by stride. Defaults to 32.
+            mode (str, optional): Padding mode. One of constant, reflect,
+                replicate or circular. Defaults to "constant".
+            value (float, optional): Value for constant padding.
+                Defaults to 0.0.
+            shape (tuple[int, int], optional): Shape of the padded image
+                (H, W). Defaults to None.
+            pad2square (bool, optional): Pad to square. Defaults to False.
+        """
+        self.stride = stride
+        self.mode = mode
+        self.value = value
+        self.update_input_hw = update_input_hw
+        self.shape = shape
+        self.pad2square = pad2square
+
+    def __call__(
+        self, images: list[NDArrayF32], input_hw: list[tuple[int, int]]
+    ) -> tuple[list[NDArrayF32], list[PadParam], list[tuple[int, int]]]:
+        """Pad images to consistent size."""
+        heights = [im.shape[1] for im in images]
+        widths = [im.shape[2] for im in images]
+
+        max_hw = _get_max_shape(
+            heights, widths, self.stride, self.shape, self.pad2square
+        )
+
+        # generate params for torch pad
+        pad_params = []
+        target_input_hw = []
+        paddings = []
+        for i, (image, h, w) in enumerate(zip(images, heights, widths)):
+            pad_top, pad_bottom = (max_hw[0] - h) // 2, max_hw[0] - h - (
+                max_hw[0] - h
+            ) // 2
+
+            pad_left, pad_right = (max_hw[1] - w) // 2, max_hw[1] - w - (
+                max_hw[1] - w
+            ) // 2
+
+            image_ = torch.from_numpy(image).permute(0, 3, 1, 2)
+            image_ = F.pad(
+                image_,
+                (pad_left, pad_right, pad_top, pad_bottom),
+                self.mode,
+                self.value,
+            )
+            images[i] = image_.permute(0, 2, 3, 1).numpy()
+
+            pad_params.append(
+                PadParam(
+                    pad_top=pad_top,
+                    pad_bottom=pad_bottom,
+                    pad_left=pad_left,
+                    pad_right=pad_right,
+                )
+            )
+
+            paddings.append([pad_left, pad_right, pad_top, pad_bottom])
+
+            target_input_hw.append(max_hw)
+
+        if self.update_input_hw:
+            input_hw = target_input_hw
+
+        return images, pad_params, input_hw, paddings
+
+
+@Transform([K.intrinsics, "transforms.pad"], K.intrinsics)
+class CenterPadIntrinsics:
+    """Resize Intrinsics."""
+
+    def __call__(
+        self, intrinsics: list[NDArrayF32], pad_params: list[PadParam]
+    ) -> list[NDArrayF32]:
+        """Scale camera intrinsics when resizing."""
+        for i, intrinsic in enumerate(intrinsics):
+            intrinsic[0, 2] += pad_params[i]["pad_left"]
+            intrinsic[1, 2] += pad_params[i]["pad_top"]
+
+            intrinsics[i] = intrinsic
+        return intrinsics
+
+
+@Transform([K.boxes2d, "transforms.pad"], K.boxes2d)
+class CenterPadBoxes2D:
+    """Pad batch of depth maps at the bottom right."""
+
+    def __call__(
+        self, boxes_list: list[NDArrayF32], pad_params: list[PadParam]
+    ) -> list[NDArrayF32]:
+        """Scale camera intrinsics when resizing."""
+        for i, boxes in enumerate(boxes_list):
+            boxes[:, 0] += pad_params[i]["pad_left"]
+            boxes[:, 1] += pad_params[i]["pad_top"]
+            boxes[:, 2] += pad_params[i]["pad_left"]
+            boxes[:, 3] += pad_params[i]["pad_top"]
+
+            boxes_list[i] = boxes
+
+        return boxes_list
+
+
+@Transform([K.depth_maps, "transforms.pad"], K.depth_maps)
+class CenterPadDepthMaps:
+    """Pad batch of depth maps at the bottom right."""
+
+    def __init__(self, mode: str = "constant", value: int = 0) -> None:
+        """Creates an instance."""
+        self.mode = mode
+        self.value = value
+
+    def __call__(
+        self, depth_maps: list[NDArrayF32], pad_params: list[PadParam]
+    ) -> list[NDArrayF32]:
+        """Pad images to consistent size."""
+
+        # generate params for torch pad
+        for i, (depth, pad_param_dict) in enumerate(
+            zip(depth_maps, pad_params)
+        ):
+            pad_param = (
+                pad_param_dict["pad_left"],
+                pad_param_dict["pad_right"],
+                pad_param_dict["pad_top"],
+                pad_param_dict["pad_bottom"],
+            )
+
+            depth_ = torch.from_numpy(depth).unsqueeze(0).unsqueeze(0)
+            depth_ = F.pad(depth_, pad_param, self.mode, self.value)
+            depth_maps[i] = depth_.squeeze(0).squeeze(0).numpy()
+
+        return depth_maps

opendet3d/data/transforms/resize.py

@@ -0,0 +1,121 @@
+"""Resize transformation."""
+
+from __future__ import annotations
+
+import math
+
+import numpy as np
+import torch
+from vis4d.common.typing import NDArrayF32, NDArrayI64
+from vis4d.data.const import CommonKeys as K
+from vis4d.data.transforms.base import Transform
+from vis4d.data.transforms.resize import ResizeParam, resize_tensor
+
+
+@Transform(K.images, ["transforms.resize", K.input_hw])
+class GenResizeParameters:
+    """Generate the parameters for a resize operation."""
+
+    def __init__(
+        self, shape: tuple[int, int], scales: tuple[float, float] | float = 1.0
+    ) -> None:
+        """Create a new instance of the class."""
+        self.shape = shape
+        self.scales = scales
+
+    def __call__(
+        self, images: list[NDArrayF32]
+    ) -> tuple[list[ResizeParam], list[tuple[int, int]]]:
+        """Compute the parameters and put them in the data dict."""
+        if isinstance(self.scales, float):
+            random_scale = self.scales
+        else:
+            random_scale = np.random.uniform(self.scales[0], self.scales[1])
+
+        shape = (
+            math.ceil(self.shape[0] * random_scale - 0.5),
+            math.ceil(self.shape[1] * random_scale - 0.5),
+        )
+
+        output_ratio = shape[1] / shape[0]
+
+        image = images[0]
+
+        input_h, input_w = (image.shape[1], image.shape[2])
+        input_ratio = input_w / input_h
+
+        if output_ratio > input_ratio:
+            scale = shape[0] / input_h
+        else:
+            scale = shape[1] / input_w
+
+        target_shape = (
+            math.ceil(input_h * scale - 0.5),
+            math.ceil(input_w * scale - 0.5),
+        )
+
+        scale_factor = (target_shape[0] / input_h, target_shape[1] / input_w)
+
+        resize_params = [
+            ResizeParam(target_shape=target_shape, scale_factor=scale_factor)
+        ] * len(images)
+        target_shapes = [target_shape] * len(images)
+
+        return resize_params, target_shapes
+
+
+@Transform(
+    [K.panoptic_masks, "transforms.resize.target_shape"], K.panoptic_masks
+)
+class ResizePanopticMasks:
+    """Resize panoptic segmentation masks."""
+
+    def __call__(
+        self,
+        masks_list: list[NDArrayI64],
+        target_shape_list: list[tuple[int, int]],
+    ) -> list[NDArrayI64]:
+        """Resize masks."""
+        for i, (masks, target_shape) in enumerate(
+            zip(masks_list, target_shape_list)
+        ):
+            masks_ = torch.from_numpy(masks)
+            masks_ = (
+                resize_tensor(
+                    masks_.float().unsqueeze(0).unsqueeze(0),
+                    target_shape,
+                    interpolation="nearest",
+                )
+                .type(masks_.dtype)
+                .squeeze(0)
+                .squeeze(0)
+            )
+            masks_list[i] = masks_.numpy()
+        return masks_list
+
+
+@Transform([K.boxes3d, "transforms.resize.scale_factor"], K.boxes3d)
+class ResizeBoxes3D:
+    """Resize list of 2D bounding boxes."""
+
+    def __call__(
+        self,
+        boxes_list: list[NDArrayF32],
+        scale_factors: list[tuple[float, float]],
+    ) -> list[NDArrayF32]:
+        """Resize 2D bounding boxes.
+
+        Args:
+            boxes_list: (list[NDArrayF32]): The bounding boxes to be resized.
+            scale_factors (list[tuple[float, float]]): scaling factors.
+
+        Returns:
+            list[NDArrayF32]: Resized bounding boxes according to parameters in
+                resize.
+        """
+        for i, (boxes, scale_factor) in enumerate(
+            zip(boxes_list, scale_factors)
+        ):
+            boxes[:, 2] /= scale_factor[0]
+            boxes_list[i] = boxes
+        return boxes_list

opendet3d/eval/detect3d.py

@@ -0,0 +1,1249 @@
+"""3D Multiple Object Detection Evaluator."""
+
+import contextlib
+import copy
+import datetime
+import io
+import itertools
+import json
+import os
+import time
+from collections import defaultdict
+
+import numpy as np
+import pycocotools.mask as maskUtils
+import torch
+from pycocotools.cocoeval import COCOeval
+from scipy.spatial.distance import cdist
+from terminaltables import AsciiTable
+from vis4d.common.array import array_to_numpy
+from vis4d.common.distributed import all_gather_object_cpu
+from vis4d.common.typing import (
+    ArrayLike,
+    DictStrAny,
+    GenericFunc,
+    MetricLogs,
+    NDArrayF32,
+    NDArrayI64,
+)
+from vis4d.data.const import AxisMode
+from vis4d.eval.base import Evaluator
+from vis4d.eval.coco.detect import xyxy_to_xywh
+from vis4d.op.box.box3d import boxes3d_to_corners
+from vis4d.op.geometry.rotation import quaternion_to_matrix
+
+from opendet3d.data.datasets.coco3d import COCO3D
+from opendet3d.op.box.box3d import box3d_overlap
+from opendet3d.op.geometric.rotation import so3_relative_angle
+
+
+class Detect3DEvaluator(Evaluator):
+    """3D object detection evaluation with COCO format."""
+
+    def __init__(
+        self,
+        det_map: dict[str, int],
+        cat_map: dict[str, int],
+        annotation: str,
+        id2name: dict[int, str] | None = None,
+        per_class_eval: bool = True,
+        eval_prox: bool = False,
+        iou_type: str = "bbox",
+        num_columns: int = 6,
+        base_classes: list[str] | None = None,
+    ) -> None:
+        """Create an instance of the class."""
+        if id2name is None:
+            self.id2name = {v: k for k, v in det_map.items()}
+        else:
+            self.id2name = id2name
+
+        self.annotation = annotation
+        self.per_class_eval = per_class_eval
+        self.eval_prox = eval_prox
+        self.iou_type = iou_type
+        self.num_columns = num_columns
+        self.base_classes = base_classes
+
+        self.tp_errors = ["ATE", "AOE", "ASE"]
+
+        category_names = sorted(det_map, key=det_map.get)
+
+        with contextlib.redirect_stdout(io.StringIO()):
+            self._coco_gt = COCO3D([annotation], category_names)
+
+        self.cat_map = cat_map
+
+        self.bbox_2D_evals_per_cat_area: DictStrAny = {}
+        self.bbox_3D_evals_per_cat_area: DictStrAny = {}
+        self._predictions: list[DictStrAny] = []
+
+    def __repr__(self) -> str:
+        """Returns the string representation of the object."""
+        return f"3D Object Detection Evaluator with {self.annotation}"
+
+    @property
+    def metrics(self) -> list[str]:
+        """Supported metrics.
+
+        Returns:
+            list[str]: Metrics to evaluate.
+        """
+        return ["2D", "3D"]
+
+    def gather(self, gather_func: GenericFunc) -> None:
+        """Accumulate predictions across processes."""
+        all_preds = all_gather_object_cpu(
+            self._predictions, use_system_tmp=False
+        )
+        if all_preds is not None:
+            self._predictions = list(itertools.chain(*all_preds))
+
+    def reset(self) -> None:
+        """Reset the saved predictions to start new round of evaluation."""
+        self._predictions.clear()
+        self.bbox_2D_evals_per_cat_area.clear()
+        self.bbox_3D_evals_per_cat_area.clear()
+
+    def process_batch(
+        self,
+        coco_image_id: list[int],
+        pred_boxes: list[ArrayLike],
+        pred_scores: list[ArrayLike],
+        pred_classes: list[ArrayLike],
+        pred_boxes3d: list[ArrayLike] | None = None,
+    ) -> None:
+        """Process sample and convert detections to coco format."""
+        for i, image_id in enumerate(coco_image_id):
+            boxes = array_to_numpy(
+                pred_boxes[i].to(torch.float32), n_dims=None, dtype=np.float32
+            )
+            scores = array_to_numpy(
+                pred_scores[i].to(torch.float32), n_dims=None, dtype=np.float32
+            )
+            classes = array_to_numpy(
+                pred_classes[i], n_dims=None, dtype=np.int64
+            )
+
+            if pred_boxes3d is not None:
+                boxes3d = array_to_numpy(
+                    pred_boxes3d[i].to(torch.float32),
+                    n_dims=None,
+                    dtype=np.float32,
+                )
+            else:
+                boxes3d = None
+
+            self._predictions_to_coco(
+                image_id, boxes, boxes3d, scores, classes
+            )
+
+    def _predictions_to_coco(
+        self,
+        img_id: int,
+        boxes: NDArrayF32,
+        boxes3d: NDArrayF32 | None,
+        scores: NDArrayF32,
+        classes: NDArrayI64,
+    ) -> None:
+        """Convert predictions to COCO format."""
+        boxes_xyxy = copy.deepcopy(boxes)
+        boxes_xywh = xyxy_to_xywh(boxes_xyxy)
+
+        if boxes3d is not None:
+            # FIXME: Make axismode configurable
+            corners_3d = boxes3d_to_corners(
+                torch.from_numpy(boxes3d), AxisMode.OPENCV
+            )
+
+        for i, (box, box_score, box_class) in enumerate(
+            zip(boxes_xywh, scores, classes)
+        ):
+            xywh = box.tolist()
+
+            result = {
+                "image_id": img_id,
+                "bbox": xywh,
+                "category_id": self.cat_map[self.id2name[box_class.item()]],
+                "score": box_score.item(),
+            }
+
+            # mapping to Omni3D format
+            if boxes3d is not None:
+                result["center_cam"] = boxes3d[i][:3].tolist()
+
+                # wlh to whl
+                result["dimensions"] = boxes3d[i][[3, 5, 4]].tolist()
+
+                result["R_cam"] = (
+                    quaternion_to_matrix(torch.from_numpy(boxes3d[i][6:10]))
+                    .numpy()
+                    .tolist()
+                )
+
+                corners = corners_3d[i].numpy().tolist()
+
+                result["bbox3D"] = [
+                    corners[6],
+                    corners[4],
+                    corners[0],
+                    corners[2],
+                    corners[7],
+                    corners[5],
+                    corners[1],
+                    corners[3],
+                ]
+
+                result["depth"] = boxes3d[i][2].item()
+
+            self._predictions.append(result)
+
+    def evaluate(self, metric: str) -> tuple[MetricLogs, str]:
+        """Evaluate predictions."""
+        if metric == "2D":
+            metrics = ["AP", "AP50", "AP75", "AP95", "APs", "APm", "APl"]
+        else:
+            if self.iou_type == "bbox":
+                metrics = ["AP", "AP15", "AP25", "AP50", "APn", "APm", "APf"]
+                main_metric = "AP"
+            else:
+                metrics = ["AP", "ATE", "ASE", "AOE", "ODS"]
+                main_metric = "ODS"
+
+            if self.base_classes is not None:
+                metrics += [f"{main_metric}_Base", f"{main_metric}_Novel"]
+
+        if len(self._predictions) == 0:
+            return {m: 0.0 for m in metrics}, "No predictions to evaluate."
+
+        with contextlib.redirect_stdout(io.StringIO()):
+            coco_dt = self._coco_gt.loadRes(self._predictions)
+
+            assert coco_dt is not None
+            evaluator = Detect3Deval(
+                self._coco_gt,
+                coco_dt,
+                mode=metric,
+                eval_prox=self.eval_prox,
+                iou_type=self.iou_type,
+            )
+            evaluator.evaluate()
+            evaluator.accumulate()
+
+        if self.iou_type == "bbox":
+            log_str = "\n" + evaluator.summarize()
+
+        # precision: (iou, recall, cls, area range, max dets)
+        precisions = evaluator.eval["precision"]
+        assert len(self._coco_gt.getCatIds()) == precisions.shape[2]
+
+        if metric == "2D":
+            self.bbox_2D_evals_per_cat_area = evaluator.evals_per_cat_area
+
+            score_dict = dict(zip(metrics, evaluator.stats))
+        else:
+            if self.iou_type == "bbox":
+                self.bbox_3D_evals_per_cat_area = evaluator.evals_per_cat_area
+
+                score_dict = dict(zip(metrics, evaluator.stats))
+            else:
+                trans_tp_errors = evaluator.eval["trans_tp_errors"]
+                rot_tp_errors = evaluator.eval["rot_tp_errors"]
+                scale_tp_errors = evaluator.eval["scale_tp_errors"]
+
+                precision = precisions[:, :, :, 0, -1]
+                precision = precision[precision > -1]
+                if precision.size:
+                    mAP = np.mean(precision).item()
+                else:
+                    mAP = float("nan")
+
+                trans_tp = trans_tp_errors[:, :, :, 0, -1]
+                trans_tp = trans_tp[trans_tp > -1]
+
+                rot_tp = rot_tp_errors[:, :, :, 0, -1]
+                rot_tp = rot_tp[rot_tp > -1]
+
+                scale_tp = scale_tp_errors[:, :, :, 0, -1]
+                scale_tp = scale_tp[scale_tp > -1]
+
+                if trans_tp.size:
+                    mATE = np.mean(trans_tp).item()
+                    mAOE = np.mean(rot_tp).item()
+                    mASE = np.mean(scale_tp).item()
+
+                    mODS = (
+                        np.sum(mAP * 3 + (1 - mATE) + (1 - mAOE) + (1 - mASE))
+                        / 6
+                    )
+
+                else:
+                    mATE = float("nan")
+                    mAOE = float("nan")
+                    mASE = float("nan")
+                    mODS = float("nan")
+
+                score_dict = {
+                    "AP": mAP,
+                    "ATE": mATE,
+                    "ASE": mASE,
+                    "AOE": mAOE,
+                    "ODS": mODS,
+                }
+
+                log_str = "\nHigh-level metrics:"
+                for k, v in score_dict.items():
+                    log_str += f"\n{k}: {v:.4f}"
+
+        if self.per_class_eval:
+            results_per_category = []
+            score_base_list = []
+            score_novel_list = []
+
+            for idx, cat_id in enumerate(self._coco_gt.getCatIds()):
+                # area range index 0: all area ranges
+                # max dets index -1: typically 100 per image
+                nm = self._coco_gt.loadCats(cat_id)[0]
+                precision = precisions[:, :, idx, 0, -1]
+                precision = precision[precision > -1]
+                if precision.size:
+                    ap = np.mean(precision).item()
+                else:
+                    ap = float("nan")
+
+                if self.iou_type == "dist":
+                    trans_tp = trans_tp_errors[:, :, idx, 0, -1]
+                    trans_tp = trans_tp[trans_tp > -1]
+
+                    rot_tp = rot_tp_errors[:, :, idx, 0, -1]
+                    rot_tp = rot_tp[rot_tp > -1]
+
+                    scale_tp = scale_tp_errors[:, :, idx, 0, -1]
+                    scale_tp = scale_tp[scale_tp > -1]
+
+                    if trans_tp.size:
+                        ate = np.mean(trans_tp).item()
+                        aoe = np.mean(rot_tp).item()
+                        ase = np.mean(scale_tp).item()
+
+                        ods = (
+                            np.sum(ap * 3 + (1 - ate) + (1 - aoe) + (1 - ase))
+                            / 6
+                        )
+
+                    else:
+                        ate = float("nan")
+                        aoe = float("nan")
+                        ase = float("nan")
+                        ods = float("nan")
+
+                    results_per_category.append(
+                        (
+                            f'{nm["name"]}',
+                            f"{ap:0.3f}",
+                            f"{ate:0.3f}",
+                            f"{ase:0.3f}",
+                            f"{aoe:0.3f}",
+                            f"{ods:0.3f}",
+                        )
+                    )
+                else:
+                    results_per_category.append(
+                        (f'{nm["name"]}', f"{ap:0.3f}")
+                    )
+
+                if self.base_classes is not None:
+                    if self.iou_type == "dist":
+                        score = ods
+                    else:
+                        score = ap
+
+                    if nm["name"] in self.base_classes:
+                        score_base_list.append(score)
+                    else:
+                        score_novel_list.append(score)
+
+            results_flatten = list(itertools.chain(*results_per_category))
+
+            if self.iou_type == "dist":
+                num_columns = 6
+                headers = ["category", "AP", "ATE", "ASE", "AOE", "ODS"]
+            else:
+                num_columns = min(
+                    self.num_columns, len(results_per_category) * 2
+                )
+                headers = ["category", "AP"] * (num_columns // 2)
+            results = itertools.zip_longest(
+                *[results_flatten[i::num_columns] for i in range(num_columns)]
+            )
+            table_data = [headers] + list(results)
+            table = AsciiTable(table_data)
+            log_str = f"\n{table.table}\n{log_str}"
+
+        if self.base_classes is not None:
+            score_dict[f"{main_metric}_Base"] = np.mean(score_base_list).item()
+            score_dict[f"{main_metric}_Novel"] = np.mean(
+                score_novel_list
+            ).item()
+
+        return score_dict, log_str
+
+    def save(
+        self, metric: str, output_dir: str, prefix: str | None = None
+    ) -> None:
+        """Save the results to json files."""
+        assert metric in self.metrics
+
+        if prefix is not None:
+            result_folder = os.path.join(output_dir, prefix)
+            os.makedirs(result_folder, exist_ok=True)
+        else:
+            result_folder = output_dir
+
+        result_file = os.path.join(
+            result_folder, f"detect_{metric}_results.json"
+        )
+
+        with open(result_file, mode="w", encoding="utf-8") as f:
+            json.dump(self._predictions, f)
+
+
+class Detect3Deval(COCOeval):
+    """COCOeval Wrapper for 2D and 3D box evaluation.
+
+    Now it support bbox IoU matching only.
+    """
+
+    def __init__(
+        self,
+        cocoGt=None,
+        cocoDt=None,
+        mode: str = "2D",
+        iou_type: str = "bbox",
+        eval_prox: bool = False,
+    ):
+        """Initialize Detect3Deval using coco APIs for Gt and Dt.
+
+        Args:
+            cocoGt: COCO object with ground truth annotations
+            cocoDt: COCO object with detection results
+            mode: (str) defines whether to evaluate 2D or 3D performance.
+                One of {"2D", "3D"}
+            eval_prox: (bool) if True, performs "Proximity Evaluation", i.e.
+                evaluates detections in the proximity of the ground truth2D
+                boxes. This is used for datasets which are not exhaustively
+                annotated.
+        """
+        if mode not in {"2D", "3D"}:
+            raise Exception(f"{mode} mode is not supported")
+        self.mode = mode
+        self.iou_type = iou_type
+        self.eval_prox = eval_prox
+
+        self.cocoGt = cocoGt  # ground truth COCO API
+        self.cocoDt = cocoDt  # detections COCO API
+
+        # per-image per-category evaluation results [KxAxI] elements
+        self.evalImgs = defaultdict(list)
+
+        self.eval = {}  # accumulated evaluation results
+        self._gts = defaultdict(list)  # gt for evaluation
+        self._dts = defaultdict(list)  # dt for evaluation
+        self.params = Detect3DParams(mode=mode, iouType=iou_type)  # parameters
+        self._paramsEval = {}  # parameters for evaluation
+        self.stats = []  # result summarization
+        self.ious = {}  # ious between all gts and dts
+
+        if cocoGt is not None:
+            self.params.imgIds = sorted(cocoGt.getImgIds())
+            self.params.catIds = sorted(cocoGt.getCatIds())
+
+        self.evals_per_cat_area = None
+
+    def _prepare(self) -> None:
+        """Prepare ._gts and ._dts for evaluation based on params."""
+        p = self.params
+
+        if p.useCats:
+            gts = self.cocoGt.loadAnns(
+                self.cocoGt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)
+            )
+            dts = self.cocoDt.loadAnns(
+                self.cocoDt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds)
+            )
+
+        else:
+            gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds))
+            dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds))
+
+        # set ignore flag
+        ignore_flag = "ignore2D" if self.mode == "2D" else "ignore3D"
+        for gt in gts:
+            gt[ignore_flag] = gt[ignore_flag] if ignore_flag in gt else 0
+
+        self._gts = defaultdict(list)  # gt for evaluation
+        self._dts = defaultdict(list)  # dt for evaluation
+
+        for gt in gts:
+            self._gts[gt["image_id"], gt["category_id"]].append(gt)
+
+        for dt in dts:
+            self._dts[dt["image_id"], dt["category_id"]].append(dt)
+
+        self.evalImgs = defaultdict(
+            list
+        )  # per-image per-category evaluation results
+        self.eval = {}  # accumulated evaluation results
+
+    def accumulate(self, p=None) -> None:
+        """Accumulate per image evaluation and store the result in self.eval.
+
+        Args:
+            p: input params for evaluation
+        """
+        print("Accumulating evaluation results...")
+        assert self.evalImgs, "Please run evaluate() first"
+
+        tic = time.time()
+
+        # allows input customized parameters
+        if p is None:
+            p = self.params
+
+        p.catIds = p.catIds if p.useCats == 1 else [-1]
+
+        T = len(p.iouThrs)
+        R = len(p.recThrs)
+        K = len(p.catIds) if p.useCats else 1
+        A = len(p.areaRng)
+        M = len(p.maxDets)
+
+        precision = -np.ones(
+            (T, R, K, A, M)
+        )  # -1 for the precision of absent categories
+        trans_tp_errors = -np.ones((T, R, K, A, M))
+        rot_tp_errors = -np.ones((T, R, K, A, M))
+        scale_tp_errors = -np.ones((T, R, K, A, M))
+        recall = -np.ones((T, K, A, M))
+        scores = -np.ones((T, R, K, A, M))
+
+        # create dictionary for future indexing
+        _pe = self._paramsEval
+
+        catIds = _pe.catIds if _pe.useCats else [-1]
+        setK = set(catIds)
+        setA = set(map(tuple, _pe.areaRng))
+        setM = set(_pe.maxDets)
+        setI = set(_pe.imgIds)
+
+        # get inds to evaluate
+        catid_list = [k for n, k in enumerate(p.catIds) if k in setK]
+        k_list = [n for n, k in enumerate(p.catIds) if k in setK]
+        m_list = [m for n, m in enumerate(p.maxDets) if m in setM]
+        a_list = [
+            n
+            for n, a in enumerate(map(lambda x: tuple(x), p.areaRng))
+            if a in setA
+        ]
+        i_list = [n for n, i in enumerate(p.imgIds) if i in setI]
+
+        I0 = len(_pe.imgIds)
+        A0 = len(_pe.areaRng)
+
+        has_precomputed_evals = not (self.evals_per_cat_area is None)
+
+        if has_precomputed_evals:
+            evals_per_cat_area = self.evals_per_cat_area
+        else:
+            evals_per_cat_area = {}
+
+        # retrieve E at each category, area range, and max number of detections
+        for k, (k0, catId) in enumerate(zip(k_list, catid_list)):
+            Nk = k0 * A0 * I0
+            for a, a0 in enumerate(a_list):
+                Na = a0 * I0
+
+                if has_precomputed_evals:
+                    E = evals_per_cat_area[(catId, a)]
+
+                else:
+                    E = [self.evalImgs[Nk + Na + i] for i in i_list]
+                    E = [e for e in E if not e is None]
+                    evals_per_cat_area[(catId, a)] = E
+
+                if len(E) == 0:
+                    continue
+
+                for m, maxDet in enumerate(m_list):
+
+                    dtScores = np.concatenate(
+                        [e["dtScores"][0:maxDet] for e in E]
+                    )
+
+                    # different sorting method generates slightly different results.
+                    # mergesort is used to be consistent as Matlab implementation.
+                    inds = np.argsort(-dtScores, kind="mergesort")
+                    dtScoresSorted = dtScores[inds]
+
+                    dtm = np.concatenate(
+                        [e["dtMatches"][:, 0:maxDet] for e in E], axis=1
+                    )[:, inds]
+                    dtIg = np.concatenate(
+                        [e["dtIgnore"][:, 0:maxDet] for e in E], axis=1
+                    )[:, inds]
+                    gtIg = np.concatenate([e["gtIgnore"] for e in E])
+                    npig = np.count_nonzero(gtIg == 0)
+
+                    if npig == 0:
+                        continue
+
+                    tps = np.logical_and(dtm, np.logical_not(dtIg))
+                    fps = np.logical_and(
+                        np.logical_not(dtm), np.logical_not(dtIg)
+                    )
+
+                    tp_sum = np.cumsum(tps, axis=1).astype(dtype=np.float64)
+                    fp_sum = np.cumsum(fps, axis=1).astype(dtype=np.float64)
+
+                    # Compute TP error
+                    if self.iou_type == "dist":
+                        tems = np.concatenate(
+                            [e["dtTranslationError"][:, 0:maxDet] for e in E],
+                            axis=1,
+                        )[:, inds]
+
+                        oems = np.concatenate(
+                            [e["dtOrientationError"][:, 0:maxDet] for e in E],
+                            axis=1,
+                        )[:, inds]
+
+                        sems = np.concatenate(
+                            [e["dtScaleError"][:, 0:maxDet] for e in E], axis=1
+                        )[:, inds]
+
+                    for t, (tp, fp) in enumerate(zip(tp_sum, fp_sum)):
+                        tp = np.array(tp)
+                        fp = np.array(fp)
+                        nd = len(tp)
+                        rc = tp / npig
+                        pr = tp / (fp + tp + np.spacing(1))
+
+                        q = np.zeros((R,))
+                        ss = np.zeros((R,))
+                        tran_tp_error = np.ones((R,))
+                        rot_tp_error = np.ones((R,))
+                        scale_tp_error = np.ones((R,))
+
+                        if nd:
+                            recall[t, k, a, m] = rc[-1]
+
+                        else:
+                            recall[t, k, a, m] = 0
+
+                        # numpy is slow without cython optimization for accessing elements
+                        # use python array gets significant speed improvement
+                        pr = pr.tolist()
+                        q = q.tolist()
+                        tran_tp_error = tran_tp_error.tolist()
+                        rot_tp_error = rot_tp_error.tolist()
+                        scale_tp_error = scale_tp_error.tolist()
+
+                        for i in range(nd - 1, 0, -1):
+                            if pr[i] > pr[i - 1]:
+                                pr[i - 1] = pr[i]
+
+                        inds = np.searchsorted(rc, p.recThrs, side="left")
+
+                        try:
+                            for ri, pi in enumerate(inds):
+                                q[ri] = pr[pi]
+                                ss[ri] = dtScoresSorted[pi]
+                                if self.iou_type == "dist":
+                                    tran_tp_error[ri] = tems[t][pi]
+                                    rot_tp_error[ri] = oems[t][pi]
+                                    scale_tp_error[ri] = sems[t][pi]
+                        except:
+                            pass
+
+                        precision[t, :, k, a, m] = np.array(q)
+                        scores[t, :, k, a, m] = np.array(ss)
+
+                        if self.iou_type == "dist":
+                            trans_tp_errors[t, :, k, a, m] = np.array(
+                                tran_tp_error
+                            )
+                            rot_tp_errors[t, :, k, a, m] = np.array(
+                                rot_tp_error
+                            )
+                            scale_tp_errors[t, :, k, a, m] = np.array(
+                                scale_tp_error
+                            )
+
+        self.evals_per_cat_area = evals_per_cat_area
+
+        self.eval = {
+            "params": p,
+            "counts": [T, R, K, A, M],
+            "date": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            "precision": precision,
+            "recall": recall,
+            "scores": scores,
+            "trans_tp_errors": trans_tp_errors,
+            "rot_tp_errors": rot_tp_errors,
+            "scale_tp_errors": scale_tp_errors,
+        }
+
+        toc = time.time()
+        print("DONE (t={:0.2f}s).".format(toc - tic))
+
+    def evaluate(self) -> None:
+        """Run per image evaluation on given images.
+
+        It will store results (a list of dict) in self.evalImgs
+        """
+        print("Running per image evaluation...")
+
+        p = self.params
+        print(f"Evaluate annotation type *{p.iouType}*")
+
+        tic = time.time()
+
+        p.imgIds = list(np.unique(p.imgIds))
+        if p.useCats:
+            p.catIds = list(np.unique(p.catIds))
+
+        p.maxDets = sorted(p.maxDets)
+        self.params = p
+
+        self._prepare()
+
+        catIds = p.catIds if p.useCats else [-1]
+
+        # loop through images, area range, max detection number
+        self.ious = {
+            (imgId, catId): self.computeIoU(imgId, catId)
+            for imgId in p.imgIds
+            for catId in catIds
+        }
+
+        maxDet = p.maxDets[-1]
+
+        self.evalImgs = [
+            self.evaluateImg(imgId, catId, areaRng, maxDet)
+            for catId in catIds
+            for areaRng in p.areaRng
+            for imgId in p.imgIds
+        ]
+
+        self._paramsEval = copy.deepcopy(self.params)
+
+        toc = time.time()
+        print("DONE (t={:0.2f}s).".format(toc - tic))
+
+    def computeIoU(self, imgId, catId) -> tuple[NDArrayF32, NDArrayF32]:
+        """Computes the IoUs by sorting based on score"""
+        p = self.params
+
+        if p.useCats:
+            gt = self._gts[imgId, catId]
+            dt = self._dts[imgId, catId]
+        else:
+            gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
+            dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
+
+        if len(gt) == 0 and len(dt) == 0:
+            return []
+
+        inds = np.argsort([-d["score"] for d in dt], kind="mergesort")
+        dt = [dt[i] for i in inds]
+        if len(dt) > p.maxDets[-1]:
+            dt = dt[0 : p.maxDets[-1]]
+
+        if self.mode == "2D":
+            g = [g["bbox"] for g in gt]
+            d = [d["bbox"] for d in dt]
+        elif self.mode == "3D":
+            g = [g["bbox3D"] for g in gt]
+            d = [d["bbox3D"] for d in dt]
+
+        # compute iou between each dt and gt region
+        # iscrowd is required in builtin maskUtils so we
+        # use a dummy buffer for it
+        iscrowd = [0 for _ in gt]
+        if self.mode == "2D":
+            ious = maskUtils.iou(d, g, iscrowd)
+        elif len(d) > 0 and len(g) > 0:
+            if p.iouType == "bbox":
+                dd = torch.tensor(d, dtype=torch.float32)
+                gg = torch.tensor(g, dtype=torch.float32)
+
+                ious = box3d_overlap(dd, gg).cpu().numpy()
+            else:
+                ious = np.zeros((len(d), len(g)))
+
+                dd = [d["center_cam"] for d in dt]
+                gg = [g["center_cam"] for g in gt]
+
+                ious = cdist(dd, gg, metric="euclidean")
+        else:
+            ious = []
+
+        in_prox = None
+
+        if self.eval_prox:
+            g = [g["bbox"] for g in gt]
+            d = [d["bbox"] for d in dt]
+            iscrowd = [0 for o in gt]
+            ious2d = maskUtils.iou(d, g, iscrowd)
+
+            if type(ious2d) == list:
+                in_prox = []
+
+            else:
+                in_prox = ious2d > p.proximity_thresh
+
+        return ious, in_prox
+
+    def evaluateImg(self, imgId, catId, aRng, maxDet):
+        """
+        Perform evaluation for single category and image
+        Returns:
+            dict (single image results)
+        """
+
+        p = self.params
+        if p.useCats:
+            gt = self._gts[imgId, catId]
+            dt = self._dts[imgId, catId]
+
+        else:
+            gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
+            dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
+
+        if len(gt) == 0 and len(dt) == 0:
+            return None
+
+        flag_range = "area" if self.mode == "2D" else "depth"
+        flag_ignore = "ignore2D" if self.mode == "2D" else "ignore3D"
+
+        for g in gt:
+            if g[flag_ignore] or (
+                g[flag_range] < aRng[0] or g[flag_range] > aRng[1]
+            ):
+                g["_ignore"] = 1
+            else:
+                g["_ignore"] = 0
+
+        # sort dt highest score first, sort gt ignore last
+        gtind = np.argsort([g["_ignore"] for g in gt], kind="mergesort")
+        gt = [gt[i] for i in gtind]
+        dtind = np.argsort([-d["score"] for d in dt], kind="mergesort")
+        dt = [dt[i] for i in dtind[0:maxDet]]
+
+        # load computed ious
+        ious = (
+            self.ious[imgId, catId][0][:, gtind]
+            if len(self.ious[imgId, catId][0]) > 0
+            else self.ious[imgId, catId][0]
+        )
+
+        if self.eval_prox:
+            in_prox = (
+                self.ious[imgId, catId][1][:, gtind]
+                if len(self.ious[imgId, catId][1]) > 0
+                else self.ious[imgId, catId][1]
+            )
+
+        T = len(p.iouThrs)
+        G = len(gt)
+        D = len(dt)
+        gtm = np.zeros((T, G))
+        dtm = np.zeros((T, D))
+        tem = np.ones((T, D))  # Translation Error
+        sem = np.ones((T, D))  # Scale Error
+        oem = np.ones((T, D))  # Oritentation Error
+        gtIg = np.array([g["_ignore"] for g in gt])
+        dtIg = np.zeros((T, D))
+
+        dist_thres = 1
+        if not len(ious) == 0:
+            for tind, t in enumerate(p.iouThrs):
+                for dind, d in enumerate(dt):
+
+                    # information about best match so far (m=-1 -> unmatched)
+                    iou = min([t, 1 - 1e-10])
+                    m = -1
+
+                    for gind, g in enumerate(gt):
+                        # in case of proximity evaluation, if not in proximity continue
+                        if self.eval_prox and not in_prox[dind, gind]:
+                            continue
+
+                        # if this gt already matched, continue
+                        if gtm[tind, gind] > 0:
+                            continue
+
+                        # if dt matched to reg gt, and on ignore gt, stop
+                        if m > -1 and gtIg[m] == 0 and gtIg[gind] == 1:
+                            break
+
+                        # continue to next gt unless better match made
+                        if p.iouType == "bbox" and ious[dind, gind] < iou:
+                            continue
+
+                        if p.iouType == "dist":
+                            # Compute Object Radius
+                            gt_obj_radius = (
+                                np.linalg.norm(np.array(g["dimensions"])) / 2
+                            )
+                            if ious[dind, gind] > gt_obj_radius * iou:
+                                continue
+                            else:
+                                dist_thres = gt_obj_radius * iou
+
+                        # if match successful and best so far, store appropriately
+                        iou = ious[dind, gind]
+                        m = gind
+
+                    # if match made store id of match for both dt and gt
+                    if m == -1:
+                        continue
+
+                    dtIg[tind, dind] = gtIg[m]
+                    dtm[tind, dind] = gt[m]["id"]
+                    gtm[tind, m] = d["id"]
+
+                    if p.iouType == "dist":
+                        # Translation Error
+                        tem[tind, dind] = np.linalg.norm(
+                            np.array(d["center_cam"])
+                            - np.array(gt[m]["center_cam"])
+                        ) / (dist_thres)
+
+                        # Orientation Error
+                        oem[tind, dind] = (
+                            so3_relative_angle(
+                                torch.tensor(d["R_cam"])[None],
+                                torch.tensor(gt[m]["R_cam"])[None],
+                                cos_bound=1e-2,
+                                eps=1e-2,
+                            ).item()
+                            / np.pi
+                        )
+
+                        # Scale Error
+                        min_whl = np.minimum(
+                            d["dimensions"], gt[m]["dimensions"]
+                        )
+                        volume_annotation = np.prod(gt[m]["dimensions"])
+                        volume_result = np.prod(d["dimensions"])
+
+                        intersection = np.prod(min_whl)
+                        union = (
+                            volume_annotation + volume_result - intersection
+                        )
+                        scale_iou = intersection / union
+
+                        sem[tind, dind] = 1 - scale_iou
+
+        # set unmatched detections outside of area range to ignore
+        a = np.array(
+            [d[flag_range] < aRng[0] or d[flag_range] > aRng[1] for d in dt]
+        ).reshape((1, len(dt)))
+
+        dtIg = np.logical_or(
+            dtIg, np.logical_and(dtm == 0, np.repeat(a, T, 0))
+        )
+
+        # in case of proximity evaluation, ignore detections which are far from gt regions
+        if self.eval_prox and len(in_prox) > 0:
+            dt_far = in_prox.any(1) == 0
+            dtIg = np.logical_or(
+                dtIg, np.repeat(dt_far.reshape((1, len(dt))), T, 0)
+            )
+
+        # store results for given image and category
+        return {
+            "image_id": imgId,
+            "category_id": catId,
+            "aRng": aRng,
+            "maxDet": maxDet,
+            "dtIds": [d["id"] for d in dt],
+            "gtIds": [g["id"] for g in gt],
+            "dtMatches": dtm,
+            "gtMatches": gtm,
+            "dtScores": [d["score"] for d in dt],
+            "gtIgnore": gtIg,
+            "dtIgnore": dtIg,
+            "dtTranslationError": tem,
+            "dtScaleError": sem,
+            "dtOrientationError": oem,
+        }
+
+    def summarize(self):
+        """
+        Compute and display summary metrics for evaluation results.
+        Note this functin can *only* be applied on the default parameter setting
+        """
+
+        def _summarize(
+            mode, ap=1, iouThr=None, areaRng="all", maxDets=100, log_str=""
+        ):
+            p = self.params
+            eval = self.eval
+
+            if mode == "2D":
+                if self.iou_type == "bbox":
+                    iStr = " {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}"
+                else:
+                    iStr = " {:<18} {} @[ Dist={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}"
+
+            elif mode == "3D":
+                if self.iou_type == "bbox":
+                    iStr = " {:<18} {} @[ IoU={:<9} | depth={:>6s} | maxDets={:>3d} ] = {:0.3f}"
+                else:
+                    iStr = " {:<18} {} @[ Dist={:<9} | depth={:>6s} | maxDets={:>3d} ] = {:0.3f}"
+
+            titleStr = "Average Precision" if ap == 1 else "Average Recall"
+            typeStr = "(AP)" if ap == 1 else "(AR)"
+
+            iouStr = (
+                "{:0.2f}:{:0.2f}".format(p.iouThrs[0], p.iouThrs[-1])
+                if iouThr is None
+                else "{:0.2f}".format(iouThr)
+            )
+
+            aind = [
+                i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng
+            ]
+            mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets]
+
+            if ap == 1:
+
+                # dimension of precision: [TxRxKxAxM]
+                s = eval["precision"]
+
+                # IoU
+                if iouThr is not None:
+                    t = np.where(np.isclose(iouThr, p.iouThrs.astype(float)))[
+                        0
+                    ]
+                    s = s[t]
+
+                s = s[:, :, :, aind, mind]
+
+            else:
+                # dimension of recall: [TxKxAxM]
+                s = eval["recall"]
+                if iouThr is not None:
+                    t = np.where(iouThr == p.iouThrs)[0]
+                    s = s[t]
+                s = s[:, :, aind, mind]
+
+            if len(s[s > -1]) == 0:
+                mean_s = -1
+
+            else:
+                mean_s = np.mean(s[s > -1])
+
+            if log_str != "":
+                log_str += "\n"
+
+            log_str += "mode={} ".format(mode) + iStr.format(
+                titleStr, typeStr, iouStr, areaRng, maxDets, mean_s
+            )
+
+            return mean_s, log_str
+
+        def _summarizeDets(mode):
+
+            params = self.params
+
+            # Define the thresholds to be printed
+            if mode == "2D":
+                thres = [0.5, 0.75, 0.95]
+            else:
+                if self.iou_type == "bbox":
+                    thres = [0.15, 0.25, 0.50]
+                else:
+                    thres = [0.5, 0.75, 1.0]
+
+            stats = np.zeros((13,))
+            stats[0], log_str = _summarize(mode, 1)
+
+            stats[1], log_str = _summarize(
+                mode,
+                1,
+                iouThr=thres[0],
+                maxDets=params.maxDets[2],
+                log_str=log_str,
+            )
+
+            stats[2], log_str = _summarize(
+                mode,
+                1,
+                iouThr=thres[1],
+                maxDets=params.maxDets[2],
+                log_str=log_str,
+            )
+
+            stats[3], log_str = _summarize(
+                mode,
+                1,
+                iouThr=thres[2],
+                maxDets=params.maxDets[2],
+                log_str=log_str,
+            )
+
+            stats[4], log_str = _summarize(
+                mode,
+                1,
+                areaRng=params.areaRngLbl[1],
+                maxDets=params.maxDets[2],
+                log_str=log_str,
+            )
+
+            stats[5], log_str = _summarize(
+                mode,
+                1,
+                areaRng=params.areaRngLbl[2],
+                maxDets=params.maxDets[2],
+                log_str=log_str,
+            )
+
+            stats[6], log_str = _summarize(
+                mode,
+                1,
+                areaRng=params.areaRngLbl[3],
+                maxDets=params.maxDets[2],
+                log_str=log_str,
+            )
+
+            stats[7], log_str = _summarize(
+                mode, 0, maxDets=params.maxDets[0], log_str=log_str
+            )
+
+            stats[8], log_str = _summarize(
+                mode, 0, maxDets=params.maxDets[1], log_str=log_str
+            )
+
+            stats[9], log_str = _summarize(
+                mode, 0, maxDets=params.maxDets[2], log_str=log_str
+            )
+
+            stats[10], log_str = _summarize(
+                mode,
+                0,
+                areaRng=params.areaRngLbl[1],
+                maxDets=params.maxDets[2],
+                log_str=log_str,
+            )
+
+            stats[11], log_str = _summarize(
+                mode,
+                0,
+                areaRng=params.areaRngLbl[2],
+                maxDets=params.maxDets[2],
+                log_str=log_str,
+            )
+
+            stats[12], log_str = _summarize(
+                mode,
+                0,
+                areaRng=params.areaRngLbl[3],
+                maxDets=params.maxDets[2],
+                log_str=log_str,
+            )
+
+            return stats, log_str
+
+        if not self.eval:
+            raise Exception("Please run accumulate() first")
+
+        stats, log_str = _summarizeDets(self.mode)
+        self.stats = stats
+
+        return log_str
+
+
+class Detect3DParams:
+    """Params for the 3d detection evaluation API."""
+
+    def __init__(
+        self,
+        mode: str = "2D",
+        iouType: str = "bbox",
+        proximity_thresh: float = 0.3,
+    ) -> None:
+        """Create an instance of Detect3DParams.
+
+        Args:
+            mode: (str) defines whether to evaluate 2D or 3D performance.
+            iouType: (str) defines the type of IoU to be used for evaluation.
+            proximity_thresh (float): It defines the neighborhood when
+                evaluating on non-exhaustively annotated datasets.
+        """
+        assert iouType in {"bbox", "dist"}, f"Invalid iouType {iouType}."
+        self.iouType = iouType
+
+        if mode == "2D":
+            self.setDet2DParams()
+        elif mode == "3D":
+            self.setDet3DParams()
+        else:
+            raise Exception(f"{mode} mode is not supported")
+        self.mode = mode
+        self.proximity_thresh = proximity_thresh
+
+    def setDet2DParams(self) -> None:
+        """Set parameters for 2D detection evaluation."""
+        self.imgIds = []
+        self.catIds = []
+
+        # np.arange causes trouble.  the data point on arange is slightly larger than the true value
+        self.iouThrs = np.linspace(
+            0.5, 0.95, int(np.round((0.95 - 0.5) / 0.05)) + 1, endpoint=True
+        )
+
+        self.recThrs = np.linspace(
+            0.0, 1.00, int(np.round((1.00 - 0.0) / 0.01)) + 1, endpoint=True
+        )
+
+        self.maxDets = [1, 10, 100]
+        self.areaRng = [
+            [0**2, 1e5**2],
+            [0**2, 32**2],
+            [32**2, 96**2],
+            [96**2, 1e5**2],
+        ]
+
+        self.areaRngLbl = ["all", "small", "medium", "large"]
+        self.useCats = 1
+
+    def setDet3DParams(self) -> None:
+        """Set parameters for 3D detection evaluation."""
+        self.imgIds = []
+        self.catIds = []
+
+        # np.arange causes trouble. The data point on arange is slightly
+        # larger than the true value
+        if self.iouType == "bbox":
+            self.iouThrs = np.linspace(
+                0.05,
+                0.5,
+                int(np.round((0.5 - 0.05) / 0.05)) + 1,
+                endpoint=True,
+            )
+        else:
+            self.iouThrs = np.linspace(
+                0.5, 1.0, int(np.round((1.00 - 0.5) / 0.05)) + 1, endpoint=True
+            )
+
+        self.recThrs = np.linspace(
+            0.0, 1.00, int(np.round((1.00 - 0.0) / 0.01)) + 1, endpoint=True
+        )
+
+        self.maxDets = [1, 10, 100]
+        self.areaRng = [[0, 1e5], [0, 10], [10, 35], [35, 1e5]]
+        self.areaRngLbl = ["all", "near", "medium", "far"]
+        self.useCats = 1

opendet3d/eval/omni3d.py

@@ -0,0 +1,285 @@
+"""Omni3D 3D detection evaluation."""
+
+import contextlib
+import copy
+import io
+import itertools
+import os
+from collections.abc import Sequence
+
+import numpy as np
+from terminaltables import AsciiTable
+from vis4d.common.logging import rank_zero_info
+from vis4d.common.typing import GenericFunc, MetricLogs, NDArrayNumber
+from vis4d.eval.base import Evaluator
+
+from opendet3d.data.datasets.omni3d.omni3d_classes import omni3d_class_map
+from opendet3d.data.datasets.omni3d.util import get_dataset_det_map
+
+from .detect3d import Detect3Deval, Detect3DEvaluator
+
+omni3d_in = {
+    "stationery",
+    "sink",
+    "table",
+    "floor mat",
+    "bottle",
+    "bookcase",
+    "bin",
+    "blinds",
+    "pillow",
+    "bicycle",
+    "refrigerator",
+    "night stand",
+    "chair",
+    "sofa",
+    "books",
+    "oven",
+    "towel",
+    "cabinet",
+    "window",
+    "curtain",
+    "bathtub",
+    "laptop",
+    "desk",
+    "television",
+    "clothes",
+    "stove",
+    "cup",
+    "shelves",
+    "box",
+    "shoes",
+    "mirror",
+    "door",
+    "picture",
+    "lamp",
+    "machine",
+    "counter",
+    "bed",
+    "toilet",
+}
+
+omni3d_out = {
+    "cyclist",
+    "pedestrian",
+    "trailer",
+    "bus",
+    "motorcycle",
+    "car",
+    "barrier",
+    "truck",
+    "van",
+    "traffic cone",
+    "bicycle",
+}
+
+
+class Omni3DEvaluator(Evaluator):
+    """Omni3D 3D detection evaluator."""
+
+    def __init__(
+        self,
+        data_root: str = "data/omni3d",
+        omni3d50: bool = True,
+        datasets: Sequence[str] = (
+            "KITTI_test",
+            "nuScenes_test",
+            "SUNRGBD_test",
+            "Hypersim_test",
+            "ARKitScenes_test",
+            "Objectron_test",
+        ),
+        per_class_eval: bool = True,
+    ) -> None:
+        """Initialize the evaluator."""
+        super().__init__()
+        self.id_to_name = {v: k for k, v in omni3d_class_map.items()}
+        self.dataset_names = datasets
+        self.per_class_eval = per_class_eval
+
+        # Each dataset evaluator is stored here
+        self.evaluators: dict[str, Detect3DEvaluator] = {}
+
+        # These store the evaluations for each category and area,
+        # concatenated from ALL evaluated datasets. Doing so avoids
+        # the need to re-compute them when accumulating results.
+        self.evals_per_cat_area2D = {}
+        self.evals_per_cat_area3D = {}
+
+        self.overall_imgIds = set()
+        self.overall_catIds = set()
+
+        for dataset_name in self.dataset_names:
+            annotation = os.path.join(
+                data_root, "annotations", f"{dataset_name}.json"
+            )
+
+            det_map = get_dataset_det_map(
+                dataset_name=dataset_name, omni3d50=omni3d50
+            )
+
+            # create an individual dataset evaluator
+            self.evaluators[dataset_name] = Detect3DEvaluator(
+                det_map,
+                cat_map=omni3d_class_map,
+                annotation=annotation,
+                eval_prox=(
+                    "Objectron" in dataset_name or "SUNRGBD" in dataset_name
+                ),
+            )
+
+            self.overall_imgIds.update(
+                set(self.evaluators[dataset_name]._coco_gt.getImgIds())
+            )
+            self.overall_catIds.update(
+                set(self.evaluators[dataset_name]._coco_gt.getCatIds())
+            )
+
+    def __repr__(self) -> str:
+        """Returns the string representation of the object."""
+        datasets_str = ", ".join(self.dataset_names)
+        return f"Omni3DEvaluator ({datasets_str})"
+
+    @property
+    def metrics(self) -> list[str]:
+        """Supported metrics.
+
+        Returns:
+            list[str]: Metrics to evaluate.
+        """
+        return ["2D", "3D"]
+
+    def reset(self) -> None:
+        """Reset the saved predictions to start new round of evaluation."""
+        for dataset_name in self.dataset_names:
+            self.evaluators[dataset_name].reset()
+        self.evals_per_cat_area2D.clear()
+        self.evals_per_cat_area3D.clear()
+
+    def gather(self, gather_func: GenericFunc) -> None:
+        """Accumulate predictions across processes."""
+        for dataset_name in self.dataset_names:
+            self.evaluators[dataset_name].gather(gather_func)
+
+    def process_batch(
+        self,
+        coco_image_id: list[int],
+        dataset_names: list[str],
+        pred_boxes: list[NDArrayNumber],
+        pred_scores: list[NDArrayNumber],
+        pred_classes: list[NDArrayNumber],
+        pred_boxes3d: list[NDArrayNumber] | None = None,
+    ) -> None:
+        """Process sample and convert detections to coco format."""
+        for i, dataset_name in enumerate(dataset_names):
+            self.evaluators[dataset_name].process_batch(
+                [coco_image_id[i]],
+                [pred_boxes[i]],
+                [pred_scores[i]],
+                [pred_classes[i]],
+                pred_boxes3d=[pred_boxes3d[i]] if pred_boxes3d else None,
+            )
+
+    def evaluate(self, metric: str) -> tuple[MetricLogs, str]:
+        """Evaluate predictions and return the results."""
+        assert metric in self.metrics, f"Unsupported metric: {metric}"
+
+        log_dict = {}
+
+        for dataset_name in self.dataset_names:
+            rank_zero_info(f"Evaluating {dataset_name}...")
+            per_dataset_log_dict, dataset_log_str = self.evaluators[
+                dataset_name
+            ].evaluate(metric)
+
+            log_dict[f"AP_{dataset_name}"] = per_dataset_log_dict["AP"]
+
+            rank_zero_info(dataset_log_str + "\n")
+
+            # store the partially accumulated evaluations per category per area
+            if metric == "2D":
+                for key, item in self.evaluators[
+                    dataset_name
+                ].bbox_2D_evals_per_cat_area.items():
+                    if not key in self.evals_per_cat_area2D:
+                        self.evals_per_cat_area2D[key] = []
+                    self.evals_per_cat_area2D[key] += item
+            else:
+                for key, item in self.evaluators[
+                    dataset_name
+                ].bbox_3D_evals_per_cat_area.items():
+                    if not key in self.evals_per_cat_area3D:
+                        self.evals_per_cat_area3D[key] = []
+                    self.evals_per_cat_area3D[key] += item
+
+        results_per_category_dict = {}
+        results_per_category = []
+
+        rank_zero_info(f"Evaluating Omni3D for {metric} Detection...")
+
+        evaluator = Detect3Deval(mode=metric)
+        evaluator.params.catIds = list(self.overall_catIds)
+        evaluator.params.imgIds = list(self.overall_imgIds)
+        evaluator.evalImgs = True
+
+        if metric == "2D":
+            evaluator.evals_per_cat_area = self.evals_per_cat_area2D
+            metrics = ["AP", "AP50", "AP75", "AP95", "APs", "APm", "APl"]
+        else:
+            evaluator.evals_per_cat_area = self.evals_per_cat_area3D
+            metrics = ["AP", "AP15", "AP25", "AP50", "APn", "APm", "APf"]
+
+        evaluator._paramsEval = copy.deepcopy(evaluator.params)
+
+        with contextlib.redirect_stdout(io.StringIO()):
+            evaluator.accumulate()
+            log_str = "\n" + evaluator.summarize()
+
+        log_dict.update(dict(zip(metrics, evaluator.stats)))
+
+        if self.per_class_eval:
+            precisions = evaluator.eval["precision"]
+            for idx, cat_id in enumerate(self.overall_catIds):
+                cat_name = self.id_to_name[cat_id]
+                precision = precisions[:, :, idx, 0, -1]
+                precision = precision[precision > -1]
+                if precision.size:
+                    ap = float(np.mean(precision).item())
+                else:
+                    ap = float("nan")
+
+                results_per_category_dict[cat_name] = ap
+                results_per_category.append((f"{cat_name}", f"{ap:0.3f}"))
+
+            num_columns = min(6, len(results_per_category) * 2)
+            results_flatten = list(itertools.chain(*results_per_category))
+            headers = ["category", "AP"] * (num_columns // 2)
+            results_2d = itertools.zip_longest(
+                *[results_flatten[i::num_columns] for i in range(num_columns)]
+            )
+            table_data = [headers] + list(results_2d)
+            table = AsciiTable(table_data)
+            log_str = f"\n{table.table}\n{log_str}"
+
+        # Omni3D Outdoor performance
+        ap_out_lst = []
+        for cat in omni3d_out:
+            ap_out_lst.append(results_per_category_dict.get(cat, 0.0))
+
+        log_dict["Omni3D_Out"] = np.mean(ap_out_lst).item()
+
+        # Omni3D Indoor performance
+        ap_in_lst = []
+        for cat in omni3d_in:
+            ap_in_lst.append(results_per_category_dict.get(cat, 0.0))
+
+        log_dict["Omni3D_In"] = np.mean(ap_in_lst).item()
+
+        return log_dict, log_str
+
+    def save(self, metric: str, output_dir: str) -> None:
+        """Save the results to json files."""
+        for dataset_name in self.dataset_names:
+            self.evaluators[dataset_name].save(
+                metric, output_dir, prefix=dataset_name
+            )

opendet3d/eval/open.py

@@ -0,0 +1,140 @@
+"""Multi-data 3D detection evaluation."""
+
+from collections.abc import Sequence
+
+from vis4d.common.logging import rank_zero_info
+from vis4d.common.typing import GenericFunc, MetricLogs, NDArrayNumber
+from vis4d.eval.base import Evaluator
+
+from .detect3d import Detect3DEvaluator
+from .omni3d import Omni3DEvaluator
+
+
+class OpenDetect3DEvaluator(Evaluator):
+    """Multi-data 3D detection evaluator."""
+
+    def __init__(
+        self,
+        datasets: Sequence[str],
+        evaluators: Sequence[Detect3DEvaluator],
+        omni3d_evaluator: Omni3DEvaluator | None = None,
+    ) -> None:
+        """Initialize the evaluator."""
+        super().__init__()
+        self.dataset_names = datasets
+        self.evaluators = {
+            name: evaluator for name, evaluator in zip(datasets, evaluators)
+        }
+
+        self.omni3d_evaluator = omni3d_evaluator
+
+    def __repr__(self) -> str:
+        """Returns the string representation of the object."""
+        datasets_str = ", ".join(self.dataset_names)
+        return f"Open 3D Object Detection Evaluator ({datasets_str})"
+
+    @property
+    def metrics(self) -> list[str]:
+        """Supported metrics.
+
+        Returns:
+            list[str]: Metrics to evaluate.
+        """
+        return ["2D", "3D"]
+
+    def reset(self) -> None:
+        """Reset the saved predictions to start new round of evaluation."""
+        for dataset_name in self.dataset_names:
+            self.evaluators[dataset_name].reset()
+
+        if self.omni3d_evaluator is not None:
+            self.omni3d_evaluator.reset()
+
+    def gather(self, gather_func: GenericFunc) -> None:
+        """Accumulate predictions across processes."""
+        for dataset_name in self.dataset_names:
+            self.evaluators[dataset_name].gather(gather_func)
+
+        if self.omni3d_evaluator is not None:
+            self.omni3d_evaluator.gather(gather_func)
+
+    def process_batch(
+        self,
+        coco_image_id: list[int],
+        dataset_names: list[str],
+        pred_boxes: list[NDArrayNumber],
+        pred_scores: list[NDArrayNumber],
+        pred_classes: list[NDArrayNumber],
+        pred_boxes3d: list[NDArrayNumber] | None = None,
+    ) -> None:
+        """Process sample and convert detections to coco format."""
+        for i, dataset_name in enumerate(dataset_names):
+            if (
+                self.omni3d_evaluator is not None
+                and dataset_name in self.omni3d_evaluator.dataset_names
+            ):
+                self.omni3d_evaluator.process_batch(
+                    [coco_image_id[i]],
+                    [dataset_name],
+                    [pred_boxes[i]],
+                    [pred_scores[i]],
+                    [pred_classes[i]],
+                    pred_boxes3d=[pred_boxes3d[i]] if pred_boxes3d else None,
+                )
+            else:
+                self.evaluators[dataset_name].process_batch(
+                    [coco_image_id[i]],
+                    [pred_boxes[i]],
+                    [pred_scores[i]],
+                    [pred_classes[i]],
+                    pred_boxes3d=[pred_boxes3d[i]] if pred_boxes3d else None,
+                )
+
+    def evaluate(self, metric: str) -> tuple[MetricLogs, str]:
+        """Evaluate predictions and return the results."""
+        assert metric in self.metrics, f"Unsupported metric: {metric}"
+
+        log_dict = {}
+        log_str = ""
+
+        if self.omni3d_evaluator is not None:
+            log_dict_omni3d, omni3d_log_str = self.omni3d_evaluator.evaluate(
+                metric
+            )
+
+            log_dict.update(log_dict_omni3d)
+            log_str += omni3d_log_str
+
+        for dataset_name in self.dataset_names:
+            rank_zero_info(f"Evaluating {dataset_name}...")
+            per_dataset_log_dict, dataset_log_str = self.evaluators[
+                dataset_name
+            ].evaluate(metric)
+
+            if "ODS" in per_dataset_log_dict:
+                score = "ODS"
+            else:
+                score = "AP"
+
+            log_dict[f"{score}_{dataset_name}"] = per_dataset_log_dict[score]
+
+            if self.evaluators[dataset_name].base_classes is not None:
+                log_dict[f"{score}_Base_{dataset_name}"] = (
+                    per_dataset_log_dict[f"{score}_Base"]
+                )
+                log_dict[f"{score}_Novel_{dataset_name}"] = (
+                    per_dataset_log_dict[f"{score}_Novel"]
+                )
+
+            log_str += f"\nCheck {dataset_name} results in log dict."
+
+            rank_zero_info(dataset_log_str + "\n")
+
+        return log_dict, log_str
+
+    def save(self, metric: str, output_dir: str) -> None:
+        """Save the results to json files."""
+        for dataset_name in self.dataset_names:
+            self.evaluators[dataset_name].save(
+                metric, output_dir, prefix=dataset_name
+            )

opendet3d/model/detect/grounding_dino.py

@@ -0,0 +1,1050 @@
+"""Grounding DINO model.
+
+modified from mmdetection.
+"""
+
+from collections.abc import Sequence
+from typing import NamedTuple
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+from transformers import BatchEncoding
+from vis4d.common.ckpt import load_model_checkpoint
+from vis4d.common.logging import rank_zero_warn
+from vis4d.op.base import BaseModel
+from vis4d.op.layer.positional_encoding import SinePositionalEncoding
+
+from opendet3d.model.language.mm_bert import BertModel
+from opendet3d.op.detect.deformable_detr import get_valid_ratio
+from opendet3d.op.detect.dino import CdnQueryGenerator
+from opendet3d.op.detect.grounding_dino import (
+    GroundingDINOHead,
+    GroundingDinoTransformerDecoder,
+    GroundingDinoTransformerEncoder,
+    RoI2Det,
+)
+from opendet3d.op.fpp.channel_mapper import ChannelMapper
+from opendet3d.op.language.grounding import (
+    chunks,
+    clean_label_name,
+    create_positive_map,
+    create_positive_map_label_to_token,
+    run_ner,
+)
+
+REV_KEYS = [
+    (r"\.conv.weight", ".weight"),
+    (r"\.conv.bias", ".bias"),
+    (r"\.gn", ".norm"),
+]
+
+
+class GroundingDINOOut(NamedTuple):
+    """Output of the Grounding DINO model."""
+
+    all_layers_cls_scores: list[Tensor]
+    all_layers_bbox_preds: list[Tensor]
+    enc_outputs_class: Tensor
+    enc_outputs_coord: Tensor
+    text_token_mask: Tensor
+    dn_meta: dict[str, Tensor]
+    positive_maps: list[Tensor]
+
+
+class DetOut(NamedTuple):
+    """Output of the Grounding DINO model."""
+
+    boxes: list[Tensor]
+    scores: list[Tensor]
+    class_ids: list[Tensor]
+    categories: list[list[str]] | None = None
+
+
+class GroundingDINO(nn.Module):
+    """Grounding DINO."""
+
+    def __init__(
+        self,
+        basemodel: BaseModel,
+        neck: ChannelMapper,
+        texts: list[str] | None = None,
+        custom_entities: bool = True,
+        chunked_size: int = -1,
+        num_queries: int = 900,
+        num_feature_levels: int = 4,
+        use_checkpoint: bool = False,
+        bbox_head: GroundingDINOHead | None = None,
+        language_model: BertModel | None = None,
+        roi2det: RoI2Det | None = None,
+        weights: str | None = None,
+    ) -> None:
+        """Create the Grounding DINO model."""
+        super().__init__()
+        self.texts = texts
+        self.custom_entities = custom_entities
+        self.chunked_size = chunked_size
+
+        self.num_queries = num_queries
+
+        self.backbone = basemodel
+        self.neck = neck
+
+        # Encoder
+        self.encoder = GroundingDinoTransformerEncoder(
+            num_levels=num_feature_levels, use_checkpoint=use_checkpoint
+        )
+
+        self.embed_dims = self.encoder.embed_dims
+        self.positional_encoding = SinePositionalEncoding(
+            num_feats=128, normalize=True, offset=0.0, temperature=20
+        )
+
+        num_feats = self.positional_encoding.num_feats
+        assert num_feats * 2 == self.embed_dims, (
+            f"embed_dims should be exactly 2 times of num_feats. "
+            f"Found {self.embed_dims} and {num_feats}."
+        )
+
+        self.level_embed = nn.Parameter(
+            torch.Tensor(num_feature_levels, self.embed_dims)
+        )
+
+        self.memory_trans_fc = nn.Linear(self.embed_dims, self.embed_dims)
+        self.memory_trans_norm = nn.LayerNorm(self.embed_dims)
+
+        # Decoder
+        self.decoder = GroundingDinoTransformerDecoder(
+            num_levels=num_feature_levels
+        )
+
+        self.query_embedding = nn.Embedding(self.num_queries, self.embed_dims)
+
+        # Grounding DINO head
+        self.bbox_head = bbox_head or GroundingDINOHead(
+            num_classes=256, num_decoder_layer=self.decoder.num_layers
+        )
+
+        self.roi2det = roi2det or RoI2Det()
+
+        self.dn_query_generator = CdnQueryGenerator(
+            num_classes=self.bbox_head.num_classes,
+            embed_dims=self.embed_dims,
+            num_matching_queries=self.num_queries,
+            label_noise_scale=0.5,
+            box_noise_scale=1.0,  # 0.4 for DN-DETR
+            dynamic=True,
+            num_groups=None,
+            num_dn_queries=100,
+        )
+
+        # Language model configuration
+        self._special_tokens = ". "
+
+        # text modules
+        self.language_model = language_model or BertModel(
+            name="bert-base-uncased",
+            max_tokens=256,
+            pad_to_max=False,
+            use_sub_sentence_represent=True,
+            special_tokens_list=["[CLS]", "[SEP]", ".", "?"],
+            add_pooling_layer=False,
+            use_checkpoint=use_checkpoint,
+        )
+
+        self.text_feat_map = nn.Linear(
+            self.language_model.language_backbone.body.language_dim,
+            self.embed_dims,
+            bias=True,
+        )
+
+        self._init_weights()
+
+        if weights is not None:
+            load_model_checkpoint(self, weights, rev_keys=REV_KEYS)
+
+    def _init_weights(self) -> None:
+        """Initialize weights."""
+        # DINO
+        for coder in self.encoder, self.decoder:
+            for p in coder.parameters():
+                if p.dim() > 1:
+                    nn.init.xavier_uniform_(p)
+
+        nn.init.xavier_uniform_(self.memory_trans_fc.weight)
+        nn.init.xavier_uniform_(self.query_embedding.weight)
+        nn.init.normal_(self.level_embed)
+
+        # G-DINO
+        nn.init.constant_(self.text_feat_map.bias.data, 0)
+        nn.init.xavier_uniform_(self.text_feat_map.weight.data)
+
+    def get_captions_and_tokens_positive(
+        self,
+        text_prompt: list[str],
+        text_prompt_mapping: dict[str, dict[str, str]] | None = None,
+    ) -> tuple[str, list[list[int]]]:
+        """Enhance the text prompts with the text mapping."""
+        captions = ""
+        tokens_positive = []
+        for word in text_prompt:
+            if text_prompt_mapping is not None and word in text_prompt_mapping:
+                enhanced_text_dict = text_prompt_mapping[word]
+                if "prefix" in enhanced_text_dict:
+                    captions += enhanced_text_dict["prefix"]
+
+                start_i = len(captions)
+                if "name" in enhanced_text_dict:
+                    captions += enhanced_text_dict["name"]
+                else:
+                    captions += word
+                end_i = len(captions)
+
+                tokens_positive.append([[start_i, end_i]])
+
+                if "suffix" in enhanced_text_dict:
+                    captions += enhanced_text_dict["suffix"]
+            else:
+                tokens_positive.append(
+                    [[len(captions), len(captions) + len(word)]]
+                )
+                captions += word
+            captions += self._special_tokens
+        return captions, tokens_positive
+
+    def get_tokens_and_prompts(
+        self,
+        text_prompt: str | list[str],
+        text_prompt_mapping: dict[str, dict[str, str]] | None = None,
+    ) -> tuple[BatchEncoding, str, list[list[int]], list[str]]:
+        """Get the tokens positive and prompts for the caption."""
+        if isinstance(text_prompt, list):
+            captions, tokens_positive = self.get_captions_and_tokens_positive(
+                text_prompt, text_prompt_mapping
+            )
+
+            tokenized = self.language_model.tokenizer(
+                [captions], padding="longest", return_tensors="pt"
+            )
+            entities = text_prompt
+        else:
+            if not text_prompt.endswith("."):
+                captions = text_prompt + self._special_tokens
+            else:
+                captions = text_prompt
+
+            tokenized = self.language_model.tokenizer(
+                [captions], padding="longest", return_tensors="pt"
+            )
+            tokens_positive, entities = run_ner(captions)
+
+        return tokenized, captions, tokens_positive, entities
+
+    def get_positive_map(
+        self, tokenized: BatchEncoding, tokens_positive: list[list[int]]
+    ) -> tuple[dict, Tensor]:
+        """Get the positive map and label to token."""
+        positive_map = create_positive_map(
+            tokenized,
+            tokens_positive,
+            max_num_entities=self.bbox_head.cls_branches[
+                self.decoder.num_layers
+            ].max_text_len,
+        )
+        positive_map_label_to_token = create_positive_map_label_to_token(
+            positive_map, plus=1
+        )
+        return positive_map_label_to_token, positive_map
+
+    def get_tokens_positive_and_prompts(
+        self,
+        text_prompt: str | list[str],
+        custom_entities: bool = False,
+        tokens_positive: list[list[int, int]] | int | None = None,
+        text_prompt_mapping: dict[str, dict[str, str]] | None = None,
+    ) -> tuple[dict, str, Tensor, list]:
+        """Get the tokens positive and prompts for the caption.
+
+        Args:
+            original_caption (str): The original caption, e.g. 'bench . car .'
+            custom_entities (bool, optional): Whether to use custom entities.
+                If ``True``, the ``original_caption`` should be a list of
+                strings, each of which is a word. Defaults to False.
+
+        Returns:
+            Tuple[dict, str, dict, str]: The dict is a mapping from each entity
+            id, which is numbered from 1, to its positive token id.
+            The str represents the prompts.
+        """
+        if tokens_positive is not None:
+            assert isinstance(
+                text_prompt, str
+            ), "Text prompt should be a string with given positive tokens."
+
+            if not text_prompt.endswith("."):
+                captions = text_prompt + self._special_tokens
+            else:
+                captions = text_prompt
+
+            if tokens_positive == -1:
+                return None, captions, None, captions
+            else:
+                assert isinstance(
+                    tokens_positive, list
+                ), "Positive tokens should be a list of list[int] if not -1."
+                tokenized = self.language_model.tokenizer(
+                    [captions], padding="longest", return_tensors="pt"
+                )
+                positive_map_label_to_token, positive_map = (
+                    self.get_positive_map(tokenized, tokens_positive)
+                )
+
+                entities = []
+                for token_positive in tokens_positive:
+                    instance_entities = []
+                    for t in token_positive:
+                        instance_entities.append(captions[t[0] : t[1]])
+                    entities.append(" / ".join(instance_entities))
+
+                return (
+                    positive_map_label_to_token,
+                    captions,
+                    positive_map,
+                    entities,
+                )
+
+        if custom_entities:
+            if isinstance(text_prompt, str):
+                text_prompt = text_prompt.strip(self._special_tokens)
+                text_prompt = text_prompt.split(self._special_tokens)
+                text_prompt = list(filter(lambda x: len(x) > 0, text_prompt))
+            text_prompt = [clean_label_name(i) for i in text_prompt]
+
+        if self.chunked_size > 0:
+            assert not self.training, "Chunked size is only for testing."
+            (
+                positive_map_label_to_token,
+                captions,
+                positive_map,
+                entities,
+            ) = self.get_tokens_positive_and_prompts_chunked(
+                text_prompt, text_prompt_mapping
+            )
+        else:
+            tokenized, captions, tokens_positive, entities = (
+                self.get_tokens_and_prompts(text_prompt, text_prompt_mapping)
+            )
+            positive_map_label_to_token, positive_map = self.get_positive_map(
+                tokenized, tokens_positive
+            )
+
+        return positive_map_label_to_token, captions, positive_map, entities
+
+    def get_tokens_positive_and_prompts_chunked(
+        self,
+        text_prompt: list[str],
+        text_prompt_mapping: dict[str, dict[str, str]] | None = None,
+    ):
+        """Get the tokens positive and prompts for the caption."""
+        text_prompt_chunked = chunks(text_prompt, self.chunked_size)
+        ids_chunked = chunks(
+            list(range(1, len(text_prompt) + 1)), self.chunked_size
+        )
+
+        positive_map_label_to_token_chunked = []
+        captions_chunked = []
+        positive_map_chunked = []
+        entities_chunked = []
+        for i in range(len(ids_chunked)):
+            captions, tokens_positive = self.get_captions_and_tokens_positive(
+                text_prompt_chunked[i], text_prompt_mapping
+            )
+
+            tokenized = self.language_model.tokenizer(
+                [captions], padding="longest", return_tensors="pt"
+            )
+            if tokenized.input_ids.shape[1] > self.language_model.max_tokens:
+                rank_zero_warn(
+                    "Caption is too long will result in poor performance. "
+                    "Please reduce the chunked size."
+                )
+
+            positive_map_label_to_token, positive_map = self.get_positive_map(
+                tokenized, tokens_positive
+            )
+
+            captions_chunked.append(captions)
+            positive_map_label_to_token_chunked.append(
+                positive_map_label_to_token
+            )
+            positive_map_chunked.append(positive_map)
+            entities_chunked.append(text_prompt_chunked[i])
+
+        return (
+            positive_map_label_to_token_chunked,
+            captions_chunked,
+            positive_map_chunked,
+            entities_chunked,
+        )
+
+    # TODO: Move this to deformable DETR
+    def pre_transformer(
+        self,
+        feats: list[Tensor],
+        input_hw: list[tuple[int, int]],
+        batch_input_shape: tuple[int, int],
+        padding: list[list[int]] | None = None,
+    ) -> tuple[Tensor, Tensor, Tensor | None, Tensor, Tensor, Tensor]:
+        """Process image features before transformer."""
+        batch_size = feats[0].size(0)
+
+        # construct binary masks for the transformer.
+        batch_input_img_h, batch_input_img_w = batch_input_shape
+        same_shape_flag = all(
+            [
+                s[0] == batch_input_img_h and s[1] == batch_input_img_w
+                for s in input_hw
+            ]
+        )
+
+        if same_shape_flag:
+            mlvl_masks = []
+            mlvl_pos_embeds = []
+            for feat in feats:
+                mlvl_masks.append(None)
+                mlvl_pos_embeds.append(
+                    self.positional_encoding(None, inputs=feat)
+                )
+        else:
+            check_center = not (padding is None)
+            masks = feats[0].new_ones(
+                (batch_size, batch_input_img_h, batch_input_img_w)
+            )
+            for img_id in range(batch_size):
+                img_h, img_w = input_hw[img_id]
+
+                if padding is None:
+                    masks[img_id, :img_h, :img_w] = 0
+                else:
+                    pad_left, pad_right, pad_top, pad_bottom = padding[img_id]
+                    masks[
+                        img_id,
+                        pad_top : batch_input_img_h - pad_bottom,
+                        pad_left : batch_input_img_w - pad_right,
+                    ] = 0
+
+            # NOTE following the official DETR repo, non-zero
+            # values representing ignored positions, while
+            # zero values means valid positions.
+            mlvl_masks = []
+            mlvl_pos_embeds = []
+            for feat in feats:
+                mlvl_masks.append(
+                    F.interpolate(masks[None], size=feat.shape[-2:])
+                    .to(torch.bool)
+                    .squeeze(0)
+                )
+                mlvl_pos_embeds.append(
+                    self.positional_encoding(mlvl_masks[-1])
+                )
+
+        feat_flatten = []
+        lvl_pos_embed_flatten = []
+        mask_flatten = []
+        spatial_shapes = []
+        for lvl, (feat, mask, pos_embed) in enumerate(
+            zip(feats, mlvl_masks, mlvl_pos_embeds)
+        ):
+            batch_size, c, _, _ = feat.shape
+            spatial_shape = torch._shape_as_tensor(feat)[2:].to(feat.device)
+            # [bs, c, h_lvl, w_lvl] -> [bs, h_lvl*w_lvl, c]
+            feat = feat.view(batch_size, c, -1).permute(0, 2, 1)
+            pos_embed = pos_embed.view(batch_size, c, -1).permute(0, 2, 1)
+            lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1)
+            # [bs, h_lvl, w_lvl] -> [bs, h_lvl*w_lvl]
+            if mask is not None:
+                mask = mask.flatten(1)
+
+            feat_flatten.append(feat)
+            lvl_pos_embed_flatten.append(lvl_pos_embed)
+            mask_flatten.append(mask)
+            spatial_shapes.append(spatial_shape)
+
+        # (bs, num_feat_points, dim)
+        feat_flatten = torch.cat(feat_flatten, 1)
+        lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)
+        # (bs, num_feat_points), where num_feat_points = sum_lvl(h_lvl*w_lvl)
+        if mask_flatten[0] is not None:
+            mask_flatten = torch.cat(mask_flatten, 1)
+        else:
+            mask_flatten = None
+
+        # (num_level, 2)
+        spatial_shapes = torch.cat(spatial_shapes).view(-1, 2)
+        level_start_index = torch.cat(
+            (
+                spatial_shapes.new_zeros((1,)),  # (num_level)
+                spatial_shapes.prod(1).cumsum(0)[:-1],
+            )
+        )
+        if mlvl_masks[0] is not None:
+            valid_ratios = torch.stack(  # (bs, num_level, 2)
+                [get_valid_ratio(m, check_center) for m in mlvl_masks],
+                1,
+            )
+        else:
+            valid_ratios = feats[0].new_ones(batch_size, len(feats), 2)
+
+        return (
+            feat_flatten,
+            lvl_pos_embed_flatten,
+            mask_flatten,
+            spatial_shapes,
+            level_start_index,
+            valid_ratios,
+        )
+
+    def forward_transformer(
+        self,
+        feat_flatten: Tensor,
+        lvl_pos_embed_flatten: Tensor,
+        memory_mask: Tensor | None,
+        spatial_shapes: Tensor,
+        level_start_index: Tensor,
+        valid_ratios: Tensor,
+        text_dict: dict[str, Tensor],
+        boxes: Tensor | None = None,
+        class_ids: Tensor | None = None,
+        input_hw: list[tuple[int, int]] | None = None,
+    ) -> tuple[Tensor, Tensor, Tensor, list[Tensor]]:
+        """Forward function for the transformer."""
+        text_token_mask = text_dict["text_token_mask"]
+
+        memory, memory_text = self.encoder(
+            query=feat_flatten,
+            query_pos=lvl_pos_embed_flatten,
+            key_padding_mask=memory_mask,
+            spatial_shapes=spatial_shapes,
+            level_start_index=level_start_index,
+            valid_ratios=valid_ratios,
+            memory_text=text_dict["embedded"],
+            text_attention_mask=~text_token_mask,
+            position_ids=text_dict["position_ids"],
+            text_self_attention_masks=text_dict["masks"],
+        )
+
+        bs = memory.shape[0]
+
+        output_memory, output_proposals = self.gen_encoder_output_proposals(
+            memory, memory_mask, spatial_shapes
+        )
+
+        enc_outputs_class = self.bbox_head.cls_branches[
+            self.decoder.num_layers
+        ](output_memory, memory_text, text_token_mask)
+        cls_out_features = self.bbox_head.cls_branches[
+            self.decoder.num_layers
+        ].max_text_len
+        enc_outputs_coord_unact = (
+            self.bbox_head.reg_branches[self.decoder.num_layers](output_memory)
+            + output_proposals
+        )
+
+        # NOTE The DINO selects top-k proposals according to scores of
+        # multi-class classification, while DeformDETR, where the input
+        # is `enc_outputs_class[..., 0]` selects according to scores of
+        # binary classification.
+        topk_indices = torch.topk(
+            enc_outputs_class.max(-1)[0], k=self.num_queries, dim=1
+        )[1]
+
+        topk_score = torch.gather(
+            enc_outputs_class,
+            1,
+            topk_indices.unsqueeze(-1).repeat(1, 1, cls_out_features),
+        )
+        topk_coords_unact = torch.gather(
+            enc_outputs_coord_unact,
+            1,
+            topk_indices.unsqueeze(-1).repeat(1, 1, 4),
+        )
+        topk_coords = topk_coords_unact.sigmoid()
+        topk_coords_unact = topk_coords_unact.detach()
+
+        query = self.query_embedding.weight[:, None, :]
+        query = query.repeat(1, bs, 1).transpose(0, 1)
+
+        if self.training:
+            dn_label_query, dn_bbox_query, dn_mask, dn_meta = (
+                self.dn_query_generator(boxes, class_ids, input_hw)
+            )
+            query = torch.cat([dn_label_query, query], dim=1)
+            reference_points = torch.cat(
+                [dn_bbox_query, topk_coords_unact], dim=1
+            )
+        else:
+            reference_points = topk_coords_unact
+            dn_mask, dn_meta = None, None
+
+        reference_points = reference_points.sigmoid()
+
+        # NOTE DINO calculates encoder losses on scores and coordinates
+        # of selected top-k encoder queries, while DeformDETR is of all
+        # encoder queries.
+        if self.training:
+            enc_outputs_class = topk_score
+            enc_outputs_coord = topk_coords
+
+        hidden_states, references = self.decoder(
+            query=query,
+            value=memory,
+            key_padding_mask=memory_mask,
+            self_attn_mask=dn_mask,
+            reference_points=reference_points,
+            spatial_shapes=spatial_shapes,
+            level_start_index=level_start_index,
+            valid_ratios=valid_ratios,
+            reg_branches=self.bbox_head.reg_branches,
+            memory_text=memory_text,
+            text_attention_mask=~text_token_mask,
+        )
+
+        if len(query) == self.num_queries:
+            # NOTE: This is to make sure label_embeding can be involved to
+            # produce loss even if there is no denoising query (no ground truth
+            # target in this GPU), otherwise, this will raise runtime error in
+            # distributed training.
+            hidden_states[0] += (
+                self.dn_query_generator.label_embedding.weight[0, 0] * 0.0
+            )
+
+        if self.training:
+            return (
+                memory_text,
+                text_token_mask,
+                hidden_states,
+                list(references),
+                enc_outputs_class,
+                enc_outputs_coord,
+                dn_meta,
+            )
+
+        return memory_text, text_token_mask, hidden_states, list(references)
+
+    # TODO: Move this to deformable DETR
+    def gen_encoder_output_proposals(
+        self,
+        memory: Tensor,
+        memory_mask: Tensor | None,
+        spatial_shapes: Tensor,
+    ) -> tuple[Tensor, Tensor]:
+        """Generate proposals from encoded memory. The function will only be
+        used when `as_two_stage` is `True`.
+
+        Args:
+            memory (Tensor): The output embeddings of the Transformer encoder,
+                has shape (bs, num_feat_points, dim).
+            memory_mask (Tensor): ByteTensor, the padding mask of the memory,
+                has shape (bs, num_feat_points).
+            spatial_shapes (Tensor): Spatial shapes of features in all levels,
+                has shape (num_levels, 2), last dimension represents (h, w).
+
+        Returns:
+            tuple: A tuple of transformed memory and proposals.
+
+            - output_memory (Tensor): The transformed memory for obtaining
+              top-k proposals, has shape (bs, num_feat_points, dim).
+            - output_proposals (Tensor): The inverse-normalized proposal, has
+              shape (batch_size, num_keys, 4) with the last dimension arranged
+              as (cx, cy, w, h).
+        """
+
+        bs = memory.size(0)
+        proposals = []
+        _cur = 0  # start index in the sequence of the current level
+        for lvl, HW in enumerate(spatial_shapes):
+            H, W = HW
+
+            if memory_mask is not None:
+                mask_flatten_ = memory_mask[:, _cur : (_cur + H * W)].view(
+                    bs, H, W, 1
+                )
+                valid_H = torch.sum(~mask_flatten_[:, :, 0, 0], 1).unsqueeze(
+                    -1
+                )
+                valid_W = torch.sum(~mask_flatten_[:, 0, :, 0], 1).unsqueeze(
+                    -1
+                )
+                scale = torch.cat([valid_W, valid_H], 1).view(bs, 1, 1, 2)
+            else:
+                if not isinstance(HW, Tensor):
+                    HW = memory.new_tensor(HW)
+                scale = HW.unsqueeze(0).flip(dims=[0, 1]).view(1, 1, 1, 2)
+
+            grid_y, grid_x = torch.meshgrid(
+                torch.linspace(
+                    0, H - 1, H, dtype=torch.float32, device=memory.device
+                ),
+                torch.linspace(
+                    0, W - 1, W, dtype=torch.float32, device=memory.device
+                ),
+                indexing="ij",
+            )
+            grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1)
+            grid = (grid.unsqueeze(0).expand(bs, -1, -1, -1) + 0.5) / scale
+            wh = torch.ones_like(grid) * 0.05 * (2.0**lvl)
+            proposal = torch.cat((grid, wh), -1).view(bs, -1, 4)
+            proposals.append(proposal)
+            _cur += H * W
+
+        output_proposals = torch.cat(proposals, 1)
+
+        # do not use `all` to make it exportable to onnx
+        output_proposals_valid = (
+            (output_proposals > 0.01) & (output_proposals < 0.99)
+        ).sum(-1, keepdim=True) == output_proposals.shape[-1]
+
+        # inverse_sigmoid
+        output_proposals = torch.log(output_proposals / (1 - output_proposals))
+        if memory_mask is not None:
+            output_proposals = output_proposals.masked_fill(
+                memory_mask.unsqueeze(-1), float("inf")
+            )
+
+        output_proposals = output_proposals.masked_fill(
+            ~output_proposals_valid, float("inf")
+        )
+
+        if memory_mask is not None:
+            output_memory = memory.masked_fill(
+                memory_mask.unsqueeze(-1), float(0)
+            )
+        else:
+            output_memory = memory
+
+        # [bs, sum(hw), 2]
+        output_memory = output_memory.masked_fill(
+            ~output_proposals_valid, float(0)
+        )
+        output_memory = self.memory_trans_fc(output_memory)
+        output_memory = self.memory_trans_norm(output_memory)
+
+        return output_memory, output_proposals
+
+    def _forward_train(
+        self,
+        images: Tensor,
+        input_texts: list[list[str]] | None,
+        boxes2d: Tensor,
+        boxes2d_classes: Tensor,
+        input_hw: list[tuple[int, int]],
+        input_tokens_positive: list[list[int, int]] | None = None,
+    ) -> GroundingDINOOut:
+        """Forward train."""
+        batch_size = images.shape[0]
+
+        # if "tokens_positive" in batch_data_samples[0]:
+        if input_tokens_positive is not None:
+            positive_maps = []
+            for tokens_positive_dict, text_prompt, gt_label in zip(
+                input_tokens_positive, input_texts, boxes2d_classes
+            ):
+                tokenized = self.language_model.tokenizer(
+                    [text_prompt], padding="longest", return_tensors="pt"
+                )
+                new_tokens_positive = [
+                    tokens_positive_dict[label.item()] for label in gt_label
+                ]
+                _, positive_map = self.get_positive_map(
+                    tokenized, new_tokens_positive
+                )
+                positive_maps.append(positive_map)
+            new_text_prompts = input_texts
+        else:
+            new_text_prompts = []
+            positive_maps = []
+
+            # All the text prompts are the same, so there is no need to
+            # calculate them multiple times.
+            if (
+                input_texts is None
+                or len(set(["".join(t) for t in input_texts])) == 1
+            ):
+                if input_texts is None:
+                    text_prompt = self.texts
+                else:
+                    text_prompt = input_texts[0]
+
+                tokenized, caption_string, tokens_positive, _ = (
+                    self.get_tokens_and_prompts(text_prompt)
+                )
+                new_text_prompts = [caption_string] * batch_size
+                for gt_label in boxes2d_classes:
+                    new_tokens_positive = [
+                        tokens_positive[label] for label in gt_label
+                    ]
+                    _, positive_map = self.get_positive_map(
+                        tokenized, new_tokens_positive
+                    )
+                    positive_maps.append(positive_map)
+            else:
+                for text_prompt, gt_label in zip(input_texts, boxes2d_classes):
+                    tokenized, caption_string, tokens_positive, _ = (
+                        self.get_tokens_and_prompts(text_prompt)
+                    )
+
+                    new_text_prompts.append(caption_string)
+
+                    new_tokens_positive = [
+                        tokens_positive[label] for label in gt_label
+                    ]
+                    _, positive_map = self.get_positive_map(
+                        tokenized, new_tokens_positive
+                    )
+
+                    positive_maps.append(positive_map)
+
+        for i in range(batch_size):
+            positive_maps[i] = (
+                positive_maps[i].to(images.device).bool().float()
+            )
+
+        text_dict = self.language_model(new_text_prompts)
+
+        if self.text_feat_map is not None:
+            text_dict["embedded"] = self.text_feat_map(text_dict["embedded"])
+
+        text_token_masks = []
+        for i in range(batch_size):
+            text_token_masks.append(
+                text_dict["text_token_mask"][i]
+                .unsqueeze(0)
+                .repeat(len(positive_map), 1)
+            )
+
+        visual_feats = self.backbone(images)[2:]
+        visual_feats = self.neck(visual_feats)
+
+        batch_input_img_h, batch_input_img_w = images.shape[-2:]
+        batch_input_shape = (batch_input_img_h, batch_input_img_w)
+
+        (
+            feat_flatten,
+            lvl_pos_embed_flatten,
+            memory_mask,
+            spatial_shapes,
+            level_start_index,
+            valid_ratios,
+        ) = self.pre_transformer(visual_feats, input_hw, batch_input_shape)
+
+        (
+            memory_text,
+            text_token_mask,
+            hidden_states,
+            references,
+            enc_outputs_class,
+            enc_outputs_coord,
+            dn_meta,
+        ) = self.forward_transformer(
+            feat_flatten,
+            lvl_pos_embed_flatten,
+            memory_mask,
+            spatial_shapes,
+            level_start_index,
+            valid_ratios,
+            text_dict,
+            boxes2d,
+            boxes2d_classes,
+            input_hw,
+        )
+
+        all_layers_cls_scores, all_layers_bbox_preds = self.bbox_head(
+            hidden_states, references, memory_text, text_token_mask
+        )
+
+        return GroundingDINOOut(
+            all_layers_cls_scores,
+            all_layers_bbox_preds,
+            enc_outputs_class,
+            enc_outputs_coord,
+            text_token_mask,
+            dn_meta,
+            positive_maps,
+        )
+
+    def _forward_test(
+        self,
+        images: Tensor,
+        input_texts: list[str] | None,
+        text_prompt_mapping: list[dict[str, dict[str, str]]] | None,
+        input_hw: list[tuple[int, int]],
+        original_hw: list[tuple[int, int]],
+    ) -> DetOut:
+        """Forward."""
+        batch_size = images.shape[0]
+
+        token_positive_maps = []
+        text_prompts = []
+        entities = []
+        for i in range(batch_size):
+            if self.texts is not None:
+                text_prompt = self.texts
+            else:
+                text_prompt = input_texts[i]
+
+            if text_prompt_mapping is not None:
+                prompt_mapping = text_prompt_mapping[i]
+            else:
+                prompt_mapping = None
+
+            token_positive_map, captions, _, _entities = (
+                self.get_tokens_positive_and_prompts(
+                    text_prompt,
+                    self.custom_entities,
+                    text_prompt_mapping=prompt_mapping,
+                )
+            )
+
+            token_positive_maps.append(token_positive_map)
+            text_prompts.append(captions)
+            entities.append(_entities)
+
+        # image feature extraction
+        batch_input_img_h, batch_input_img_w = images.shape[-2:]
+
+        visual_feats = self.backbone(images)[2:]
+        visual_feats = self.neck(visual_feats)
+
+        if isinstance(text_prompts[0], list):
+            # TODO: Support chunked text prompts in the future.
+            pass
+            # assert batch_size == 1, "Batch size should be 1 for chunked text."
+            # count = 0
+            # results_list = []
+
+            # entities = [[item for lst in entities[0] for item in lst]]
+
+            # for i, captions in enumerate(text_prompts[0]):
+            #     token_positive_map = token_positive_maps[0][i]
+
+            #     text_dict = self.language_model(captions)
+
+            #     # text feature map layer
+            #     if self.text_feat_map is not None:
+            #         text_dict["embedded"] = self.text_feat_map(
+            #             text_dict["embedded"]
+            #         )
+
+            #     head_inputs_dict = self.forward_transformer(
+            #         copy.deepcopy(visual_feats), text_dict, input_hw
+            #     )
+            #     pred_instances = self.bbox_head.predict(
+            #         **head_inputs_dict,
+            #         batch_token_positive_maps=token_positive_map,
+            #     )[0]
+
+            #     if len(pred_instances) > 0:
+            #         pred_instances.labels += count
+            #     count += len(token_positive_maps_once)
+            #     results_list.append(pred_instances)
+            # results_list = [results_list[0].cat(results_list)]
+        else:
+            # extract text feats
+            text_dict = self.language_model(list(text_prompts))
+
+            # text feature map layer
+            if self.text_feat_map is not None:
+                text_dict["embedded"] = self.text_feat_map(
+                    text_dict["embedded"]
+                )
+
+            batch_input_shape = (batch_input_img_h, batch_input_img_w)
+
+            (
+                feat_flatten,
+                lvl_pos_embed_flatten,
+                memory_mask,
+                spatial_shapes,
+                level_start_index,
+                valid_ratios,
+            ) = self.pre_transformer(visual_feats, input_hw, batch_input_shape)
+
+            (
+                memory_text,
+                text_token_mask,
+                hidden_states,
+                references,
+            ) = self.forward_transformer(
+                feat_flatten,
+                lvl_pos_embed_flatten,
+                memory_mask,
+                spatial_shapes,
+                level_start_index,
+                valid_ratios,
+                text_dict,
+            )
+
+            all_layers_cls_scores, all_layers_bbox_preds = self.bbox_head(
+                hidden_states, references, memory_text, text_token_mask
+            )
+
+            cls_scores = all_layers_cls_scores[-1]
+            bbox_preds = all_layers_bbox_preds[-1]
+
+            boxes = []
+            scores = []
+            class_ids = []
+            categories = []
+            for i, bbox_pred in enumerate(bbox_preds):
+                cls_score = cls_scores[i]
+                det_bboxes, det_scores, det_labels = self.roi2det(
+                    cls_score,
+                    bbox_pred,
+                    token_positive_maps[i],
+                    input_hw[i],
+                    original_hw[i],
+                )
+                boxes.append(det_bboxes)
+                scores.append(det_scores)
+                class_ids.append(det_labels)
+
+                # Get the categories text
+                cur_categories = []
+                for label in det_labels:
+                    cur_categories.append(entities[i][label])
+
+                categories.append(cur_categories)
+
+        return DetOut(boxes, scores, class_ids, categories=categories)
+
+    def forward(
+        self,
+        images: Tensor,
+        input_hw: list[tuple[int, int]],
+        boxes2d: Tensor | None = None,
+        boxes2d_classes: Tensor | None = None,
+        original_hw: list[tuple[int, int]] | None = None,
+        input_texts: Sequence[str] | str | None = None,
+        input_tokens_positive: list[dict[int, list[int, int]]] | None = None,
+        text_prompt_mapping: dict[str, dict[str, str]] | None = None,
+    ) -> GroundingDINOOut | DetOut:
+        """Forward function."""
+        if self.training:
+            assert boxes2d is not None and boxes2d_classes is not None
+            return self._forward_train(
+                images,
+                input_texts,
+                boxes2d,
+                boxes2d_classes,
+                input_hw,
+                input_tokens_positive=input_tokens_positive,
+            )
+
+        assert original_hw is not None
+        return self._forward_test(
+            images,
+            input_texts,
+            text_prompt_mapping,
+            input_hw,
+            original_hw,
+        )

opendet3d/model/detect3d/grounding_dino_3d.py

@@ -0,0 +1,812 @@
+"""3D-MOOD."""
+
+from __future__ import annotations
+
+import copy
+from collections.abc import Sequence
+from typing import NamedTuple
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+from vis4d.op.base import BaseModel
+from vis4d.op.fpp.fpn import FPN
+
+from opendet3d.model.detect.grounding_dino import GroundingDINO
+from opendet3d.model.language.mm_bert import BertModel
+from opendet3d.op.detect3d.grounding_dino_3d import (
+    GroundingDINO3DHead,
+    RoI2Det3D,
+)
+from opendet3d.op.detect.grounding_dino import GroundingDINOHead, RoI2Det
+from opendet3d.op.fpp.channel_mapper import ChannelMapper
+
+
+class Det3DOut(NamedTuple):
+    """Output of the detection model.
+
+    boxes (list[Tensor]): 2D bounding boxes of shape [N, 4] in xyxy format.
+    boxes3d (list[Tensor]): 3D bounding boxes of shape [N, 10].
+    scores (list[Tensor]): confidence scores of shape [N,].
+    class_ids (list[Tensor]): class ids of shape [N,].
+    """
+
+    boxes: list[Tensor]
+    boxes3d: list[Tensor]
+    scores: list[Tensor]
+    class_ids: list[Tensor]
+    depth_maps: list[Tensor] | None
+    categories: list[list[str]] | None = None
+
+
+class GroundingDINO3DOut(NamedTuple):
+    """Output of the Grounding DINO model."""
+
+    all_layers_cls_scores: list[Tensor]
+    all_layers_bbox_preds: list[Tensor]
+    all_layers_bbox_3d_preds: list[Tensor]
+    enc_outputs_class: Tensor
+    enc_outputs_coord: Tensor
+    enc_outputs_3d: Tensor
+    text_token_mask: Tensor
+    dn_meta: dict[str, Tensor]
+    positive_maps: list[Tensor]
+    depth_maps: Tensor | None
+
+
+class GroundingDINO3D(GroundingDINO):
+    """Grounding DINO."""
+
+    def __init__(
+        self,
+        basemodel: BaseModel,
+        neck: ChannelMapper,
+        texts: list[str] | None = None,
+        custom_entities: bool = True,
+        chunked_size: int = -1,
+        num_queries: int = 900,
+        num_feature_levels: int = 4,
+        use_checkpoint: bool = False,
+        bbox_head: GroundingDINOHead | None = None,
+        language_model: BertModel | None = None,
+        roi2det: RoI2Det | None = None,
+        bbox3d_head: GroundingDINO3DHead | None = None,
+        roi2det3d: RoI2Det3D | None = None,
+        depth_head: nn.Module | None = None,
+        fpn: FPN | None = None,
+        freeze_detector: bool = False,
+        weights: str | None = None,
+        cat_mapping: dict[str, int] | None = None,
+    ) -> None:
+        """Create the Grounding DINO model."""
+        super().__init__(
+            basemodel=basemodel,
+            neck=neck,
+            texts=texts,
+            custom_entities=custom_entities,
+            chunked_size=chunked_size,
+            num_queries=num_queries,
+            num_feature_levels=num_feature_levels,
+            use_checkpoint=use_checkpoint,
+            bbox_head=bbox_head,
+            roi2det=roi2det,
+            language_model=language_model,
+            weights=weights,
+        )
+
+        self.bbox3d_head = bbox3d_head or GroundingDINO3DHead()
+        self.roi2det3d = roi2det3d or RoI2Det3D()
+
+        # Depth Head
+        self.fpn = fpn
+        self.depth_head = depth_head
+
+        if freeze_detector:
+            self._freeze_detector()
+
+        self.cat_mapping = cat_mapping
+
+    def _freeze_detector(self):
+        """Freeze the detector."""
+        for model in [
+            self.backbone,
+            self.neck,
+            self.encoder,
+            self.positional_encoding,
+            self.memory_trans_fc,
+            self.memory_trans_norm,
+            self.decoder,
+            self.bbox_head,
+            self.dn_query_generator,
+            self.language_model,
+            self.text_feat_map,
+        ]:
+            model.eval()
+            for param in model.parameters():
+                param.requires_grad = False
+
+        self.level_embed.requires_grad = False
+        self.query_embedding.requires_grad = False
+
+    def forward_transformer(
+        self,
+        feat_flatten: Tensor,
+        lvl_pos_embed_flatten: Tensor,
+        memory_mask: Tensor | None,
+        spatial_shapes: Tensor,
+        level_start_index: Tensor,
+        valid_ratios: Tensor,
+        text_dict: dict[str, Tensor],
+        boxes: Tensor | None = None,
+        class_ids: Tensor | None = None,
+        input_hw: list[tuple[int, int]] | None = None,
+        ray_embeddings: Tensor | None = None,
+        depth_latents: Tensor | None = None,
+    ) -> tuple[Tensor, Tensor, Tensor, list[Tensor]]:
+        """Forward function for the transformer."""
+        text_token_mask = text_dict["text_token_mask"]
+
+        memory, memory_text = self.encoder(
+            query=feat_flatten,
+            query_pos=lvl_pos_embed_flatten,
+            key_padding_mask=memory_mask,
+            spatial_shapes=spatial_shapes,
+            level_start_index=level_start_index,
+            valid_ratios=valid_ratios,
+            memory_text=text_dict["embedded"],
+            text_attention_mask=~text_token_mask,
+            position_ids=text_dict["position_ids"],
+            text_self_attention_masks=text_dict["masks"],
+        )
+
+        bs = memory.shape[0]
+
+        output_memory, output_proposals = self.gen_encoder_output_proposals(
+            memory, memory_mask, spatial_shapes
+        )
+
+        enc_outputs_class = self.bbox_head.cls_branches[
+            self.decoder.num_layers
+        ](output_memory, memory_text, text_token_mask)
+        cls_out_features = self.bbox_head.cls_branches[
+            self.decoder.num_layers
+        ].max_text_len
+        enc_outputs_coord_unact = (
+            self.bbox_head.reg_branches[self.decoder.num_layers](output_memory)
+            + output_proposals
+        )
+
+        # NOTE The DINO selects top-k proposals according to scores of
+        # multi-class classification, while DeformDETR, where the input
+        # is `enc_outputs_class[..., 0]` selects according to scores of
+        # binary classification.
+        topk_indices = torch.topk(
+            enc_outputs_class.max(-1)[0], k=self.num_queries, dim=1
+        )[1]
+
+        topk_score = torch.gather(
+            enc_outputs_class,
+            1,
+            topk_indices.unsqueeze(-1).repeat(1, 1, cls_out_features),
+        )
+        topk_coords_unact = torch.gather(
+            enc_outputs_coord_unact,
+            1,
+            topk_indices.unsqueeze(-1).repeat(1, 1, 4),
+        )
+        topk_coords = topk_coords_unact.sigmoid()
+        topk_coords_unact = topk_coords_unact.detach()
+
+        # Top-k 3D proposals
+        topk_output_memory = torch.gather(
+            output_memory, 1, topk_indices.unsqueeze(-1).repeat(1, 1, 256)
+        )
+
+        topk_output_3d = self.bbox3d_head.single_forward(
+            self.decoder.num_layers,
+            topk_output_memory,
+            ray_embeddings,
+            depth_latents,
+        )
+
+        query = self.query_embedding.weight[:, None, :]
+        query = query.repeat(1, bs, 1).transpose(0, 1)
+
+        if self.training:
+            dn_label_query, dn_bbox_query, dn_mask, dn_meta = (
+                self.dn_query_generator(boxes, class_ids, input_hw)
+            )
+            query = torch.cat([dn_label_query, query], dim=1)
+            reference_points = torch.cat(
+                [dn_bbox_query, topk_coords_unact], dim=1
+            )
+        else:
+            reference_points = topk_coords_unact
+            dn_mask, dn_meta = None, None
+
+        reference_points = reference_points.sigmoid()
+
+        # NOTE DINO calculates encoder losses on scores and coordinates
+        # of selected top-k encoder queries, while DeformDETR is of all
+        # encoder queries.
+        if self.training:
+            enc_outputs_class = topk_score
+            enc_outputs_coord = topk_coords
+
+        hidden_states, references = self.decoder(
+            query=query,
+            value=memory,
+            key_padding_mask=memory_mask,
+            self_attn_mask=dn_mask,
+            reference_points=reference_points,
+            spatial_shapes=spatial_shapes,
+            level_start_index=level_start_index,
+            valid_ratios=valid_ratios,
+            reg_branches=self.bbox_head.reg_branches,
+            memory_text=memory_text,
+            text_attention_mask=~text_token_mask,
+        )
+
+        if len(query) == self.num_queries:
+            # NOTE: This is to make sure label_embeding can be involved to
+            # produce loss even if there is no denoising query (no ground truth
+            # target in this GPU), otherwise, this will raise runtime error in
+            # distributed training.
+            hidden_states[0] += (
+                self.dn_query_generator.label_embedding.weight[0, 0] * 0.0
+            )
+
+        if self.training:
+            return (
+                memory_text,
+                text_token_mask,
+                hidden_states,
+                list(references),
+                enc_outputs_class,
+                enc_outputs_coord,
+                topk_output_3d,
+                dn_meta,
+            )
+
+        return (
+            memory_text,
+            text_token_mask,
+            hidden_states,
+            list(references),
+        )
+
+    def _extract_image_features(
+        self, images: Tensor
+    ) -> tuple[list[Tensor], list[Tensor] | None]:
+        """Extract image features."""
+        visual_feats = self.backbone(images)[2:]
+
+        if self.fpn is not None:
+            depth_feats = self.fpn(visual_feats)
+
+        if len(visual_feats) > len(self.neck.in_channels):
+            start_index = len(visual_feats) - len(self.neck.in_channels)
+            # NOTE: This is to make sure the number of input channels is the
+            # same as the number of input channels of the neck.
+            visual_feats = visual_feats[start_index:]
+
+        visual_feats = self.neck(visual_feats)
+
+        if self.fpn is None:
+            depth_feats = visual_feats
+
+        return visual_feats, depth_feats
+
+    def _forward_train(
+        self,
+        images: Tensor,
+        input_texts: list[list[str]] | None,
+        boxes2d: list[Tensor],
+        boxes2d_classes: list[Tensor],
+        input_hw: list[tuple[int, int]],
+        intrinsics: Tensor,
+        input_tokens_positive: list[list[int, int]] | list[None] | None = None,
+        padding: list[list[int]] | None = None,
+    ) -> GroundingDINO3DOut:
+        """Forward function for training."""
+        batch_size = images.shape[0]
+
+        new_text_prompts = []
+        positive_maps = []
+        if input_tokens_positive is not None:
+            for tokens_positive_dict, text_prompt, gt_label in zip(
+                input_tokens_positive, input_texts, boxes2d_classes
+            ):
+                if tokens_positive_dict is not None:
+                    tokenized = self.language_model.tokenizer(
+                        [text_prompt], padding="longest", return_tensors="pt"
+                    )
+
+                    new_text_prompts.append(text_prompt)
+
+                    new_tokens_positive = [
+                        tokens_positive_dict[label.item()]
+                        for label in gt_label
+                    ]
+                else:
+                    tokenized, caption_string, tokens_positive, _ = (
+                        self.get_tokens_and_prompts(text_prompt)
+                    )
+
+                    new_text_prompts.append(caption_string)
+
+                    new_tokens_positive = [
+                        tokens_positive[label] for label in gt_label
+                    ]
+
+                _, positive_map = self.get_positive_map(
+                    tokenized, new_tokens_positive
+                )
+                positive_maps.append(positive_map)
+        else:
+            # All the text prompts are the same or using the self.texts,
+            # so there is no need to calculate them multiple times.
+            if (
+                input_texts is None
+                or len(set(["".join(t) for t in input_texts])) == 1
+            ):
+                if input_texts is None:
+                    assert self.texts is not None, "Texts should be provided."
+                    text_prompt = self.texts
+                else:
+                    text_prompt = input_texts[0]
+
+                tokenized, caption_string, tokens_positive, _ = (
+                    self.get_tokens_and_prompts(text_prompt)
+                )
+                new_text_prompts = [caption_string] * batch_size
+                for gt_label in boxes2d_classes:
+                    new_tokens_positive = [
+                        tokens_positive[label] for label in gt_label
+                    ]
+                    _, positive_map = self.get_positive_map(
+                        tokenized, new_tokens_positive
+                    )
+                    positive_maps.append(positive_map)
+            else:
+                for text_prompt, gt_label in zip(input_texts, boxes2d_classes):
+                    tokenized, caption_string, tokens_positive, _ = (
+                        self.get_tokens_and_prompts(text_prompt)
+                    )
+
+                    new_text_prompts.append(caption_string)
+
+                    new_tokens_positive = [
+                        tokens_positive[label] for label in gt_label
+                    ]
+                    _, positive_map = self.get_positive_map(
+                        tokenized, new_tokens_positive
+                    )
+
+                    positive_maps.append(positive_map)
+
+        for i in range(batch_size):
+            positive_maps[i] = (
+                positive_maps[i].to(images.device).bool().float()
+            )
+
+        text_dict = self.language_model(new_text_prompts)
+
+        if self.text_feat_map is not None:
+            text_dict["embedded"] = self.text_feat_map(text_dict["embedded"])
+
+        text_token_masks = []
+        for i in range(batch_size):
+            text_token_masks.append(
+                text_dict["text_token_mask"][i]
+                .unsqueeze(0)
+                .repeat(len(positive_map), 1)
+            )
+
+        visual_feats, depth_feats = self._extract_image_features(images)
+
+        batch_input_img_h, batch_input_img_w = images.shape[-2:]
+        batch_input_shape = (batch_input_img_h, batch_input_img_w)
+
+        (
+            feat_flatten,
+            lvl_pos_embed_flatten,
+            memory_mask,
+            spatial_shapes,
+            level_start_index,
+            valid_ratios,
+        ) = self.pre_transformer(
+            visual_feats, input_hw, batch_input_shape, padding=padding
+        )
+
+        ray_embeddings = self.bbox3d_head.get_camera_embeddings(
+            intrinsics, batch_input_shape
+        )
+
+        # Depth Head
+        depth_preds, depth_latents = self.depth_head(
+            depth_feats, intrinsics, batch_input_shape
+        )
+
+        (
+            memory_text,
+            text_token_mask,
+            hidden_states,
+            references,
+            enc_outputs_class,
+            enc_outputs_coord,
+            enc_outputs_3d,
+            dn_meta,
+        ) = self.forward_transformer(
+            feat_flatten,
+            lvl_pos_embed_flatten,
+            memory_mask,
+            spatial_shapes,
+            level_start_index,
+            valid_ratios,
+            text_dict,
+            boxes2d,
+            boxes2d_classes,
+            input_hw,
+            ray_embeddings,
+            depth_latents,
+        )
+
+        all_layers_cls_scores, all_layers_bbox_preds = self.bbox_head(
+            hidden_states, references, memory_text, text_token_mask
+        )
+
+        # Not using denoising for 3D
+        hidden_states_3d = hidden_states[
+            :, :, dn_meta["num_denoising_queries"] :, :
+        ]
+
+        all_layers_outputs_3d = self.bbox3d_head(
+            hidden_states_3d, ray_embeddings, depth_latents
+        )
+
+        return GroundingDINO3DOut(
+            all_layers_cls_scores,
+            all_layers_bbox_preds,
+            all_layers_outputs_3d,
+            enc_outputs_class,
+            enc_outputs_coord,
+            enc_outputs_3d,
+            text_token_mask,
+            dn_meta,
+            positive_maps,
+            depth_maps=depth_preds,
+        )
+
+    def _forward_test(
+        self,
+        images: Tensor,
+        input_texts: list[str] | None,
+        text_prompt_mapping: list[dict[str, dict[str, str]]] | None,
+        input_hw: list[tuple[int, int]],
+        original_hw: list[tuple[int, int]],
+        intrinsics: list[Tensor] | None,
+        padding: list[list[int]] | None,
+    ) -> Det3DOut:
+        """Forward."""
+        batch_size = images.shape[0]
+
+        token_positive_maps = []
+        text_prompts = []
+        entities = []
+        for i in range(batch_size):
+            if self.texts is not None:
+                text_prompt = self.texts
+            else:
+                text_prompt = input_texts[i]
+
+            if text_prompt_mapping is not None:
+                prompt_mapping = text_prompt_mapping[i]
+            else:
+                prompt_mapping = None
+
+            token_positive_map, captions, _, _entities = (
+                self.get_tokens_positive_and_prompts(
+                    text_prompt,
+                    self.custom_entities,
+                    text_prompt_mapping=prompt_mapping,
+                )
+            )
+
+            token_positive_maps.append(token_positive_map)
+            text_prompts.append(captions)
+            entities.append(_entities)
+
+        # image feature extraction
+        batch_input_img_h, batch_input_img_w = images.shape[-2:]
+
+        visual_feats, depth_feats = self._extract_image_features(images)
+
+        batch_input_shape = (batch_input_img_h, batch_input_img_w)
+
+        if isinstance(text_prompts[0], list):
+            assert batch_size == 1, "Batch size should be 1 for chunked text."
+            assert (
+                self.cat_mapping is not None
+            ), "Category mapping should be provided."
+
+            boxes = []
+            boxes3d = []
+            scores = []
+            class_ids = []
+            categories = []
+            for i, captions in enumerate(text_prompts[0]):
+                token_positive_map = token_positive_maps[0][i]
+                cur_entities = entities[0][i]
+
+                text_dict = self.language_model([captions])
+
+                # text feature map layer
+                if self.text_feat_map is not None:
+                    text_dict["embedded"] = self.text_feat_map(
+                        text_dict["embedded"]
+                    )
+
+                (
+                    feat_flatten,
+                    lvl_pos_embed_flatten,
+                    memory_mask,
+                    spatial_shapes,
+                    level_start_index,
+                    valid_ratios,
+                ) = self.pre_transformer(
+                    copy.deepcopy(visual_feats), input_hw, batch_input_shape
+                )
+
+                ray_embeddings = self.bbox3d_head.get_camera_embeddings(
+                    intrinsics, batch_input_shape
+                )
+
+                # Depth Head
+                depth_preds, depth_latents = self.depth_head(
+                    depth_feats, intrinsics, batch_input_shape
+                )
+
+                depth_maps = []
+                for i, depth_pred in enumerate(depth_preds):
+                    if padding is not None:
+                        pad_left, pad_right, pad_top, pad_bottom = padding[i]
+
+                        depth_pred = depth_pred[
+                            pad_top : batch_input_img_h - pad_bottom,
+                            pad_left : batch_input_img_w - pad_right,
+                        ]
+
+                    depth_maps.append(
+                        F.interpolate(
+                            depth_pred.unsqueeze(0).unsqueeze(0),
+                            size=original_hw[i],
+                            mode="bilinear",
+                            align_corners=False,
+                            antialias=True,
+                        )
+                        .squeeze(0)
+                        .squeeze(0)
+                    )
+
+                (
+                    memory_text,
+                    text_token_mask,
+                    hidden_states,
+                    references,
+                ) = self.forward_transformer(
+                    feat_flatten,
+                    lvl_pos_embed_flatten,
+                    memory_mask,
+                    spatial_shapes,
+                    level_start_index,
+                    valid_ratios,
+                    text_dict,
+                    ray_embeddings=ray_embeddings,
+                    depth_latents=depth_latents,
+                )
+
+                all_layers_cls_scores, all_layers_bbox_preds = self.bbox_head(
+                    hidden_states, references, memory_text, text_token_mask
+                )
+
+                all_layers_outputs_3d = self.bbox3d_head(
+                    hidden_states, ray_embeddings, depth_latents=depth_latents
+                )
+
+                cls_scores = all_layers_cls_scores[-1]
+                bbox_preds = all_layers_bbox_preds[-1]
+                boxes3d_preds = all_layers_outputs_3d[-1]
+
+                cls_score = cls_scores[0]
+                det_bboxes, det_scores, det_labels, det_bboxes3d = (
+                    self.roi2det3d(
+                        cls_score,
+                        bbox_preds[0],
+                        token_positive_map,
+                        input_hw[0],
+                        original_hw[0],
+                        boxes3d_preds[0],
+                        intrinsics[0],
+                        padding[0] if padding is not None else None,
+                    )
+                )
+
+                boxes.append(det_bboxes)
+                scores.append(det_scores)
+                boxes3d.append(det_bboxes3d)
+
+                # Get the categories text and class ids
+                cur_class_ids = []
+                cur_categories = []
+                for label in det_labels:
+                    cur_class_ids.append(self.cat_mapping[cur_entities[label]])
+                    cur_categories.append(cur_entities[label])
+
+                class_ids.append(cur_class_ids)
+                categories.append(cur_categories)
+
+            boxes = [torch.cat([b for b in boxes])]
+            boxes3d = [torch.cat([b for b in boxes3d])]
+            scores = [torch.cat([s for s in scores])]
+            class_ids = [sum(class_ids, [])]
+            categories = [sum(categories, [])]
+        else:
+            # extract text feats
+            text_dict = self.language_model(list(text_prompts))
+
+            # text feature map layer
+            if self.text_feat_map is not None:
+                text_dict["embedded"] = self.text_feat_map(
+                    text_dict["embedded"]
+                )
+
+            (
+                feat_flatten,
+                lvl_pos_embed_flatten,
+                memory_mask,
+                spatial_shapes,
+                level_start_index,
+                valid_ratios,
+            ) = self.pre_transformer(visual_feats, input_hw, batch_input_shape)
+
+            ray_embeddings = self.bbox3d_head.get_camera_embeddings(
+                intrinsics, batch_input_shape
+            )
+
+            # Depth Head
+            depth_preds, depth_latents = self.depth_head(
+                depth_feats, intrinsics, batch_input_shape
+            )
+
+            depth_maps = []
+            for i, depth_pred in enumerate(depth_preds):
+                if padding is not None:
+                    pad_left, pad_right, pad_top, pad_bottom = padding[i]
+
+                    depth_pred = depth_pred[
+                        pad_top : batch_input_img_h - pad_bottom,
+                        pad_left : batch_input_img_w - pad_right,
+                    ]
+
+                depth_maps.append(
+                    F.interpolate(
+                        depth_pred.unsqueeze(0).unsqueeze(0),
+                        size=original_hw[i],
+                        mode="bilinear",
+                        align_corners=False,
+                        antialias=True,
+                    )
+                    .squeeze(0)
+                    .squeeze(0)
+                )
+
+            (
+                memory_text,
+                text_token_mask,
+                hidden_states,
+                references,
+            ) = self.forward_transformer(
+                feat_flatten,
+                lvl_pos_embed_flatten,
+                memory_mask,
+                spatial_shapes,
+                level_start_index,
+                valid_ratios,
+                text_dict,
+                ray_embeddings=ray_embeddings,
+                depth_latents=depth_latents,
+            )
+
+            all_layers_cls_scores, all_layers_bbox_preds = self.bbox_head(
+                hidden_states, references, memory_text, text_token_mask
+            )
+
+            all_layers_outputs_3d = self.bbox3d_head(
+                hidden_states, ray_embeddings, depth_latents=depth_latents
+            )
+
+            cls_scores = all_layers_cls_scores[-1]
+            bbox_preds = all_layers_bbox_preds[-1]
+            boxes3d_preds = all_layers_outputs_3d[-1]
+
+            boxes = []
+            boxes3d = []
+            scores = []
+            class_ids = []
+            categories = []
+            for i, bbox_pred in enumerate(bbox_preds):
+                cls_score = cls_scores[i]
+                det_bboxes, det_scores, det_labels, det_bboxes3d = (
+                    self.roi2det3d(
+                        cls_score,
+                        bbox_pred,
+                        token_positive_maps[i],
+                        input_hw[i],
+                        original_hw[i],
+                        boxes3d_preds[i],
+                        intrinsics[i],
+                        padding[i] if padding is not None else None,
+                    )
+                )
+
+                boxes.append(det_bboxes)
+                scores.append(det_scores)
+                class_ids.append(det_labels)
+                boxes3d.append(det_bboxes3d)
+
+                # Get the categories text
+                cur_categories = []
+                for label in det_labels:
+                    cur_categories.append(entities[i][label])
+
+                categories.append(cur_categories)
+
+        return Det3DOut(
+            boxes,
+            boxes3d,
+            scores,
+            class_ids,
+            depth_maps=depth_maps,
+            categories=categories,
+        )
+
+    def forward(
+        self,
+        images: Tensor,
+        input_hw: list[tuple[int, int]],
+        intrinsics: Tensor,
+        boxes2d: Tensor | None = None,
+        boxes2d_classes: Tensor | None = None,
+        original_hw: list[tuple[int, int]] | None = None,
+        input_texts: Sequence[str] | str | None = None,
+        input_tokens_positive: list[dict[int, list[int, int]]] | None = None,
+        text_prompt_mapping: dict[str, dict[str, str]] | None = None,
+        padding: list[list[int]] | None = None,
+        **kwargs,
+    ) -> GroundingDINO3DOut | Det3DOut:
+        """Forward function."""
+        if self.training:
+            assert boxes2d is not None and boxes2d_classes is not None
+            return self._forward_train(
+                images,
+                input_texts,
+                boxes2d,
+                boxes2d_classes,
+                input_hw,
+                intrinsics,
+                input_tokens_positive=input_tokens_positive,
+                padding=padding,
+                **kwargs,
+            )
+
+        assert original_hw is not None
+        return self._forward_test(
+            images,
+            input_texts,
+            text_prompt_mapping,
+            input_hw,
+            original_hw,
+            intrinsics,
+            padding,
+        )

opendet3d/model/language/mm_bert.py

@@ -0,0 +1,255 @@
+"""BERT model from mmdetection."""
+
+import logging
+import os
+from collections import OrderedDict
+from collections.abc import Sequence
+
+import torch
+from torch import nn
+from transformers import AutoTokenizer, BertConfig
+from transformers import BertModel as HFBertModel
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+
+def generate_masks_with_special_tokens_and_transfer_map(
+    tokenized, special_tokens_list
+):
+    """Generate attention mask between each pair of special tokens.
+
+    Only token pairs in between two special tokens are attended to
+    and thus the attention mask for these pairs is positive.
+
+    Args:
+        input_ids (torch.Tensor): input ids. Shape: [bs, num_token]
+        special_tokens_mask (list): special tokens mask.
+
+    Returns:
+        Tuple(Tensor, Tensor):
+        - attention_mask is the attention mask between each tokens.
+          Only token pairs in between two special tokens are positive.
+          Shape: [bs, num_token, num_token].
+        - position_ids is the position id of tokens within each valid sentence.
+          The id starts from 0 whenenver a special token is encountered.
+          Shape: [bs, num_token]
+    """
+    input_ids = tokenized["input_ids"]
+    bs, num_token = input_ids.shape
+    # special_tokens_mask:
+    # bs, num_token. 1 for special tokens. 0 for normal tokens
+    special_tokens_mask = torch.zeros(
+        (bs, num_token), device=input_ids.device
+    ).bool()
+
+    for special_token in special_tokens_list:
+        special_tokens_mask |= input_ids == special_token
+
+    # idxs: each row is a list of indices of special tokens
+    idxs = torch.nonzero(special_tokens_mask)
+
+    # generate attention mask and positional ids
+    attention_mask = (
+        torch.eye(num_token, device=input_ids.device)
+        .bool()
+        .unsqueeze(0)
+        .repeat(bs, 1, 1)
+    )
+    position_ids = torch.zeros((bs, num_token), device=input_ids.device)
+    previous_col = 0
+    for i in range(idxs.shape[0]):
+        row, col = idxs[i]
+        if (col == 0) or (col == num_token - 1):
+            attention_mask[row, col, col] = True
+            position_ids[row, col] = 0
+        else:
+            attention_mask[
+                row, previous_col + 1 : col + 1, previous_col + 1 : col + 1
+            ] = True
+            position_ids[row, previous_col + 1 : col + 1] = torch.arange(
+                0, col - previous_col, device=input_ids.device
+            )
+        previous_col = col
+
+    return attention_mask, position_ids.to(torch.long)
+
+
+class BertModel(nn.Module):
+    """BERT model for language embedding only encoder.
+
+    Args:
+        name (str, optional): name of the pretrained BERT model from
+            HuggingFace. Defaults to bert-base-uncased.
+        max_tokens (int, optional): maximum number of tokens to be
+            used for BERT. Defaults to 256.
+        pad_to_max (bool, optional): whether to pad the tokens to max_tokens.
+             Defaults to True.
+        use_sub_sentence_represent (bool, optional): whether to use sub
+            sentence represent introduced in `Grounding DINO
+            <https://arxiv.org/abs/2303.05499>`. Defaults to False.
+        special_tokens_list (list, optional): special tokens used to split
+            subsentence. It cannot be None when `use_sub_sentence_represent`
+            is True. Defaults to None.
+        add_pooling_layer (bool, optional): whether to adding pooling
+            layer in bert encoder. Defaults to False.
+        num_layers_of_embedded (int, optional): number of layers of
+            the embedded model. Defaults to 1.
+        use_checkpoint (bool, optional): whether to use gradient checkpointing.
+             Defaults to False.
+    """
+
+    def __init__(
+        self,
+        name: str = "bert-base-uncased",
+        max_tokens: int = 256,
+        pad_to_max: bool = True,
+        use_sub_sentence_represent: bool = False,
+        special_tokens_list: list = None,
+        add_pooling_layer: bool = False,
+        num_layers_of_embedded: int = 1,
+        use_checkpoint: bool = False,
+    ) -> None:
+        """Create the BERT model."""
+        super().__init__()
+        self.max_tokens = max_tokens
+        self.pad_to_max = pad_to_max
+
+        self.tokenizer = AutoTokenizer.from_pretrained(name)
+        self.language_backbone = nn.Sequential(
+            OrderedDict(
+                [
+                    (
+                        "body",
+                        BertEncoder(
+                            name,
+                            add_pooling_layer=add_pooling_layer,
+                            num_layers_of_embedded=num_layers_of_embedded,
+                            use_checkpoint=use_checkpoint,
+                        ),
+                    )
+                ]
+            )
+        )
+
+        self.use_sub_sentence_represent = use_sub_sentence_represent
+        if self.use_sub_sentence_represent:
+            assert (
+                special_tokens_list is not None
+            ), "special_tokens should not be None \
+                    if use_sub_sentence_represent is True"
+
+            self.special_tokens = self.tokenizer.convert_tokens_to_ids(
+                special_tokens_list
+            )
+
+    def forward(self, captions: Sequence[str]) -> dict:
+        """Forward function."""
+        device = next(self.language_backbone.parameters()).device
+        tokenized = self.tokenizer.batch_encode_plus(
+            captions,
+            max_length=self.max_tokens,
+            padding="max_length" if self.pad_to_max else "longest",
+            return_special_tokens_mask=True,
+            return_tensors="pt",
+            truncation=True,
+        ).to(device)
+        input_ids = tokenized.input_ids
+        if self.use_sub_sentence_represent:
+            attention_mask, position_ids = (
+                generate_masks_with_special_tokens_and_transfer_map(
+                    tokenized, self.special_tokens
+                )
+            )
+            token_type_ids = tokenized["token_type_ids"]
+
+        else:
+            attention_mask = tokenized.attention_mask
+            position_ids = None
+            token_type_ids = None
+
+        tokenizer_input = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "position_ids": position_ids,
+            "token_type_ids": token_type_ids,
+        }
+        language_dict_features = self.language_backbone(tokenizer_input)
+        if self.use_sub_sentence_represent:
+            language_dict_features["position_ids"] = position_ids
+            language_dict_features["text_token_mask"] = (
+                tokenized.attention_mask.bool()
+            )
+        return language_dict_features
+
+
+class BertEncoder(nn.Module):
+    """BERT encoder for language embedding.
+
+    Args:
+        name (str): name of the pretrained BERT model from HuggingFace.
+                Defaults to bert-base-uncased.
+        add_pooling_layer (bool): whether to add a pooling layer.
+        num_layers_of_embedded (int): number of layers of the embedded model.
+                Defaults to 1.
+        use_checkpoint (bool): whether to use gradient checkpointing.
+                Defaults to False.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        add_pooling_layer: bool = False,
+        num_layers_of_embedded: int = 1,
+        use_checkpoint: bool = False,
+    ):
+        super().__init__()
+        config = BertConfig.from_pretrained(name)
+        config.gradient_checkpointing = use_checkpoint
+
+        loggers = [
+            logging.getLogger(name) for name in logging.root.manager.loggerDict
+        ]
+        for logger in loggers:
+            if "transformers" in logger.name.lower():
+                logger.setLevel(logging.ERROR)
+
+        # only encoder
+        self.model = HFBertModel.from_pretrained(
+            name,
+            add_pooling_layer=add_pooling_layer,
+            config=config,
+            attn_implementation="eager",
+        )
+
+        self.language_dim = config.hidden_size
+        self.num_layers_of_embedded = num_layers_of_embedded
+
+    def forward(self, x) -> dict:
+        mask = x["attention_mask"]
+
+        outputs = self.model(
+            input_ids=x["input_ids"],
+            attention_mask=mask,
+            position_ids=x["position_ids"],
+            token_type_ids=x["token_type_ids"],
+            output_hidden_states=True,
+        )
+
+        # outputs has 13 layers, 1 input layer and 12 hidden layers
+        encoded_layers = outputs.hidden_states[1:]
+        features = torch.stack(
+            encoded_layers[-self.num_layers_of_embedded :], 1
+        ).mean(1)
+        # language embedding has shape [len(phrase), seq_len, language_dim]
+        features = features / self.num_layers_of_embedded
+        if mask.dim() == 2:
+            embedded = features * mask.unsqueeze(-1).float()
+        else:
+            embedded = features
+
+        results = {
+            "embedded": embedded,
+            "masks": mask,
+            "hidden": encoded_layers[-1],
+        }
+        return results