diff --git a/OpenCOOD/opencood/data_utils/datasets/basedataset.py b/OpenCOOD/opencood/data_utils/datasets/basedataset.py index a61543d5..084d6b4e 100644 --- a/OpenCOOD/opencood/data_utils/datasets/basedataset.py +++ b/OpenCOOD/opencood/data_utils/datasets/basedataset.py @@ -2,10 +2,13 @@ Basedataset class for all kinds of fusion. """ +from __future__ import annotations + import os import math import logging from collections import OrderedDict +from typing import TYPE_CHECKING import torch import numpy as np @@ -17,6 +20,9 @@ from opencood.utils.pcd_utils import downsample_lidar_minimum from opencood.utils.transformation_utils import x1_to_x2 +if TYPE_CHECKING: + from opencood.models.communication_adapters.base import ModelCommunicationAdapter + logger = logging.getLogger("cavise.opencda.OpenCOOD.opencood.data_utils.datasets.basedataset") @@ -54,12 +60,16 @@ class BaseDataset(Dataset): data_augmentor : opencood.data_augmentor Used to augment data. + communication_adapter : ModelCommunicationAdapter | None + Model-owned wire adapter attached by the CAPI model manager. + """ def __init__(self, params, visualize, train=True): self.params = params self.visualize = visualize self.train = train + self.communication_adapter: ModelCommunicationAdapter | None = None self.pre_processor = None self.post_processor = None diff --git a/OpenCOOD/opencood/data_utils/datasets/early_fusion_dataset.py b/OpenCOOD/opencood/data_utils/datasets/early_fusion_dataset.py index 682c5cc8..bcda0107 100644 --- a/OpenCOOD/opencood/data_utils/datasets/early_fusion_dataset.py +++ b/OpenCOOD/opencood/data_utils/datasets/early_fusion_dataset.py @@ -5,6 +5,7 @@ import math import logging from collections import OrderedDict +from typing import Any import numpy as np import torch @@ -13,6 +14,7 @@ from opencood.utils import box_utils from opencood.data_utils.post_processor import build_postprocessor from opencood.data_utils.datasets import basedataset +from opencood.models.communication_adapters import EarlyFusionWirePayload from opencood.data_utils.pre_processor import build_preprocessor from opencood.utils.pcd_utils import downsample_lidar_minimum from opencood.utils.transformation_utils import x1_to_x2 @@ -57,60 +59,72 @@ def extract_data(self, idx): if self.payload_handler is not None: for cav_id, selected_cav_base in base_data_dict.items(): selected_cav_processed = self.get_item_single_car(selected_cav_base, ego_lidar_pose) + payload = self.build_wire_payload(selected_cav_processed) + self.payload_handler.set_opencda_payload(cav_id, self.module_name, payload) - with self.payload_handler.handle_opencda_payload(cav_id, self.module_name) as msg: - msg["object_ids"] = selected_cav_processed["object_ids"] # list - msg["object_bbx_center"] = selected_cav_processed["object_bbx_center"] - msg["projected_lidar"] = selected_cav_processed["projected_lidar"] - msg["projected_lidar_spoofing_mask"] = selected_cav_processed["projected_lidar_spoofing_mask"] + @staticmethod + def build_wire_payload(selected_cav_processed: dict[str, Any]) -> EarlyFusionWirePayload: + """ + Build the early-fusion payload sent over V2X. + + Parameters + ---------- + selected_cav_processed : dict[str, Any] + Processed data for one sending agent. + + Returns + ------- + EarlyFusionWirePayload + Typed payload containing only the remote model input. + """ + return EarlyFusionWirePayload( + projected_lidar=selected_cav_processed["projected_lidar"], + ) + + @staticmethod + def decode_wire_payload(payload: object) -> EarlyFusionWirePayload: + """ + Validate and decode an early-fusion payload received over V2X. + + Parameters + ---------- + payload : object + Deserialized module payload received over V2X. + + Returns + ------- + EarlyFusionWirePayload + Validated early-fusion payload. + + Raises + ------ + TypeError + If the received payload has an unexpected type. + """ + if not isinstance(payload, EarlyFusionWirePayload): + raise TypeError(f"Expected EarlyFusionWirePayload, got {type(payload).__name__}") + return payload def __process_with_messages(self, ego_id, ego_lidar_pose, base_data_dict): - object_stack = [] - object_id_stack = [] projected_lidar_stack = [] - projected_lidar_roles = [] - projected_lidar_agent_ids = [] - projected_lidar_spoofing_masks = [] ego_cav_base = base_data_dict.get(ego_id) ego_cav_processed = self.get_item_single_car(ego_cav_base, ego_lidar_pose) - object_id_stack += ego_cav_processed["object_ids"] - object_stack.append(ego_cav_processed["object_bbx_center"]) projected_lidar_stack.append(ego_cav_processed["projected_lidar"]) - projected_lidar_roles.append("ego") - projected_lidar_agent_ids.append(ego_id) - projected_lidar_spoofing_masks.append(ego_cav_processed["projected_lidar_spoofing_mask"]) if ego_id in self.payload_handler.current_artery_payload: for cav_id, _ in base_data_dict.items(): - if cav_id in self.payload_handler.current_artery_payload[ego_id]: - with self.payload_handler.handle_artery_payload(ego_id, cav_id, self.module_name) as msg: - object_id_stack += msg["object_ids"] - object_stack.append(msg["object_bbx_center"]) - projected_lidar_stack.append(msg["projected_lidar"]) - projected_lidar_roles.append("other") - projected_lidar_agent_ids.append(cav_id) - projected_lidar_spoofing_masks.append( - np.asarray(msg.get("projected_lidar_spoofing_mask", np.zeros((msg["projected_lidar"].shape[0],), dtype=np.bool_))) - ) + raw_payload = self.payload_handler.get_artery_payload(ego_id, cav_id, self.module_name) + if raw_payload is None: + continue + payload = self.decode_wire_payload(raw_payload) + projected_lidar_stack.append(payload.projected_lidar) - return { - "object_stack": object_stack, - "object_id_stack": object_id_stack, - "projected_lidar_stack": projected_lidar_stack, - "projected_lidar_roles": projected_lidar_roles, - "projected_lidar_agent_ids": projected_lidar_agent_ids, - "projected_lidar_spoofing_masks": projected_lidar_spoofing_masks, - } + return {"projected_lidar_stack": projected_lidar_stack} def __process_without_messages(self, ego_lidar_pose, base_data_dict): projected_lidar_stack = [] - object_stack = [] - object_id_stack = [] - projected_lidar_roles = [] - projected_lidar_agent_ids = [] - projected_lidar_spoofing_masks = [] # loop over all CAVs to process information for cav_id, selected_cav_base in base_data_dict.items(): @@ -124,55 +138,121 @@ def __process_without_messages(self, ego_lidar_pose, base_data_dict): selected_cav_processed = self.get_item_single_car(selected_cav_base, ego_lidar_pose) projected_lidar_stack.append(selected_cav_processed["projected_lidar"]) - object_stack.append(selected_cav_processed["object_bbx_center"]) - object_id_stack += selected_cav_processed["object_ids"] - projected_lidar_roles.append("ego" if selected_cav_base["ego"] else "other") + + return {"projected_lidar_stack": projected_lidar_stack} + + def build_local_supervision(self, base_data_dict: dict[str, Any], ego_lidar_pose: list[float]) -> dict[str, Any]: + """ + Build ego-frame ground truth from the local scene snapshot. + + Parameters + ---------- + base_data_dict : dict[str, Any] + Complete local dataset snapshot for the current scene. + ego_lidar_pose : list[float] + Current ego LiDAR pose in world coordinates. + + Returns + ------- + dict[str, Any] + Ground-truth boxes, mask, and object identifiers. + """ + object_bbx_center, object_bbx_mask, object_ids = self.post_processor.generate_object_center(list(base_data_dict.values()), ego_lidar_pose) + return { + "object_bbx_center": object_bbx_center, + "object_bbx_mask": object_bbx_mask, + "object_ids": object_ids, + } + + def build_visualization_context( + self, + ego_id: str, + ego_lidar_pose: list[float], + base_data_dict: dict[str, Any], + ) -> dict[str, Any]: + """ + Build point-cloud visualization data from the local scene. + + Parameters + ---------- + ego_id : str + Identifier of the receiving ego agent. + ego_lidar_pose : list[float] + Current ego LiDAR pose in world coordinates. + base_data_dict : dict[str, Any] + Local dataset snapshot, independent of delivered messages. + + Returns + ------- + dict[str, Any] + Per-agent point clouds, identities, roles, and spoofing masks. + """ + if not self.visualize: + return {} + + projected_lidar_stack = [] + projected_lidar_roles = [] + projected_lidar_agent_ids = [] + projected_lidar_spoofing_masks = [] + lidar_range = self.params["preprocess"]["cav_lidar_range"] + + for cav_id, cav_base in base_data_dict.items(): + cav_processed = self.get_item_single_car(cav_base, ego_lidar_pose) + projected_lidar = cav_processed["projected_lidar"] + range_mask = ( + (projected_lidar[:, 0] > lidar_range[0]) + & (projected_lidar[:, 0] < lidar_range[3]) + & (projected_lidar[:, 1] > lidar_range[1]) + & (projected_lidar[:, 1] < lidar_range[4]) + & (projected_lidar[:, 2] > lidar_range[2]) + & (projected_lidar[:, 2] < lidar_range[5]) + ) + projected_lidar_stack.append(projected_lidar[range_mask]) + projected_lidar_roles.append("ego" if cav_id == ego_id else "other") projected_lidar_agent_ids.append(cav_id) - projected_lidar_spoofing_masks.append(selected_cav_processed["projected_lidar_spoofing_mask"]) + projected_lidar_spoofing_masks.append(cav_processed["projected_lidar_spoofing_mask"][range_mask]) return { - "object_stack": object_stack, - "object_id_stack": object_id_stack, - "projected_lidar_stack": projected_lidar_stack, - "projected_lidar_roles": projected_lidar_roles, - "projected_lidar_agent_ids": projected_lidar_agent_ids, - "projected_lidar_spoofing_masks": projected_lidar_spoofing_masks, + "origin_lidar": np.vstack(projected_lidar_stack), + "origin_lidar_by_agent": projected_lidar_stack, + "origin_lidar_roles": projected_lidar_roles, + "origin_lidar_agent_ids": projected_lidar_agent_ids, + "origin_lidar_spoofing_masks": projected_lidar_spoofing_masks, } - def __getitem__(self, idx): - base_data_dict = self.retrieve_base_data(idx) - processed_data_dict = OrderedDict() - processed_data_dict["ego"] = {} + def assemble_inference_sample( + self, + inference_input: dict[str, Any], + local_supervision: dict[str, Any], + visualization_context: dict[str, Any], + ) -> OrderedDict[str, dict[str, Any]]: + """ + Combine delivered LiDAR with local supervision and visualization. - ego_id, ego_lidar_pose = self.__find_ego_vehicle(base_data_dict) + Parameters + ---------- + inference_input : dict[str, Any] + Ego LiDAR plus point clouds from messages that were delivered. + local_supervision : dict[str, Any] + Ground truth produced solely from the local dataset snapshot. + visualization_context : dict[str, Any] + Visualization data produced solely from the local snapshot. - if self.payload_handler is not None: - data = self.__process_with_messages(ego_id, ego_lidar_pose, base_data_dict) - else: - data = self.__process_without_messages(ego_lidar_pose, base_data_dict) - - # exclude all repetitive objects - unique_indices = [data["object_id_stack"].index(x) for x in set(data["object_id_stack"])] - object_stack = np.vstack(data["object_stack"]) - object_stack = object_stack[unique_indices] - - # make sure bounding boxes across all frames have the same number - object_bbx_center = np.zeros((self.params["postprocess"]["max_num"], 7)) - mask = np.zeros(self.params["postprocess"]["max_num"]) - object_bbx_center[: object_stack.shape[0], :] = object_stack - mask[: object_stack.shape[0]] = 1 - - # convert list to numpy array, (N, 4) - projected_lidar_stack = np.vstack(data["projected_lidar_stack"]) - point_source = np.concatenate( - [np.full(points.shape[0], idx, dtype=np.int32) for idx, points in enumerate(data["projected_lidar_stack"])], axis=0 + Returns + ------- + collections.OrderedDict + Early-fusion sample ready for collation. + """ + projected_lidar_stack = np.vstack(inference_input["projected_lidar_stack"]) + object_bbx_center = local_supervision["object_bbx_center"] + object_bbx_mask = local_supervision["object_bbx_mask"] + + projected_lidar_stack, object_bbx_center, object_bbx_mask = self.augment( + projected_lidar_stack, + object_bbx_center, + object_bbx_mask, ) - point_spoofing = np.concatenate([np.asarray(mask, dtype=np.bool_) for mask in data["projected_lidar_spoofing_masks"]], axis=0) - # data augmentation - projected_lidar_stack, object_bbx_center, mask = self.augment(projected_lidar_stack, object_bbx_center, mask) - - # we do lidar filtering in the stacked lidar lidar_range = self.params["preprocess"]["cav_lidar_range"] lidar_mask = ( (projected_lidar_stack[:, 0] > lidar_range[0]) @@ -183,55 +263,53 @@ def __getitem__(self, idx): & (projected_lidar_stack[:, 2] < lidar_range[5]) ) projected_lidar_stack = projected_lidar_stack[lidar_mask] - point_source = point_source[lidar_mask] - point_spoofing = point_spoofing[lidar_mask] - # augmentation may remove some of the bbx out of range - object_bbx_center_valid = object_bbx_center[mask == 1] - object_bbx_center_valid, range_mask = box_utils.mask_boxes_outside_range_numpy( - object_bbx_center_valid, lidar_range, self.params["postprocess"]["order"], return_mask=True - ) - mask[object_bbx_center_valid.shape[0] :] = 0 - object_bbx_center[: object_bbx_center_valid.shape[0]] = object_bbx_center_valid - object_bbx_center[object_bbx_center_valid.shape[0] :] = 0 - unique_indices = list(np.array(unique_indices)[range_mask]) - # pre-process the lidar to voxel/bev/downsampled lidar - lidar_dict = self.pre_processor.preprocess(projected_lidar_stack) + valid_indices = np.flatnonzero(object_bbx_mask == 1) + valid_object_ids = [local_supervision["object_ids"][index] for index in valid_indices] + valid_boxes, range_mask = box_utils.mask_boxes_outside_range_numpy( + object_bbx_center[valid_indices], + lidar_range, + self.params["postprocess"]["order"], + return_mask=True, + ) + object_bbx_center.fill(0) + object_bbx_mask.fill(0) + object_bbx_center[: valid_boxes.shape[0]] = valid_boxes + object_bbx_mask[: valid_boxes.shape[0]] = 1 - # generate the anchor boxes anchor_box = self.post_processor.generate_anchor_box() - - # generate targets label - label_dict = self.post_processor.generate_label(gt_box_center=object_bbx_center, anchors=anchor_box, mask=mask) - - processed_data_dict["ego"].update( - { - "object_bbx_center": object_bbx_center, - "object_bbx_mask": mask, - "object_ids": [data["object_id_stack"][i] for i in unique_indices], - "anchor_box": anchor_box, - "processed_lidar": lidar_dict, - "label_dict": label_dict, - } + label_dict = self.post_processor.generate_label( + gt_box_center=object_bbx_center, + anchors=anchor_box, + mask=object_bbx_mask, ) + ego_sample = { + "object_bbx_center": object_bbx_center, + "object_bbx_mask": object_bbx_mask, + "object_ids": [object_id for object_id, keep in zip(valid_object_ids, range_mask) if keep], + "anchor_box": anchor_box, + "processed_lidar": self.pre_processor.preprocess(projected_lidar_stack), + "label_dict": label_dict, + **visualization_context, + } + return OrderedDict({"ego": ego_sample}) - if self.visualize: - valid_agent_indices = [idx for idx in range(len(data["projected_lidar_roles"])) if np.any(point_source == idx)] - processed_data_dict["ego"].update( - { - "origin_lidar": projected_lidar_stack, - "origin_lidar_by_agent": [projected_lidar_stack[point_source == idx] for idx in valid_agent_indices], - "origin_lidar_roles": [data["projected_lidar_roles"][idx] for idx in valid_agent_indices], - "origin_lidar_agent_ids": [data["projected_lidar_agent_ids"][idx] for idx in valid_agent_indices], - "origin_lidar_spoofing_masks": [point_spoofing[point_source == idx] for idx in valid_agent_indices], - } - ) + def __getitem__(self, idx): + base_data_dict = self.retrieve_base_data(idx) + ego_id, ego_lidar_pose = self.__find_ego_vehicle(base_data_dict) - return processed_data_dict + if self.payload_handler is not None: + inference_input = self.__process_with_messages(ego_id, ego_lidar_pose, base_data_dict) + else: + inference_input = self.__process_without_messages(ego_lidar_pose, base_data_dict) + + local_supervision = self.build_local_supervision(base_data_dict, ego_lidar_pose) + visualization_context = self.build_visualization_context(ego_id, ego_lidar_pose, base_data_dict) + return self.assemble_inference_sample(inference_input, local_supervision, visualization_context) def get_item_single_car(self, selected_cav_base, ego_pose): """ - Project the lidar and bbx to ego space first, and then do clipping. + Project one agent's LiDAR into ego coordinates. Parameters ---------- @@ -250,9 +328,6 @@ def get_item_single_car(self, selected_cav_base, ego_pose): # calculate the transformation matrix transformation_matrix = x1_to_x2(selected_cav_base["params"]["lidar_pose"], ego_pose) - # retrieve objects under ego coordinates - object_bbx_center, object_bbx_mask, object_ids = self.post_processor.generate_object_center([selected_cav_base], ego_pose) - # filter lidar lidar_np = selected_cav_base["lidar_np"] spoofing_mask = np.asarray(selected_cav_base.get("spoofing_mask", np.zeros((lidar_np.shape[0],), dtype=np.bool_)), dtype=np.bool_) @@ -270,8 +345,6 @@ def get_item_single_car(self, selected_cav_base, ego_pose): selected_cav_processed.update( { - "object_bbx_center": object_bbx_center[object_bbx_mask == 1], - "object_ids": object_ids, "projected_lidar": lidar_np, "projected_lidar_spoofing_mask": spoofing_mask, } diff --git a/OpenCOOD/opencood/data_utils/datasets/intermediate_fusion_dataset.py b/OpenCOOD/opencood/data_utils/datasets/intermediate_fusion_dataset.py index 365b8ac8..545f8e7f 100644 --- a/OpenCOOD/opencood/data_utils/datasets/intermediate_fusion_dataset.py +++ b/OpenCOOD/opencood/data_utils/datasets/intermediate_fusion_dataset.py @@ -5,6 +5,7 @@ import math import logging from collections import OrderedDict +from typing import Any import numpy as np import torch @@ -13,11 +14,30 @@ import opencood.data_utils.post_processor as post_processor from opencood.utils import box_utils from opencood.data_utils.datasets import basedataset +from opencood.models.communication_adapters import ( + InferenceInput, + IntermediateFusionWirePayload, + IntermediateMetadata, + PoseFrameMetadata, + V2XViTMetadata, + build_inference_input, + merge_inference_inputs, +) from opencood.data_utils.pre_processor import build_preprocessor from opencood.utils.pcd_utils import pcd_to_np, mask_points_by_range, mask_ego_points, shuffle_points, downsample_lidar_minimum +from opencood.utils.transformation_utils import x1_to_x2 logger = logging.getLogger("cavise.opencda.OpenCOOD.opencood.data_utils.datasets.intermediate_fusion_dataset") +PAIRWISE_METADATA_MODELS = frozenset( + { + "point_pillar_coalign", + "point_pillar_v2vnet", + "point_pillar_where2comm", + } +) +V2X_VIT_METADATA_MODELS = frozenset({"point_pillar_transformer"}) + class IntermediateFusionDataset(basedataset.BaseDataset): """ @@ -47,6 +67,7 @@ def __init__(self, params, visualize, train=True, payload_handler=None): self.payload_handler = payload_handler self.module_name = "OpenCOOD.IntermediateFusionDataset" + self.model_name = params["model"]["core_method"] def extract_data(self, idx): base_data_dict = self.retrieve_base_data(idx, cur_ego_pose_flag=self.cur_ego_pose_flag) @@ -55,17 +76,102 @@ def extract_data(self, idx): if self.payload_handler is not None: for cav_id, selected_cav_base in base_data_dict.items(): selected_cav_processed = self.get_item_single_car(selected_cav_base, ego_lidar_pose) + payload = self.build_wire_payload(cav_id, selected_cav_base, selected_cav_processed, idx) + self.payload_handler.set_opencda_payload(cav_id, self.module_name, payload) + + def build_wire_payload( + self, + cav_id: str, + selected_cav_base: dict[str, Any], + selected_cav_processed: dict[str, Any], + receive_frame: int, + ) -> IntermediateFusionWirePayload: + """ + Build the intermediate-fusion payload sent over V2X. + + Parameters + ---------- + cav_id : str + Identifier of the sending CAV or RSU. + selected_cav_base : dict[str, Any] + Raw dataset entry for the sending agent. + selected_cav_processed : dict[str, Any] + Processed data for the same sending agent. + receive_frame : int + Dataset frame in which the payload is sent. + + Returns + ------- + IntermediateFusionWirePayload + Typed payload containing only remote model-input data. + """ + return IntermediateFusionWirePayload( + inference_input=selected_cav_processed["inference_input"], + metadata=self.build_model_metadata(cav_id, selected_cav_base, receive_frame), + ) + + def build_model_metadata( + self, + cav_id: str, + selected_cav_base: dict[str, Any], + receive_frame: int, + ) -> IntermediateMetadata | None: + """ + Build only the metadata consumed by the configured model. - with self.payload_handler.handle_opencda_payload(cav_id, self.module_name) as msg: - msg["infra"] = 1 if "rsu" in cav_id else 0 - msg["velocity"] = selected_cav_processed["velocity"] - msg["time_delay"] = selected_cav_base["time_delay"] - msg["object_ids"] = selected_cav_processed["object_ids"] - msg["object_bbx_center"] = selected_cav_processed["object_bbx_center"] - msg["spatial_correction_matrix"] = selected_cav_base["params"]["spatial_correction_matrix"] - msg["voxel_num_points"] = selected_cav_processed["processed_features"]["voxel_num_points"] - msg["voxel_features"] = selected_cav_processed["processed_features"]["voxel_features"] - msg["voxel_coords"] = selected_cav_processed["processed_features"]["voxel_coords"] + Parameters + ---------- + cav_id : str + Identifier of the sending CAV or RSU. + selected_cav_base : dict[str, Any] + Local dataset entry for the sending agent. + receive_frame : int + Dataset frame in which the model input is received. + + Returns + ------- + IntermediateMetadata | None + Model-specific metadata, or ``None`` for models that consume only + the preprocessed inference input. + """ + capture_frame = receive_frame - int(selected_cav_base["time_delay"]) + if self.model_name in PAIRWISE_METADATA_MODELS: + return PoseFrameMetadata( + lidar_pose=selected_cav_base["params"]["lidar_pose"], + capture_frame=capture_frame, + ) + if self.model_name in V2X_VIT_METADATA_MODELS: + return V2XViTMetadata( + velocity=float(selected_cav_base["params"]["ego_speed"] / 30), + agent_type=1 if "rsu" in cav_id else 0, + capture_frame=capture_frame, + spatial_correction_matrix=selected_cav_base["params"]["spatial_correction_matrix"], + ) + return None + + @staticmethod + def decode_wire_payload(payload: object) -> IntermediateFusionWirePayload: + """ + Validate an intermediate-fusion payload received over V2X. + + Parameters + ---------- + payload : object + Deserialized module payload received over V2X. + + Returns + ------- + IntermediateFusionWirePayload + Validated intermediate-fusion payload. + + Raises + ------ + TypeError + If the received payload has an unexpected type. + """ + if not isinstance(payload, IntermediateFusionWirePayload): + raise TypeError(f"Expected IntermediateFusionWirePayload, got {type(payload).__name__}") + return payload def __find_ego_vehicle(self, base_data_dict): ego_id = -1 @@ -114,22 +220,6 @@ def __retrieve_visualization_base_data(self, idx): data[cav_id]["spoofing_mask"] = cav_snapshot["spoofing_mask"] return data - def __prepare_object_stack(self, object_stack, object_id_stack): - # exclude all repetitive objects - unique_indices = [object_id_stack.index(x) for x in set(object_id_stack)] - object_stack = np.vstack(object_stack) - object_stack = object_stack[unique_indices] - - # make sure bounding boxes across all frames have the same number - object_bbx_center = np.zeros((self.params["postprocess"]["max_num"], 7)) - mask = np.zeros(self.params["postprocess"]["max_num"]) - object_bbx_center[: object_stack.shape[0], :] = object_stack - mask[: object_stack.shape[0]] = 1 - - filtered_object_ids = [object_id_stack[i] for i in unique_indices] - - return object_bbx_center, mask, filtered_object_ids - def __pad_spatial_matrix(self, matrix_list): matrix_list = np.stack(matrix_list) padding_eye = np.tile(np.eye(4)[None], (self.max_cav - len(matrix_list), 1, 1)) @@ -138,43 +228,97 @@ def __pad_spatial_matrix(self, matrix_list): def __pad_to_max(self, lst, pad_value): return lst + (self.max_cav - len(lst)) * [pad_value] + def __build_pairwise_transformation(self, metadata: list[PoseFrameMetadata]) -> np.ndarray: + """ + Build pairwise transformations for the received agents. + + Parameters + ---------- + metadata : list[PoseFrameMetadata] + Pose metadata ordered exactly like the received model inputs. + + Returns + ------- + numpy.ndarray + Pairwise transformation matrix padded to ``max_cav``. + """ + pairwise_t_matrix = np.zeros((self.max_cav, self.max_cav, 4, 4)) + if self.proj_first: + pairwise_t_matrix[:, :] = np.identity(4) + return pairwise_t_matrix + + for source_index, source_metadata in enumerate(metadata): + for target_index, target_metadata in enumerate(metadata): + pairwise_t_matrix[source_index, target_index] = x1_to_x2( + source_metadata.lidar_pose, + target_metadata.lidar_pose, + ) + return pairwise_t_matrix + + def __assemble_model_metadata( + self, + metadata: list[IntermediateMetadata | None], + receive_frame: int, + ) -> dict[str, Any]: + """ + Derive model tensors from received metadata. + + Parameters + ---------- + metadata : list[IntermediateMetadata | None] + Metadata ordered exactly like the available inference inputs. + receive_frame : int + Current receiver dataset frame. + + Returns + ------- + dict[str, Any] + Model fields derived from metadata. + + Raises + ------ + TypeError + If the received metadata does not match the configured model. + """ + if self.model_name in PAIRWISE_METADATA_MODELS: + if not all(isinstance(agent_metadata, PoseFrameMetadata) for agent_metadata in metadata): + raise TypeError(f"{self.model_name} requires PoseFrameMetadata for every agent") + pose_metadata = [agent_metadata for agent_metadata in metadata if isinstance(agent_metadata, PoseFrameMetadata)] + return {"pairwise_t_matrix": self.__build_pairwise_transformation(pose_metadata)} + + if self.model_name in V2X_VIT_METADATA_MODELS: + if not all(isinstance(agent_metadata, V2XViTMetadata) for agent_metadata in metadata): + raise TypeError(f"{self.model_name} requires V2XViTMetadata for every agent") + v2x_metadata = [agent_metadata for agent_metadata in metadata if isinstance(agent_metadata, V2XViTMetadata)] + return { + "velocity": self.__pad_to_max([agent_metadata.velocity for agent_metadata in v2x_metadata], 0.0), + "time_delay": self.__pad_to_max( + [float(receive_frame - agent_metadata.capture_frame) for agent_metadata in v2x_metadata], + 0.0, + ), + "infra": self.__pad_to_max([agent_metadata.agent_type for agent_metadata in v2x_metadata], 0.0), + "spatial_correction_matrix": self.__pad_spatial_matrix([agent_metadata.spatial_correction_matrix for agent_metadata in v2x_metadata]), + } + + if any(agent_metadata is not None for agent_metadata in metadata): + raise TypeError(f"{self.model_name} does not consume wire metadata") + return {} + @staticmethod def __build_model_data(): return { - "processed_features": [], - "object_stack": [], - "object_id_stack": [], - "velocity": [], - "time_delay": [], - "infra": [], - "spatial_correction_matrix": [], + "inference_inputs": [], + "metadata": [], } - @staticmethod - def __append_processed_model_data(data, cav_id, cav_base, cav_processed): - data["infra"].append(1 if "rsu" in cav_id else 0) - data["velocity"].append(cav_processed["velocity"]) - data["time_delay"].append(float(cav_base["time_delay"])) - data["object_id_stack"] += cav_processed["object_ids"] - data["object_stack"].append(cav_processed["object_bbx_center"]) - data["spatial_correction_matrix"].append(cav_base["params"]["spatial_correction_matrix"]) - data["processed_features"].append(cav_processed["processed_features"]) + def __append_processed_model_data(self, data, cav_id, cav_base, cav_processed, receive_frame): + data["inference_inputs"].append(cav_processed["inference_input"]) + data["metadata"].append(self.build_model_metadata(cav_id, cav_base, receive_frame)) @staticmethod - def __append_message_model_data(data, msg): - data["infra"].append(msg["infra"]) - data["velocity"].append(msg["velocity"]) - data["time_delay"].append(msg["time_delay"]) - data["object_id_stack"] += msg["object_ids"] - data["object_stack"].append(msg["object_bbx_center"]) - data["spatial_correction_matrix"].append(msg["spatial_correction_matrix"]) - data["processed_features"].append( - { - "voxel_num_points": msg["voxel_num_points"], - "voxel_features": msg["voxel_features"], - "voxel_coords": msg["voxel_coords"], - } - ) + def __append_message_model_data(data: dict[str, Any], payload: IntermediateFusionWirePayload) -> None: + data["inference_inputs"].append(payload.inference_input) + data["metadata"].append(payload.metadata) def __agent_in_communication_range(self, cav_base, ego_lidar_pose): dx = cav_base["params"]["lidar_pose"][0] - ego_lidar_pose[0] @@ -189,7 +333,22 @@ def __project_lidar_for_visualization(self, cav_base): lidar_np[:, :3] = box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix) return mask_points_by_range(lidar_np, self.params["preprocess"]["cav_lidar_range"]) - def __build_visualization_data(self, ego_id, base_data_dict): + def build_visualization_context(self, ego_id: str, base_data_dict: dict[str, Any]) -> dict[str, Any]: + """ + Build visualization data from the local scene snapshot. + + Parameters + ---------- + ego_id : str + Identifier of the receiving ego agent. + base_data_dict : dict[str, Any] + Local dataset snapshot, independent of delivered messages. + + Returns + ------- + dict[str, Any] + Per-agent projected point clouds and their local identities. + """ if not self.visualize: return { "projected_lidar_stack": [], @@ -212,92 +371,141 @@ def __build_visualization_data(self, ego_id, base_data_dict): "projected_lidar_agent_ids": projected_lidar_agent_ids, } - def __append_visualization_data(self, data, ego_id, base_data_dict): - data.update(self.__build_visualization_data(ego_id, base_data_dict)) - return data - - def __process_with_messages(self, ego_id, ego_lidar_pose, base_data_dict, visualization_base_data_dict): + def __process_with_messages(self, ego_id, ego_lidar_pose, base_data_dict, receive_frame): data = self.__build_model_data() ego_cav_base = base_data_dict.get(ego_id) ego_cav_processed = self.get_item_single_car(ego_cav_base, ego_lidar_pose) - self.__append_processed_model_data(data, ego_id, ego_cav_base, ego_cav_processed) + self.__append_processed_model_data(data, ego_id, ego_cav_base, ego_cav_processed, receive_frame) if ego_id in self.payload_handler.current_artery_payload: for cav_id, _ in base_data_dict.items(): - if cav_id in self.payload_handler.current_artery_payload[ego_id]: - with self.payload_handler.handle_artery_payload(ego_id, cav_id, self.module_name) as msg: - self.__append_message_model_data(data, msg) + raw_payload = self.payload_handler.get_artery_payload(ego_id, cav_id, self.module_name) + if raw_payload is None: + continue + payload = self.decode_wire_payload(raw_payload) + self.__append_message_model_data(data, payload) - return self.__append_visualization_data(data, ego_id, visualization_base_data_dict) + return data - def __process_without_messages(self, ego_id, ego_lidar_pose, base_data_dict, visualization_base_data_dict): + def __process_without_messages(self, ego_id, ego_lidar_pose, base_data_dict, receive_frame): data = self.__build_model_data() for cav_id, selected_cav_base in base_data_dict.items(): if not self.__agent_in_communication_range(selected_cav_base, ego_lidar_pose): continue selected_cav_processed = self.get_item_single_car(selected_cav_base, ego_lidar_pose) - self.__append_processed_model_data(data, cav_id, selected_cav_base, selected_cav_processed) - - return self.__append_visualization_data(data, ego_id, visualization_base_data_dict) - - def __getitem__(self, idx): - base_data_dict = self.retrieve_base_data(idx, cur_ego_pose_flag=self.cur_ego_pose_flag) - processed_data_dict = OrderedDict() - processed_data_dict["ego"] = {} + self.__append_processed_model_data(data, cav_id, selected_cav_base, selected_cav_processed, receive_frame) - ego_id, ego_lidar_pose = self.__find_ego_vehicle(base_data_dict) - pairwise_t_matrix = self.get_pairwise_transformation(base_data_dict, self.max_cav) - visualization_base_data_dict = self.__retrieve_visualization_base_data(idx) if self.visualize else base_data_dict + return data - if self.payload_handler is not None: - data = self.__process_with_messages(ego_id, ego_lidar_pose, base_data_dict, visualization_base_data_dict) - else: - data = self.__process_without_messages(ego_id, ego_lidar_pose, base_data_dict, visualization_base_data_dict) + def build_local_supervision(self, base_data_dict: dict[str, Any], ego_lidar_pose: list[float]) -> dict[str, Any]: + """ + Build ground truth and training targets from the local scene. - object_bbx_center, mask, object_ids = self.__prepare_object_stack(data["object_stack"], data["object_id_stack"]) + Parameters + ---------- + base_data_dict : dict[str, Any] + Complete local dataset snapshot for the current scene. + ego_lidar_pose : list[float] + Current ego LiDAR pose in world coordinates. - merged_feature_dict = self.merge_features_to_dict(data["processed_features"]) + Returns + ------- + dict[str, Any] + Ground-truth boxes, object identifiers, anchors, and labels. + """ + object_bbx_center, object_bbx_mask, object_ids = self.post_processor.generate_object_center(list(base_data_dict.values()), ego_lidar_pose) anchor_box = self.post_processor.generate_anchor_box() - label_dict = self.post_processor.generate_label(gt_box_center=object_bbx_center, anchors=anchor_box, mask=mask) + label_dict = self.post_processor.generate_label( + gt_box_center=object_bbx_center, + anchors=anchor_box, + mask=object_bbx_mask, + ) + return { + "object_bbx_center": object_bbx_center, + "object_bbx_mask": object_bbx_mask, + "object_ids": object_ids, + "anchor_box": anchor_box, + "label_dict": label_dict, + } - spatial_correction_matrix = self.__pad_spatial_matrix(data["spatial_correction_matrix"]) - velocity = self.__pad_to_max(data["velocity"], 0.0) - time_delay = self.__pad_to_max(data["time_delay"], 0.0) - infra = self.__pad_to_max(data["infra"], 0.0) + def assemble_inference_sample( + self, + inference_input: dict[str, Any], + local_supervision: dict[str, Any], + visualization_context: dict[str, Any], + receive_frame: int, + ) -> OrderedDict[str, dict[str, Any]]: + """ + Assemble model input with local supervision and visualization. - processed_data_dict["ego"].update( - { - "object_bbx_center": object_bbx_center, - "object_bbx_mask": mask, - "object_ids": object_ids, - "anchor_box": anchor_box, - "processed_lidar": merged_feature_dict, - "label_dict": label_dict, - "cav_num": len(data["processed_features"]), - "velocity": velocity, - "time_delay": time_delay, - "infra": infra, - "spatial_correction_matrix": spatial_correction_matrix, - "pairwise_t_matrix": pairwise_t_matrix, - } - ) + Parameters + ---------- + inference_input : dict[str, Any] + Ego features plus features from messages that were delivered. + local_supervision : dict[str, Any] + Ground truth produced solely from the local dataset snapshot. + visualization_context : dict[str, Any] + Visualization data produced solely from the local snapshot. + receive_frame : int + Current receiver dataset frame. + + Returns + ------- + collections.OrderedDict + Dataset sample consumed by the collate and inference pipeline. + Raises + ------ + ValueError + If model inputs and metadata entries are not aligned. + """ + typed_inputs: list[InferenceInput] = inference_input["inference_inputs"] + merged_feature_dict = merge_inference_inputs(typed_inputs) + metadata: list[IntermediateMetadata | None] = inference_input["metadata"] + if len(metadata) != len(typed_inputs): + raise ValueError("Every inference input must have a matching metadata entry") + + ego_sample = { + **local_supervision, + "processed_lidar": merged_feature_dict, + "cav_num": len(typed_inputs), + **self.__assemble_model_metadata(metadata, receive_frame), + } if self.visualize: - processed_data_dict["ego"].update( + ego_sample.update( { - "origin_lidar": np.vstack(data["projected_lidar_stack"]), - "origin_lidar_by_agent": data["projected_lidar_stack"], - "origin_lidar_roles": data["projected_lidar_roles"], - "origin_lidar_agent_ids": data["projected_lidar_agent_ids"], + "origin_lidar": np.vstack(visualization_context["projected_lidar_stack"]), + "origin_lidar_by_agent": visualization_context["projected_lidar_stack"], + "origin_lidar_roles": visualization_context["projected_lidar_roles"], + "origin_lidar_agent_ids": visualization_context["projected_lidar_agent_ids"], } ) - return processed_data_dict + return OrderedDict({"ego": ego_sample}) + + def __getitem__(self, idx): + base_data_dict = self.retrieve_base_data(idx, cur_ego_pose_flag=self.cur_ego_pose_flag) + ego_id, ego_lidar_pose = self.__find_ego_vehicle(base_data_dict) + visualization_base_data_dict = self.__retrieve_visualization_base_data(idx) if self.visualize else base_data_dict + + if self.payload_handler is not None: + inference_input = self.__process_with_messages(ego_id, ego_lidar_pose, base_data_dict, idx) + else: + inference_input = self.__process_without_messages(ego_id, ego_lidar_pose, base_data_dict, idx) + + local_supervision = self.build_local_supervision(base_data_dict, ego_lidar_pose) + visualization_context = self.build_visualization_context(ego_id, visualization_base_data_dict) + return self.assemble_inference_sample( + inference_input, + local_supervision, + visualization_context, + idx, + ) def get_item_single_car(self, selected_cav_base, ego_pose): """ - Project the lidar and bbx to ego space first, and then do clipping. + Project and preprocess one agent's LiDAR input. Parameters ---------- @@ -316,9 +524,6 @@ def get_item_single_car(self, selected_cav_base, ego_pose): # calculate the transformation matrix transformation_matrix = selected_cav_base["params"]["transformation_matrix"] - # retrieve objects under ego coordinates - object_bbx_center, object_bbx_mask, object_ids = self.post_processor.generate_object_center([selected_cav_base], ego_pose) - # filter lidar lidar_np = selected_cav_base["lidar_np"] lidar_np = shuffle_points(lidar_np) @@ -328,20 +533,12 @@ def get_item_single_car(self, selected_cav_base, ego_pose): if self.proj_first: lidar_np[:, :3] = box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix) lidar_np = mask_points_by_range(lidar_np, self.params["preprocess"]["cav_lidar_range"]) - processed_lidar = self.pre_processor.preprocess(lidar_np) - - # velocity - velocity = selected_cav_base["params"]["ego_speed"] - # normalize veloccity by average speed 30 km/h - velocity = velocity / 30 + inference_input = build_inference_input(self.pre_processor.preprocess(lidar_np)) selected_cav_processed.update( { - "object_bbx_center": object_bbx_center[object_bbx_mask == 1], - "object_ids": object_ids, "projected_lidar": lidar_np, - "processed_features": processed_lidar, - "velocity": velocity, + "inference_input": inference_input, } ) @@ -414,12 +611,14 @@ def collate_batch_train(self, batch): processed_lidar_list.append(ego_dict["processed_lidar"]) record_len.append(ego_dict["cav_num"]) label_dict_list.append(ego_dict["label_dict"]) - pairwise_t_matrix_list.append(ego_dict["pairwise_t_matrix"]) - velocity.append(ego_dict["velocity"]) - time_delay.append(ego_dict["time_delay"]) - infra.append(ego_dict["infra"]) - spatial_correction_matrix_list.append(ego_dict["spatial_correction_matrix"]) + if self.model_name in PAIRWISE_METADATA_MODELS: + pairwise_t_matrix_list.append(ego_dict["pairwise_t_matrix"]) + elif self.model_name in V2X_VIT_METADATA_MODELS: + velocity.append(ego_dict["velocity"]) + time_delay.append(ego_dict["time_delay"]) + infra.append(ego_dict["infra"]) + spatial_correction_matrix_list.append(ego_dict["spatial_correction_matrix"]) if self.visualize: origin_lidar.append(ego_dict["origin_lidar"]) @@ -427,24 +626,13 @@ def collate_batch_train(self, batch): object_bbx_center = torch.from_numpy(np.array(object_bbx_center)) object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask)) - # example: {'voxel_features':[np.array([1,2,3]]), - # np.array([3,5,6]), ...]} + # Merge per-sample agent inputs before preprocessor-specific collation. merged_feature_dict = self.merge_features_to_dict(processed_lidar_list) processed_lidar_torch_dict = self.pre_processor.collate_batch(merged_feature_dict) # [2, 3, 4, ..., M], M <= max_cav record_len = torch.from_numpy(np.array(record_len, dtype=int)) label_torch_dict = self.post_processor.collate_batch(label_dict_list) - # (B, max_cav) - velocity = torch.from_numpy(np.array(velocity)) - time_delay = torch.from_numpy(np.array(time_delay)) - infra = torch.from_numpy(np.array(infra)) - spatial_correction_matrix_list = torch.from_numpy(np.array(spatial_correction_matrix_list)) - # (B, max_cav, 3) - prior_encoding = torch.stack([velocity, time_delay, infra], dim=-1).float() - # (B, max_cav) - pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list)) - # object id is only used during inference, where batch size is 1. # so here we only get the first element. output_dict["ego"].update( @@ -455,12 +643,22 @@ def collate_batch_train(self, batch): "record_len": record_len, "label_dict": label_torch_dict, "object_ids": object_ids[0], - "prior_encoding": prior_encoding, - "spatial_correction_matrix": spatial_correction_matrix_list, - "pairwise_t_matrix": pairwise_t_matrix, } ) + if self.model_name in PAIRWISE_METADATA_MODELS: + output_dict["ego"]["pairwise_t_matrix"] = torch.from_numpy(np.array(pairwise_t_matrix_list)) + elif self.model_name in V2X_VIT_METADATA_MODELS: + velocity_tensor = torch.from_numpy(np.array(velocity)) + time_delay_tensor = torch.from_numpy(np.array(time_delay)) + infra_tensor = torch.from_numpy(np.array(infra)) + output_dict["ego"].update( + { + "prior_encoding": torch.stack([velocity_tensor, time_delay_tensor, infra_tensor], dim=-1).float(), + "spatial_correction_matrix": torch.from_numpy(np.array(spatial_correction_matrix_list)), + } + ) + if self.visualize: origin_lidar = np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar)) origin_lidar = torch.from_numpy(origin_lidar) @@ -514,47 +712,3 @@ def post_process(self, data_dict, output_dict): gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict) return pred_box_tensor, pred_score, gt_box_tensor - - def get_pairwise_transformation(self, base_data_dict, max_cav): - """ - Get pair-wise transformation matrix accross different agents. - - Parameters - ---------- - base_data_dict : dict - Key : cav id, item: transformation matrix to ego, lidar points. - - max_cav : int - The maximum number of cav, default 5 - - Return - ------ - pairwise_t_matrix : np.array - The pairwise transformation matrix across each cav. - shape: (L, L, 4, 4) - """ - pairwise_t_matrix = np.zeros((max_cav, max_cav, 4, 4)) - - if self.proj_first: - # if lidar projected to ego first, then the pairwise matrix - # becomes identity - pairwise_t_matrix[:, :] = np.identity(4) - else: - t_list = [] - - # save all transformation matrix in a list in order first. - for cav_id, cav_content in base_data_dict.items(): - t_list.append(cav_content["params"]["transformation_matrix"]) - - for i in range(len(t_list)): - for j in range(len(t_list)): - # identity matrix to self - if i == j: - t_matrix = np.eye(4) - pairwise_t_matrix[i, j] = t_matrix - continue - # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj - t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i]) - pairwise_t_matrix[i, j] = t_matrix - - return pairwise_t_matrix diff --git a/OpenCOOD/opencood/data_utils/datasets/intermediate_fusion_dataset_v2.py b/OpenCOOD/opencood/data_utils/datasets/intermediate_fusion_dataset_v2.py index 081a1974..c21f2e72 100644 --- a/OpenCOOD/opencood/data_utils/datasets/intermediate_fusion_dataset_v2.py +++ b/OpenCOOD/opencood/data_utils/datasets/intermediate_fusion_dataset_v2.py @@ -4,6 +4,8 @@ import math from collections import OrderedDict +from collections.abc import Callable +from typing import Any import numpy as np import torch @@ -12,6 +14,7 @@ import opencood.data_utils.post_processor as post_processor from opencood.utils import box_utils from opencood.data_utils.datasets import basedataset +from opencood.models.communication_adapters import PoseFrameMetadata from opencood.data_utils.pre_processor import build_preprocessor from opencood.utils.pcd_utils import pcd_to_np, mask_points_by_range, mask_ego_points, shuffle_points from opencood.utils.transformation_utils import x1_to_x2 @@ -46,20 +49,42 @@ def __init__(self, params, visualize, train=True, payload_handler=None): self.payload_handler = payload_handler self.module_name = "OpenCOOD.IntermediateFusionDatasetV2" - def extract_data(self, idx): + def extract_data( + self, + idx: int, + agent_payload_builder: Callable[[dict[str, Any], PoseFrameMetadata], object] | None = None, + ) -> None: + """ + Build model-specific payloads from local agent observations. + + Parameters + ---------- + idx : int + Dataset frame to publish. + agent_payload_builder : Callable[[dict[str, Any], PoseFrameMetadata], object] | None + Model adapter callback that converts one local agent observation + and its metadata into a typed wire payload. + + Raises + ------ + NotImplementedError + If the configured model has no communication adapter for this + dataset contract. + """ base_data_dict = self.retrieve_base_data(idx, cur_ego_pose_flag=self.cur_ego_pose_flag) - _, ego_lidar_pose = self.__find_ego_vehicle(base_data_dict) if self.payload_handler is not None: + if agent_payload_builder is None: + raise NotImplementedError("IntermediateFusionDatasetV2 requires a model-specific communication adapter") for cav_id, selected_cav_base in base_data_dict.items(): - selected_cav_processed = self.get_item_single_car(selected_cav_base, ego_lidar_pose) - - with self.payload_handler.handle_opencda_payload(cav_id, self.module_name) as msg: - msg["object_ids"] = selected_cav_processed["object_ids"] # list - msg["object_bbx_center"] = selected_cav_processed["object_bbx_center"] - msg["voxel_num_points"] = selected_cav_processed["processed_features"]["voxel_num_points"] - msg["voxel_features"] = selected_cav_processed["processed_features"]["voxel_features"] - msg["voxel_coords"] = selected_cav_processed["processed_features"]["voxel_coords"] + sender_pose = selected_cav_base["params"]["lidar_pose"] + selected_cav_processed = self.build_agent_model_input(selected_cav_base, sender_pose) + metadata = PoseFrameMetadata( + lidar_pose=tuple(float(value) for value in sender_pose), + capture_frame=int(idx - selected_cav_base["time_delay"]), + ) + payload = agent_payload_builder(selected_cav_processed, metadata) + self.payload_handler.set_opencda_payload(cav_id, self.module_name, payload) def __find_ego_vehicle(self, base_data_dict): ego_id = -1 @@ -115,7 +140,25 @@ def __project_lidar_for_visualization(self, cav_base, ego_lidar_pose): lidar_np[:, :3] = box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix) return mask_points_by_range(lidar_np, self.params["preprocess"]["cav_lidar_range"]) - def __build_visualization_data(self, ego_id, ego_lidar_pose, base_data_dict): + def build_visualization_context(self, ego_id, ego_lidar_pose, base_data_dict): + """ + Build local visualization data independently of wire delivery. + + Parameters + ---------- + ego_id : str + Identifier of the receiving agent. + ego_lidar_pose : list[float] + Receiver pose used to project visualization point clouds. + base_data_dict : dict[str, Any] + Complete local scene snapshot. + + Returns + ------- + dict[str, Any] + Per-agent point clouds, roles, and identifiers used only by + visualization code. + """ if not self.visualize: return { "projected_lidar_stack": [], @@ -140,51 +183,75 @@ def __build_visualization_data(self, ego_id, ego_lidar_pose, base_data_dict): def __process_with_messages(self, ego_id, ego_lidar_pose, base_data_dict, visualization_ego_lidar_pose, visualization_base_data_dict): processed_features = [] - object_stack = [] - object_id_stack = [] model_lidar_stack = [] + remote_agent_outputs = [] ego_cav_base = base_data_dict.get(ego_id) - ego_cav_processed = self.get_item_single_car(ego_cav_base, ego_lidar_pose) + ego_cav_processed = self.build_agent_model_input(ego_cav_base, ego_lidar_pose) - object_id_stack += ego_cav_processed["object_ids"] - object_stack.append(ego_cav_processed["object_bbx_center"]) processed_features.append(ego_cav_processed["processed_features"]) model_lidar_stack.append(ego_cav_processed["projected_lidar"]) + if self.communication_adapter is None: + raise RuntimeError("IntermediateFusionDatasetV2 requires an attached communication adapter") + if ego_id in self.payload_handler.current_artery_payload: for cav_id, _ in base_data_dict.items(): - if cav_id in self.payload_handler.current_artery_payload[ego_id]: - with self.payload_handler.handle_artery_payload(ego_id, cav_id, self.module_name) as msg: - cav_processed = self.get_item_single_car(base_data_dict[cav_id], ego_lidar_pose) - projected = cav_processed["projected_lidar"] - - if len(projected) > 10: - model_lidar_stack.append(projected) - object_id_stack += msg["object_ids"] - object_stack.append(msg["object_bbx_center"]) - - voxel_num_points = msg["voxel_num_points"] - voxel_features = msg["voxel_features"] - voxel_coords = msg["voxel_coords"] - - processed_features.append( - {"voxel_num_points": voxel_num_points, "voxel_features": voxel_features, "voxel_coords": voxel_coords} - ) + if cav_id == ego_id: + continue + raw_payload = self.payload_handler.get_artery_payload(ego_id, cav_id, self.module_name) + if raw_payload is None: + continue + remote_agent_outputs.append(self.communication_adapter.decode_received_payload(raw_payload, ego_lidar_pose)) return { "processed_features": processed_features, - "object_stack": object_stack, - "object_id_stack": object_id_stack, "model_lidar_stack": model_lidar_stack, - **self.__build_visualization_data(ego_id, visualization_ego_lidar_pose, visualization_base_data_dict), + "remote_agent_outputs": remote_agent_outputs, + **self.build_visualization_context(ego_id, visualization_ego_lidar_pose, visualization_base_data_dict), + } + + def build_local_supervision(self, base_data_dict: dict[str, Any], ego_lidar_pose: list[float]) -> dict[str, Any]: + """ + Build two-stage ground truth from the complete local scene. + + Parameters + ---------- + base_data_dict : dict[str, Any] + Complete local dataset snapshot, independent of packet delivery. + ego_lidar_pose : list[float] + Current receiver LiDAR pose. + + Returns + ------- + dict[str, Any] + Ground-truth boxes, masks, identifiers, and anchors, plus + stage-two targets when building a training sample. + """ + object_bbx_center, object_bbx_mask, object_ids = self.post_processor.generate_object_center( + list(base_data_dict.values()), + ego_lidar_pose, + ) + anchor_box = self.post_processor.generate_anchor_box() + supervision: dict[str, Any] = { + "object_bbx_center": object_bbx_center, + "object_bbx_mask": object_bbx_mask, + "object_ids": object_ids, + "anchor_box": anchor_box, } + if self.payload_handler is None: + supervision["stage2_label"] = self.post_processor.generate_label( + gt_box_center=object_bbx_center, + anchors=anchor_box, + mask=object_bbx_mask, + ) + return supervision def __process_without_messages(self, ego_id, ego_lidar_pose, base_data_dict, visualization_ego_lidar_pose, visualization_base_data_dict): processed_features = [] object_stack = [] - object_id_stack = [] model_lidar_stack = [] + agent_to_ego = [] for cav_id, selected_cav_base in base_data_dict.items(): dx = selected_cav_base["params"]["lidar_pose"][0] - ego_lidar_pose[0] @@ -194,25 +261,26 @@ def __process_without_messages(self, ego_id, ego_lidar_pose, base_data_dict, vis if distance > opencood.data_utils.datasets.COM_RANGE: continue - selected_cav_processed = self.get_item_single_car(selected_cav_base, ego_lidar_pose) + agent_pose = selected_cav_base["params"]["lidar_pose"] + selected_cav_processed = self.get_item_single_car(selected_cav_base, agent_pose) if len(selected_cav_processed["projected_lidar"]) > 10: object_stack.append(selected_cav_processed["object_bbx_center"]) - object_id_stack += selected_cav_processed["object_ids"] processed_features.append(selected_cav_processed["processed_features"]) model_lidar_stack.append(selected_cav_processed["projected_lidar"]) + agent_to_ego.append(x1_to_x2(agent_pose, ego_lidar_pose)) return { "processed_features": processed_features, "object_stack": object_stack, - "object_id_stack": object_id_stack, "model_lidar_stack": model_lidar_stack, - **self.__build_visualization_data(ego_id, visualization_ego_lidar_pose, visualization_base_data_dict), + "agent_to_ego": agent_to_ego, + **self.build_visualization_context(ego_id, visualization_ego_lidar_pose, visualization_base_data_dict), } def __getitem__(self, idx): # put here to avoid initialization error - base_data_dict = self.retrieve_base_data(idx) + base_data_dict = self.retrieve_base_data(idx, cur_ego_pose_flag=self.cur_ego_pose_flag) processed_data_dict = OrderedDict() processed_data_dict["ego"] = {} @@ -238,71 +306,54 @@ def __getitem__(self, idx): visualization_base_data_dict, ) - # exclude all repetitive objects - unique_indices = [data["object_id_stack"].index(x) for x in set(data["object_id_stack"])] - object_stack_all = np.vstack(data["object_stack"]) - object_stack_all = object_stack_all[unique_indices] - - # make sure bounding boxes across all frames have the same number - object_bbx_center = np.zeros((self.params["postprocess"]["max_num"], 7)) - mask = np.zeros(self.params["postprocess"]["max_num"]) - object_bbx_center[: object_stack_all.shape[0], :] = object_stack_all - mask[: object_stack_all.shape[0]] = 1 - # merge preprocessed features from different cavs into the same dict cav_num = len(data["processed_features"]) merged_feature_dict = self.merge_features_to_dict(data["processed_features"]) - # generate the anchor boxes - anchor_box = self.post_processor.generate_anchor_box() + local_supervision = self.build_local_supervision(base_data_dict, ego_lidar_pose) + anchor_box = local_supervision["anchor_box"] - # generate targets label - label_dict = self.post_processor.generate_label( - gt_box_center=object_bbx_center, # hwl - anchors=anchor_box, - mask=mask, - ) - - # Filter empty boxes - object_stack_filtered = [] - label_dict_no_coop = [] - for boxes, points in zip(data["object_stack"], data["model_lidar_stack"]): - point_indices = points_in_boxes_cpu(points[:, :3], boxes[:, [0, 1, 2, 5, 4, 3, 6]]) - cur_mask = point_indices.sum(axis=1) > 0 - if cur_mask.sum() == 0: - label_dict_no_coop.append( - { - "pos_equal_one": np.zeros((*anchor_box.shape[:2], self.post_processor.anchor_num)), - "neg_equal_one": np.ones((*anchor_box.shape[:2], self.post_processor.anchor_num)), - "targets": np.zeros((*anchor_box.shape[:2], self.post_processor.anchor_num * 7)), - } - ) - continue - object_stack_filtered.append(boxes[cur_mask]) - bbx_center = np.zeros((self.params["postprocess"]["max_num"], 7)) - bbx_mask = np.zeros(self.params["postprocess"]["max_num"]) - bbx_center[: boxes[cur_mask].shape[0], :] = boxes[cur_mask] - bbx_mask[: boxes[cur_mask].shape[0]] = 1 - label_dict_no_coop.append( - self.post_processor.generate_label( - gt_box_center=bbx_center, # hwl - anchors=anchor_box, - mask=bbx_mask, - ) - ) - label_dict = {"stage1": label_dict_no_coop, "stage2": label_dict} processed_data_dict["ego"].update( { - "object_bbx_center": object_bbx_center, - "object_bbx_mask": mask, - "object_ids": [data["object_id_stack"][i] for i in unique_indices], + "object_bbx_center": local_supervision["object_bbx_center"], + "object_bbx_mask": local_supervision["object_bbx_mask"], + "object_ids": local_supervision["object_ids"], "anchor_box": anchor_box, "processed_lidar": merged_feature_dict, - "label_dict": label_dict, "cav_num": cav_num, } ) + if "stage2_label" in local_supervision: + label_dict_no_coop = [] + for boxes, points in zip(data["object_stack"], data["model_lidar_stack"]): + point_indices = points_in_boxes_cpu(points[:, :3], boxes[:, [0, 1, 2, 5, 4, 3, 6]]) + cur_mask = point_indices.sum(axis=1) > 0 + if cur_mask.sum() == 0: + label_dict_no_coop.append( + { + "pos_equal_one": np.zeros((*anchor_box.shape[:2], self.post_processor.anchor_num)), + "neg_equal_one": np.ones((*anchor_box.shape[:2], self.post_processor.anchor_num)), + "targets": np.zeros((*anchor_box.shape[:2], self.post_processor.anchor_num * 7)), + } + ) + continue + bbx_center = np.zeros((self.params["postprocess"]["max_num"], 7)) + bbx_mask = np.zeros(self.params["postprocess"]["max_num"]) + bbx_center[: boxes[cur_mask].shape[0], :] = boxes[cur_mask] + bbx_mask[: boxes[cur_mask].shape[0]] = 1 + label_dict_no_coop.append( + self.post_processor.generate_label( + gt_box_center=bbx_center, + anchors=anchor_box, + mask=bbx_mask, + ) + ) + processed_data_dict["ego"]["label_dict"] = { + "stage1": label_dict_no_coop, + "stage2": local_supervision["stage2_label"], + } + processed_data_dict["ego"].update( { "origin_lidar": data["model_lidar_stack"], @@ -311,51 +362,59 @@ def __getitem__(self, idx): "origin_lidar_agent_ids": data["projected_lidar_agent_ids"], } ) + if "remote_agent_outputs" in data: + processed_data_dict["ego"]["remote_agent_outputs"] = data["remote_agent_outputs"] + if "agent_to_ego" in data: + processed_data_dict["ego"]["agent_to_ego"] = data["agent_to_ego"] return processed_data_dict - def get_item_single_car(self, selected_cav_base, ego_pose): + def build_agent_model_input(self, selected_cav_base, target_pose): """ - Project the lidar and bbx to ego space first, and then do clipping. + Build one agent's local two-stage input without supervision data. Parameters ---------- selected_cav_base : dict - The dictionary contains a single CAV's raw information. - ego_pose : list - The ego vehicle lidar pose under world coordinate. + Raw and metadata fields for one agent. + target_pose : list + Pose of the coordinate frame used for this model input. Returns ------- - selected_cav_processed : dict - The dictionary contains the cav's processed information. + dict + Projected raw points and their voxelized representation. """ - selected_cav_processed = {} - - # calculate the transformation matrix - transformation_matrix = x1_to_x2(selected_cav_base["params"]["lidar_pose"], ego_pose) - - # retrieve objects under ego coordinates - object_bbx_center, object_bbx_mask, object_ids = self.post_processor.generate_object_center([selected_cav_base], ego_pose) + transformation_matrix = x1_to_x2(selected_cav_base["params"]["lidar_pose"], target_pose) - # filter lidar lidar_np = selected_cav_base["lidar_np"] lidar_np = shuffle_points(lidar_np) - # remove points that hit itself lidar_np = mask_ego_points(lidar_np) - # project the lidar to ego space lidar_np[:, :3] = box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix) lidar_np = mask_points_by_range(lidar_np, self.params["preprocess"]["cav_lidar_range"]) - processed_lidar = self.pre_processor.preprocess(lidar_np) + return { + "projected_lidar": lidar_np, + "processed_features": self.pre_processor.preprocess(lidar_np), + } - selected_cav_processed.update( - { - "object_bbx_center": object_bbx_center[object_bbx_mask == 1], - "object_ids": object_ids, - "projected_lidar": lidar_np, - "processed_features": processed_lidar, - } - ) + def get_item_single_car(self, selected_cav_base, ego_pose): + """ + Build one local training agent with its stage-one ground truth. + + Parameters + ---------- + selected_cav_base : dict + Raw and metadata fields for one agent. + ego_pose : list + Pose of the coordinate frame used for model input and targets. + Returns + ------- + dict + Agent model input plus local stage-one ground-truth boxes. + """ + selected_cav_processed = self.build_agent_model_input(selected_cav_base, ego_pose) + object_bbx_center, object_bbx_mask, _ = self.post_processor.generate_object_center([selected_cav_base], ego_pose) + selected_cav_processed["object_bbx_center"] = object_bbx_center[object_bbx_mask == 1] return selected_cav_processed @staticmethod @@ -401,8 +460,10 @@ def collate_batch_train(self, batch): record_len = [] label_dict_list = [] label_dict_no_coop_list = [] + has_labels = "label_dict" in batch[0]["ego"] origin_lidar = [] + agent_to_ego = [] # added by yys, fpvrcnn needs anchors for # first stage proposal generation @@ -417,10 +478,13 @@ def collate_batch_train(self, batch): processed_lidar_list.append(ego_dict["processed_lidar"]) record_len.append(ego_dict["cav_num"]) - label_dict_no_coop_list.append(ego_dict["label_dict"]["stage1"]) - label_dict_list.append(ego_dict["label_dict"]["stage2"]) + if has_labels: + label_dict_no_coop_list.append(ego_dict["label_dict"]["stage1"]) + label_dict_list.append(ego_dict["label_dict"]["stage2"]) origin_lidar.append(ego_dict["origin_lidar"]) + if "agent_to_ego" in ego_dict: + agent_to_ego.extend(ego_dict["agent_to_ego"]) # convert to numpy, (B, max_num, 7) object_bbx_center = torch.from_numpy(np.array(object_bbx_center)) @@ -432,12 +496,6 @@ def collate_batch_train(self, batch): processed_lidar_torch_dict = self.pre_processor.collate_batch(merged_feature_dict) # [2, 3, 4, ..., M], M <= 5 record_len = torch.from_numpy(np.array(record_len, dtype=int)) - label_torch_dict = self.post_processor.collate_batch(label_dict_list) - label_dict_no_coop_list_ = [label_dict for label_list in label_dict_no_coop_list for label_dict in label_list] - for i in range(len(label_dict_no_coop_list_)): - if isinstance(label_dict_no_coop_list_[i], list): - print("debug") - label_no_coop_torch_dict = self.post_processor.collate_batch(label_dict_no_coop_list_) # object id is only used during inference, where batch size is 1. # so here we only get the first element. output_dict["ego"].update( @@ -446,10 +504,16 @@ def collate_batch_train(self, batch): "object_bbx_mask": object_bbx_mask, "processed_lidar": processed_lidar_torch_dict, "record_len": record_len, - "label_dict": {"stage1": label_no_coop_torch_dict, "stage2": label_torch_dict}, "object_ids": object_ids[0], } ) + if has_labels: + label_torch_dict = self.post_processor.collate_batch(label_dict_list) + label_dict_no_coop = [label_dict for label_list in label_dict_no_coop_list for label_dict in label_list] + output_dict["ego"]["label_dict"] = { + "stage1": self.post_processor.collate_batch(label_dict_no_coop), + "stage2": label_torch_dict, + } coords = [] idx = 0 @@ -464,6 +528,17 @@ def collate_batch_train(self, batch): origin_lidar = torch.from_numpy(origin_lidar) output_dict["ego"].update({"origin_lidar": origin_lidar}) + if agent_to_ego: + output_dict["ego"]["agent_to_ego"] = torch.from_numpy(np.asarray(agent_to_ego)).float() + + if "remote_agent_outputs" in batch[0]["ego"]: + if len(batch) != 1: + raise ValueError("Distributed two-stage inference requires batch size one") + output_dict["ego"]["remote_agent_outputs"] = [ + {field: torch.from_numpy(np.asarray(value)) for field, value in agent_output.items()} + for agent_output in batch[0]["ego"]["remote_agent_outputs"] + ] + return output_dict def collate_batch_test(self, batch): diff --git a/OpenCOOD/opencood/data_utils/datasets/late_fusion_dataset.py b/OpenCOOD/opencood/data_utils/datasets/late_fusion_dataset.py index 5d08d897..b04fade0 100644 --- a/OpenCOOD/opencood/data_utils/datasets/late_fusion_dataset.py +++ b/OpenCOOD/opencood/data_utils/datasets/late_fusion_dataset.py @@ -6,13 +6,21 @@ import random import logging from collections import OrderedDict +from typing import Any import numpy as np +import numpy.typing as npt import torch import opencood.data_utils.datasets from opencood.data_utils.post_processor import build_postprocessor from opencood.data_utils.datasets import basedataset +from opencood.models.communication_adapters import ( + LateFusionWirePayload, + PoseFrameMetadata, + build_inference_input, + inference_input_to_dict, +) from opencood.data_utils.pre_processor import build_preprocessor from opencood.utils import box_utils from opencood.utils.pcd_utils import pcd_to_np, mask_points_by_range, mask_ego_points, shuffle_points, downsample_lidar_minimum @@ -49,20 +57,64 @@ def extract_data(self, idx): if self.payload_handler is not None: for cav_id, selected_cav_base in base_data_dict.items(): - selected_cav_processed = self.get_item_single_car(selected_cav_base) - - with self.payload_handler.handle_opencda_payload(cav_id, self.module_name) as msg: - msg["object_ids"] = selected_cav_processed["object_ids"] - msg["lidar_pose"] = selected_cav_base["params"]["lidar_pose"] - msg["object_bbx_center"] = selected_cav_processed["object_bbx_center"] - msg["object_bbx_mask"] = selected_cav_processed["object_bbx_mask"] - msg["anchor_box"] = selected_cav_processed["anchor_box"] - msg["pos_equal_one"] = selected_cav_processed["label_dict"]["pos_equal_one"] - msg["neg_equal_one"] = selected_cav_processed["label_dict"]["neg_equal_one"] - msg["targets"] = selected_cav_processed["label_dict"]["targets"] - msg["voxel_features"] = selected_cav_processed["processed_lidar"]["voxel_features"] - msg["voxel_coords"] = selected_cav_processed["processed_lidar"]["voxel_coords"] - msg["voxel_num_points"] = selected_cav_processed["processed_lidar"]["voxel_num_points"] + selected_cav_processed = self.__build_model_input(selected_cav_base) + payload = self.build_wire_payload(selected_cav_base, selected_cav_processed, idx) + self.payload_handler.set_opencda_payload(cav_id, self.module_name, payload) + + @staticmethod + def build_wire_payload( + selected_cav_base: dict[str, Any], + selected_cav_processed: dict[str, Any], + receive_frame: int, + ) -> LateFusionWirePayload: + """ + Build the late-fusion payload sent over V2X. + + Parameters + ---------- + selected_cav_base : dict[str, Any] + Raw dataset entry for one sending agent. + selected_cav_processed : dict[str, Any] + Processed data for the same sending agent. + receive_frame : int + Dataset frame in which the payload is sent. + + Returns + ------- + LateFusionWirePayload + Typed payload containing only remote model-input data. + """ + return LateFusionWirePayload( + inference_input=selected_cav_processed["inference_input"], + metadata=PoseFrameMetadata( + lidar_pose=selected_cav_base["params"]["lidar_pose"], + capture_frame=receive_frame - int(selected_cav_base["time_delay"]), + ), + ) + + @staticmethod + def decode_wire_payload(payload: object) -> LateFusionWirePayload: + """ + Validate and decode a late-fusion payload received over V2X. + + Parameters + ---------- + payload : object + Deserialized module payload received over V2X. + + Returns + ------- + LateFusionWirePayload + Validated late-fusion payload. + + Raises + ------ + TypeError + If the received payload has an unexpected type. + """ + if not isinstance(payload, LateFusionWirePayload): + raise TypeError(f"Expected LateFusionWirePayload, got {type(payload).__name__}") + return payload def __find_ego_vehicle(self, base_data_dict): ego_id = -1 @@ -117,7 +169,29 @@ def __project_lidar_for_visualization(self, cav_base, ego_lidar_pose): lidar_np[:, :3] = box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix) return mask_points_by_range(lidar_np, self.params["preprocess"]["cav_lidar_range"]) - def __build_visualization_data(self, ego_id, ego_lidar_pose, base_data_dict): + def build_visualization_context( + self, + ego_id: str, + ego_lidar_pose: list[float], + base_data_dict: dict[str, Any], + ) -> dict[str, Any]: + """ + Build visualization data from the local scene snapshot. + + Parameters + ---------- + ego_id : str + Identifier of the receiving ego agent. + ego_lidar_pose : list[float] + Current ego LiDAR pose in world coordinates. + base_data_dict : dict[str, Any] + Local dataset snapshot, independent of delivered messages. + + Returns + ------- + dict[str, Any] + Per-agent projected point clouds and their local identities. + """ if not self.visualize: return {} @@ -136,52 +210,36 @@ def __build_visualization_data(self, ego_id, ego_lidar_pose, base_data_dict): "origin_lidar_agent_ids": origin_lidar_agent_ids, } - def __append_visualization_data(self, processed_data_dict, ego_id, ego_lidar_pose, base_data_dict): - if "ego" in processed_data_dict: - processed_data_dict["ego"].update(self.__build_visualization_data(ego_id, ego_lidar_pose, base_data_dict)) - return processed_data_dict - - def __process_with_messages(self, ego_id, ego_lidar_pose, base_data_dict, visualization_ego_lidar_pose, visualization_base_data_dict): + def __process_with_messages(self, ego_id, ego_lidar_pose, base_data_dict): processed_data_dict = OrderedDict() ego_cav_base = base_data_dict.get(ego_id) - ego_cav_processed = self.get_item_single_car(ego_cav_base) + ego_cav_processed = self.__build_model_input(ego_cav_base) transformation_matrix_info = x1_to_x2(ego_lidar_pose, ego_lidar_pose) - ego_cav_processed["transformation_matrix"] = transformation_matrix_info - - processed_data_dict.update({"ego": ego_cav_processed}) + processed_data_dict["ego"] = { + "inference_input": ego_cav_processed["inference_input"], + "transformation_matrix": transformation_matrix_info, + } if ego_id in self.payload_handler.current_artery_payload: for cav_id, _ in base_data_dict.items(): - if cav_id in self.payload_handler.current_artery_payload[ego_id]: - with self.payload_handler.handle_artery_payload(ego_id, cav_id, self.module_name) as msg: - cav_lidar_pose = msg["lidar_pose"] - transformation_matrix_info = x1_to_x2(cav_lidar_pose, ego_lidar_pose) - - selected_cav_processed = { - "object_bbx_center": msg["object_bbx_center"], - "object_bbx_mask": msg["object_bbx_mask"], - "object_ids": msg["object_ids"], - "anchor_box": msg["anchor_box"], - "label_dict": { - "pos_equal_one": msg["pos_equal_one"], - "neg_equal_one": msg["neg_equal_one"], - "targets": msg["targets"], - }, - "processed_lidar": { - "voxel_features": msg["voxel_features"], - "voxel_coords": msg["voxel_coords"], - "voxel_num_points": msg["voxel_num_points"], - }, - "transformation_matrix": transformation_matrix_info, - } - - processed_data_dict.update({cav_id: selected_cav_processed}) - - return self.__append_visualization_data(processed_data_dict, ego_id, visualization_ego_lidar_pose, visualization_base_data_dict) - - def __process_without_messages(self, ego_id, ego_lidar_pose, base_data_dict, visualization_ego_lidar_pose, visualization_base_data_dict): + raw_payload = self.payload_handler.get_artery_payload(ego_id, cav_id, self.module_name) + if raw_payload is None: + continue + payload = self.decode_wire_payload(raw_payload) + transformation_matrix_info = x1_to_x2(payload.metadata.lidar_pose, ego_lidar_pose) + + selected_cav_processed = { + "inference_input": payload.inference_input, + "transformation_matrix": transformation_matrix_info, + } + + processed_data_dict.update({cav_id: selected_cav_processed}) + + return processed_data_dict + + def __process_without_messages(self, ego_id, ego_lidar_pose, base_data_dict): processed_data_dict = OrderedDict() for cav_id, selected_cav_base in base_data_dict.items(): @@ -196,12 +254,111 @@ def __process_without_messages(self, ego_id, ego_lidar_pose, base_data_dict, vis cav_lidar_pose = selected_cav_base["params"]["lidar_pose"] transformation_matrix = x1_to_x2(cav_lidar_pose, ego_lidar_pose) - selected_cav_processed = self.get_item_single_car(selected_cav_base) - selected_cav_processed.update({"transformation_matrix": transformation_matrix}) + selected_cav_processed = self.__build_model_input(selected_cav_base) update_cav = "ego" if cav_id == ego_id else cav_id - processed_data_dict.update({update_cav: selected_cav_processed}) + processed_data_dict[update_cav] = { + "inference_input": selected_cav_processed["inference_input"], + "transformation_matrix": transformation_matrix, + } - return self.__append_visualization_data(processed_data_dict, ego_id, visualization_ego_lidar_pose, visualization_base_data_dict) + return processed_data_dict + + def build_local_supervision(self, base_data_dict: dict[str, Any], ego_lidar_pose: list[float]) -> dict[str, Any]: + """ + Build ego-frame ground truth and labels from the local scene. + + Parameters + ---------- + base_data_dict : dict[str, Any] + Complete local dataset snapshot for the current scene. + ego_lidar_pose : list[float] + Current ego LiDAR pose in world coordinates. + + Returns + ------- + dict[str, Any] + Ego-frame ground truth, anchors, and target labels. + """ + object_bbx_center, object_bbx_mask, object_ids = self.post_processor.generate_object_center(list(base_data_dict.values()), ego_lidar_pose) + anchor_box = self.post_processor.generate_anchor_box() + label_dict = self.post_processor.generate_label( + gt_box_center=object_bbx_center, + anchors=anchor_box, + mask=object_bbx_mask, + ) + return { + "object_bbx_center": object_bbx_center, + "object_bbx_mask": object_bbx_mask, + "object_ids": object_ids, + "anchor_box": anchor_box, + "label_dict": label_dict, + } + + def assemble_inference_sample( + self, + inference_input: OrderedDict[str, dict[str, Any]], + local_supervision: dict[str, Any], + visualization_context: dict[str, Any], + ) -> OrderedDict[str, dict[str, Any]]: + """ + Combine delivered model inputs with local-only dataset data. + + Parameters + ---------- + inference_input : collections.OrderedDict + Ego model input plus model inputs from delivered messages. + local_supervision : dict[str, Any] + Ground truth and labels produced from the local scene. + visualization_context : dict[str, Any] + Visualization data produced from the local scene. + + Returns + ------- + collections.OrderedDict + Late-fusion sample ready for collation. + """ + anchor_box = local_supervision["anchor_box"] + for cav_content in inference_input.values(): + cav_content["anchor_box"] = anchor_box + + inference_input["ego"].update(local_supervision) + inference_input["ego"].update(visualization_context) + return inference_input + + def __prepare_lidar(self, selected_cav_base: dict[str, Any]) -> npt.NDArray[Any]: + """ + Filter one agent's local LiDAR point cloud. + + Parameters + ---------- + selected_cav_base : dict[str, Any] + Raw dataset entry for one agent. + + Returns + ------- + numpy.ndarray + Shuffled and range-filtered LiDAR points. + """ + lidar_np = shuffle_points(selected_cav_base["lidar_np"]) + lidar_np = mask_points_by_range(lidar_np, self.params["preprocess"]["cav_lidar_range"]) + return mask_ego_points(lidar_np) + + def __build_model_input(self, selected_cav_base: dict[str, Any]) -> dict[str, Any]: + """ + Build model input without constructing local supervision. + + Parameters + ---------- + selected_cav_base : dict[str, Any] + Raw dataset entry for one agent. + + Returns + ------- + dict[str, Any] + Preprocessed LiDAR features used by the detector. + """ + lidar_np = self.__prepare_lidar(selected_cav_base) + return {"inference_input": build_inference_input(self.pre_processor.preprocess(lidar_np))} def get_item_single_car(self, selected_cav_base): """ @@ -219,12 +376,7 @@ def get_item_single_car(self, selected_cav_base): """ selected_cav_processed = {} - # filter lidar - lidar_np = selected_cav_base["lidar_np"] - lidar_np = shuffle_points(lidar_np) - lidar_np = mask_points_by_range(lidar_np, self.params["preprocess"]["cav_lidar_range"]) - # remove points that hit ego vehicle - lidar_np = mask_ego_points(lidar_np) + lidar_np = self.__prepare_lidar(selected_cav_base) # generate the bounding box(n, 7) under the cav's space object_bbx_center, object_bbx_mask, object_ids = self.post_processor.generate_object_center( @@ -275,49 +427,22 @@ def get_item_test(self, idx, base_data_dict): _, visualization_ego_lidar_pose = self.__find_ego_vehicle(visualization_base_data_dict) if self.payload_handler is not None: - processed_data_dict = self.__process_with_messages( - ego_id, - ego_lidar_pose, - base_data_dict, - visualization_ego_lidar_pose, - visualization_base_data_dict, - ) + inference_input = self.__process_with_messages(ego_id, ego_lidar_pose, base_data_dict) else: - processed_data_dict = self.__process_without_messages( - ego_id, - ego_lidar_pose, - base_data_dict, - visualization_ego_lidar_pose, - visualization_base_data_dict, - ) - - return processed_data_dict + inference_input = self.__process_without_messages(ego_id, ego_lidar_pose, base_data_dict) - def __collate_processed_lidar(self, cav_content): - if "processed_lidar" in cav_content: - return self.pre_processor.collate_batch([cav_content["processed_lidar"]]) - - return self.pre_processor.collate_batch( - { - "voxel_features": cav_content["voxel_features"], - "voxel_coords": cav_content["voxel_coords"], - "voxel_num_points": cav_content["voxel_num_points"], - } + local_supervision = self.build_local_supervision(base_data_dict, ego_lidar_pose) + visualization_context = self.build_visualization_context( + ego_id, + visualization_ego_lidar_pose, + visualization_base_data_dict, ) + return self.assemble_inference_sample(inference_input, local_supervision, visualization_context) - def __collate_label_dict(self, cav_content): - if "label_dict" in cav_content: - return self.post_processor.collate_batch([cav_content["label_dict"]]) - - return self.post_processor.collate_batch( - [ - { - "pos_equal_one": cav_content["pos_equal_one"], - "neg_equal_one": cav_content["neg_equal_one"], - "targets": cav_content["targets"], - } - ] - ) + def __collate_processed_lidar(self, cav_content): + if "inference_input" in cav_content: + return self.pre_processor.collate_batch([inference_input_to_dict(cav_content["inference_input"])]) + return self.pre_processor.collate_batch([cav_content["processed_lidar"]]) def collate_batch_test(self, batch): """ @@ -349,10 +474,6 @@ def collate_batch_test(self, batch): for cav_id, cav_content in batch.items(): output_dict.update({cav_id: {}}) - # shape: (1, max_num, 7) - object_bbx_center = torch.from_numpy(np.array([cav_content["object_bbx_center"]])) - object_bbx_mask = torch.from_numpy(np.array([cav_content["object_bbx_mask"]])) - object_ids = cav_content["object_ids"] # the anchor box is the same for all bounding boxes usually, thus # we don't need the batch dimension. @@ -368,23 +489,27 @@ def collate_batch_test(self, batch): # processed lidar dictionary processed_lidar_torch_dict = self.__collate_processed_lidar(cav_content) - # label dictionary - label_torch_dict = self.__collate_label_dict(cav_content) # save the transformation matrix (4, 4) to ego vehicle transformation_matrix_torch = torch.from_numpy(np.array(cav_content["transformation_matrix"])).float() output_dict[cav_id].update( { - "object_bbx_center": object_bbx_center, - "object_bbx_mask": object_bbx_mask, "processed_lidar": processed_lidar_torch_dict, - "label_dict": label_torch_dict, - "object_ids": object_ids, "transformation_matrix": transformation_matrix_torch, } ) + if "object_bbx_center" in cav_content: + output_dict[cav_id].update( + { + "object_bbx_center": torch.from_numpy(np.array([cav_content["object_bbx_center"]])), + "object_bbx_mask": torch.from_numpy(np.array([cav_content["object_bbx_mask"]])), + "label_dict": self.post_processor.collate_batch([cav_content["label_dict"]]), + "object_ids": cav_content["object_ids"], + } + ) + if self.visualize: if "origin_lidar" in cav_content: origin_lidar = np.array(downsample_lidar_minimum(pcd_np_list=[local_lidar])) @@ -427,6 +552,6 @@ def post_process(self, data_dict, output_dict): The tensor of gt bounding box. """ pred_box_tensor, pred_score = self.post_processor.post_process(data_dict, output_dict) - gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict) + gt_box_tensor = self.post_processor.generate_gt_bbx({"ego": data_dict["ego"]}) return pred_box_tensor, pred_score, gt_box_tensor diff --git a/OpenCOOD/opencood/data_utils/pre_processor/base_preprocessor.py b/OpenCOOD/opencood/data_utils/pre_processor/base_preprocessor.py index 955047e9..bb170bf7 100644 --- a/OpenCOOD/opencood/data_utils/pre_processor/base_preprocessor.py +++ b/OpenCOOD/opencood/data_utils/pre_processor/base_preprocessor.py @@ -1,4 +1,5 @@ import numpy as np +import torch from opencood.utils import pcd_utils @@ -41,6 +42,35 @@ def preprocess(self, pcd_np): return data_dict + @staticmethod + def collate_batch(batch): + """Collate sampled point-cloud inputs. + + Parameters + ---------- + batch : list[dict] | dict[str, list] + Point-cloud inputs represented either as a list of samples or a + dict of agent input lists. + + Returns + ------- + dict[str, torch.Tensor] + Batched point clouds under the ``downsample_lidar`` model key. + + Raises + ------ + TypeError + If ``batch`` is neither a list nor a dictionary. + """ + if isinstance(batch, list): + point_clouds = [sample["downsample_lidar"] for sample in batch] + elif isinstance(batch, dict): + point_clouds = batch["downsample_lidar"] + else: + raise TypeError(f"Unsupported point-cloud batch type: {type(batch).__name__}") + + return {"downsample_lidar": torch.from_numpy(np.stack(point_clouds))} + def project_points_to_bev_map(self, points, ratio=0.1): """ Project points to BEV occupancy map with default ratio=0.1. diff --git a/OpenCOOD/opencood/models/communication_adapters/__init__.py b/OpenCOOD/opencood/models/communication_adapters/__init__.py new file mode 100644 index 00000000..f43d32bd --- /dev/null +++ b/OpenCOOD/opencood/models/communication_adapters/__init__.py @@ -0,0 +1,94 @@ +"""Model-owned communication adapter selection.""" + +from typing import cast + +import torch +from torch import nn + +from opencood.models.communication_adapters.base import ModelCommunicationAdapter +from opencood.models.communication_adapters.inference_input import ( + BevInferenceInput, + FpvrcnnAgentInferenceInput, + InferenceInput, + PointCloudInferenceInput, + VoxelInferenceInput, + build_inference_input, + inference_input_to_dict, + merge_inference_inputs, +) +from opencood.models.communication_adapters.inference_metadata import ( + IntermediateMetadata, + PoseFrameMetadata, + V2XViTMetadata, +) +from opencood.models.communication_adapters.pending import PendingModelCommunicationAdapter +from opencood.models.communication_adapters.wire_payload import ( + EarlyFusionWirePayload, + FpvrcnnWirePayload, + IntermediateFusionWirePayload, + LateFusionWirePayload, +) + + +def build_communication_adapter( + model: nn.Module, + device: torch.device, +) -> ModelCommunicationAdapter: + """ + Build the communication adapter declared by a perception model. + + Models without a declared adapter temporarily retain their dataset-managed + payload through :class:`PendingModelCommunicationAdapter`. + + Parameters + ---------- + model : torch.nn.Module + Cooperative perception model instance. + device : torch.device + Device used for sender-side model execution. + + Returns + ------- + ModelCommunicationAdapter + Adapter responsible for the model's network communication boundary. + + Raises + ------ + TypeError + If the model declares an invalid adapter class. + """ + # TODO(#191): Remove the fallback after every model declares the adapter + # matching its real sender/receiver communication boundary. + adapter_class: object = getattr( + type(model), + "communication_adapter_class", + PendingModelCommunicationAdapter, + ) + if not isinstance(adapter_class, type) or not issubclass( + adapter_class, + ModelCommunicationAdapter, + ): + raise TypeError("model.communication_adapter_class must derive from ModelCommunicationAdapter") + return cast(type[ModelCommunicationAdapter], adapter_class)(model, device) + + +__all__ = ( + "BevInferenceInput", + "EarlyFusionWirePayload", + "FpvrcnnAgentInferenceInput", + "FpvrcnnWirePayload", + "InferenceInput", + "IntermediateFusionWirePayload", + "IntermediateMetadata", + "LateFusionWirePayload", + "ModelCommunicationAdapter", + "PendingModelCommunicationAdapter", + "PointCloudInferenceInput", + "PoseFrameMetadata", + "V2XViTMetadata", + "VoxelInferenceInput", + "build_communication_adapter", + "build_inference_input", + "inference_input_to_dict", + "merge_inference_inputs", +) diff --git a/OpenCOOD/opencood/models/communication_adapters/base.py b/OpenCOOD/opencood/models/communication_adapters/base.py new file mode 100644 index 00000000..1899514b --- /dev/null +++ b/OpenCOOD/opencood/models/communication_adapters/base.py @@ -0,0 +1,80 @@ +"""Common interface for model-specific network communication adapters.""" + +from abc import ABC, abstractmethod +from typing import Any + +import torch +from torch import nn + + +class ModelCommunicationAdapter(ABC): + """Bridge a perception model to its sender/receiver wire contract.""" + + def __init__(self, model: nn.Module, device: torch.device) -> None: + self.model = model + self.device = device + + @abstractmethod + def prepare_transmission_payloads(self, dataset: Any, idx: int) -> None: + """ + Build and publish every sender payload for one dataset frame. + + Parameters + ---------- + dataset : Any + OpenCOOD dataset providing the local scene observations. + idx : int + Dataset frame to encode and publish. + """ + + def move_to_device(self, value: Any) -> Any: + """ + Recursively move a model input to the adapter device. + + Parameters + ---------- + value : Any + Nested tensors, mappings, or sequences forming a model input. + + Returns + ------- + Any + Input with every tensor moved to the configured device. + """ + if isinstance(value, torch.Tensor): + return value.to(self.device) + if isinstance(value, dict): + return {key: self.move_to_device(item) for key, item in value.items()} + if isinstance(value, list): + return [self.move_to_device(item) for item in value] + if isinstance(value, tuple): + return tuple(self.move_to_device(item) for item in value) + return value + + def decode_received_payload( + self, + payload: object, + receiver_pose: list[float], + ) -> dict[str, Any]: + """ + Decode one delivered model payload for receiver-side inference. + + Parameters + ---------- + payload : object + Deserialized model payload delivered over the V2X communication + channel. + receiver_pose : list[float] + Receiver pose used for model-specific coordinate conversion. + + Returns + ------- + dict[str, Any] + Model-specific remote inference input. + + Raises + ------ + NotImplementedError + If the model still uses dataset-managed payload decoding. + """ + raise NotImplementedError(f"{type(self).__name__} does not implement receiver-side payload decoding") diff --git a/OpenCOOD/opencood/models/communication_adapters/fpvrcnn.py b/OpenCOOD/opencood/models/communication_adapters/fpvrcnn.py new file mode 100644 index 00000000..823d2b7b --- /dev/null +++ b/OpenCOOD/opencood/models/communication_adapters/fpvrcnn.py @@ -0,0 +1,201 @@ +"""Communication adapter for distributed FPV-RCNN inference.""" + +from collections.abc import Callable, Mapping +from typing import Any, Protocol, runtime_checkable + +import numpy as np +import numpy.typing as npt +import torch + +from opencood.models.communication_adapters.base import ModelCommunicationAdapter +from opencood.models.communication_adapters.inference_input import ( + FpvrcnnAgentInferenceInput, +) +from opencood.models.communication_adapters.inference_metadata import PoseFrameMetadata +from opencood.models.communication_adapters.wire_payload import FpvrcnnWirePayload +from opencood.utils import box_utils +from opencood.utils.transformation_utils import x1_to_x2 + + +@runtime_checkable +class FpvrcnnCommunicationDataset(Protocol): + """Dataset capabilities required by the FPV-RCNN adapter.""" + + pre_processor: Any + post_processor: Any + + def extract_data( + self, + idx: int, + agent_payload_builder: Callable[[dict[str, Any], PoseFrameMetadata], object] | None, + ) -> None: ... + + +class FpvrcnnCommunicationAdapter(ModelCommunicationAdapter): + """Run FPV-RCNN stage one and VSA before serializing an agent payload.""" + + def prepare_transmission_payloads(self, dataset: Any, idx: int) -> None: + """ + Encode and publish FPV-RCNN outputs for every local agent. + + Parameters + ---------- + dataset : Any + Dataset exposing the FPV-RCNN agent-sample capabilities. + idx : int + Dataset frame to encode and publish. + + Raises + ------ + TypeError + If the selected dataset cannot prepare FPV-RCNN agent samples. + """ + if not isinstance(dataset, FpvrcnnCommunicationDataset): + raise TypeError("FPV-RCNN communication requires a dataset capable of building local agent observations") + dataset.extract_data( + idx, + agent_payload_builder=lambda sample, metadata: self._build_wire_payload( + dataset, + sample, + metadata, + ), + ) + + def _build_wire_payload( + self, + dataset: FpvrcnnCommunicationDataset, + sample: dict[str, Any], + metadata: PoseFrameMetadata, + ) -> FpvrcnnWirePayload: + """ + Run sender-side FPV-RCNN and build its typed wire payload. + + Parameters + ---------- + dataset : FpvrcnnCommunicationDataset + Dataset providing the FPV-RCNN preprocessor and anchors. + sample : dict[str, Any] + Local sender model input without supervision fields. + metadata : PoseFrameMetadata + Sender pose and feature capture frame. + + Returns + ------- + FpvrcnnWirePayload + Typed payload containing proposals and VSA keypoint features. + + Raises + ------ + TypeError + If the model output is not the expected tensor mapping. + """ + encode_agent = getattr(self.model, "encode_agent", None) + if encode_agent is None: + raise TypeError("FpvrcnnCommunicationAdapter requires a model with encode_agent()") + + processed_lidar = dataset.pre_processor.collate_batch([sample["processed_features"]]) + origin_lidar = np.pad(sample["projected_lidar"], ((0, 0), (1, 0)), mode="constant") + batch = self.move_to_device( + { + "processed_lidar": processed_lidar, + "origin_lidar": torch.from_numpy(origin_lidar), + "anchor_box": torch.from_numpy(np.asarray(dataset.post_processor.generate_anchor_box())), + "record_len": torch.ones(1, dtype=torch.long), + } + ) + self.model.eval() + with torch.no_grad(): + output = encode_agent(batch) + if not isinstance(output, Mapping): + raise TypeError("FPV-RCNN encode_agent() must return a tensor mapping") + + return FpvrcnnWirePayload( + inference_input=FpvrcnnAgentInferenceInput( + proposals=self._tensor_to_numpy(output, "proposals"), + proposal_scores=self._tensor_to_numpy(output, "proposal_scores"), + point_coords=self._tensor_to_numpy(output, "point_coords"), + point_features=self._tensor_to_numpy(output, "point_features"), + ), + metadata=metadata, + ) + + def decode_received_payload( + self, + payload: object, + receiver_pose: list[float], + ) -> dict[str, Any]: + """ + Decode and project one FPV-RCNN payload into the receiver frame. + + Parameters + ---------- + payload : object + Deserialized FPV-RCNN payload delivered over V2X. + receiver_pose : list[float] + Receiver LiDAR pose for the current inference frame. + + Returns + ------- + dict[str, Any] + Remote proposals and VSA keypoint data in the receiver frame. + + Raises + ------ + TypeError + If a payload belonging to another model contract is received. + """ + if not isinstance(payload, FpvrcnnWirePayload): + raise TypeError(f"Expected FpvrcnnWirePayload, got {type(payload).__name__}") + + transformation_matrix = x1_to_x2(payload.metadata.lidar_pose, receiver_pose) + point_coords = np.array(payload.inference_input.point_coords, copy=True) + if len(point_coords) > 0: + point_coords[:, :3] = box_utils.project_points_by_matrix_torch( + point_coords[:, :3], + transformation_matrix, + ) + + proposals = np.array(payload.inference_input.proposals, copy=True) + if len(proposals) > 0: + proposals[:, :3] = box_utils.project_points_by_matrix_torch( + proposals[:, :3], + transformation_matrix, + ) + proposals[:, 6] += np.arctan2( + transformation_matrix[1, 0], + transformation_matrix[0, 0], + ) + + return { + "proposals": proposals, + "proposal_scores": np.array(payload.inference_input.proposal_scores, copy=True), + "point_coords": point_coords, + "point_features": np.array(payload.inference_input.point_features, copy=True), + } + + @staticmethod + def _tensor_to_numpy(output: Mapping[str, Any], field: str) -> npt.NDArray[Any]: + """ + Detach one required model output as a NumPy array. + + Parameters + ---------- + output : collections.abc.Mapping[str, Any] + Sender-side model output. + field : str + Required tensor field to convert. + + Returns + ------- + numpy.typing.NDArray[Any] + Detached CPU representation of the tensor. + + Raises + ------ + TypeError + If the required field is not a tensor. + """ + value = output.get(field) + if not isinstance(value, torch.Tensor): + raise TypeError(f'FPV-RCNN output field "{field}" must be a tensor') + return value.detach().cpu().numpy() diff --git a/OpenCOOD/opencood/models/communication_adapters/inference_input.py b/OpenCOOD/opencood/models/communication_adapters/inference_input.py new file mode 100644 index 00000000..430d0254 --- /dev/null +++ b/OpenCOOD/opencood/models/communication_adapters/inference_input.py @@ -0,0 +1,190 @@ +"""Typed model inputs exchanged through communication adapters.""" + +from dataclasses import dataclass +from typing import Any, Mapping, TypeAlias + +import numpy.typing as npt + + +VOXEL_INPUT_FIELDS = frozenset({"voxel_features", "voxel_coords"}) +SPARSE_VOXEL_INPUT_FIELDS = VOXEL_INPUT_FIELDS | {"voxel_num_points"} +BEV_INPUT_FIELDS = frozenset({"bev_input"}) +POINT_CLOUD_INPUT_FIELDS = frozenset({"downsample_lidar"}) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class VoxelInferenceInput: + """ + Voxelized LiDAR input. + + Attributes + ---------- + voxel_features : numpy.typing.NDArray[Any] + Features stored for every non-empty voxel. + voxel_coords : numpy.typing.NDArray[Any] + Coordinates of the non-empty voxels. + voxel_num_points : numpy.typing.NDArray[Any] | None + Number of points in every voxel when supplied by the preprocessor. + """ + + voxel_features: npt.NDArray[Any] + voxel_coords: npt.NDArray[Any] + voxel_num_points: npt.NDArray[Any] | None = None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class BevInferenceInput: + """ + Bird's-eye-view model input. + + Attributes + ---------- + bev_input : numpy.typing.NDArray[Any] + Dense BEV feature map produced by the preprocessor. + """ + + bev_input: npt.NDArray[Any] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class PointCloudInferenceInput: + """ + Point-cloud model input. + + Attributes + ---------- + points : numpy.typing.NDArray[Any] + Sampled LiDAR points produced by the preprocessor. + """ + + points: npt.NDArray[Any] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class FpvrcnnAgentInferenceInput: + """ + Agent-side FPV-RCNN output consumed by cooperative stage two. + + Attributes + ---------- + proposals : numpy.typing.NDArray[Any] + Stage-one bounding-box proposals in the sender coordinate frame. + proposal_scores : numpy.typing.NDArray[Any] + Confidence scores corresponding to ``proposals``. + point_coords : numpy.typing.NDArray[Any] + Coordinates of VSA keypoints in the sender coordinate frame. + point_features : numpy.typing.NDArray[Any] + VSA features corresponding to ``point_coords``. + """ + + proposals: npt.NDArray[Any] + proposal_scores: npt.NDArray[Any] + point_coords: npt.NDArray[Any] + point_features: npt.NDArray[Any] + + +InferenceInput: TypeAlias = VoxelInferenceInput | BevInferenceInput | PointCloudInferenceInput + + +def build_inference_input(preprocessor_output: Mapping[str, Any]) -> InferenceInput: + """ + Convert a preprocessor result into a typed inference input. + + The concrete type is selected from the fields actually returned by the + preprocessor. Dataset and transport code therefore remain independent of + the configured preprocessor class. + + Parameters + ---------- + preprocessor_output : collections.abc.Mapping[str, Any] + Result returned by ``preprocessor.preprocess()``. + + Returns + ------- + InferenceInput + Typed representation of the preprocessor result. + + Raises + ------ + ValueError + If the returned field set does not match a supported input type. + """ + fields = frozenset(preprocessor_output) + + if fields == VOXEL_INPUT_FIELDS or fields == SPARSE_VOXEL_INPUT_FIELDS: + return VoxelInferenceInput( + voxel_features=preprocessor_output["voxel_features"], + voxel_coords=preprocessor_output["voxel_coords"], + voxel_num_points=preprocessor_output.get("voxel_num_points"), + ) + if fields == BEV_INPUT_FIELDS: + return BevInferenceInput(bev_input=preprocessor_output["bev_input"]) + if fields == POINT_CLOUD_INPUT_FIELDS: + return PointCloudInferenceInput(points=preprocessor_output["downsample_lidar"]) + + raise ValueError(f"Unsupported preprocessor output fields: {sorted(fields)}") + + +def inference_input_to_dict(inference_input: InferenceInput) -> dict[str, npt.NDArray[Any]]: + """ + Convert a typed inference input to the preprocessor collate format. + + Parameters + ---------- + inference_input : InferenceInput + Typed model input. + + Returns + ------- + dict[str, numpy.typing.NDArray[Any]] + Mapping accepted by the corresponding preprocessor's collate path. + """ + if isinstance(inference_input, VoxelInferenceInput): + output = { + "voxel_features": inference_input.voxel_features, + "voxel_coords": inference_input.voxel_coords, + } + if inference_input.voxel_num_points is not None: + output["voxel_num_points"] = inference_input.voxel_num_points + return output + if isinstance(inference_input, BevInferenceInput): + return {"bev_input": inference_input.bev_input} + if isinstance(inference_input, PointCloudInferenceInput): + return {"downsample_lidar": inference_input.points} + raise TypeError(f"Unsupported inference input type: {type(inference_input).__name__}") + + +def merge_inference_inputs(inference_inputs: list[InferenceInput]) -> dict[str, list[npt.NDArray[Any]]]: + """ + Merge homogeneous agent inputs for preprocessor collation. + + Parameters + ---------- + inference_inputs : list[InferenceInput] + Inputs for the agents participating in the current inference. + + Returns + ------- + dict[str, list[numpy.typing.NDArray[Any]]] + Dict-of-lists accepted by ``preprocessor.collate_batch()``. + + Raises + ------ + ValueError + If the input list is empty. + TypeError + If agents provide different input representations. + """ + if not inference_inputs: + raise ValueError("At least one inference input is required") + + input_type = type(inference_inputs[0]) + if any(type(inference_input) is not input_type for inference_input in inference_inputs[1:]): + raise TypeError("All agents must use the same inference input representation") + + input_dicts = [inference_input_to_dict(inference_input) for inference_input in inference_inputs] + input_fields = set(input_dicts[0]) + if any(set(input_dict) != input_fields for input_dict in input_dicts[1:]): + raise TypeError("All agents must provide the same inference input fields") + + return {field: [input_dict[field] for input_dict in input_dicts] for field in input_dicts[0]} diff --git a/OpenCOOD/opencood/models/communication_adapters/inference_metadata.py b/OpenCOOD/opencood/models/communication_adapters/inference_metadata.py new file mode 100644 index 00000000..b351b37e --- /dev/null +++ b/OpenCOOD/opencood/models/communication_adapters/inference_metadata.py @@ -0,0 +1,49 @@ +"""Typed metadata exchanged through communication adapters.""" + +from dataclasses import dataclass +from typing import Any, TypeAlias + +import numpy.typing as npt + + +@dataclass(frozen=True, slots=True, kw_only=True) +class PoseFrameMetadata: + """ + Pose and capture frame used to derive pairwise transformations. + + Attributes + ---------- + lidar_pose : list[float] | tuple[float, ...] + Agent LiDAR pose at feature capture time. + capture_frame : int + Dataset frame in which the model input was captured. + """ + + lidar_pose: list[float] | tuple[float, ...] + capture_frame: int + + +@dataclass(frozen=True, slots=True, kw_only=True) +class V2XViTMetadata: + """ + Metadata consumed by the V2X-ViT fusion model. + + Attributes + ---------- + velocity : float + Sender velocity normalized by 30 km/h. + agent_type : int + ``1`` for an infrastructure agent and ``0`` for a vehicle. + capture_frame : int + Dataset frame in which the model input was captured. + spatial_correction_matrix : numpy.typing.NDArray[Any] + Transformation correcting the captured feature to the receive frame. + """ + + velocity: float + agent_type: int + capture_frame: int + spatial_correction_matrix: npt.NDArray[Any] + + +IntermediateMetadata: TypeAlias = PoseFrameMetadata | V2XViTMetadata diff --git a/OpenCOOD/opencood/models/communication_adapters/pending.py b/OpenCOOD/opencood/models/communication_adapters/pending.py new file mode 100644 index 00000000..a058931b --- /dev/null +++ b/OpenCOOD/opencood/models/communication_adapters/pending.py @@ -0,0 +1,35 @@ +"""Temporary adapters for models awaiting a true network boundary.""" + +from typing import Any + +from opencood.models.communication_adapters.base import ModelCommunicationAdapter + + +class PendingModelCommunicationAdapter(ModelCommunicationAdapter): + """ + Preserve the current dataset-managed payload until a model is split. + + This adapter intentionally performs no model-side encoding. It exists so + models that have not been migrated yet still use their current payload + implementation without adding model-name branches to the manager. + + TODO + ---- + Replace this fallback with model-side adapters for Early, Late and the + remaining Intermediate models. Each adapter must place the wire boundary + at the model's real communication stage instead of sending preprocessor + output by default. + """ + + def prepare_transmission_payloads(self, dataset: Any, idx: int) -> None: + """ + Delegate payload construction to the existing dataset path. + + Parameters + ---------- + dataset : Any + Dataset retaining its temporary payload implementation. + idx : int + Dataset frame to publish. + """ + dataset.extract_data(idx) diff --git a/OpenCOOD/opencood/models/communication_adapters/wire_payload.py b/OpenCOOD/opencood/models/communication_adapters/wire_payload.py new file mode 100644 index 00000000..0aab313d --- /dev/null +++ b/OpenCOOD/opencood/models/communication_adapters/wire_payload.py @@ -0,0 +1,82 @@ +"""Typed cooperative perception payloads exchanged over V2X.""" + +from dataclasses import dataclass +from typing import Any + +import numpy.typing as npt + +from opencood.models.communication_adapters.inference_input import ( + FpvrcnnAgentInferenceInput, + InferenceInput, +) +from opencood.models.communication_adapters.inference_metadata import ( + IntermediateMetadata, + PoseFrameMetadata, +) + + +@dataclass(frozen=True, slots=True, kw_only=True) +class EarlyFusionWirePayload: + """ + Wire payload produced by :class:`EarlyFusionDataset`. + + Attributes + ---------- + projected_lidar : numpy.typing.NDArray[Any] + LiDAR points projected into the ego coordinate frame. + """ + + projected_lidar: npt.NDArray[Any] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class LateFusionWirePayload: + """ + Wire payload produced by :class:`LateFusionDataset`. + + Attributes + ---------- + inference_input : InferenceInput + Typed model input produced by the sender's preprocessor. + metadata : PoseFrameMetadata + Pose and capture frame required for late-fusion post-processing. + """ + + inference_input: InferenceInput + metadata: PoseFrameMetadata + + +@dataclass(frozen=True, slots=True, kw_only=True) +class IntermediateFusionWirePayload: + """ + Wire payload produced by :class:`IntermediateFusionDataset`. + + Attributes + ---------- + inference_input : InferenceInput + Typed model input produced by the sender's preprocessor. + metadata : IntermediateMetadata | None + Metadata required by the configured model, or ``None`` when the model + only consumes the inference input. + """ + + inference_input: InferenceInput + metadata: IntermediateMetadata | None + + +@dataclass(frozen=True, slots=True, kw_only=True) +class FpvrcnnWirePayload: + """ + Agent-side FPV-RCNN output exchanged over V2X. + + Attributes + ---------- + inference_input : FpvrcnnAgentInferenceInput + Stage-one proposals and VSA keypoint features computed by the sender. + metadata : PoseFrameMetadata + Sender pose and capture frame used to transform the output at the + receiver. + """ + + inference_input: FpvrcnnAgentInferenceInput + metadata: PoseFrameMetadata diff --git a/OpenCOOD/opencood/models/fpvrcnn.py b/OpenCOOD/opencood/models/fpvrcnn.py index 0ac0b22f..f3e71b3c 100644 --- a/OpenCOOD/opencood/models/fpvrcnn.py +++ b/OpenCOOD/opencood/models/fpvrcnn.py @@ -1,6 +1,8 @@ -from torch import nn import numpy as np +import torch +from torch import nn +from opencood.models.communication_adapters.fpvrcnn import FpvrcnnCommunicationAdapter from opencood.models.sub_modules.mean_vfe import MeanVFE from opencood.models.sub_modules.sparse_backbone_3d import VoxelBackBone8x from opencood.models.sub_modules.height_compression import HeightCompression @@ -12,6 +14,10 @@ class FPVRCNN(nn.Module): + # The model declares its own wire boundary; neither the manager nor the + # dataset needs to infer it from the selected fusion dataset. + communication_adapter_class = FpvrcnnCommunicationAdapter + def __init__(self, args): super(FPVRCNN, self).__init__() lidar_range = np.array(args["lidar_range"]) @@ -27,7 +33,20 @@ def __init__(self, args): self.roi_head = RoIHead(args["roi_head"]) self.train_stage2 = args["activate_stage2"] - def forward(self, batch_dict): + def _forward_stage_one(self, batch_dict): + """ + Run the per-agent backbone and proposal head. + + Parameters + ---------- + batch_dict : dict + Collated voxel input for one or more independent agents. + + Returns + ------- + tuple[dict, list[torch.Tensor] | None, list[torch.Tensor] | None] + Updated model state, proposal boxes, and proposal scores. + """ voxel_features = batch_dict["processed_lidar"]["voxel_features"] voxel_coords = batch_dict["processed_lidar"]["voxel_coords"] voxel_num_points = batch_dict["processed_lidar"]["voxel_num_points"] @@ -57,12 +76,151 @@ def forward(self, batch_dict): batch_dict["det_boxes"] = pred_box3d_list batch_dict["det_scores"] = scores_list + return batch_dict, pred_box3d_list, scores_list + + def encode_agent(self, batch_dict): + """ + Build the FPV-RCNN representation transmitted by one agent. + + Parameters + ---------- + batch_dict : dict + Single-agent voxel input, raw LiDAR points, and anchors. + + Returns + ------- + dict[str, torch.Tensor] + Stage-one proposals and VSA keypoint coordinates/features in the + sender coordinate frame. + + Raises + ------ + RuntimeError + If cooperative stage two is disabled in the model configuration. + """ + if not self.train_stage2: + raise RuntimeError("Distributed FPV-RCNN inference requires activate_stage2=true") + batch_dict, proposals, proposal_scores = self._forward_stage_one(batch_dict) + if proposals is None or proposal_scores is None: + device = batch_dict["voxel_features"].device + return { + "proposals": torch.empty((0, 7), device=device), + "proposal_scores": torch.empty((0,), device=device), + "point_coords": torch.empty((0, 4), device=device), + "point_features": torch.empty((0, self.vsa.num_point_features), device=device), + } + + batch_dict = self.vsa(batch_dict) + return { + "proposals": proposals[0], + "proposal_scores": proposal_scores[0], + "point_coords": batch_dict["point_coords"][0], + "point_features": batch_dict["point_features"][0], + } + + def _forward_distributed(self, batch_dict): + """ + Fuse local FPV-RCNN output with delivered agent outputs. + + Parameters + ---------- + batch_dict : dict + Ego input containing ``remote_agent_outputs`` decoded from wire + payloads. + + Returns + ------- + dict + Cooperative stage-two model output. + """ + remote_outputs = batch_dict.pop("remote_agent_outputs") + local_output = self.encode_agent(batch_dict) + agent_outputs = [local_output, *remote_outputs] + + batch_dict["det_boxes"] = [output["proposals"] for output in agent_outputs] + batch_dict["det_scores"] = [output["proposal_scores"] for output in agent_outputs] + batch_dict["point_coords"] = [output["point_coords"] for output in agent_outputs] + batch_dict["point_features"] = [output["point_features"] for output in agent_outputs] + batch_dict["record_len"] = torch.tensor( + [len(agent_outputs)], + dtype=torch.long, + device=batch_dict["voxel_features"].device, + ) + + if not any(len(boxes) > 0 for boxes in batch_dict["det_boxes"]): + batch_dict["fpvrcnn_out"] = None + return batch_dict + + batch_dict = self.matcher(batch_dict) + return self.roi_head(batch_dict) + + @staticmethod + def _transform_agent_outputs(batch_dict): + """ + Transform local proposals and keypoints into each scene's ego frame. + + Parameters + ---------- + batch_dict : dict + Model state containing per-agent detections, keypoints, and one + ``agent_to_ego`` matrix per flattened agent. + + Returns + ------- + dict + Model state with proposal centers, headings, and keypoint + coordinates expressed in their corresponding ego frames. + """ + transformations = batch_dict.pop("agent_to_ego") + if len(transformations) != len(batch_dict["det_boxes"]) or len(transformations) != len(batch_dict["point_coords"]): + raise ValueError("Every FPV-RCNN agent output requires one agent-to-ego transformation") + transformed_boxes = [] + transformed_points = [] + + for boxes, points, transformation in zip( + batch_dict["det_boxes"], + batch_dict["point_coords"], + transformations, + ): + boxes = boxes.clone() + if len(boxes) > 0: + box_transformation = transformation.to(dtype=boxes.dtype) + homogeneous_centers = torch.cat( + [boxes[:, :3], torch.ones((len(boxes), 1), dtype=boxes.dtype, device=boxes.device)], + dim=1, + ) + boxes[:, :3] = (box_transformation @ homogeneous_centers.T).T[:, :3] + boxes[:, 6] += torch.atan2(box_transformation[1, 0], box_transformation[0, 0]) + transformed_boxes.append(boxes) + + points = points.clone() + if len(points) > 0: + point_transformation = transformation.to(dtype=points.dtype) + homogeneous_points = torch.cat( + [points[:, :3], torch.ones((len(points), 1), dtype=points.dtype, device=points.device)], + dim=1, + ) + points[:, :3] = (point_transformation @ homogeneous_points.T).T[:, :3] + transformed_points.append(points) + + batch_dict["det_boxes"] = transformed_boxes + batch_dict["point_coords"] = transformed_points + return batch_dict + + def forward(self, batch_dict): + if "remote_agent_outputs" in batch_dict: + return self._forward_distributed(batch_dict) + + batch_dict, pred_box3d_list, _ = self._forward_stage_one(batch_dict) + if pred_box3d_list is not None and self.train_stage2: batch_dict = self.vsa(batch_dict) + if "agent_to_ego" in batch_dict: + batch_dict = self._transform_agent_outputs(batch_dict) batch_dict = self.matcher(batch_dict) batch_dict = self.roi_head(batch_dict) if pred_box3d_list is None: - data_dict["ego"]["fpvrcnn_out"] = None + batch_dict["fpvrcnn_out"] = None return batch_dict diff --git a/OpenCOOD/opencood/models/point_pillar_cobevt.py b/OpenCOOD/opencood/models/point_pillar_cobevt.py index 95baf99d..6e526a0a 100644 --- a/OpenCOOD/opencood/models/point_pillar_cobevt.py +++ b/OpenCOOD/opencood/models/point_pillar_cobevt.py @@ -68,7 +68,6 @@ def forward(self, data_dict): voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] record_len = data_dict["record_len"] - _ = data_dict["spatial_correction_matrix"] # spatial_correction_matrix batch_dict = {"voxel_features": voxel_features, "voxel_coords": voxel_coords, "voxel_num_points": voxel_num_points, "record_len": record_len} # n, 4 -> n, c diff --git a/OpenCOOD/opencood/models/sub_modules/roi_head.py b/OpenCOOD/opencood/models/sub_modules/roi_head.py index 06faa502..28e4f0e0 100644 --- a/OpenCOOD/opencood/models/sub_modules/roi_head.py +++ b/OpenCOOD/opencood/models/sub_modules/roi_head.py @@ -180,6 +180,33 @@ def assign_targets(self, batch_dict): return batch_dict + def prepare_rois(self, batch_dict): + """ + Prepare fused proposals for inference without ground truth. + + Parameters + ---------- + batch_dict : dict + Model state containing one ``boxes_fused`` tensor per scene. + + Returns + ------- + dict + Model state with the proposal fields required by RoI pooling and + stage-two box decoding. + """ + rois_by_scene = batch_dict["boxes_fused"] + rois = torch.cat(rois_by_scene, dim=0) + rois_anchor = rois.clone().detach().view(-1, self.code_size) + rois_anchor[:, 0:3] = 0 + rois_anchor[:, 6] = 0 + batch_dict["rcnn_label_dict"] = { + "rois": rois, + "rois_anchor": rois_anchor, + "record_len": [rois.shape[0] for rois in rois_by_scene], + } + return batch_dict + def roi_grid_pool(self, batch_dict): batch_size = len(batch_dict["record_len"]) rois = batch_dict["rcnn_label_dict"]["rois"] @@ -189,7 +216,7 @@ def roi_grid_pool(self, batch_dict): point_features = torch.cat(point_features, dim=0) # (BxN, 6x6x6, 3) - global_roi_grid_points, local_roi_grid_points = self.get_global_grid_points_of_roi(rois) + global_roi_grid_points, _ = self.get_global_grid_points_of_roi(rois) # (B, Nx6x6x6, 3) global_roi_grid_points = global_roi_grid_points.view(batch_size, -1, 3) @@ -215,7 +242,7 @@ def roi_grid_pool(self, batch_dict): return pooled_features def forward(self, batch_dict): - batch_dict = self.assign_targets(batch_dict) + batch_dict = self.assign_targets(batch_dict) if self.training else self.prepare_rois(batch_dict) # RoI aware pooling pooled_features = self.roi_grid_pool(batch_dict) # (BxN, 6x6x6, C) diff --git a/OpenCOOD/test/conftest.py b/OpenCOOD/test/conftest.py new file mode 100644 index 00000000..91bd8e96 --- /dev/null +++ b/OpenCOOD/test/conftest.py @@ -0,0 +1,9 @@ +"""Pytest configuration for OpenCOOD unit tests.""" + +import sys +from pathlib import Path + + +OPENCOOD_ROOT = Path(__file__).resolve().parents[1] +if str(OPENCOOD_ROOT) not in sys.path: + sys.path.insert(0, str(OPENCOOD_ROOT)) diff --git a/OpenCOOD/test/test_communication_adapters.py b/OpenCOOD/test/test_communication_adapters.py new file mode 100644 index 00000000..d07a43f3 --- /dev/null +++ b/OpenCOOD/test/test_communication_adapters.py @@ -0,0 +1,222 @@ +"""Tests for model-owned communication adapters.""" + +import sys +from typing import Any +from unittest.mock import MagicMock + +import numpy as np +import pytest + +mocked_opencood = sys.modules.get("opencood") +if mocked_opencood is not None and not hasattr(mocked_opencood, "__path__"): + pytest.skip( + "full OpenCOOD tests are disabled when the lightweight OpenCDA test doubles are active", + allow_module_level=True, + ) + +torch = pytest.importorskip("torch") +from torch import nn # noqa: E402 + +from opencood.models.communication_adapters import ( # noqa: E402 + FpvrcnnAgentInferenceInput, + FpvrcnnWirePayload, + ModelCommunicationAdapter, + PendingModelCommunicationAdapter, + PoseFrameMetadata, + build_communication_adapter, +) +from opencood.models.communication_adapters import fpvrcnn as fpvrcnn_module # noqa: E402 +from opencood.models.communication_adapters.fpvrcnn import FpvrcnnCommunicationAdapter # noqa: E402 + + +class DeclaredAdapter(ModelCommunicationAdapter): + """Minimal declared adapter used to verify adapter selection.""" + + def prepare_transmission_payloads(self, dataset: Any, idx: int) -> None: + dataset.extract_data(idx) + + +class DeclaredModel(nn.Module): + communication_adapter_class = DeclaredAdapter + + +class PlainModel(nn.Module): + pass + + +class InvalidAdapterModel(nn.Module): + communication_adapter_class = object + + +class FakePreprocessor: + def __init__(self) -> None: + self.received_features = None + + def collate_batch(self, features): + self.received_features = features + return { + "voxel_features": torch.ones((2, 4), dtype=torch.float32), + "voxel_coords": torch.zeros((2, 4), dtype=torch.int32), + "voxel_num_points": torch.ones(2, dtype=torch.int32), + } + + +class FakePostprocessor: + @staticmethod + def generate_anchor_box(): + return np.ones((2, 7), dtype=np.float32) + + +class FakeFpvrcnnDataset: + def __init__(self) -> None: + self.pre_processor = FakePreprocessor() + self.post_processor = FakePostprocessor() + self.sample = { + "processed_features": { + "voxel_features": np.ones((2, 4), dtype=np.float32), + "voxel_coords": np.zeros((2, 4), dtype=np.int32), + "voxel_num_points": np.ones(2, dtype=np.int32), + }, + "projected_lidar": np.array( + [[1.0, 2.0, 3.0, 0.5], [4.0, 5.0, 6.0, 0.8]], + dtype=np.float32, + ), + } + self.metadata = PoseFrameMetadata( + lidar_pose=(10.0, 20.0, 0.0, 0.0, 0.0, 0.0), + capture_frame=12, + ) + self.published_payload = None + + def extract_data(self, idx, agent_payload_builder=None) -> None: + assert idx == 12 + assert agent_payload_builder is not None + self.published_payload = agent_payload_builder(self.sample, self.metadata) + + +class FakeFpvrcnnModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.received_batch = None + + def encode_agent(self, batch): + self.received_batch = batch + return { + "proposals": torch.tensor( + [[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.25]], + dtype=torch.float32, + ), + "proposal_scores": torch.tensor([0.75], dtype=torch.float32), + "point_coords": torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32), + "point_features": torch.tensor([[0.1, 0.2]], dtype=torch.float32), + } + + +def test_build_communication_adapter_uses_model_declaration(): + model = DeclaredModel() + + adapter = build_communication_adapter(model, torch.device("cpu")) + + assert isinstance(adapter, DeclaredAdapter) + assert adapter.model is model + + +def test_build_communication_adapter_uses_pending_fallback(): + adapter = build_communication_adapter(PlainModel(), torch.device("cpu")) + + assert isinstance(adapter, PendingModelCommunicationAdapter) + + +def test_build_communication_adapter_rejects_invalid_declaration(): + with pytest.raises(TypeError, match="must derive from ModelCommunicationAdapter"): + build_communication_adapter(InvalidAdapterModel(), torch.device("cpu")) + + +def test_pending_adapter_delegates_to_dataset_payload_path(): + dataset = MagicMock() + adapter = PendingModelCommunicationAdapter(PlainModel(), torch.device("cpu")) + + adapter.prepare_transmission_payloads(dataset, idx=4) + + dataset.extract_data.assert_called_once_with(4) + + +def test_fpvrcnn_adapter_builds_payload_from_sender_model_outputs(): + model = FakeFpvrcnnModel() + dataset = FakeFpvrcnnDataset() + adapter = FpvrcnnCommunicationAdapter(model, torch.device("cpu")) + + adapter.prepare_transmission_payloads(dataset, idx=12) + + payload = dataset.published_payload + assert isinstance(payload, FpvrcnnWirePayload) + assert payload.metadata is dataset.metadata + np.testing.assert_array_equal( + payload.inference_input.proposals, + np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.25]], dtype=np.float32), + ) + np.testing.assert_array_equal( + payload.inference_input.point_features, + np.array([[0.1, 0.2]], dtype=np.float32), + ) + assert dataset.pre_processor.received_features[0] is dataset.sample["processed_features"] + assert model.training is False + np.testing.assert_array_equal( + model.received_batch["origin_lidar"].cpu().numpy(), + np.array( + [[0.0, 1.0, 2.0, 3.0, 0.5], [0.0, 4.0, 5.0, 6.0, 0.8]], + dtype=np.float32, + ), + ) + + +def test_fpvrcnn_adapter_projects_received_outputs_without_mutating_payload(monkeypatch): + transformation = np.array( + [ + [0.0, -1.0, 0.0, 10.0], + [1.0, 0.0, 0.0, 20.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ], + dtype=np.float32, + ) + monkeypatch.setattr(fpvrcnn_module, "x1_to_x2", lambda _sender, _receiver: transformation) + original_proposals = np.array( + [[1.0, 0.0, 2.0, 4.0, 5.0, 6.0, 0.25]], + dtype=np.float32, + ) + original_points = np.array([[1.0, 0.0, 2.0]], dtype=np.float32) + payload = FpvrcnnWirePayload( + inference_input=FpvrcnnAgentInferenceInput( + proposals=original_proposals, + proposal_scores=np.array([0.9], dtype=np.float32), + point_coords=original_points, + point_features=np.array([[0.1, 0.2]], dtype=np.float32), + ), + metadata=PoseFrameMetadata( + lidar_pose=(1.0, 2.0, 0.0, 0.0, 0.0, 0.0), + capture_frame=5, + ), + ) + adapter = FpvrcnnCommunicationAdapter(FakeFpvrcnnModel(), torch.device("cpu")) + + decoded = adapter.decode_received_payload( + payload, + receiver_pose=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ) + + np.testing.assert_allclose(decoded["proposals"][0, :3], [10.0, 21.0, 2.0]) + np.testing.assert_allclose(decoded["proposals"][0, 6], 0.25 + np.pi / 2) + np.testing.assert_allclose(decoded["point_coords"][0], [10.0, 21.0, 2.0]) + np.testing.assert_array_equal(payload.inference_input.proposals, original_proposals) + np.testing.assert_array_equal(payload.inference_input.point_coords, original_points) + + +def test_fpvrcnn_adapter_rejects_payload_from_another_contract(): + adapter = FpvrcnnCommunicationAdapter(FakeFpvrcnnModel(), torch.device("cpu")) + + with pytest.raises(TypeError, match="Expected FpvrcnnWirePayload"): + adapter.decode_received_payload( + object(), + receiver_pose=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ) diff --git a/OpenCOOD/test/test_communication_contracts.py b/OpenCOOD/test/test_communication_contracts.py new file mode 100644 index 00000000..d64b37d8 --- /dev/null +++ b/OpenCOOD/test/test_communication_contracts.py @@ -0,0 +1,162 @@ +"""Tests for typed cooperative-perception communication contracts.""" + +from dataclasses import fields +import pickle +import sys + +import numpy as np +import pytest + +mocked_opencood = sys.modules.get("opencood") +if mocked_opencood is not None and not hasattr(mocked_opencood, "__path__"): + pytest.skip( + "full OpenCOOD tests are disabled when the lightweight OpenCDA test doubles are active", + allow_module_level=True, + ) + +pytest.importorskip("torch") + +from opencood.models.communication_adapters import ( # noqa: E402 + BevInferenceInput, + EarlyFusionWirePayload, + FpvrcnnAgentInferenceInput, + FpvrcnnWirePayload, + IntermediateFusionWirePayload, + LateFusionWirePayload, + PointCloudInferenceInput, + PoseFrameMetadata, + VoxelInferenceInput, + build_inference_input, + inference_input_to_dict, + merge_inference_inputs, +) + + +@pytest.mark.parametrize( + ("preprocessor_output", "expected_type"), + [ + ( + { + "voxel_features": np.ones((2, 4), dtype=np.float32), + "voxel_coords": np.zeros((2, 4), dtype=np.int32), + }, + VoxelInferenceInput, + ), + ( + { + "voxel_features": np.ones((2, 4), dtype=np.float32), + "voxel_coords": np.zeros((2, 4), dtype=np.int32), + "voxel_num_points": np.ones(2, dtype=np.int32), + }, + VoxelInferenceInput, + ), + ({"bev_input": np.ones((2, 3, 4), dtype=np.float32)}, BevInferenceInput), + ( + {"downsample_lidar": np.ones((8, 4), dtype=np.float32)}, + PointCloudInferenceInput, + ), + ], +) +def test_build_inference_input_uses_actual_preprocessor_fields( + preprocessor_output, + expected_type, +): + inference_input = build_inference_input(preprocessor_output) + + assert isinstance(inference_input, expected_type) + converted = inference_input_to_dict(inference_input) + assert converted.keys() == preprocessor_output.keys() + for field_name, value in preprocessor_output.items(): + np.testing.assert_array_equal(converted[field_name], value) + + +def test_build_inference_input_rejects_unknown_or_mixed_fields(): + with pytest.raises(ValueError, match="Unsupported preprocessor output fields"): + build_inference_input( + { + "bev_input": np.ones((2, 2), dtype=np.float32), + "voxel_features": np.ones((2, 2), dtype=np.float32), + } + ) + + +def test_merge_inference_inputs_preserves_all_agent_values(): + first = BevInferenceInput(bev_input=np.full((2, 2), 1.0, dtype=np.float32)) + second = BevInferenceInput(bev_input=np.full((2, 2), 2.0, dtype=np.float32)) + + merged = merge_inference_inputs([first, second]) + + assert list(merged) == ["bev_input"] + np.testing.assert_array_equal(merged["bev_input"][0], first.bev_input) + np.testing.assert_array_equal(merged["bev_input"][1], second.bev_input) + + +def test_merge_inference_inputs_rejects_heterogeneous_representations(): + with pytest.raises(TypeError, match="same inference input representation"): + merge_inference_inputs( + [ + BevInferenceInput(bev_input=np.ones((2, 2), dtype=np.float32)), + PointCloudInferenceInput(points=np.ones((2, 4), dtype=np.float32)), + ] + ) + + +def test_wire_payloads_expose_only_network_contract_fields(): + forbidden_fields = { + "anchor_box", + "label_dict", + "neg_equal_one", + "object_bbx_center", + "object_bbx_mask", + "object_ids", + "pos_equal_one", + "schema_version", + "spoofing_mask", + "targets", + } + expected_fields = { + EarlyFusionWirePayload: {"projected_lidar"}, + LateFusionWirePayload: {"inference_input", "metadata"}, + IntermediateFusionWirePayload: {"inference_input", "metadata"}, + FpvrcnnWirePayload: {"inference_input", "metadata"}, + } + + for payload_type, expected in expected_fields.items(): + payload_fields = {field.name for field in fields(payload_type)} + assert payload_fields == expected + assert payload_fields.isdisjoint(forbidden_fields) + + +def test_fpvrcnn_wire_payload_contains_stage_one_and_vsa_outputs(): + metadata = PoseFrameMetadata(lidar_pose=(1.0, 2.0, 3.0, 0.0, 0.0, 0.0), capture_frame=17) + inference_input = FpvrcnnAgentInferenceInput( + proposals=np.ones((2, 7), dtype=np.float32), + proposal_scores=np.array([0.9, 0.8], dtype=np.float32), + point_coords=np.ones((3, 3), dtype=np.float32), + point_features=np.ones((3, 8), dtype=np.float32), + ) + + payload = FpvrcnnWirePayload(inference_input=inference_input, metadata=metadata) + + assert payload.metadata.capture_frame == 17 + assert payload.inference_input.proposals.shape == (2, 7) + assert payload.inference_input.point_features.shape == (3, 8) + + +def test_typed_wire_payload_survives_v2x_pickle_round_trip(): + payload = LateFusionWirePayload( + inference_input=BevInferenceInput(bev_input=np.arange(12, dtype=np.float32).reshape(3, 4)), + metadata=PoseFrameMetadata( + lidar_pose=(1.0, 2.0, 3.0, 0.0, 0.0, 0.0), + capture_frame=21, + ), + ) + + restored = pickle.loads(pickle.dumps(payload)) + + assert isinstance(restored, LateFusionWirePayload) + assert restored.metadata == payload.metadata + np.testing.assert_array_equal( + restored.inference_input.bev_input, + payload.inference_input.bev_input, + ) diff --git a/OpenCOOD/test/test_intermediate_fusion_v2_contract.py b/OpenCOOD/test/test_intermediate_fusion_v2_contract.py new file mode 100644 index 00000000..81963bbe --- /dev/null +++ b/OpenCOOD/test/test_intermediate_fusion_v2_contract.py @@ -0,0 +1,197 @@ +"""Tests for the model-neutral IntermediateFusionDatasetV2 boundary.""" + +from collections import OrderedDict +import sys +from unittest.mock import MagicMock + +import numpy as np +import pytest + +mocked_opencood = sys.modules.get("opencood") +if mocked_opencood is not None and not hasattr(mocked_opencood, "__path__"): + pytest.skip( + "full OpenCOOD tests are disabled when the lightweight OpenCDA test doubles are active", + allow_module_level=True, + ) + +pytest.importorskip("torch") +pytest.importorskip("open3d") + +from opencood.data_utils.datasets.intermediate_fusion_dataset_v2 import ( # noqa: E402 + IntermediateFusionDatasetV2, +) +from opencood.models.communication_adapters import PoseFrameMetadata # noqa: E402 + + +class RecordingPayloadHandler: + def __init__(self) -> None: + self.published = [] + + def set_opencda_payload(self, cav_id, module_name, payload) -> None: + self.published.append((cav_id, module_name, payload)) + + +def _build_dataset_for_extraction(): + dataset = IntermediateFusionDatasetV2.__new__(IntermediateFusionDatasetV2) + dataset.cur_ego_pose_flag = True + dataset.module_name = "OpenCOOD.IntermediateFusionDatasetV2" + dataset.payload_handler = RecordingPayloadHandler() + dataset.retrieve_base_data = MagicMock( + return_value=OrderedDict( + { + "cav-1": { + "params": {"lidar_pose": [1.0, 2.0, 0.0, 0.0, 0.0, 0.0]}, + "time_delay": 0, + }, + "cav-2": { + "params": {"lidar_pose": [3.0, 4.0, 0.0, 0.0, 0.0, 0.0]}, + "time_delay": 2, + }, + } + ) + ) + dataset.build_agent_model_input = MagicMock( + side_effect=lambda cav, pose: { + "agent_pose": tuple(pose), + "source": cav, + } + ) + return dataset + + +def test_extract_data_delegates_payload_construction_to_model_adapter(): + dataset = _build_dataset_for_extraction() + builder = MagicMock(side_effect=lambda sample, metadata: (sample, metadata)) + + dataset.extract_data(idx=10, agent_payload_builder=builder) + + assert builder.call_count == 2 + assert [entry[0] for entry in dataset.payload_handler.published] == ["cav-1", "cav-2"] + assert all(entry[1] == "OpenCOOD.IntermediateFusionDatasetV2" for entry in dataset.payload_handler.published) + first_metadata = dataset.payload_handler.published[0][2][1] + second_metadata = dataset.payload_handler.published[1][2][1] + assert isinstance(first_metadata, PoseFrameMetadata) + assert first_metadata.capture_frame == 10 + assert second_metadata.capture_frame == 8 + assert first_metadata.lidar_pose == (1.0, 2.0, 0.0, 0.0, 0.0, 0.0) + + +def test_extract_data_does_not_assume_fpvrcnn_without_adapter(): + dataset = _build_dataset_for_extraction() + + with pytest.raises(NotImplementedError, match="model-specific communication adapter"): + dataset.extract_data(idx=10) + + +def test_local_supervision_uses_complete_scene_and_stays_outside_wire_path(): + dataset = IntermediateFusionDatasetV2.__new__(IntermediateFusionDatasetV2) + dataset.payload_handler = object() + dataset.post_processor = MagicMock() + object_centers = np.ones((4, 7), dtype=np.float32) + object_mask = np.array([1.0, 1.0, 0.0, 0.0], dtype=np.float32) + object_ids = ["vehicle-1", "vehicle-2"] + anchors = np.ones((2, 2, 7), dtype=np.float32) + dataset.post_processor.generate_object_center.return_value = ( + object_centers, + object_mask, + object_ids, + ) + dataset.post_processor.generate_anchor_box.return_value = anchors + base_data = OrderedDict( + { + "cav-1": {"objects": ["vehicle-1"]}, + "cav-2": {"objects": ["vehicle-2"]}, + } + ) + ego_pose = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + + supervision = dataset.build_local_supervision(base_data, ego_pose) + + dataset.post_processor.generate_object_center.assert_called_once_with( + list(base_data.values()), + ego_pose, + ) + assert supervision["object_ids"] == object_ids + np.testing.assert_array_equal(supervision["object_bbx_center"], object_centers) + np.testing.assert_array_equal(supervision["object_bbx_mask"], object_mask) + np.testing.assert_array_equal(supervision["anchor_box"], anchors) + assert "stage2_label" not in supervision + dataset.post_processor.generate_label.assert_not_called() + + +def test_local_supervision_builds_stage_two_labels_only_without_capi_transport(): + dataset = IntermediateFusionDatasetV2.__new__(IntermediateFusionDatasetV2) + dataset.payload_handler = None + dataset.post_processor = MagicMock() + object_centers = np.ones((2, 7), dtype=np.float32) + object_mask = np.ones(2, dtype=np.float32) + anchors = np.ones((2, 2, 7), dtype=np.float32) + stage_two_label = {"targets": np.ones((2, 2, 7), dtype=np.float32)} + dataset.post_processor.generate_object_center.return_value = ( + object_centers, + object_mask, + ["vehicle-1", "vehicle-2"], + ) + dataset.post_processor.generate_anchor_box.return_value = anchors + dataset.post_processor.generate_label.return_value = stage_two_label + + supervision = dataset.build_local_supervision( + OrderedDict({"cav-1": {}, "cav-2": {}}), + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ) + + assert supervision["stage2_label"] is stage_two_label + dataset.post_processor.generate_label.assert_called_once_with( + gt_box_center=object_centers, + anchors=anchors, + mask=object_mask, + ) + + +def test_received_payloads_are_decoded_by_attached_model_adapter(): + dataset = IntermediateFusionDatasetV2.__new__(IntermediateFusionDatasetV2) + dataset.communication_adapter = MagicMock() + dataset.communication_adapter.decode_received_payload.return_value = {"model_specific_feature": np.array([1.0], dtype=np.float32)} + dataset.payload_handler = MagicMock() + dataset.payload_handler.current_artery_payload = {"cav-1": {"cav-2": object()}} + raw_payload = object() + dataset.payload_handler.get_artery_payload.side_effect = [raw_payload, None] + local_input = { + "processed_features": {"voxel_features": np.ones((1, 4), dtype=np.float32)}, + "projected_lidar": np.ones((2, 4), dtype=np.float32), + } + dataset.build_agent_model_input = MagicMock(return_value=local_input) + dataset.build_visualization_context = MagicMock( + return_value={ + "projected_lidar_stack": [], + "projected_lidar_roles": [], + "projected_lidar_agent_ids": [], + } + ) + base_data = OrderedDict( + { + "cav-1": {"ego": True}, + "cav-2": {"ego": False}, + "cav-3": {"ego": False}, + } + ) + ego_pose = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + + result = dataset._IntermediateFusionDatasetV2__process_with_messages( + "cav-1", + ego_pose, + base_data, + ego_pose, + base_data, + ) + + dataset.communication_adapter.decode_received_payload.assert_called_once_with( + raw_payload, + ego_pose, + ) + assert result["processed_features"][0] is local_input["processed_features"] + assert len(result["remote_agent_outputs"]) == 1 + np.testing.assert_array_equal( + result["remote_agent_outputs"][0]["model_specific_feature"], + np.array([1.0], dtype=np.float32), + ) diff --git a/OpenCOOD/test/test_roi_head_inference.py b/OpenCOOD/test/test_roi_head_inference.py new file mode 100644 index 00000000..1e9c3a55 --- /dev/null +++ b/OpenCOOD/test/test_roi_head_inference.py @@ -0,0 +1,86 @@ +"""Tests for GT-independent FPV-RCNN RoI inference.""" + +import sys +from unittest.mock import MagicMock + +import pytest + +mocked_opencood = sys.modules.get("opencood") +if mocked_opencood is not None and not hasattr(mocked_opencood, "__path__"): + pytest.skip( + "full OpenCOOD tests are disabled when the lightweight OpenCDA test doubles are active", + allow_module_level=True, + ) + +torch = pytest.importorskip("torch") +from torch import nn # noqa: E402 + +from opencood.models.sub_modules.roi_head import RoIHead # noqa: E402 + + +def _build_minimal_roi_head(proposal_count: int) -> RoIHead: + head = RoIHead.__new__(RoIHead) + nn.Module.__init__(head) + head.code_size = 7 + head.grid_size = 1 + head.roi_grid_pool = MagicMock(return_value=torch.ones((proposal_count, 1, 1), dtype=torch.float32)) + head.shared_fc_layers = nn.Identity() + head.cls_layers = nn.Identity() + head.iou_layers = nn.Identity() + head.reg_layers = nn.Identity() + return head + + +def test_roi_head_inference_prepares_rois_without_ground_truth_assignment(): + boxes_by_scene = [ + torch.tensor( + [ + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.25], + [7.0, 8.0, 9.0, 2.0, 3.0, 4.0, -0.5], + ], + dtype=torch.float32, + ), + torch.tensor( + [[10.0, 11.0, 12.0, 1.0, 2.0, 3.0, 1.25]], + dtype=torch.float32, + ), + ] + head = _build_minimal_roi_head(proposal_count=3) + head.assign_targets = MagicMock(side_effect=AssertionError("assign_targets must not run during inference")) + head.eval() + batch = { + "boxes_fused": boxes_by_scene, + "record_len": torch.tensor([1, 1]), + } + + output = head(batch) + + head.assign_targets.assert_not_called() + assert output["rcnn_label_dict"]["record_len"] == [2, 1] + assert output["rcnn_label_dict"]["rois"].shape == (3, 7) + anchors = output["rcnn_label_dict"]["rois_anchor"] + torch.testing.assert_close(anchors[:, :3], torch.zeros((3, 3))) + torch.testing.assert_close(anchors[:, 6], torch.zeros(3)) + torch.testing.assert_close(anchors[:, 3:6], output["rcnn_label_dict"]["rois"][:, 3:6]) + assert set(output["fpvrcnn_out"]) == {"rcnn_cls", "rcnn_iou", "rcnn_reg"} + + +def test_roi_head_training_keeps_target_assignment_path(): + head = _build_minimal_roi_head(proposal_count=1) + head.assign_targets = MagicMock( + side_effect=lambda batch: { + **batch, + "rcnn_label_dict": { + "rois": torch.ones((1, 7), dtype=torch.float32), + "rois_anchor": torch.ones((1, 7), dtype=torch.float32), + "record_len": [1], + }, + } + ) + head.prepare_rois = MagicMock(side_effect=AssertionError("prepare_rois must not replace training targets")) + head.train() + + head({"record_len": torch.tensor([1])}) + + head.assign_targets.assert_called_once() + head.prepare_rois.assert_not_called() diff --git a/opencda/core/common/communication/payload_handler.py b/opencda/core/common/communication/payload_handler.py index a869c3ff..0ea6bbf6 100644 --- a/opencda/core/common/communication/payload_handler.py +++ b/opencda/core/common/communication/payload_handler.py @@ -1,9 +1,8 @@ import sys -import pickle # TODO: In the future pickle module will be replaced with our own safe implementation import logging import pathlib -from contextlib import contextmanager -from typing import Any, Generator +import pickle # TODO: In the future pickle module will be replaced with our own safe implementation +from typing import Any sys.path.append(str((pathlib.Path(__file__).resolve().parent / "protos" / "cavise").resolve())) @@ -20,17 +19,38 @@ def __init__(self) -> None: self.current_opencda_payload: dict[str, dict[str, Any]] = {} self.current_artery_payload: dict[str, dict[str, dict[str, Any]]] = {} - @contextmanager - def handle_opencda_payload(self, id: str, module: str) -> Generator[dict[str, Any], None, None]: - self.current_opencda_payload.setdefault(id, {}).setdefault(module, {}) - - yield self.current_opencda_payload[id][module] - - @contextmanager - def handle_artery_payload(self, ego_id: str, id: str, module: str) -> Generator[dict[str, Any], None, None]: - self.current_artery_payload.setdefault(ego_id, {}).setdefault(id, {}).setdefault(module, {}) - - yield self.current_artery_payload[ego_id][id][module] + def set_opencda_payload(self, entity_id: str, module: str, payload: object) -> None: + """Store a module payload for the next OpenCDA message. + + Parameters + ---------- + entity_id : str + Identifier of the sending CAV or RSU. + module : str + Name of the module that owns the payload contract. + payload : object + Pickle-serializable module payload. + """ + self.current_opencda_payload.setdefault(entity_id, {})[module] = payload + + def get_artery_payload(self, ego_id: str, entity_id: str, module: str) -> object | None: + """Return a received module payload if Artery delivered it. + + Parameters + ---------- + ego_id : str + Identifier of the receiving ego agent. + entity_id : str + Identifier of the sending CAV or RSU. + module : str + Name of the module that owns the payload contract. + + Returns + ------- + object | None + Received payload, or ``None`` when no matching payload exists. + """ + return self.current_artery_payload.get(ego_id, {}).get(entity_id, {}).get(module) def make_opencda_message(self) -> proto_opencda.OpenCDAMessage: opencda_message = proto_opencda.OpenCDAMessage() diff --git a/opencda/core/common/communication/test/test_payload_handler.py b/opencda/core/common/communication/test/test_payload_handler.py index 72067860..6eae0299 100644 --- a/opencda/core/common/communication/test/test_payload_handler.py +++ b/opencda/core/common/communication/test/test_payload_handler.py @@ -90,80 +90,33 @@ def _make_entity(entity_id: str, payload: dict) -> FakeEntity: class TestPayloadHandler: - def test_handle_opencda_payload_creates_and_persists_mutations(self, payload_handler_mod): + def test_set_opencda_payload_stores_module_payload(self, payload_handler_mod): ph = payload_handler_mod.PayloadHandler() + payload = {"x": 1} - with ph.handle_opencda_payload("ego1", "loc") as payload: - assert isinstance(payload, dict) - assert payload == {} - payload["x"] = 1 + ph.set_opencda_payload("ego1", "loc", payload) - assert ph.current_opencda_payload["ego1"]["loc"] == {"x": 1} + assert ph.current_opencda_payload["ego1"]["loc"] is payload - with ph.handle_opencda_payload("ego1", "loc") as payload2: - assert payload2 is ph.current_opencda_payload["ego1"]["loc"] - assert payload2 == {"x": 1} - payload2["y"] = 2 - - assert ph.current_opencda_payload["ego1"]["loc"] == {"x": 1, "y": 2} - - def test_handle_opencda_payload_new_module_under_existing_id(self, payload_handler_mod): + def test_set_opencda_payload_replaces_existing_module_payload(self, payload_handler_mod): ph = payload_handler_mod.PayloadHandler() + ph.set_opencda_payload("ego1", "loc", {"x": 1}) - with ph.handle_opencda_payload("ego1", "loc") as loc_payload: - loc_payload["x"] = 1 + ph.set_opencda_payload("ego1", "loc", {"y": 2}) - with ph.handle_opencda_payload("ego1", "perc") as perc_payload: - assert isinstance(perc_payload, dict) - assert perc_payload == {} - assert perc_payload is not loc_payload - perc_payload["z"] = 3 + assert ph.current_opencda_payload["ego1"]["loc"] == {"y": 2} - assert ph.current_opencda_payload["ego1"]["loc"] == {"x": 1} - assert ph.current_opencda_payload["ego1"]["perc"] == {"z": 3} - - def test_handle_artery_payload_creates_and_persists_mutations(self, payload_handler_mod): + def test_get_artery_payload_returns_matching_module(self, payload_handler_mod): ph = payload_handler_mod.PayloadHandler() + payload = {"data": 100} + ph.current_artery_payload = {"ego1": {"ent1": {"module": payload}}} - with ph.handle_artery_payload("ego1", "ent1", "loc") as payload: - assert isinstance(payload, dict) - assert payload == {} - payload["a"] = 10 - - assert ph.current_artery_payload["ego1"]["ent1"]["loc"] == {"a": 10} + assert ph.get_artery_payload("ego1", "ent1", "module") is payload - with ph.handle_artery_payload("ego1", "ent1", "loc") as payload2: - assert payload2 is ph.current_artery_payload["ego1"]["ent1"]["loc"] - payload2["b"] = 20 - - assert ph.current_artery_payload["ego1"]["ent1"]["loc"] == {"a": 10, "b": 20} - - def test_handle_artery_payload_isolation_between_ego_ids(self, payload_handler_mod): + def test_get_artery_payload_returns_none_when_missing(self, payload_handler_mod): ph = payload_handler_mod.PayloadHandler() - with ph.handle_artery_payload("ego1", "ent1", "module") as p1: - p1["data"] = 100 - - with ph.handle_artery_payload("ego2", "ent1", "module") as p2: - p2["data"] = 200 - - assert ph.current_artery_payload["ego1"]["ent1"]["module"]["data"] == 100 - assert ph.current_artery_payload["ego2"]["ent1"]["module"]["data"] == 200 - - def test_handle_artery_payload_new_module_under_existing_entity(self, payload_handler_mod): - ph = payload_handler_mod.PayloadHandler() - - with ph.handle_artery_payload("ego1", "ent1", "loc") as p_loc: - p_loc["a"] = 1 - - with ph.handle_artery_payload("ego1", "ent1", "perc") as p_perc: - assert isinstance(p_perc, dict) - assert p_perc == {} - assert p_perc is not p_loc - p_perc["b"] = 2 - - assert ph.current_artery_payload["ego1"]["ent1"]["loc"] == {"a": 1} - assert ph.current_artery_payload["ego1"]["ent1"]["perc"] == {"b": 2} + assert ph.get_artery_payload("ego1", "ent1", "module") is None def test_make_opencda_message_empty(self, payload_handler_mod): ph = payload_handler_mod.PayloadHandler() @@ -176,10 +129,8 @@ def test_make_opencda_message_empty(self, payload_handler_mod): def test_make_opencda_message_single_entity_round_trip(self, payload_handler_mod): ph = payload_handler_mod.PayloadHandler() - with ph.handle_opencda_payload("ego1", "loc") as loc: - loc["val"] = 1 - with ph.handle_opencda_payload("ego1", "perc") as perc: - perc["score"] = 0.5 + ph.set_opencda_payload("ego1", "loc", {"val": 1}) + ph.set_opencda_payload("ego1", "perc", {"score": 0.5}) msg = ph.make_opencda_message() assert len(msg.entity) == 1 diff --git a/opencda/core/common/coperception_model_manager.py b/opencda/core/common/coperception_model_manager.py index e332c6ec..2f688f4c 100644 --- a/opencda/core/common/coperception_model_manager.py +++ b/opencda/core/common/coperception_model_manager.py @@ -16,6 +16,7 @@ import opencood.hypes_yaml.yaml_utils as yaml_utils from opencood.tools import train_utils, inference_utils from opencood.data_utils.datasets import build_dataset +from opencood.models.communication_adapters import build_communication_adapter from opencood.visualization import vis_utils from opencda.metrics_tools.collection_models import MetricCollection from opencda.metrics_tools.config import resolve_metric_collector_config @@ -505,6 +506,7 @@ def __init__( self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.saved_path = self.opt.model_dir self.model = self._init_model() + self.communication_adapter = build_communication_adapter(self.model, self.device) self.opencood_dataset: DatasetOpenCOOD | None = None self.data_loader: DataLoader[Any] | None = None self.current_memory_data = None @@ -569,6 +571,7 @@ def _init_model(self): def _init_dataset(self) -> None: logger.info("Initial Dataset Building") self.opencood_dataset = cast(DatasetOpenCOOD, build_dataset(self.hypes, visualize=True, train=False, payload_handler=self.payload_handler)) + self.opencood_dataset.communication_adapter = self.communication_adapter self.data_loader = self._create_data_loader(self.opencood_dataset) @staticmethod @@ -591,6 +594,25 @@ def update_dataset(self, data=None): if len(self.opencood_dataset) == 0: logger.warning("No samples found in dataset after update.") + def prepare_transmission_payloads(self, idx: int) -> None: + """ + Build cooperative perception payloads for the current frame. + + Parameters + ---------- + idx : int + Dataset frame to encode and publish through the payload handler. + + Raises + ------ + RuntimeError + If the cooperative perception dataset is unavailable. + """ + if self.opencood_dataset is None: + raise RuntimeError("Coperception dataset is missing; payload extraction cannot continue") + + self.communication_adapter.prepare_transmission_payloads(self.opencood_dataset, idx) + def _resolve_inference_callable(self) -> Callable[..., CoperceptionInferenceResult]: core_method = self.hypes.get("fusion", {}).get("core_method") return InferenceMapper.resolve(type(self), core_method) diff --git a/opencda/core/common/test/conftest.py b/opencda/core/common/test/conftest.py index 46b39adb..74916f7c 100644 --- a/opencda/core/common/test/conftest.py +++ b/opencda/core/common/test/conftest.py @@ -24,6 +24,8 @@ "opencood.tools.inference_utils", "opencood.data_utils", "opencood.data_utils.datasets", + "opencood.models", + "opencood.models.communication_adapters", "opencood.visualization", "opencood.visualization.simple_vis", "opencood.visualization.vis_utils", @@ -131,6 +133,13 @@ def __len__(self): data_utils.datasets = datasets opencood.data_utils = data_utils + models = types.ModuleType("opencood.models") + communication_adapters = types.ModuleType("opencood.models.communication_adapters") + communication_adapter = MagicMock(name="communication_adapter") + communication_adapters.build_communication_adapter = Mock(return_value=communication_adapter) + models.communication_adapters = communication_adapters + opencood.models = models + visualization = types.ModuleType("opencood.visualization") simple_vis = types.ModuleType("opencood.visualization.simple_vis") vis_utils = types.ModuleType("opencood.visualization.vis_utils") @@ -206,6 +215,8 @@ def __len__(self): "opencood.tools.inference_utils": inference_utils, "opencood.data_utils": data_utils, "opencood.data_utils.datasets": datasets, + "opencood.models": models, + "opencood.models.communication_adapters": communication_adapters, "opencood.visualization": visualization, "opencood.visualization.simple_vis": simple_vis, "opencood.visualization.vis_utils": vis_utils, @@ -250,4 +261,8 @@ def fake_heavy_deps(): """ Returns the currently mocked modules from sys.modules for use in tests. """ - return {"torch": sys.modules["torch"], "opencood": sys.modules["opencood"], "open3d": sys.modules["open3d"]} + return { + "torch": sys.modules["torch"], + "opencood": sys.modules["opencood"], + "open3d": sys.modules["open3d"], + } diff --git a/opencda/core/common/test/test_coperception_model_manager.py b/opencda/core/common/test/test_coperception_model_manager.py index 241ede7c..95512c15 100644 --- a/opencda/core/common/test/test_coperception_model_manager.py +++ b/opencda/core/common/test/test_coperception_model_manager.py @@ -65,6 +65,8 @@ def manager_deps(self, fake_heavy_deps): opencood.tools.inference_utils.inference_intermediate_fusion, opencood.tools.inference_utils.save_prediction_gt, opencood.data_utils.datasets.build_dataset, + opencood.models.communication_adapters.build_communication_adapter, + opencood.models.communication_adapters.build_communication_adapter.return_value, opencood.visualization.simple_vis.visualize, opencood.visualization.vis_utils.visualize_inference_sample_dataloader, opencood.visualization.vis_utils.linset_assign_list, @@ -92,6 +94,8 @@ def manager_deps(self, fake_heavy_deps): opencood.hypes_yaml.yaml_utils.load_yaml.return_value = hypes model = MagicMock() + opencood.data_utils.datasets.build_dataset.return_value = MagicMock() + opencood.models.communication_adapters.build_communication_adapter.return_value = MagicMock() opencood.tools.train_utils.create_model.return_value = model opencood.tools.train_utils.load_saved_model.return_value = (None, model) @@ -105,6 +109,7 @@ def manager_deps(self, fake_heavy_deps): coperception_model_manager_module.train_utils = opencood.tools.train_utils coperception_model_manager_module.inference_utils = opencood.tools.inference_utils coperception_model_manager_module.build_dataset = opencood.data_utils.datasets.build_dataset + coperception_model_manager_module.build_communication_adapter = opencood.models.communication_adapters.build_communication_adapter coperception_model_manager_module.vis_utils = opencood.visualization.vis_utils # Return dict for easy access @@ -116,6 +121,7 @@ def manager_deps(self, fake_heavy_deps): "simple_vis": opencood.visualization.simple_vis, "eval_utils": opencood.utils.eval_utils, "build_dataset": opencood.data_utils.datasets.build_dataset, + "build_communication_adapter": (opencood.models.communication_adapters.build_communication_adapter), "Visualizer": open3d.visualization.Visualizer, "torch": torch, "model": model, @@ -139,6 +145,35 @@ def test_init_cuda(self, manager_deps): assert manager.device == "device(cuda)" manager_deps["model"].cuda.assert_called_once() + def test_init_attaches_model_communication_adapter_to_dataset(self, manager_deps): + dataset = DummyDataset() + adapter = MagicMock() + manager_deps["build_dataset"].return_value = dataset + manager_deps["build_communication_adapter"].return_value = adapter + + manager = CoperceptionModelManager(DummyOpt(), "2023_01_01") + + manager_deps["build_communication_adapter"].assert_called_once_with(manager.model, manager.device) + assert manager.communication_adapter is adapter + assert dataset.communication_adapter is adapter + + def test_prepare_transmission_payloads_delegates_to_adapter(self, manager_deps): + manager = CoperceptionModelManager(DummyOpt(), "2023_01_01") + + manager.prepare_transmission_payloads(idx=7) + + manager.communication_adapter.prepare_transmission_payloads.assert_called_once_with( + manager.opencood_dataset, + 7, + ) + + def test_prepare_transmission_payloads_requires_dataset(self, manager_deps): + manager = CoperceptionModelManager(DummyOpt(), "2023_01_01") + manager.opencood_dataset = None + + with pytest.raises(RuntimeError, match="dataset is missing"): + manager.prepare_transmission_payloads(idx=0) + def test_coperception_metrics_use_scenario_config(self, manager_deps): opt = DummyOpt() manager = CoperceptionModelManager( diff --git a/opencda/scenario_testing/scenario.py b/opencda/scenario_testing/scenario.py index 978392de..67063ef5 100644 --- a/opencda/scenario_testing/scenario.py +++ b/opencda/scenario_testing/scenario.py @@ -505,7 +505,7 @@ def capi_loop(self, opt: argparse.Namespace) -> None: opencood_dataset = self.coperception_model_manager.opencood_dataset if opencood_dataset is None: self._abort_simulation("Coperception dataset is missing; prediction pipeline cannot continue.") - opencood_dataset.extract_data(idx=0) + self.coperception_model_manager.prepare_transmission_payloads(idx=0) can_predict_current_tick = True opencda_message = payload_handler.make_opencda_message()