"""CA1M dataset.""" from __future__ import annotations import os import pickle from collections import defaultdict import numpy as np import torch import torch.nn.functional as F from mapdet3d.common.typing import ArgsType, DictStrAny from mapdet3d.data.const import AxisMode from mapdet3d.data.const import CommonKeys as K from mapdet3d.op.mapanything.image import preprocess_inputs from .base import VideoDataset, VideoMapping from .util import CacheMappingMixin, im_decode def select_training_frame_ids( key_fid: int, max_fid: int, num_views: int, max_sampling_rate: int = 1, ) -> list[int]: """Select future training frame ids with a random or fixed stride.""" if num_views < 1: raise ValueError("num_views must be greater than 0.") if max_sampling_rate < 1: raise ValueError("max_sampling_rate must be greater than 0.") sampling_rate = int(np.random.randint(1, max_sampling_rate + 1)) selected_fids = list(range(key_fid, max_fid + 1, sampling_rate))[ :num_views ] if len(selected_fids) < num_views: selected_fids = [key_fid] * ( num_views - len(selected_fids) ) + selected_fids return selected_fids class CA1M(CacheMappingMixin, VideoDataset): """CA1M dataset.""" def __init__( self, data_root: str, split: str = "train", max_depth: float = 10.0, depth_scale: float = 1000.0, remove_empty: bool = False, valid_scenes: list[str] | None = None, cache_as_binary: bool = False, cached_dir: str = "cache", use_arkit_depth: bool = True, max_sampling_rate: int = 5, **kwargs: ArgsType, ) -> None: """Init.""" super().__init__(**kwargs) self.data_root = data_root self.split = split self.max_depth = max_depth self.depth_scale = depth_scale self.use_arkit_depth = use_arkit_depth self.max_sampling_rate = max_sampling_rate self.remove_empty = remove_empty self.valid_scenes = valid_scenes self.cache_as_binary = cache_as_binary self.cached_dir = cached_dir self.cached_file_path = os.path.join( self.data_root, self.cached_dir, f"{self.split}.pkl" ) # Load annotations self.samples, _ = self._load_mapping( self._generate_data_mapping, self._filter_data, cache_as_binary=cache_as_binary, cached_file_path=self.cached_file_path, ) # Generate video mapping self.video_mapping = self._generate_video_mapping() def __repr__(self) -> str: """Concise representation of the dataset.""" return f"CA1M {self.split}" def _filter_data(self, data: list[DictStrAny]) -> list[DictStrAny]: """Remove empty samples.""" if not self.remove_empty: return data samples = [] for sample in data: if sample["sequence_name"] in self.valid_scenes: samples.append(sample) return samples def _generate_video_mapping(self) -> VideoMapping: """Group dataset sample indices by their associated video ID. The sample index is an integer while video IDs are string. Returns: VideoMapping: Mapping of video IDs to sample indices and frame IDs. """ video_to_indices: dict[str, list[int]] = defaultdict(list) video_to_frame_ids: dict[str, list[int]] = defaultdict(list) for i, sample in enumerate(self.samples): seq = sample["sequence_name"] fid = sample["frame_id"] video_to_indices[seq].append(i) video_to_frame_ids[seq].append(fid) return self._sort_video_mapping( { "video_to_indices": video_to_indices, "video_to_frame_ids": video_to_frame_ids, } ) def _generate_data_mapping(self) -> list[DictStrAny]: """Generates the data mapping.""" with open(self.cached_file_path, "rb") as file: data = pickle.loads(file.read()) return data def __len__(self): return len(self.samples) def _get_sample_data( self, sample: DictStrAny, sample_data: list[DictStrAny] ) -> DictStrAny: """Get single sample from raw data.""" data_dict = {} data_dict[K.sample_names] = sample["timestamp"] data_dict["image_ids"] = int(sample["timestamp"]) data_dict[K.timestamp] = int(sample["timestamp"]) / 1e9 data_dict[K.sequence_names] = sample["sequence_name"] data_dict[K.frame_ids] = sample["frame_id"] # Load image im_bytes = self.data_backend.get(sample_data["image_file_path"]) image = np.ascontiguousarray( im_decode(im_bytes, mode=self.image_channel_mode), dtype=np.float32, )[None] intrinsics = sample_data["intrinsics"] 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.intrinsics] = intrinsics data_dict["original_intrinsics"] = intrinsics data_dict[K.extrinsics] = sample_data["extrinsics"] data_dict["T_gravity"] = sample_data["T_gravity"] # Load annotations data_dict[K.boxes2d] = sample_data["boxes2d"] data_dict[K.boxes2d_names] = sample_data["categories"] data_dict[K.boxes2d_classes] = np.zeros( len(sample_data["categories"]), dtype=np.int64 ) data_dict[K.boxes2d_track_ids] = sample_data["track_ids"] data_dict[K.boxes3d] = sample_data["boxes3d"] data_dict["boxes3d_cam_from_world"] = sample_data[ "boxes3d_cam_from_world" ] data_dict[K.boxes3d_classes] = np.zeros( len(sample_data["categories"]), dtype=np.int64 ) data_dict[K.boxes3d_names] = sample_data["categories"] data_dict[K.boxes3d_track_ids] = sample_data["track_ids"] data_dict[K.axis_mode] = AxisMode.OPENCV # Load depth depth_bytes = self.data_backend.get( sample_data["arkit_depth_file_path"] if self.use_arkit_depth else sample_data["depth_file_path"] ) depth_array = im_decode(depth_bytes) depth = np.ascontiguousarray(depth_array, dtype=np.float32) depth = depth / self.depth_scale depth[depth > self.max_depth] = 0 depth = F.interpolate( torch.from_numpy(depth)[None, None, ...], size=(image.shape[1], image.shape[2]), mode="nearest", )[0, 0].numpy() data_dict[K.depth_maps] = depth # Load global annotations if "boxes3d_world" in sample: data_dict["boxes3d_world"] = sample["boxes3d_world"] # Mesh path for visualization (optional) data_dict["mesh_path"] = os.path.join( self.data_root, "mesh", sample["sequence_name"], "mesh.ply", ) return data_dict def __getitem__(self, idx: int | tuple[int, int, float]) -> DictStrAny: """Get single sample.""" # For training if isinstance(idx, tuple): key_sample = self.samples[idx[0]] num_views = idx[1] aspect_ratio = idx[2] seq_name = key_sample["sequence_name"] frame_ids = self.video_mapping["video_to_frame_ids"][seq_name] video_indices = self.video_mapping["video_to_indices"][seq_name] with open( os.path.join( self.data_root, self.cached_dir, self.split, f"{seq_name}.pkl", ), "rb", ) as file: data = pickle.loads(file.read()) key_fid = key_sample["frame_id"] max_fid = frame_ids[-1] selected_fids = select_training_frame_ids( key_fid, max_fid, num_views, self.max_sampling_rate ) # Get data for all selected frames seq = [] for fid in selected_fids: frame_data = self._get_sample_data( self.samples[video_indices[fid]], data["seq_data"][fid] ) seq.append(frame_data) self.data_backend.close() views = [ { "img": s["pil_image"], "intrinsics": s[K.intrinsics], "camera_poses": s[K.extrinsics], "depth_z": s[K.depth_maps], "is_metric_scale": torch.tensor([True]), "boxes2d": s[K.boxes2d], } for s in seq ] processed_views = preprocess_inputs( views, padding_mode=True, aspect_ratio=aspect_ratio, ) pad_info_list = [v.pop("pad_info", None) for v in processed_views] input_hw_list = [ [v["img"].shape[-2], v["img"].shape[-1]] for v in processed_views ] intrinsics_list = [ v["intrinsics"][0].numpy().copy() for v in processed_views ] boxes2d_list = [ v.pop("boxes2d")[0].numpy() for v in processed_views ] track_ids_list = [s[K.boxes3d_track_ids] for s in seq] class_ids_list = [s[K.boxes3d_classes] for s in seq] boxes3d_list = [s[K.boxes3d] for s in seq] sample_names_list = [s[K.sample_names] for s in seq] categories_list = [s[K.boxes3d_names] for s in seq] return { "processed_views": processed_views, "sample_names": sample_names_list, "input_hw": input_hw_list, "intrinsics": intrinsics_list, "boxes2d": boxes2d_list, "boxes3d": boxes3d_list, "track_ids": track_ids_list, "class_ids": class_ids_list, "categories": categories_list, "pad_info": pad_info_list, } sample = self.samples[idx] seq_name = sample["sequence_name"] with open( os.path.join( self.data_root, self.cached_dir, self.split, f"{seq_name}.pkl" ), "rb", ) as file: data = pickle.loads(file.read()) self.data_backend.close() return self._get_sample_data( sample, data["seq_data"][sample["frame_id"]] )