From 3d8cf00b907d0a28d2eb65f1a7e97b57eb0412ec Mon Sep 17 00:00:00 2001 From: Grigoriy Khoroshilov Date: Fri, 4 Sep 2026 15:55:58 +0300 Subject: [PATCH 1/2] [OpenCOOD] Implement model-specific V2X payloads for all fusion strategies --- .../datasets/early_fusion_dataset.py | 141 ++-- .../datasets/intermediate_fusion_dataset.py | 213 +++--- .../datasets/late_fusion_dataset.py | 149 ++-- .../post_processor/bev_postprocessor.py | 83 ++- .../post_processor/ciassd_postprocessor.py | 45 ++ .../post_processor/voxel_postprocessor.py | 133 ++-- OpenCOOD/opencood/models/ciassd.py | 25 +- .../models/communication_adapters/__init__.py | 53 +- .../models/communication_adapters/base.py | 29 +- .../models/communication_adapters/early.py | 110 +++ .../communication_adapters/inference_input.py | 63 ++ .../communication_adapters/intermediate.py | 260 +++++++ .../models/communication_adapters/late.py | 169 +++++ .../models/communication_adapters/pending.py | 35 - .../communication_adapters/where2comm.py | 225 ++++++ .../communication_adapters/wire_payload.py | 72 +- .../models/fuse_modules/where2comm_fuse.py | 85 ++- .../opencood/models/pixor_intermediate.py | 69 +- .../opencood/models/point_pillar_coalign.py | 88 ++- .../opencood/models/point_pillar_cobevt.py | 70 ++ .../opencood/models/point_pillar_fcooper.py | 62 ++ .../models/point_pillar_intermediate.py | 55 +- .../models/point_pillar_intermediate_V2VAM.py | 64 +- .../models/point_pillar_transformer.py | 77 ++ .../opencood/models/point_pillar_v2vnet.py | 63 ++ .../models/point_pillar_where2comm.py | 118 +++ OpenCOOD/opencood/models/second.py | 2 +- .../opencood/models/second_intermediate.py | 59 ++ .../models/sub_modules/att_bev_backbone.py | 97 ++- .../models/sub_modules/auto_encoder.py | 32 +- .../models/sub_modules/naive_compress.py | 35 +- OpenCOOD/opencood/models/voxel_net.py | 39 +- .../opencood/models/voxel_net_intermediate.py | 98 ++- OpenCOOD/opencood/tools/inference_utils.py | 3 +- OpenCOOD/test/test_communication_adapters.py | 684 +++++++++++++++++- OpenCOOD/test/test_communication_contracts.py | 32 +- .../test/test_intermediate_fusion_contract.py | 244 +++++++ .../test/test_intermediate_model_split.py | 94 +++ .../test/test_late_fusion_postprocessors.py | 123 ++++ .../core/common/coperception_model_manager.py | 6 +- .../test/test_coperception_model_manager.py | 6 +- 41 files changed, 3633 insertions(+), 477 deletions(-) create mode 100644 OpenCOOD/opencood/models/communication_adapters/early.py create mode 100644 OpenCOOD/opencood/models/communication_adapters/intermediate.py create mode 100644 OpenCOOD/opencood/models/communication_adapters/late.py delete mode 100644 OpenCOOD/opencood/models/communication_adapters/pending.py create mode 100644 OpenCOOD/opencood/models/communication_adapters/where2comm.py create mode 100644 OpenCOOD/test/test_intermediate_fusion_contract.py create mode 100644 OpenCOOD/test/test_intermediate_model_split.py create mode 100644 OpenCOOD/test/test_late_fusion_postprocessors.py diff --git a/OpenCOOD/opencood/data_utils/datasets/early_fusion_dataset.py b/OpenCOOD/opencood/data_utils/datasets/early_fusion_dataset.py index bcda0107..fbc408ca 100644 --- a/OpenCOOD/opencood/data_utils/datasets/early_fusion_dataset.py +++ b/OpenCOOD/opencood/data_utils/datasets/early_fusion_dataset.py @@ -5,16 +5,18 @@ import math import logging from collections import OrderedDict +from collections.abc import Callable from typing import Any import numpy as np +import numpy.typing as npt import torch import opencood.data_utils.datasets 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.models.communication_adapters import PoseFrameMetadata 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 @@ -52,58 +54,35 @@ def __find_ego_vehicle(self, base_data_dict): return ego_id, ego_lidar_pose - def extract_data(self, idx): - base_data_dict = self.retrieve_base_data(idx) - _, ego_lidar_pose = self.__find_ego_vehicle(base_data_dict) - - 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) - - @staticmethod - def build_wire_payload(selected_cav_processed: dict[str, Any]) -> EarlyFusionWirePayload: + def extract_data( + self, + idx: int, + agent_payload_builder: Callable[[npt.NDArray[Any], PoseFrameMetadata], object], + ) -> None: """ - Build the early-fusion payload sent over V2X. + Publish sender-local observations through the communication adapter. Parameters ---------- - selected_cav_processed : dict[str, Any] - Processed data for one sending agent. - - Returns - ------- - EarlyFusionWirePayload - Typed payload containing only the remote model input. + idx : int + Dataset frame to publish. + agent_payload_builder : Callable + Adapter callback converting one local observation into its wire + representation. """ - 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. + base_data_dict = self.retrieve_base_data(idx) - 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 + if self.payload_handler is not None: + for cav_id, selected_cav_base in base_data_dict.items(): + lidar_points, _ = self.__prepare_local_lidar(selected_cav_base) + payload = agent_payload_builder( + lidar_points, + PoseFrameMetadata( + lidar_pose=selected_cav_base["params"]["lidar_pose"], + capture_frame=idx - int(selected_cav_base["time_delay"]), + ), + ) + self.payload_handler.set_opencda_payload(cav_id, self.module_name, payload) def __process_with_messages(self, ego_id, ego_lidar_pose, base_data_dict): projected_lidar_stack = [] @@ -115,11 +94,18 @@ def __process_with_messages(self, ego_id, ego_lidar_pose, base_data_dict): if ego_id in self.payload_handler.current_artery_payload: for cav_id, _ in base_data_dict.items(): + 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 - payload = self.decode_wire_payload(raw_payload) - projected_lidar_stack.append(payload.projected_lidar) + if self.communication_adapter is None: + raise RuntimeError("Early-fusion payload decoding requires a communication adapter") + decoded_payload = self.communication_adapter.decode_received_payload( + raw_payload, + ego_lidar_pose, + ) + projected_lidar_stack.append(decoded_payload["projected_lidar"]) return {"projected_lidar_stack": projected_lidar_stack} @@ -323,34 +309,51 @@ def get_item_single_car(self, selected_cav_base, ego_pose): selected_cav_processed : dict The dictionary contains the cav's processed information. """ - selected_cav_processed = {} - # calculate the transformation matrix transformation_matrix = x1_to_x2(selected_cav_base["params"]["lidar_pose"], 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_) - - shuffle_idx = np.random.permutation(lidar_np.shape[0]) - lidar_np = lidar_np[shuffle_idx] - spoofing_mask = spoofing_mask[shuffle_idx] - - # remove points that hit itself - ego_mask = (lidar_np[:, 0] >= -1.95) & (lidar_np[:, 0] <= 2.95) & (lidar_np[:, 1] >= -1.1) & (lidar_np[:, 1] <= 1.1) - lidar_np = lidar_np[np.logical_not(ego_mask)] - spoofing_mask = spoofing_mask[np.logical_not(ego_mask)] + lidar_np, spoofing_mask = self.__prepare_local_lidar(selected_cav_base) # project the lidar to ego space lidar_np[:, :3] = box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix) - selected_cav_processed.update( - { - "projected_lidar": lidar_np, - "projected_lidar_spoofing_mask": spoofing_mask, - } + return { + "projected_lidar": lidar_np, + "projected_lidar_spoofing_mask": spoofing_mask, + } + + @staticmethod + def __prepare_local_lidar( + selected_cav_base: dict[str, Any], + ) -> tuple[np.ndarray, np.ndarray]: + """ + Shuffle and remove ego returns without changing coordinate frames. + + Parameters + ---------- + selected_cav_base : dict[str, Any] + Raw dataset entry for one sending agent. + + Returns + ------- + tuple[numpy.ndarray, numpy.ndarray] + Sender-local LiDAR points and their aligned spoofing mask. + """ + 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_, ) - return selected_cav_processed + shuffle_idx = np.random.permutation(lidar_np.shape[0]) + lidar_np = np.array(lidar_np[shuffle_idx], copy=True) + spoofing_mask = np.array(spoofing_mask[shuffle_idx], copy=True) + + ego_mask = (lidar_np[:, 0] >= -1.95) & (lidar_np[:, 0] <= 2.95) & (lidar_np[:, 1] >= -1.1) & (lidar_np[:, 1] <= 1.1) + keep_mask = np.logical_not(ego_mask) + return lidar_np[keep_mask], spoofing_mask[keep_mask] def collate_batch_test(self, batch): """ diff --git a/OpenCOOD/opencood/data_utils/datasets/intermediate_fusion_dataset.py b/OpenCOOD/opencood/data_utils/datasets/intermediate_fusion_dataset.py index 545f8e7f..392e1072 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 collections.abc import Callable from typing import Any import numpy as np @@ -16,7 +17,6 @@ from opencood.data_utils.datasets import basedataset from opencood.models.communication_adapters import ( InferenceInput, - IntermediateFusionWirePayload, IntermediateMetadata, PoseFrameMetadata, V2XViTMetadata, @@ -69,47 +69,37 @@ def __init__(self, params, visualize, train=True, payload_handler=None): self.module_name = "OpenCOOD.IntermediateFusionDataset" self.model_name = params["model"]["core_method"] - def extract_data(self, idx): + def extract_data( + self, + idx: int, + agent_payload_builder: Callable[ + [InferenceInput, IntermediateMetadata | None], + object, + ], + ) -> None: + """ + Publish individual agent inputs through the selected model adapter. + + Parameters + ---------- + idx : int + Dataset frame to publish. + agent_payload_builder : Callable + Model adapter callback that encodes the sender-local input at the + configured model's communication boundary. + """ 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: 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) + payload = agent_payload_builder( + selected_cav_processed["inference_input"], + self.build_model_metadata(cav_id, selected_cav_base, 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, @@ -149,30 +139,6 @@ def build_model_metadata( ) 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 ego_lidar_pose = [] @@ -315,11 +281,6 @@ def __append_processed_model_data(self, data, cav_id, cav_base, cav_processed, r 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: 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] dy = cav_base["params"]["lidar_pose"][1] - ego_lidar_pose[1] @@ -372,20 +333,74 @@ def build_visualization_context(self, ego_id: str, base_data_dict: dict[str, Any } 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) + return self.__process_encoded_messages( + ego_id, + ego_lidar_pose, + base_data_dict, + receive_frame, + ) + + def __process_encoded_messages( + self, + ego_id: str, + ego_lidar_pose: list[float], + base_data_dict: dict[str, Any], + receive_frame: int, + ) -> dict[str, Any]: + """ + Assemble learned local and delivered features for model-side fusion. + + Parameters + ---------- + ego_id : str + Identifier of the receiving agent. + ego_lidar_pose : list[float] + Current receiver LiDAR pose. + base_data_dict : dict[str, Any] + Local scene observations. + receive_frame : int + Current dataset frame. + + Returns + ------- + dict[str, Any] + Successfully available learned features and their metadata. + """ + if self.communication_adapter is None: + raise RuntimeError("Intermediate feature decoding requires a communication adapter") + + ego_cav_base = base_data_dict[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, receive_frame) + features = [ + self.communication_adapter.encode_local_intermediate_input( + self, + ego_cav_processed["inference_input"], + ) + ] + metadata = [self.build_model_metadata(ego_id, ego_cav_base, receive_frame)] if ego_id in self.payload_handler.current_artery_payload: - for cav_id, _ in base_data_dict.items(): - raw_payload = self.payload_handler.get_artery_payload(ego_id, cav_id, self.module_name) + for cav_id in base_data_dict: + 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 - payload = self.decode_wire_payload(raw_payload) - self.__append_message_model_data(data, payload) + decoded = self.communication_adapter.decode_received_payload( + raw_payload, + ego_lidar_pose, + ) + metadata.append(decoded.pop("metadata")) + features.append(decoded) - return data + return { + "intermediate_features": features, + "metadata": metadata, + } def __process_without_messages(self, ego_id, ego_lidar_pose, base_data_dict, receive_frame): data = self.__build_model_data() @@ -460,16 +475,35 @@ def assemble_inference_sample( 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") + if "intermediate_features" in inference_input: + intermediate_features = inference_input["intermediate_features"] + if len(metadata) != len(intermediate_features): + raise ValueError("Every intermediate feature must have a matching metadata entry") + feature_fields = set(intermediate_features[0]) + if any(set(features) != feature_fields for features in intermediate_features[1:]): + raise TypeError("All agents must provide the same intermediate feature fields") + merged_feature_dict = { + field: ( + np.concatenate([features[field] for features in intermediate_features], axis=0) + if isinstance(intermediate_features[0][field], np.ndarray) + else np.asarray([features[field] for features in intermediate_features]) + ) + for field in feature_fields + } + cav_num = len(intermediate_features) + model_input = {"intermediate_features": merged_feature_dict} + else: + typed_inputs: list[InferenceInput] = inference_input["inference_inputs"] + if len(metadata) != len(typed_inputs): + raise ValueError("Every inference input must have a matching metadata entry") + model_input = {"processed_lidar": merge_inference_inputs(typed_inputs)} + cav_num = len(typed_inputs) ego_sample = { **local_supervision, - "processed_lidar": merged_feature_dict, - "cav_num": len(typed_inputs), + **model_input, + "cav_num": cav_num, **self.__assemble_model_metadata(metadata, receive_frame), } if self.visualize: @@ -583,6 +617,7 @@ def collate_batch_train(self, batch): object_bbx_mask = [] object_ids = [] processed_lidar_list = [] + intermediate_feature_list = [] # used to record different scenario record_len = [] label_dict_list = [] @@ -608,7 +643,10 @@ def collate_batch_train(self, batch): object_bbx_mask.append(ego_dict["object_bbx_mask"]) object_ids.append(ego_dict["object_ids"]) - processed_lidar_list.append(ego_dict["processed_lidar"]) + if "intermediate_features" in ego_dict: + intermediate_feature_list.append(ego_dict["intermediate_features"]) + else: + processed_lidar_list.append(ego_dict["processed_lidar"]) record_len.append(ego_dict["cav_num"]) label_dict_list.append(ego_dict["label_dict"]) @@ -626,9 +664,16 @@ 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)) - # 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) + if intermediate_feature_list: + if processed_lidar_list: + raise TypeError("A batch cannot mix preprocessor inputs and intermediate features") + intermediate_features_torch = { + field: torch.from_numpy(np.concatenate([features[field] for features in intermediate_feature_list], axis=0)) + for field in intermediate_feature_list[0] + } + else: + 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) @@ -639,13 +684,17 @@ def collate_batch_train(self, batch): { "object_bbx_center": object_bbx_center, "object_bbx_mask": object_bbx_mask, - "processed_lidar": processed_lidar_torch_dict, "record_len": record_len, "label_dict": label_torch_dict, "object_ids": object_ids[0], } ) + if intermediate_feature_list: + output_dict["ego"]["intermediate_features"] = intermediate_features_torch + else: + output_dict["ego"]["processed_lidar"] = processed_lidar_torch_dict + 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: diff --git a/OpenCOOD/opencood/data_utils/datasets/late_fusion_dataset.py b/OpenCOOD/opencood/data_utils/datasets/late_fusion_dataset.py index b04fade0..9a4b0464 100644 --- a/OpenCOOD/opencood/data_utils/datasets/late_fusion_dataset.py +++ b/OpenCOOD/opencood/data_utils/datasets/late_fusion_dataset.py @@ -6,6 +6,7 @@ import random import logging from collections import OrderedDict +from collections.abc import Callable from typing import Any import numpy as np @@ -16,7 +17,7 @@ from opencood.data_utils.post_processor import build_postprocessor from opencood.data_utils.datasets import basedataset from opencood.models.communication_adapters import ( - LateFusionWirePayload, + InferenceInput, PoseFrameMetadata, build_inference_input, inference_input_to_dict, @@ -31,7 +32,7 @@ class LateFusionDataset(basedataset.BaseDataset): """ - This class is for intermediate fusion where each vehicle transmit the + This class is for late fusion where each vehicle transmits its detection outputs to ego. """ @@ -52,69 +53,35 @@ def __getitem__(self, idx): return reformat_data_dict - def extract_data(self, idx): - base_data_dict = self.retrieve_base_data(idx) - - if self.payload_handler is not None: - for cav_id, selected_cav_base in base_data_dict.items(): - 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: + def extract_data( + self, + idx: int, + agent_payload_builder: Callable[[InferenceInput, PoseFrameMetadata], object], + ) -> None: """ - Build the late-fusion payload sent over V2X. + Publish local detector inputs through the late-fusion adapter. 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: + idx : int + Dataset frame to publish. + agent_payload_builder : Callable + Adapter callback that runs the sender-side detector and builds its + wire representation. """ - 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. + base_data_dict = self.retrieve_base_data(idx) - 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 + if self.payload_handler is not None: + for cav_id, selected_cav_base in base_data_dict.items(): + selected_cav_processed = self.__build_model_input(selected_cav_base) + payload = agent_payload_builder( + selected_cav_processed["inference_input"], + PoseFrameMetadata( + lidar_pose=selected_cav_base["params"]["lidar_pose"], + capture_frame=idx - int(selected_cav_base["time_delay"]), + ), + ) + self.payload_handler.set_opencda_payload(cav_id, self.module_name, payload) def __find_ego_vehicle(self, base_data_dict): ego_id = -1 @@ -224,18 +191,17 @@ def __process_with_messages(self, ego_id, ego_lidar_pose, base_data_dict): if ego_id in self.payload_handler.current_artery_payload: for cav_id, _ in base_data_dict.items(): + 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 - 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}) + if self.communication_adapter is None: + raise RuntimeError("Late-fusion payload decoding requires a communication adapter") + processed_data_dict[cav_id] = self.communication_adapter.decode_received_payload( + raw_payload, + ego_lidar_pose, + ) return processed_data_dict @@ -319,7 +285,8 @@ def assemble_inference_sample( """ anchor_box = local_supervision["anchor_box"] for cav_content in inference_input.values(): - cav_content["anchor_box"] = anchor_box + if "inference_input" in cav_content: + cav_content["anchor_box"] = anchor_box inference_input["ego"].update(local_supervision) inference_input["ego"].update(visualization_context) @@ -475,6 +442,15 @@ def collate_batch_test(self, batch): for cav_id, cav_content in batch.items(): output_dict.update({cav_id: {}}) + if "pred_box_tensor" in cav_content: + output_dict[cav_id].update( + { + "pred_box_tensor": torch.from_numpy(np.asarray(cav_content["pred_box_tensor"])), + "pred_score": torch.from_numpy(np.asarray(cav_content["pred_score"])), + } + ) + continue + # the anchor box is the same for all bounding boxes usually, thus # we don't need the batch dimension. if cav_content["anchor_box"] is not None: @@ -496,6 +472,7 @@ def collate_batch_test(self, batch): output_dict[cav_id].update( { "processed_lidar": processed_lidar_torch_dict, + "batch_size": 1, "transformation_matrix": transformation_matrix_torch, } ) @@ -551,7 +528,39 @@ def post_process(self, data_dict, output_dict): gt_box_tensor : torch.Tensor The tensor of gt bounding box. """ - pred_box_tensor, pred_score = self.post_processor.post_process(data_dict, output_dict) + pred_box_list = [] + pred_score_list = [] + for cav_id, cav_content in data_dict.items(): + if "pred_box_tensor" in cav_content: + boxes = cav_content["pred_box_tensor"] + scores = cav_content["pred_score"] + else: + boxes, scores = self.post_processor.decode_agent_predictions( + cav_content, + output_dict[cav_id], + ) + if boxes.shape[0] > 0: + transformation_matrix = cav_content["transformation_matrix"] + if boxes.shape[-1] == 3: + boxes = box_utils.project_box3d(boxes, transformation_matrix) + else: + box3d = torch.nn.functional.pad(boxes, (0, 1)) + boxes = box_utils.project_points_by_matrix_torch( + box3d.reshape(-1, 3), + transformation_matrix, + )[:, :2].reshape(-1, 4, 2) + + if boxes.shape[0] > 0: + pred_box_list.append(boxes) + pred_score_list.append(scores) + + if pred_box_list: + pred_box_tensor, pred_score = self.post_processor.post_process_detections( + torch.cat(pred_box_list, dim=0), + torch.cat(pred_score_list, dim=0), + ) + else: + pred_box_tensor, pred_score = None, None 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/post_processor/bev_postprocessor.py b/OpenCOOD/opencood/data_utils/post_processor/bev_postprocessor.py index c709f0d9..057fae72 100644 --- a/OpenCOOD/opencood/data_utils/post_processor/bev_postprocessor.py +++ b/OpenCOOD/opencood/data_utils/post_processor/bev_postprocessor.py @@ -193,6 +193,59 @@ def collate_batch(label_batch_list): } return processed_batch + def decode_agent_predictions(self, cav_content, output): + """ + Decode one agent's PIXOR output without projection or NMS. + + Parameters + ---------- + cav_content : dict + Local model input. It is accepted for a common postprocessor API. + output : dict + Raw output produced by the local detector. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + Sender-frame 2D box corners and confidence scores. + """ + del cav_content + prob = torch.sigmoid(output["cls"].squeeze(0).squeeze(0)) + reg_map = output["reg"].squeeze(0).permute(1, 2, 0) + reg_map = self.denormalize_reg_map(reg_map) + mask = torch.gt(prob, self.params["target_args"]["score_threshold"]) + if mask.sum() == 0: + return reg_map.new_empty((0, 4, 2)), prob.new_empty((0,)) + return self.reg_map_to_bbx_corners(reg_map, mask), prob[mask] + + def post_process_detections(self, boxes, scores): + """ + Apply receiver-side NMS and range filtering to 2D detections. + + Parameters + ---------- + boxes : torch.Tensor + Ego-frame boxes with shape ``(N, 4, 2)``. + scores : torch.Tensor + Confidence scores with shape ``(N,)``. + + Returns + ------- + tuple[torch.Tensor | None, torch.Tensor | None] + Final boxes and scores, or two ``None`` values when empty. + """ + if boxes.shape[0] == 0: + return None, None + keep_index = box_utils.nms_rotated(boxes, scores, self.params["nms_thresh"]) + boxes = boxes[keep_index] + scores = scores[keep_index] + + range_mask = box_utils.get_mask_for_boxes_within_range_torch(boxes) + boxes = boxes[range_mask] + scores = scores[range_mask] + assert scores.shape[0] == boxes.shape[0] + return boxes, scores + def post_process(self, data_dict, output_dict): """ Process the outputs of the model to 2D bounding box. @@ -226,18 +279,11 @@ def post_process(self, data_dict, output_dict): # the transformation matrix to ego space transformation_matrix = cav_content["transformation_matrix"] - # classification probability -- (label_shape[0], label_shape[1]) - prob = output_dict[cav_id]["cls"].squeeze(0).squeeze(0) - prob = torch.sigmoid(prob) - # regression map -- (label_shape[0], label_shape[1], 6) - reg_map = output_dict[cav_id]["reg"].squeeze(0).permute(1, 2, 0) - reg_map = self.denormalize_reg_map(reg_map) - threshold = self.params["target_args"]["score_threshold"] - mask = torch.gt(prob, threshold) - - if mask.sum() > 0: - # (number of high confidence bbx, 4, 2) - corners2d = self.reg_map_to_bbx_corners(reg_map, mask) + corners2d, scores = self.decode_agent_predictions( + cav_content, + output_dict[cav_id], + ) + if corners2d.shape[0] > 0: # assume the z-diviation in transformation_matrix is small, # thus we can pad zeros to simulate the 3d transformation. # (number of high confidence bbx, 4, 3) @@ -246,7 +292,6 @@ def post_process(self, data_dict, output_dict): projected_boxes2d = box_utils.project_points_by_matrix_torch(box3d.view(-1, 3), transformation_matrix)[:, :2] projected_boxes2d = projected_boxes2d.view(-1, 4, 2) - scores = prob[mask] pred_box2d_list.append(projected_boxes2d) pred_score_list.append(scores) @@ -256,17 +301,7 @@ def post_process(self, data_dict, output_dict): else: return None, None - keep_index = box_utils.nms_rotated(pred_box2ds, pred_scores, self.params["nms_thresh"]) - if len(keep_index): - pred_box2ds = pred_box2ds[keep_index] - pred_scores = pred_scores[keep_index] - - # filter out the prediction out of the range. - mask = box_utils.get_mask_for_boxes_within_range_torch(pred_box2ds) - pred_box2ds = pred_box2ds[mask, :, :] - pred_scores = pred_scores[mask] - assert pred_scores.shape[0] == pred_box2ds.shape[0] - return pred_box2ds, pred_scores + return self.post_process_detections(pred_box2ds, pred_scores) def reg_map_to_bbx_corners(self, reg_map, mask): """ diff --git a/OpenCOOD/opencood/data_utils/post_processor/ciassd_postprocessor.py b/OpenCOOD/opencood/data_utils/post_processor/ciassd_postprocessor.py index ee6eaf10..59487c31 100644 --- a/OpenCOOD/opencood/data_utils/post_processor/ciassd_postprocessor.py +++ b/OpenCOOD/opencood/data_utils/post_processor/ciassd_postprocessor.py @@ -15,6 +15,51 @@ def __init__(self, anchor_params, train): self.train = train self.anchor_num = self.params["anchor_args"]["num"] + def decode_agent_predictions(self, cav_content, output): + """ + Decode one agent's CIA-SSD output without projection or NMS. + + Parameters + ---------- + cav_content : dict + Local model input containing the configured anchors. + output : dict + Raw output produced by CIA-SSD. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + Sender-frame box corners and confidence scores. + """ + preds_dict = output["preds_dict_stage1"] + prob = torch.sigmoid(preds_dict["cls_preds"].permute(0, 2, 3, 1).contiguous()) + reg = preds_dict["box_preds"].permute(0, 2, 3, 1).contiguous() + iou = preds_dict["iou_preds"].permute(0, 2, 3, 1).contiguous().reshape(1, -1) + direction = preds_dict["dir_cls_preds"].permute(0, 2, 3, 1).contiguous().reshape(1, -1, 2) + + batch_box3d = self.delta_to_boxes3d(reg, cav_content["anchor_box"], False) + mask = torch.gt(prob, self.params["target_args"]["score_threshold"]).view(1, -1) + mask_reg = mask.unsqueeze(2).repeat(1, 1, 7) + assert batch_box3d.shape[0] == 1 + boxes3d = torch.masked_select(batch_box3d[0], mask_reg[0]).view(-1, 7) + scores = torch.masked_select(prob.reshape(-1), mask[0]) + if boxes3d.shape[0] == 0: + return boxes3d.new_empty((0, 8, 3)), scores + + direction_labels = torch.max(direction, dim=-1)[1][mask] + adjusted_iou = (iou + 1) * 0.5 + scores = scores * torch.pow(adjusted_iou.masked_select(mask), 4) + opposite_direction = (boxes3d[..., -1] > 0) ^ (direction_labels.byte() == 1) + boxes3d[..., -1] += torch.where( + opposite_direction, + torch.tensor(np.pi).type_as(boxes3d), + torch.tensor(0.0).type_as(boxes3d), + ) + return box_utils.boxes_to_corners_3d( + boxes3d, + order=self.params["order"], + ), scores + def post_process(self, data_dict, output_dict): """ Process the outputs of the model to 2D/3D bounding box. diff --git a/OpenCOOD/opencood/data_utils/post_processor/voxel_postprocessor.py b/OpenCOOD/opencood/data_utils/post_processor/voxel_postprocessor.py index b19f3351..beadd119 100644 --- a/OpenCOOD/opencood/data_utils/post_processor/voxel_postprocessor.py +++ b/OpenCOOD/opencood/data_utils/post_processor/voxel_postprocessor.py @@ -7,7 +7,6 @@ import numpy as np import torch -import torch.nn.functional as F from opencood.data_utils.post_processor.base_postprocessor import BasePostprocessor from opencood.utils import box_utils @@ -195,6 +194,71 @@ def collate_batch(label_batch_list): return {"targets": targets, "pos_equal_one": pos_equal_one, "neg_equal_one": neg_equal_one} + def decode_agent_predictions(self, cav_content, output): + """ + Decode one agent's detector output without projection or NMS. + + Parameters + ---------- + cav_content : dict + Local model input containing the configured anchors. + output : dict + Raw output produced by the local detector. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + Sender-frame box corners and confidence scores. + """ + anchor_box = cav_content["anchor_box"] + prob = torch.sigmoid(output["psm"].permute(0, 2, 3, 1)).reshape(1, -1) + batch_box3d = self.delta_to_boxes3d(output["rm"], anchor_box) + mask = torch.gt(prob, self.params["target_args"]["score_threshold"]).view(1, -1) + mask_reg = mask.unsqueeze(2).repeat(1, 1, 7) + + assert batch_box3d.shape[0] == 1 + boxes3d = torch.masked_select(batch_box3d[0], mask_reg[0]).view(-1, 7) + scores = torch.masked_select(prob[0], mask[0]) + if boxes3d.shape[0] == 0: + return boxes3d.new_empty((0, 8, 3)), scores + return box_utils.boxes_to_corners_3d(boxes3d, order=self.params["order"]), scores + + def post_process_detections(self, boxes, scores): + """ + Apply receiver-side filtering and NMS to decoded 3D detections. + + Parameters + ---------- + boxes : torch.Tensor + Ego-frame boxes with shape ``(N, 8, 3)``. + scores : torch.Tensor + Confidence scores with shape ``(N,)``. + + Returns + ------- + tuple[torch.Tensor | None, torch.Tensor | None] + Final boxes and scores, or two ``None`` values when empty. + """ + if boxes.shape[0] == 0: + return None, None + + keep_index = torch.logical_and( + box_utils.remove_large_pred_bbx(boxes), + box_utils.remove_bbx_abnormal_z(boxes), + ) + boxes = boxes[keep_index] + scores = scores[keep_index] + + keep_index = box_utils.nms_rotated(boxes, scores, self.params["nms_thresh"]) + boxes = boxes[keep_index] + scores = scores[keep_index] + + range_mask = box_utils.get_mask_for_boxes_within_range_torch(boxes) + boxes = boxes[range_mask] + scores = scores[range_mask] + assert scores.shape[0] == boxes.shape[0] + return boxes, scores + def post_process(self, data_dict, output_dict): """ Process the outputs of the model to 2D/3D bounding box. @@ -226,33 +290,11 @@ def post_process(self, data_dict, output_dict): # the transformation matrix to ego space transformation_matrix = cav_content["transformation_matrix"] - # (H, W, anchor_num, 7) - anchor_box = cav_content["anchor_box"] - - # classification probability - prob = output_dict[cav_id]["psm"] - prob = F.sigmoid(prob.permute(0, 2, 3, 1)) - prob = prob.reshape(1, -1) - - # regression map - reg = output_dict[cav_id]["rm"] - - # convert regression map back to bounding box - # (N, W*L*anchor_num, 7) - batch_box3d = self.delta_to_boxes3d(reg, anchor_box) - mask = torch.gt(prob, self.params["target_args"]["score_threshold"]) - mask = mask.view(1, -1) - mask_reg = mask.unsqueeze(2).repeat(1, 1, 7) - - # during validation/testing, the batch size should be 1 - assert batch_box3d.shape[0] == 1 - boxes3d = torch.masked_select(batch_box3d[0], mask_reg[0]).view(-1, 7) - scores = torch.masked_select(prob[0], mask[0]) - - # convert output to bounding box - if len(boxes3d) != 0: - # (N, 8, 3) - boxes3d_corner = box_utils.boxes_to_corners_3d(boxes3d, order=self.params["order"]) + boxes3d_corner, scores = self.decode_agent_predictions( + cav_content, + output_dict[cav_id], + ) + if len(boxes3d_corner) != 0: # (N, 8, 3) projected_boxes3d = box_utils.project_box3d(boxes3d_corner, transformation_matrix) # convert 3d bbx to 2d, (N,4) @@ -265,36 +307,11 @@ def post_process(self, data_dict, output_dict): if len(pred_box2d_list) == 0 or len(pred_box3d_list) == 0: return None, None - # shape: (N, 5) - pred_box2d_list = torch.vstack(pred_box2d_list) - # scores - scores = pred_box2d_list[:, -1] - # predicted 3d bbx - pred_box3d_tensor = torch.vstack(pred_box3d_list) - # remove large bbx - keep_index_1 = box_utils.remove_large_pred_bbx(pred_box3d_tensor) - keep_index_2 = box_utils.remove_bbx_abnormal_z(pred_box3d_tensor) - keep_index = torch.logical_and(keep_index_1, keep_index_2) - - pred_box3d_tensor = pred_box3d_tensor[keep_index] - scores = scores[keep_index] - - # nms - keep_index = box_utils.nms_rotated(pred_box3d_tensor, scores, self.params["nms_thresh"]) - - pred_box3d_tensor = pred_box3d_tensor[keep_index] - - # select cooresponding score - scores = scores[keep_index] - - # filter out the prediction out of the range. - mask = box_utils.get_mask_for_boxes_within_range_torch(pred_box3d_tensor) - pred_box3d_tensor = pred_box3d_tensor[mask, :, :] - scores = scores[mask] - - assert scores.shape[0] == pred_box3d_tensor.shape[0] - - return pred_box3d_tensor, scores + pred_box2d_tensor = torch.vstack(pred_box2d_list) + return self.post_process_detections( + torch.vstack(pred_box3d_list), + pred_box2d_tensor[:, -1], + ) @staticmethod def delta_to_boxes3d(deltas, anchors, channel_swap=True): diff --git a/OpenCOOD/opencood/models/ciassd.py b/OpenCOOD/opencood/models/ciassd.py index d9a73f24..b9f0bdbf 100644 --- a/OpenCOOD/opencood/models/ciassd.py +++ b/OpenCOOD/opencood/models/ciassd.py @@ -19,15 +19,24 @@ def __init__(self, args): self.head = Head(**args["head"]) def forward(self, batch_dict): - batch_dict["batch_size"] = batch_dict["object_bbx_center"].shape[0] - batch_dict = self.vfe(batch_dict) - batch_dict = self.spconv_block(batch_dict) - batch_dict = self.map_to_bev(batch_dict) - out = self.ssfa(batch_dict["processed_lidar"]["spatial_features"]) + processed_lidar = batch_dict["processed_lidar"] + batch_size = batch_dict.get("batch_size") + if batch_size is None: + batch_size = batch_dict["object_bbx_center"].shape[0] + model_batch = { + **processed_lidar, + "batch_size": batch_size, + } + model_batch = self.vfe(model_batch) + model_batch = self.spconv_block(model_batch) + model_batch = self.map_to_bev(model_batch) + out = self.ssfa(model_batch["spatial_features"]) out = self.head(out) - batch_dict["preds_dict_stage1"] = out - - return batch_dict + return { + "preds_dict_stage1": out, + "batch_size": batch_size, + "anchor_box": batch_dict.get("anchor_box"), + } if __name__ == "__main__": diff --git a/OpenCOOD/opencood/models/communication_adapters/__init__.py b/OpenCOOD/opencood/models/communication_adapters/__init__.py index f43d32bd..2db04053 100644 --- a/OpenCOOD/opencood/models/communication_adapters/__init__.py +++ b/OpenCOOD/opencood/models/communication_adapters/__init__.py @@ -6,39 +6,53 @@ from torch import nn from opencood.models.communication_adapters.base import ModelCommunicationAdapter +from opencood.models.communication_adapters.early import EarlyFusionCommunicationAdapter from opencood.models.communication_adapters.inference_input import ( BevInferenceInput, + DetectionInferenceInput, FpvrcnnAgentInferenceInput, InferenceInput, + MultiScaleFeatureInput, PointCloudInferenceInput, + SpatialFeatureInput, VoxelInferenceInput, + Where2CommFeatureInput, build_inference_input, inference_input_to_dict, merge_inference_inputs, ) +from opencood.models.communication_adapters.intermediate import ( + IntermediateFeatureCommunicationAdapter, + MultiScaleFeatureCommunicationAdapter, + SpatialFeatureCommunicationAdapter, +) from opencood.models.communication_adapters.inference_metadata import ( IntermediateMetadata, PoseFrameMetadata, V2XViTMetadata, ) -from opencood.models.communication_adapters.pending import PendingModelCommunicationAdapter +from opencood.models.communication_adapters.late import LateFusionCommunicationAdapter +from opencood.models.communication_adapters.where2comm import Where2CommCommunicationAdapter from opencood.models.communication_adapters.wire_payload import ( EarlyFusionWirePayload, FpvrcnnWirePayload, - IntermediateFusionWirePayload, + IntermediateFeatureWirePayload, LateFusionWirePayload, + Where2CommWirePayload, ) def build_communication_adapter( model: nn.Module, device: torch.device, + fusion_method: str | None = None, ) -> ModelCommunicationAdapter: """ Build the communication adapter declared by a perception model. - Models without a declared adapter temporarily retain their dataset-managed - payload through :class:`PendingModelCommunicationAdapter`. + Fusion-wide contracts are selected from the configured fusion method. + Intermediate models declare the adapter matching their learned-feature + communication boundary. Parameters ---------- @@ -46,6 +60,8 @@ def build_communication_adapter( Cooperative perception model instance. device : torch.device Device used for sender-side model execution. + fusion_method : str | None + Explicit fusion method selected by the model configuration. Returns ------- @@ -57,13 +73,14 @@ def build_communication_adapter( 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 fusion_method == "EarlyFusionDataset": + return EarlyFusionCommunicationAdapter(model, device) + if fusion_method == "LateFusionDataset": + return LateFusionCommunicationAdapter(model, device) + + adapter_class: object = getattr(type(model), "communication_adapter_class", None) + if adapter_class is None: + raise TypeError(f"{type(model).__name__} must declare communication_adapter_class") if not isinstance(adapter_class, type) or not issubclass( adapter_class, ModelCommunicationAdapter, @@ -74,19 +91,29 @@ def build_communication_adapter( __all__ = ( "BevInferenceInput", + "DetectionInferenceInput", "EarlyFusionWirePayload", + "EarlyFusionCommunicationAdapter", "FpvrcnnAgentInferenceInput", "FpvrcnnWirePayload", "InferenceInput", - "IntermediateFusionWirePayload", + "IntermediateFeatureCommunicationAdapter", + "IntermediateFeatureWirePayload", "IntermediateMetadata", "LateFusionWirePayload", + "LateFusionCommunicationAdapter", "ModelCommunicationAdapter", - "PendingModelCommunicationAdapter", + "MultiScaleFeatureCommunicationAdapter", + "MultiScaleFeatureInput", "PointCloudInferenceInput", "PoseFrameMetadata", + "SpatialFeatureCommunicationAdapter", + "SpatialFeatureInput", "V2XViTMetadata", "VoxelInferenceInput", + "Where2CommCommunicationAdapter", + "Where2CommFeatureInput", + "Where2CommWirePayload", "build_communication_adapter", "build_inference_input", "inference_input_to_dict", diff --git a/OpenCOOD/opencood/models/communication_adapters/base.py b/OpenCOOD/opencood/models/communication_adapters/base.py index 1899514b..137de7fa 100644 --- a/OpenCOOD/opencood/models/communication_adapters/base.py +++ b/OpenCOOD/opencood/models/communication_adapters/base.py @@ -75,6 +75,33 @@ def decode_received_payload( Raises ------ NotImplementedError - If the model still uses dataset-managed payload decoding. + If the adapter does not support receiver-side payload decoding. """ raise NotImplementedError(f"{type(self).__name__} does not implement receiver-side payload decoding") + + def encode_local_intermediate_input( + self, + dataset: Any, + inference_input: Any, + ) -> dict[str, Any]: + """ + Encode the receiver's local input at an intermediate model boundary. + + Parameters + ---------- + dataset : Any + Dataset providing model preprocessing facilities. + inference_input : Any + Receiver-local preprocessor output. + + Returns + ------- + dict[str, Any] + Model-specific learned feature input. + + Raises + ------ + NotImplementedError + If the adapter does not implement an intermediate feature split. + """ + raise NotImplementedError(f"{type(self).__name__} does not implement intermediate feature encoding") diff --git a/OpenCOOD/opencood/models/communication_adapters/early.py b/OpenCOOD/opencood/models/communication_adapters/early.py new file mode 100644 index 00000000..c55c9faa --- /dev/null +++ b/OpenCOOD/opencood/models/communication_adapters/early.py @@ -0,0 +1,110 @@ +"""Early-fusion V2X communication adapter.""" + +from collections.abc import Callable +from typing import Any, Protocol, runtime_checkable + +import numpy as np +import numpy.typing as npt + +from opencood.models.communication_adapters.base import ModelCommunicationAdapter +from opencood.models.communication_adapters.inference_metadata import PoseFrameMetadata +from opencood.models.communication_adapters.wire_payload import EarlyFusionWirePayload +from opencood.utils import box_utils +from opencood.utils.transformation_utils import x1_to_x2 + + +@runtime_checkable +class EarlyFusionCommunicationDataset(Protocol): + """Dataset capabilities required by the early-fusion adapter.""" + + def extract_data( + self, + idx: int, + agent_payload_builder: Callable[[npt.NDArray[Any], PoseFrameMetadata], object], + ) -> None: ... + + +class EarlyFusionCommunicationAdapter(ModelCommunicationAdapter): + """Exchange local LiDAR points and project them only at the receiver.""" + + def prepare_transmission_payloads(self, dataset: Any, idx: int) -> None: + """ + Build and publish an early-fusion payload for every local agent. + + Parameters + ---------- + dataset : Any + Dataset exposing local early-fusion agent observations. + idx : int + Dataset frame to publish. + + Raises + ------ + TypeError + If the selected dataset cannot build early-fusion observations. + """ + if not isinstance(dataset, EarlyFusionCommunicationDataset): + raise TypeError("Early fusion requires a dataset capable of building local LiDAR observations") + dataset.extract_data(idx, agent_payload_builder=self._build_wire_payload) + + @staticmethod + def _build_wire_payload( + lidar_points: npt.NDArray[Any], + metadata: PoseFrameMetadata, + ) -> EarlyFusionWirePayload: + """ + Build a payload without applying a receiver-specific transformation. + + Parameters + ---------- + lidar_points : numpy.typing.NDArray[Any] + LiDAR points expressed in the sender coordinate frame. + metadata : PoseFrameMetadata + Sender pose and capture frame. + + Returns + ------- + EarlyFusionWirePayload + Local LiDAR points and the metadata needed by a receiver. + """ + return EarlyFusionWirePayload( + lidar_points=lidar_points, + metadata=metadata, + ) + + def decode_received_payload( + self, + payload: object, + receiver_pose: list[float], + ) -> dict[str, Any]: + """ + Project one sender-local point cloud into the receiver frame. + + Parameters + ---------- + payload : object + Deserialized early-fusion payload delivered over V2X. + receiver_pose : list[float] + Receiver LiDAR pose for the current inference frame. + + Returns + ------- + dict[str, Any] + Point cloud expressed in the receiver coordinate frame. + + Raises + ------ + TypeError + If a payload belonging to another contract is received. + """ + if not isinstance(payload, EarlyFusionWirePayload): + raise TypeError(f"Expected EarlyFusionWirePayload, got {type(payload).__name__}") + + transformation_matrix = x1_to_x2(payload.metadata.lidar_pose, receiver_pose) + projected_lidar = np.array(payload.lidar_points, copy=True) + if len(projected_lidar) > 0: + projected_lidar[:, :3] = box_utils.project_points_by_matrix_torch( + projected_lidar[:, :3], + transformation_matrix, + ) + return {"projected_lidar": projected_lidar} diff --git a/OpenCOOD/opencood/models/communication_adapters/inference_input.py b/OpenCOOD/opencood/models/communication_adapters/inference_input.py index 430d0254..22e3962b 100644 --- a/OpenCOOD/opencood/models/communication_adapters/inference_input.py +++ b/OpenCOOD/opencood/models/communication_adapters/inference_input.py @@ -83,6 +83,69 @@ class FpvrcnnAgentInferenceInput: point_features: npt.NDArray[Any] +@dataclass(frozen=True, slots=True, kw_only=True) +class DetectionInferenceInput: + """ + Detection candidates produced by one late-fusion sender. + + Attributes + ---------- + boxes : numpy.typing.NDArray[Any] + Bounding-box corners expressed in the sender coordinate frame. + scores : numpy.typing.NDArray[Any] + Confidence scores corresponding to ``boxes``. + """ + + boxes: npt.NDArray[Any] + scores: npt.NDArray[Any] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class Where2CommFeatureInput: + """ + Learned feature data selected by a Where2Comm sender. + + Attributes + ---------- + spatial_features : numpy.typing.NDArray[Any] + Sender-side feature representation after the communication mask and + optional learned compression encoder. + communication_rate : float + Fraction of spatial locations selected for transmission. + """ + + spatial_features: npt.NDArray[Any] + communication_rate: float + + +@dataclass(frozen=True, slots=True, kw_only=True) +class SpatialFeatureInput: + """ + Single learned spatial feature map produced by an agent encoder. + + Attributes + ---------- + spatial_features : numpy.typing.NDArray[Any] + Feature map including its single-agent batch dimension. + """ + + spatial_features: npt.NDArray[Any] + + +@dataclass(frozen=True, slots=True, kw_only=True) +class MultiScaleFeatureInput: + """ + Learned feature maps produced at multiple encoder scales. + + Attributes + ---------- + feature_maps : tuple[numpy.typing.NDArray[Any], ...] + Ordered feature maps including their single-agent batch dimensions. + """ + + feature_maps: tuple[npt.NDArray[Any], ...] + + InferenceInput: TypeAlias = VoxelInferenceInput | BevInferenceInput | PointCloudInferenceInput diff --git a/OpenCOOD/opencood/models/communication_adapters/intermediate.py b/OpenCOOD/opencood/models/communication_adapters/intermediate.py new file mode 100644 index 00000000..ddbe86a9 --- /dev/null +++ b/OpenCOOD/opencood/models/communication_adapters/intermediate.py @@ -0,0 +1,260 @@ +"""Reusable adapters for learned intermediate feature maps.""" + +from collections.abc import Callable, Mapping +from typing import Any, Protocol, runtime_checkable + +import numpy as np +import torch + +from opencood.models.communication_adapters.base import ModelCommunicationAdapter +from opencood.models.communication_adapters.inference_input import ( + InferenceInput, + MultiScaleFeatureInput, + SpatialFeatureInput, + inference_input_to_dict, +) +from opencood.models.communication_adapters.inference_metadata import IntermediateMetadata +from opencood.models.communication_adapters.wire_payload import IntermediateFeatureWirePayload + + +@runtime_checkable +class IntermediateCommunicationDataset(Protocol): + """Dataset capabilities required by learned-feature adapters.""" + + pre_processor: Any + + def extract_data( + self, + idx: int, + agent_payload_builder: Callable[ + [InferenceInput, IntermediateMetadata | None], + object, + ], + ) -> None: ... + + +class IntermediateFeatureCommunicationAdapter(ModelCommunicationAdapter): + """Base adapter for models exposing an ``encode_agent`` split.""" + + def prepare_transmission_payloads(self, dataset: Any, idx: int) -> None: + """ + Encode and publish learned features for every local agent. + + Parameters + ---------- + dataset : Any + Intermediate dataset exposing individual agent inputs. + idx : int + Dataset frame to publish. + """ + communication_dataset = self._require_dataset(dataset) + communication_dataset.extract_data( + idx, + agent_payload_builder=lambda inference_input, metadata: self._build_wire_payload( + communication_dataset, + inference_input, + metadata, + ), + ) + + @staticmethod + def _require_dataset(dataset: Any) -> IntermediateCommunicationDataset: + if not isinstance(dataset, IntermediateCommunicationDataset): + raise TypeError("Intermediate feature encoding requires an intermediate communication dataset") + return dataset + + def _encode( + self, + dataset: IntermediateCommunicationDataset, + inference_input: InferenceInput, + ) -> Mapping[str, Any]: + encode_agent = getattr(self.model, "encode_agent", None) + if encode_agent is None: + raise TypeError(f"{type(self).__name__} requires a model with encode_agent()") + model_input = self.move_to_device( + { + "processed_lidar": dataset.pre_processor.collate_batch([inference_input_to_dict(inference_input)]), + "batch_size": 1, + } + ) + self.model.eval() + with torch.no_grad(): + output = encode_agent(model_input) + if not isinstance(output, Mapping): + raise TypeError("encode_agent() must return a tensor mapping") + return output + + def _build_wire_payload( + self, + dataset: IntermediateCommunicationDataset, + inference_input: InferenceInput, + metadata: IntermediateMetadata | None, + ) -> IntermediateFeatureWirePayload: + encoded = self._encode(dataset, inference_input) + return IntermediateFeatureWirePayload( + inference_input=self._serialize_features(encoded), + metadata=metadata, + ) + + def encode_local_intermediate_input( + self, + dataset: Any, + inference_input: Any, + ) -> dict[str, Any]: + """ + Encode the receiver's own observation at the model boundary. + + Parameters + ---------- + dataset : Any + Intermediate dataset providing preprocessing. + inference_input : Any + Receiver-local typed preprocessor output. + + Returns + ------- + dict[str, Any] + Learned feature maps in the receiver assembly format. + """ + communication_dataset = self._require_dataset(dataset) + encoded = self._encode(communication_dataset, inference_input) + return self._deserialize_features(self._serialize_features(encoded)) + + def decode_received_payload( + self, + payload: object, + receiver_pose: list[float], + ) -> dict[str, Any]: + """ + Decode one successfully delivered intermediate feature payload. + + Parameters + ---------- + payload : object + Deserialized model payload delivered over V2X. + receiver_pose : list[float] + Receiver pose reserved for model-side alignment. + + Returns + ------- + dict[str, Any] + Learned features and their sender metadata. + """ + del receiver_pose + if not isinstance(payload, IntermediateFeatureWirePayload): + raise TypeError(f"Expected IntermediateFeatureWirePayload, got {type(payload).__name__}") + decoded = self._deserialize_features(payload.inference_input) + decoded["metadata"] = payload.metadata + return decoded + + def _serialize_features( + self, + encoded: Mapping[str, Any], + ) -> SpatialFeatureInput | MultiScaleFeatureInput: + raise NotImplementedError + + def _deserialize_features( + self, + inference_input: SpatialFeatureInput | MultiScaleFeatureInput, + ) -> dict[str, Any]: + raise NotImplementedError + + +class SpatialFeatureCommunicationAdapter(IntermediateFeatureCommunicationAdapter): + """Exchange one dense learned feature map per agent.""" + + def _serialize_features( + self, + encoded: Mapping[str, Any], + ) -> SpatialFeatureInput: + """ + Move one encoded spatial feature map to the wire representation. + + Parameters + ---------- + encoded : collections.abc.Mapping[str, Any] + Model output containing ``spatial_features``. + + Returns + ------- + SpatialFeatureInput + CPU NumPy representation of the learned feature map. + """ + spatial_features = encoded.get("spatial_features") + if not isinstance(spatial_features, torch.Tensor): + raise TypeError('encode_agent() field "spatial_features" must be a tensor') + return SpatialFeatureInput( + spatial_features=spatial_features.detach().cpu().numpy(), + ) + + def _deserialize_features( + self, + inference_input: SpatialFeatureInput | MultiScaleFeatureInput, + ) -> dict[str, Any]: + """ + Restore one spatial feature map for receiver-side collation. + + Parameters + ---------- + inference_input : SpatialFeatureInput | MultiScaleFeatureInput + Typed feature representation received over V2X. + + Returns + ------- + dict[str, Any] + Receiver assembly mapping containing ``spatial_features``. + """ + if not isinstance(inference_input, SpatialFeatureInput): + raise TypeError(f"Expected SpatialFeatureInput, got {type(inference_input).__name__}") + return { + "spatial_features": np.array(inference_input.spatial_features, copy=True), + } + + +class MultiScaleFeatureCommunicationAdapter(IntermediateFeatureCommunicationAdapter): + """Exchange an ordered tuple of learned feature scales per agent.""" + + def _serialize_features( + self, + encoded: Mapping[str, Any], + ) -> MultiScaleFeatureInput: + """ + Move ordered model feature scales to the wire representation. + + Parameters + ---------- + encoded : collections.abc.Mapping[str, Any] + Model output containing ``feature_maps``. + + Returns + ------- + MultiScaleFeatureInput + CPU NumPy representations of all learned feature scales. + """ + feature_maps = encoded.get("feature_maps") + if not isinstance(feature_maps, (list, tuple)) or not all(isinstance(feature_map, torch.Tensor) for feature_map in feature_maps): + raise TypeError('encode_agent() field "feature_maps" must be a tensor sequence') + return MultiScaleFeatureInput( + feature_maps=tuple(feature_map.detach().cpu().numpy() for feature_map in feature_maps), + ) + + def _deserialize_features( + self, + inference_input: SpatialFeatureInput | MultiScaleFeatureInput, + ) -> dict[str, Any]: + """ + Restore ordered feature scales for receiver-side collation. + + Parameters + ---------- + inference_input : SpatialFeatureInput | MultiScaleFeatureInput + Typed feature representation received over V2X. + + Returns + ------- + dict[str, Any] + Receiver assembly mapping keyed by scale index. + """ + if not isinstance(inference_input, MultiScaleFeatureInput): + raise TypeError(f"Expected MultiScaleFeatureInput, got {type(inference_input).__name__}") + return {f"feature_{index}": np.array(feature_map, copy=True) for index, feature_map in enumerate(inference_input.feature_maps)} diff --git a/OpenCOOD/opencood/models/communication_adapters/late.py b/OpenCOOD/opencood/models/communication_adapters/late.py new file mode 100644 index 00000000..6d299094 --- /dev/null +++ b/OpenCOOD/opencood/models/communication_adapters/late.py @@ -0,0 +1,169 @@ +"""Late-fusion V2X communication adapter.""" + +from collections.abc import Callable, Mapping +from typing import Any, Protocol, runtime_checkable + +import numpy as np +import torch + +from opencood.models.communication_adapters.base import ModelCommunicationAdapter +from opencood.models.communication_adapters.inference_input import ( + DetectionInferenceInput, + InferenceInput, + inference_input_to_dict, +) +from opencood.models.communication_adapters.inference_metadata import PoseFrameMetadata +from opencood.models.communication_adapters.wire_payload import LateFusionWirePayload +from opencood.utils import box_utils +from opencood.utils.transformation_utils import x1_to_x2 + + +@runtime_checkable +class LateFusionCommunicationDataset(Protocol): + """Dataset capabilities required by the late-fusion adapter.""" + + pre_processor: Any + post_processor: Any + + def extract_data( + self, + idx: int, + agent_payload_builder: Callable[[InferenceInput, PoseFrameMetadata], object], + ) -> None: ... + + +class LateFusionCommunicationAdapter(ModelCommunicationAdapter): + """Run a complete detector locally and exchange detection candidates.""" + + def prepare_transmission_payloads(self, dataset: Any, idx: int) -> None: + """ + Run sender-side late-fusion inference for every local agent. + + Parameters + ---------- + dataset : Any + Dataset exposing local detector inputs and post-processing. + idx : int + Dataset frame to publish. + + Raises + ------ + TypeError + If the selected dataset lacks late-fusion capabilities. + """ + if not isinstance(dataset, LateFusionCommunicationDataset): + raise TypeError("Late fusion requires a dataset capable of building local detector inputs") + dataset.extract_data( + idx, + agent_payload_builder=lambda inference_input, metadata: self._build_wire_payload( + dataset, + inference_input, + metadata, + ), + ) + + def _build_wire_payload( + self, + dataset: LateFusionCommunicationDataset, + inference_input: InferenceInput, + metadata: PoseFrameMetadata, + ) -> LateFusionWirePayload: + """ + Run the local detector and serialize its decoded candidates. + + Parameters + ---------- + dataset : LateFusionCommunicationDataset + Dataset providing pre- and post-processing implementations. + inference_input : InferenceInput + Preprocessed input belonging only to the sender. + metadata : PoseFrameMetadata + Sender pose and capture frame. + + Returns + ------- + LateFusionWirePayload + Decoded sender-frame boxes and their confidence scores. + + Raises + ------ + TypeError + If the model output or postprocessor result is invalid. + """ + processed_lidar = dataset.pre_processor.collate_batch([inference_input_to_dict(inference_input)]) + anchor_box = dataset.post_processor.generate_anchor_box() + model_input = self.move_to_device( + { + "processed_lidar": processed_lidar, + "batch_size": 1, + "anchor_box": (None if anchor_box is None else torch.from_numpy(np.asarray(anchor_box))), + } + ) + + self.model.eval() + with torch.no_grad(): + output = self.model(model_input) + if not isinstance(output, Mapping): + raise TypeError("Late-fusion model must return a tensor mapping") + + decode_predictions = getattr(dataset.post_processor, "decode_agent_predictions", None) + if decode_predictions is None: + raise TypeError("Late-fusion postprocessor must implement decode_agent_predictions()") + boxes, scores = decode_predictions(model_input, output) + if not isinstance(boxes, torch.Tensor) or not isinstance(scores, torch.Tensor): + raise TypeError("decode_agent_predictions() must return box and score tensors") + + return LateFusionWirePayload( + detections=DetectionInferenceInput( + boxes=boxes.detach().cpu().numpy(), + scores=scores.detach().cpu().numpy(), + ), + metadata=metadata, + ) + + def decode_received_payload( + self, + payload: object, + receiver_pose: list[float], + ) -> dict[str, Any]: + """ + Project one sender's detections into the receiver frame. + + Parameters + ---------- + payload : object + Deserialized late-fusion payload delivered over V2X. + receiver_pose : list[float] + Receiver LiDAR pose for the current inference frame. + + Returns + ------- + dict[str, Any] + Detection candidates expressed in the receiver frame. + + Raises + ------ + TypeError + If the payload or box representation is unsupported. + """ + if not isinstance(payload, LateFusionWirePayload): + raise TypeError(f"Expected LateFusionWirePayload, got {type(payload).__name__}") + + boxes = np.array(payload.detections.boxes, copy=True) + transformation_matrix = x1_to_x2(payload.metadata.lidar_pose, receiver_pose) + if boxes.ndim != 3 or boxes.shape[-2:] not in {(4, 2), (8, 3)}: + raise TypeError(f"Unsupported late-fusion box shape: {boxes.shape}") + if boxes.shape[0] > 0: + if boxes.shape[-1] == 3: + boxes = box_utils.project_box3d(boxes, transformation_matrix) + else: + box3d = np.pad(boxes, ((0, 0), (0, 0), (0, 1))) + boxes = box_utils.project_points_by_matrix_torch( + box3d.reshape(-1, 3), + transformation_matrix, + )[:, :2].reshape(-1, 4, 2) + + return { + "pred_box_tensor": boxes, + "pred_score": np.array(payload.detections.scores, copy=True), + } diff --git a/OpenCOOD/opencood/models/communication_adapters/pending.py b/OpenCOOD/opencood/models/communication_adapters/pending.py deleted file mode 100644 index a058931b..00000000 --- a/OpenCOOD/opencood/models/communication_adapters/pending.py +++ /dev/null @@ -1,35 +0,0 @@ -"""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/where2comm.py b/OpenCOOD/opencood/models/communication_adapters/where2comm.py new file mode 100644 index 00000000..81ad0bdd --- /dev/null +++ b/OpenCOOD/opencood/models/communication_adapters/where2comm.py @@ -0,0 +1,225 @@ +"""Where2Comm intermediate-feature communication adapter.""" + +from collections.abc import Callable, Mapping +from typing import Any, Protocol, runtime_checkable + +import numpy as np +import torch + +from opencood.models.communication_adapters.base import ModelCommunicationAdapter +from opencood.models.communication_adapters.inference_input import ( + InferenceInput, + Where2CommFeatureInput, + inference_input_to_dict, +) +from opencood.models.communication_adapters.inference_metadata import PoseFrameMetadata +from opencood.models.communication_adapters.wire_payload import Where2CommWirePayload + + +@runtime_checkable +class Where2CommCommunicationDataset(Protocol): + """Dataset capabilities required by the Where2Comm adapter.""" + + pre_processor: Any + + def extract_data( + self, + idx: int, + agent_payload_builder: Callable[ + [InferenceInput, PoseFrameMetadata], + object, + ], + ) -> None: ... + + +class Where2CommCommunicationAdapter(ModelCommunicationAdapter): + """Execute the private Where2Comm encoder before V2X transmission.""" + + def prepare_transmission_payloads(self, dataset: Any, idx: int) -> None: + """ + Encode and publish learned features for every local agent. + + Parameters + ---------- + dataset : Any + Dataset exposing individual intermediate-fusion inputs. + idx : int + Dataset frame to publish. + + Raises + ------ + TypeError + If the selected dataset lacks the required capabilities. + """ + if not isinstance(dataset, Where2CommCommunicationDataset): + raise TypeError("Where2Comm requires an intermediate-fusion communication dataset") + dataset.extract_data( + idx, + agent_payload_builder=lambda inference_input, metadata: self._build_wire_payload( + dataset, + inference_input, + metadata, + ), + ) + + def _encode( + self, + dataset: Where2CommCommunicationDataset, + inference_input: InferenceInput, + *, + apply_communication_mask: bool, + ) -> tuple[np.ndarray, float]: + """ + Run the sender-private model stages for one agent. + + Parameters + ---------- + dataset : Where2CommCommunicationDataset + Dataset providing the configured preprocessor. + inference_input : InferenceInput + One agent's preprocessed LiDAR input. + apply_communication_mask : bool + Whether to apply the over-the-air feature selection mask. + + Returns + ------- + tuple[numpy.ndarray, float] + Learned feature map and selected-location ratio. + + Raises + ------ + TypeError + If the model does not expose the Where2Comm encoder contract. + """ + encode_agent = getattr(self.model, "encode_agent", None) + if encode_agent is None: + raise TypeError("Where2CommCommunicationAdapter requires a model with encode_agent()") + + model_input = self.move_to_device( + { + "processed_lidar": dataset.pre_processor.collate_batch([inference_input_to_dict(inference_input)]), + } + ) + self.model.eval() + with torch.no_grad(): + output = encode_agent( + model_input, + apply_communication_mask=apply_communication_mask, + ) + if not isinstance(output, Mapping): + raise TypeError("Where2Comm encode_agent() must return a tensor mapping") + spatial_features = output.get("spatial_features") + communication_rate = output.get("communication_rate") + if not isinstance(spatial_features, torch.Tensor): + raise TypeError('Where2Comm output field "spatial_features" must be a tensor') + if not isinstance(communication_rate, torch.Tensor): + raise TypeError('Where2Comm output field "communication_rate" must be a tensor') + return ( + spatial_features.detach().cpu().numpy(), + float(communication_rate.detach().cpu().item()), + ) + + def _build_wire_payload( + self, + dataset: Where2CommCommunicationDataset, + inference_input: InferenceInput, + metadata: PoseFrameMetadata, + ) -> Where2CommWirePayload: + """ + Build one sender's learned-feature payload. + + Parameters + ---------- + dataset : Where2CommCommunicationDataset + Dataset providing the configured preprocessor. + inference_input : InferenceInput + Sender-local preprocessor output. + metadata : PoseFrameMetadata + Sender pose and capture frame. + + Returns + ------- + Where2CommWirePayload + Masked learned features and sender metadata. + """ + spatial_features, communication_rate = self._encode( + dataset, + inference_input, + apply_communication_mask=True, + ) + return Where2CommWirePayload( + inference_input=Where2CommFeatureInput( + spatial_features=spatial_features, + communication_rate=communication_rate, + ), + metadata=metadata, + ) + + def encode_local_intermediate_input( + self, + dataset: Any, + inference_input: Any, + ) -> dict[str, Any]: + """ + Encode the receiver's own feature without an over-the-air mask. + + Parameters + ---------- + dataset : Any + Intermediate-fusion dataset providing the preprocessor. + inference_input : Any + Receiver-local typed preprocessor output. + + Returns + ------- + dict[str, Any] + Receiver-local learned feature input. + """ + if not isinstance(dataset, Where2CommCommunicationDataset): + raise TypeError("Where2Comm requires an intermediate-fusion communication dataset") + spatial_features, communication_rate = self._encode( + dataset, + inference_input, + apply_communication_mask=False, + ) + return { + "spatial_features": spatial_features, + "communication_rate": communication_rate, + } + + def decode_received_payload( + self, + payload: object, + receiver_pose: list[float], + ) -> dict[str, Any]: + """ + Decode one successfully delivered Where2Comm feature payload. + + Parameters + ---------- + payload : object + Deserialized Where2Comm payload delivered over V2X. + receiver_pose : list[float] + Receiver pose reserved for later receiver-side feature alignment. + + Returns + ------- + dict[str, Any] + Learned feature map and communication rate. + + Raises + ------ + TypeError + If a payload belonging to another contract is received. + """ + del receiver_pose + if not isinstance(payload, Where2CommWirePayload): + raise TypeError(f"Expected Where2CommWirePayload, got {type(payload).__name__}") + return { + "spatial_features": np.array( + payload.inference_input.spatial_features, + copy=True, + ), + "communication_rate": payload.inference_input.communication_rate, + "metadata": payload.metadata, + } diff --git a/OpenCOOD/opencood/models/communication_adapters/wire_payload.py b/OpenCOOD/opencood/models/communication_adapters/wire_payload.py index 0aab313d..56ee71e9 100644 --- a/OpenCOOD/opencood/models/communication_adapters/wire_payload.py +++ b/OpenCOOD/opencood/models/communication_adapters/wire_payload.py @@ -6,8 +6,11 @@ import numpy.typing as npt from opencood.models.communication_adapters.inference_input import ( + DetectionInferenceInput, FpvrcnnAgentInferenceInput, - InferenceInput, + MultiScaleFeatureInput, + SpatialFeatureInput, + Where2CommFeatureInput, ) from opencood.models.communication_adapters.inference_metadata import ( IntermediateMetadata, @@ -18,65 +21,84 @@ @dataclass(frozen=True, slots=True, kw_only=True) class EarlyFusionWirePayload: """ - Wire payload produced by :class:`EarlyFusionDataset`. + Sender-local LiDAR data exchanged for early fusion. Attributes ---------- - projected_lidar : numpy.typing.NDArray[Any] - LiDAR points projected into the ego coordinate frame. + lidar_points : numpy.typing.NDArray[Any] + LiDAR points expressed in the sender coordinate frame. + metadata : PoseFrameMetadata + Sender pose and capture frame used for receiver-side projection. """ - projected_lidar: npt.NDArray[Any] + lidar_points: npt.NDArray[Any] + metadata: PoseFrameMetadata @dataclass(frozen=True, slots=True, kw_only=True) class LateFusionWirePayload: """ - Wire payload produced by :class:`LateFusionDataset`. + Sender-side detections exchanged for late fusion. Attributes ---------- - inference_input : InferenceInput - Typed model input produced by the sender's preprocessor. + detections : DetectionInferenceInput + Bounding-box candidates and confidence scores produced by the sender. metadata : PoseFrameMetadata Pose and capture frame required for late-fusion post-processing. """ - inference_input: InferenceInput + detections: DetectionInferenceInput metadata: PoseFrameMetadata @dataclass(frozen=True, slots=True, kw_only=True) -class IntermediateFusionWirePayload: +class FpvrcnnWirePayload: """ - Wire payload produced by :class:`IntermediateFusionDataset`. + Agent-side FPV-RCNN output exchanged over V2X. 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 : 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: InferenceInput - metadata: IntermediateMetadata | None + inference_input: FpvrcnnAgentInferenceInput + metadata: PoseFrameMetadata @dataclass(frozen=True, slots=True, kw_only=True) -class FpvrcnnWirePayload: +class Where2CommWirePayload: """ - Agent-side FPV-RCNN output exchanged over V2X. + Sender-side learned features exchanged by Where2Comm. Attributes ---------- - inference_input : FpvrcnnAgentInferenceInput - Stage-one proposals and VSA keypoint features computed by the sender. + inference_input : Where2CommFeatureInput + Masked feature map and its effective communication rate. metadata : PoseFrameMetadata - Sender pose and capture frame used to transform the output at the - receiver. + Sender pose and capture frame used by the fusion pipeline. """ - inference_input: FpvrcnnAgentInferenceInput + inference_input: Where2CommFeatureInput metadata: PoseFrameMetadata + + +@dataclass(frozen=True, slots=True, kw_only=True) +class IntermediateFeatureWirePayload: + """ + Learned features exchanged by an intermediate-fusion model. + + Attributes + ---------- + inference_input : SpatialFeatureInput | MultiScaleFeatureInput + Model-specific learned features produced by the sender. + metadata : IntermediateMetadata | None + Metadata required by the receiver-side fusion implementation. + """ + + inference_input: SpatialFeatureInput | MultiScaleFeatureInput + metadata: IntermediateMetadata | None diff --git a/OpenCOOD/opencood/models/fuse_modules/where2comm_fuse.py b/OpenCOOD/opencood/models/fuse_modules/where2comm_fuse.py index f27a68eb..6d4b3665 100644 --- a/OpenCOOD/opencood/models/fuse_modules/where2comm_fuse.py +++ b/OpenCOOD/opencood/models/fuse_modules/where2comm_fuse.py @@ -35,43 +35,72 @@ def init_gaussian_filter(self, k_size=5, sigma=1.0): self.gaussian_filter.weight.data = torch.Tensor(gaussian_kernel).to(self.gaussian_filter.weight.device).unsqueeze(0).unsqueeze(0) self.gaussian_filter.bias.data.zero_() + def build_mask(self, confidence_map, *, force_first: bool): + """ + Build the spatial transmission mask for one group of agents. + + Parameters + ---------- + confidence_map : torch.Tensor + Per-agent classification logits. + force_first : bool + Whether the first feature belongs to the local receiver and must + remain unmasked. + + Returns + ------- + tuple[torch.Tensor, torch.Tensor] + Binary communication mask and selected-location ratio. + """ + _, _, height, width = confidence_map.shape + original_map, _ = confidence_map.sigmoid().max(dim=1, keepdim=True) + communication_map = self.gaussian_filter(original_map) if self.smooth else original_map + agent_count = communication_map.shape[0] + + if self.training: + selected_count = int(height * width * random.uniform(0, 1)) + flattened_map = communication_map.reshape(agent_count, height * width) + _, indices = torch.topk(flattened_map, k=selected_count, sorted=False) + communication_mask = torch.zeros_like(flattened_map) + selected_values = torch.ones( + agent_count, + selected_count, + dtype=communication_map.dtype, + device=communication_map.device, + ) + communication_mask = torch.scatter( + communication_mask, + -1, + indices, + selected_values, + ).reshape(agent_count, 1, height, width) + elif self.threshold: + communication_mask = torch.where( + communication_map > self.threshold, + torch.ones_like(communication_map), + torch.zeros_like(communication_map), + ) + else: + communication_mask = torch.ones_like(communication_map) + + communication_rate = communication_mask.sum() / (agent_count * height * width) + if force_first: + communication_mask[0] = 1 + return communication_mask, communication_rate + def forward(self, batch_confidence_maps, B): """ Args: batch_confidence_maps: [(L1, H, W), (L2, H, W), ...] """ - _, _, H, W = batch_confidence_maps[0].shape - communication_masks = [] communication_rates = [] for b in range(B): - ori_communication_maps, _ = batch_confidence_maps[b].sigmoid().max(dim=1, keepdim=True) - if self.smooth: - communication_maps = self.gaussian_filter(ori_communication_maps) - else: - communication_maps = ori_communication_maps - - L = communication_maps.shape[0] - if self.training: - # Official training proxy objective - K = int(H * W * random.uniform(0, 1)) - communication_maps = communication_maps.reshape(L, H * W) - _, indices = torch.topk(communication_maps, k=K, sorted=False) - communication_mask = torch.zeros_like(communication_maps).to(communication_maps.device) - ones_fill = torch.ones(L, K, dtype=communication_maps.dtype, device=communication_maps.device) - communication_mask = torch.scatter(communication_mask, -1, indices, ones_fill).reshape(L, 1, H, W) - elif self.threshold: - ones_mask = torch.ones_like(communication_maps).to(communication_maps.device) - zeros_mask = torch.zeros_like(communication_maps).to(communication_maps.device) - communication_mask = torch.where(communication_maps > self.threshold, ones_mask, zeros_mask) - else: - communication_mask = torch.ones_like(communication_maps).to(communication_maps.device) - - communication_rate = communication_mask.sum() / (L * H * W) - # Ego - communication_mask[0] = 1 - + communication_mask, communication_rate = self.build_mask( + batch_confidence_maps[b], + force_first=True, + ) communication_masks.append(communication_mask) communication_rates.append(communication_rate) communication_rates = sum(communication_rates) / B diff --git a/OpenCOOD/opencood/models/pixor_intermediate.py b/OpenCOOD/opencood/models/pixor_intermediate.py index b9594620..ff880c59 100644 --- a/OpenCOOD/opencood/models/pixor_intermediate.py +++ b/OpenCOOD/opencood/models/pixor_intermediate.py @@ -2,6 +2,7 @@ import torch.nn as nn +from opencood.models.communication_adapters.intermediate import MultiScaleFeatureCommunicationAdapter from opencood.models.fuse_modules.self_attn import AttFusion from opencood.models.pixor import Bottleneck, BackBone, Header @@ -15,14 +16,29 @@ def __init__(self, block, num_block, geom, use_bn=True): self.fusion_net5 = AttFusion(384) def forward(self, x, record_len): - c3, c4, c5 = self.encode(x) - + return self.fuse_agents(self.encode(x), record_len) + + def fuse_agents(self, feature_maps, record_len): + """ + Fuse encoded PIXOR scales and run the feature pyramid decoder. + + Parameters + ---------- + feature_maps : tuple[torch.Tensor, torch.Tensor, torch.Tensor] + Per-agent ``c3``, ``c4``, and ``c5`` feature maps. + record_len : torch.Tensor + Number of available agents in each sample. + + Returns + ------- + torch.Tensor + Decoded fused PIXOR feature map. + """ + c3, c4, c5 = feature_maps c5 = self.fusion_net5(c5, record_len) c4 = self.fusion_net4(c4, record_len) c3 = self.fusion_net3(c3, record_len) - - p4 = self.decode(c3, c4, c5) - return p4 + return self.decode(c3, c4, c5) class PIXORIntermediate(nn.Module): @@ -45,6 +61,8 @@ class PIXORIntermediate(nn.Module): Header used to predict the classification and coordinates. """ + communication_adapter_class = MultiScaleFeatureCommunicationAdapter + def __init__(self, args): super(PIXORIntermediate, self).__init__() geom = args["geometry_param"] @@ -67,6 +85,9 @@ def __init__(self, args): self.header.reghead.bias.data.fill_(0) def forward(self, data_dict): + if "intermediate_features" in data_dict: + return self.fuse_agents(data_dict) + bev_input = data_dict["processed_lidar"]["bev_input"] record_len = data_dict["record_len"] @@ -78,3 +99,41 @@ def forward(self, data_dict): output_dict = {"cls": cls, "reg": reg} return output_dict + + def encode_agent(self, data_dict): + """ + Encode one agent into PIXOR's three private backbone scales. + + Parameters + ---------- + data_dict : dict + Sender-local preprocessed BEV input. + + Returns + ------- + dict + Ordered ``c3``, ``c4``, and ``c5`` feature maps. + """ + return { + "feature_maps": self.backbone.encode(data_dict["processed_lidar"]["bev_input"]), + } + + def fuse_agents(self, data_dict): + """ + Fuse received PIXOR features and run its detection header. + + Parameters + ---------- + data_dict : dict + Receiver input containing learned features and ``record_len``. + + Returns + ------- + dict + PIXOR classification and regression maps. + """ + intermediate_features = data_dict["intermediate_features"] + feature_maps = tuple(intermediate_features[f"feature_{index}"] for index in range(3)) + features = self.backbone.fuse_agents(feature_maps, data_dict["record_len"]) + cls, reg = self.header(features) + return {"cls": cls, "reg": reg} diff --git a/OpenCOOD/opencood/models/point_pillar_coalign.py b/OpenCOOD/opencood/models/point_pillar_coalign.py index 7710cc19..d0193851 100644 --- a/OpenCOOD/opencood/models/point_pillar_coalign.py +++ b/OpenCOOD/opencood/models/point_pillar_coalign.py @@ -6,7 +6,7 @@ import torch.nn as nn - +from opencood.models.communication_adapters.intermediate import SpatialFeatureCommunicationAdapter from opencood.models.sub_modules.pillar_vfe import PillarVFE from opencood.models.sub_modules.point_pillar_scatter import PointPillarScatter from opencood.models.sub_modules.res_bev_backbone import ResBEVBackbone @@ -16,6 +16,8 @@ class PointPillarCoAlign(nn.Module): + communication_adapter_class = SpatialFeatureCommunicationAdapter + def __init__(self, args): super(PointPillarCoAlign, self).__init__() @@ -25,6 +27,7 @@ def __init__(self, args): self.backbone = ResBEVBackbone(args["res_bev_backbone"], 64) self.voxel_size = args["voxel_size"] + self.first_layer_stride = args["res_bev_backbone"]["layer_strides"][0] # multiscale fusion network modules self.fusion_net = nn.ModuleList() @@ -49,6 +52,9 @@ def __init__(self, args): self.reg_head = nn.Conv2d(self.out_channel, 7 * args["anchor_num"], kernel_size=1) def forward(self, data_dict): + if "intermediate_features" in data_dict: + return self.fuse_agents(data_dict) + voxel_features = data_dict["processed_lidar"]["voxel_features"] voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] @@ -89,3 +95,83 @@ def forward(self, data_dict): output_dict = {"psm": psm, "rm": rm} return output_dict + + def encode_agent(self, data_dict): + """ + Encode one agent up to CoAlign's first fusion scale. + + Parameters + ---------- + data_dict : dict + Sender-local preprocessed LiDAR input. + + Returns + ------- + dict + First-scale learned feature transmitted by CoAlign. + """ + processed_lidar = data_dict["processed_lidar"] + batch_dict = { + "voxel_features": processed_lidar["voxel_features"], + "voxel_coords": processed_lidar["voxel_coords"], + "voxel_num_points": processed_lidar["voxel_num_points"], + } + batch_dict = self.pillar_vfe(batch_dict) + batch_dict = self.scatter(batch_dict) + spatial_features = self.backbone.get_layer_i_feature( + batch_dict["spatial_features"], + layer_i=0, + ) + if self.compression: + spatial_features = self.naive_compressor.encode(spatial_features) + return {"spatial_features": spatial_features} + + def fuse_agents(self, data_dict): + """ + Fuse delivered CoAlign features and run the detection heads. + + Parameters + ---------- + data_dict : dict + Receiver input containing learned features and pairwise poses. + + Returns + ------- + dict + Classification and regression maps. + """ + first_scale = data_dict["intermediate_features"]["spatial_features"] + if self.compression: + first_scale = self.naive_compressor.decode(first_scale) + _, _, height, width = first_scale.shape + normalized_affine_matrix = normalize_pairwise_tfm( + data_dict["pairwise_t_matrix"], + height, + width, + self.voxel_size[0], + downsample_rate=self.first_layer_stride, + ) + + feature_list = [first_scale] + for level in range(1, self.backbone.num_levels): + feature_list.append( + self.backbone.get_layer_i_feature( + feature_list[level - 1], + layer_i=level, + ) + ) + fused_feature_list = [ + fuse_module( + feature_list[level], + data_dict["record_len"], + normalized_affine_matrix, + ) + for level, fuse_module in enumerate(self.fusion_net) + ] + fused_feature = self.backbone.decode_multiscale_feature(fused_feature_list) + if self.shrink_flag: + fused_feature = self.shrink_conv(fused_feature) + return { + "psm": self.cls_head(fused_feature), + "rm": self.reg_head(fused_feature), + } diff --git a/OpenCOOD/opencood/models/point_pillar_cobevt.py b/OpenCOOD/opencood/models/point_pillar_cobevt.py index 6e526a0a..7e10c49b 100644 --- a/OpenCOOD/opencood/models/point_pillar_cobevt.py +++ b/OpenCOOD/opencood/models/point_pillar_cobevt.py @@ -1,6 +1,7 @@ import torch.nn as nn from einops import repeat +from opencood.models.communication_adapters.intermediate import SpatialFeatureCommunicationAdapter from opencood.models.sub_modules.pillar_vfe import PillarVFE from opencood.models.sub_modules.point_pillar_scatter import PointPillarScatter from opencood.models.sub_modules.base_bev_backbone import BaseBEVBackbone @@ -11,6 +12,8 @@ class PointPillarCoBEVT(nn.Module): + communication_adapter_class = SpatialFeatureCommunicationAdapter + def __init__(self, args): super(PointPillarCoBEVT, self).__init__() @@ -64,6 +67,9 @@ def backbone_fix(self): p.requires_grad = False def forward(self, data_dict): + if "intermediate_features" in data_dict: + return self.fuse_agents(data_dict) + voxel_features = data_dict["processed_lidar"]["voxel_features"] voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] @@ -97,3 +103,67 @@ def forward(self, data_dict): output_dict = {"psm": psm, "rm": rm} return output_dict + + def encode_agent(self, data_dict): + """ + Encode one agent into the feature map exchanged by CoBEVT. + + Parameters + ---------- + data_dict : dict + Sender-local preprocessed LiDAR input. + + Returns + ------- + dict + Learned spatial feature map before Swap Fusion. + """ + processed_lidar = data_dict["processed_lidar"] + batch_dict = { + "voxel_features": processed_lidar["voxel_features"], + "voxel_coords": processed_lidar["voxel_coords"], + "voxel_num_points": processed_lidar["voxel_num_points"], + } + batch_dict = self.pillar_vfe(batch_dict) + batch_dict = self.scatter(batch_dict) + batch_dict = self.backbone(batch_dict) + spatial_features = batch_dict["spatial_features_2d"] + if self.shrink_flag: + spatial_features = self.shrink_conv(spatial_features) + if self.compression: + spatial_features = self.naive_compressor.encode(spatial_features) + return {"spatial_features": spatial_features} + + def fuse_agents(self, data_dict): + """ + Fuse delivered CoBEVT features and run the detection heads. + + Parameters + ---------- + data_dict : dict + Receiver input containing learned features and ``record_len``. + + Returns + ------- + dict + Classification and regression maps. + """ + spatial_features = data_dict["intermediate_features"]["spatial_features"] + if self.compression: + spatial_features = self.naive_compressor.decode(spatial_features) + regroup_feature, mask = regroup( + spatial_features, + data_dict["record_len"], + self.max_cav, + ) + communication_mask = repeat( + mask.unsqueeze(1).unsqueeze(2).unsqueeze(3), + "b h w c l -> b (h new_h) (w new_w) c l", + new_h=regroup_feature.shape[3], + new_w=regroup_feature.shape[4], + ) + fused_feature = self.fusion_net(regroup_feature, communication_mask) + return { + "psm": self.cls_head(fused_feature), + "rm": self.reg_head(fused_feature), + } diff --git a/OpenCOOD/opencood/models/point_pillar_fcooper.py b/OpenCOOD/opencood/models/point_pillar_fcooper.py index 894642a3..ff433c3a 100644 --- a/OpenCOOD/opencood/models/point_pillar_fcooper.py +++ b/OpenCOOD/opencood/models/point_pillar_fcooper.py @@ -1,5 +1,6 @@ import torch.nn as nn +from opencood.models.communication_adapters.intermediate import SpatialFeatureCommunicationAdapter from opencood.models.sub_modules.pillar_vfe import PillarVFE from opencood.models.sub_modules.point_pillar_scatter import PointPillarScatter from opencood.models.sub_modules.base_bev_backbone import BaseBEVBackbone @@ -13,6 +14,8 @@ class PointPillarFCooper(nn.Module): F-Cooper implementation with point pillar backbone. """ + communication_adapter_class = SpatialFeatureCommunicationAdapter + def __init__(self, args): super(PointPillarFCooper, self).__init__() @@ -66,6 +69,9 @@ def backbone_fix(self): p.requires_grad = False def forward(self, data_dict): + if "intermediate_features" in data_dict: + return self.fuse_agents(data_dict) + voxel_features = data_dict["processed_lidar"]["voxel_features"] voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] @@ -94,3 +100,59 @@ def forward(self, data_dict): output_dict = {"psm": psm, "rm": rm} return output_dict + + def encode_agent(self, data_dict): + """ + Encode one agent into the BEV feature exchanged by F-Cooper. + + Parameters + ---------- + data_dict : dict + Sender-local preprocessed LiDAR input. + + Returns + ------- + dict + Learned spatial feature map before F-Cooper fusion. + """ + processed_lidar = data_dict["processed_lidar"] + batch_dict = { + "voxel_features": processed_lidar["voxel_features"], + "voxel_coords": processed_lidar["voxel_coords"], + "voxel_num_points": processed_lidar["voxel_num_points"], + } + batch_dict = self.pillar_vfe(batch_dict) + batch_dict = self.scatter(batch_dict) + batch_dict = self.backbone(batch_dict) + spatial_features = batch_dict["spatial_features_2d"] + if self.shrink_flag: + spatial_features = self.shrink_conv(spatial_features) + if self.compression: + spatial_features = self.naive_compressor.encode(spatial_features) + return {"spatial_features": spatial_features} + + def fuse_agents(self, data_dict): + """ + Fuse delivered F-Cooper features and run the detection heads. + + Parameters + ---------- + data_dict : dict + Receiver input containing learned features and ``record_len``. + + Returns + ------- + dict + Classification and regression maps. + """ + spatial_features = data_dict["intermediate_features"]["spatial_features"] + if self.compression: + spatial_features = self.naive_compressor.decode(spatial_features) + fused_feature = self.fusion_net( + spatial_features, + data_dict["record_len"], + ) + return { + "psm": self.cls_head(fused_feature), + "rm": self.reg_head(fused_feature), + } diff --git a/OpenCOOD/opencood/models/point_pillar_intermediate.py b/OpenCOOD/opencood/models/point_pillar_intermediate.py index 092e28ec..b0e625d2 100644 --- a/OpenCOOD/opencood/models/point_pillar_intermediate.py +++ b/OpenCOOD/opencood/models/point_pillar_intermediate.py @@ -1,12 +1,14 @@ import torch.nn as nn - +from opencood.models.communication_adapters.intermediate import MultiScaleFeatureCommunicationAdapter from opencood.models.sub_modules.pillar_vfe import PillarVFE from opencood.models.sub_modules.point_pillar_scatter import PointPillarScatter from opencood.models.sub_modules.att_bev_backbone import AttBEVBackbone class PointPillarIntermediate(nn.Module): + communication_adapter_class = MultiScaleFeatureCommunicationAdapter + def __init__(self, args): super(PointPillarIntermediate, self).__init__() @@ -19,6 +21,9 @@ def __init__(self, args): self.reg_head = nn.Conv2d(128 * 3, 7 * args["anchor_num"], kernel_size=1) def forward(self, data_dict): + if "intermediate_features" in data_dict: + return self.fuse_agents(data_dict) + voxel_features = data_dict["processed_lidar"]["voxel_features"] voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] @@ -38,3 +43,51 @@ def forward(self, data_dict): output_dict = {"psm": psm, "rm": rm} return output_dict + + def encode_agent(self, data_dict): + """ + Encode one agent into private multi-scale PointPillar features. + + Parameters + ---------- + data_dict : dict + Sender-local preprocessed LiDAR input. + + Returns + ------- + dict + Ordered feature maps at the model's attention boundaries. + """ + processed_lidar = data_dict["processed_lidar"] + batch_dict = { + "voxel_features": processed_lidar["voxel_features"], + "voxel_coords": processed_lidar["voxel_coords"], + "voxel_num_points": processed_lidar["voxel_num_points"], + } + batch_dict = self.pillar_vfe(batch_dict) + batch_dict = self.scatter(batch_dict) + return { + "feature_maps": self.backbone.encode_agent(batch_dict["spatial_features"]), + } + + def fuse_agents(self, data_dict): + """ + Fuse received feature scales and run PointPillar detection heads. + + Parameters + ---------- + data_dict : dict + Receiver input containing learned features and ``record_len``. + + Returns + ------- + dict + Classification and regression maps. + """ + intermediate_features = data_dict["intermediate_features"] + feature_maps = tuple(intermediate_features[f"feature_{index}"] for index in range(len(self.backbone.fuse_modules))) + fused_features = self.backbone.fuse_agents(feature_maps, data_dict["record_len"]) + return { + "psm": self.cls_head(fused_features), + "rm": self.reg_head(fused_features), + } diff --git a/OpenCOOD/opencood/models/point_pillar_intermediate_V2VAM.py b/OpenCOOD/opencood/models/point_pillar_intermediate_V2VAM.py index 477c68cc..16974c33 100644 --- a/OpenCOOD/opencood/models/point_pillar_intermediate_V2VAM.py +++ b/OpenCOOD/opencood/models/point_pillar_intermediate_V2VAM.py @@ -1,5 +1,6 @@ import torch.nn as nn +from opencood.models.communication_adapters.intermediate import SpatialFeatureCommunicationAdapter from opencood.models.sub_modules.pillar_vfe import PillarVFE from opencood.models.sub_modules.point_pillar_scatter import PointPillarScatter from opencood.models.sub_modules.base_bev_backbone import BaseBEVBackbone @@ -10,9 +11,11 @@ class PointPillarintermediateV2VAM(nn.Module): """ - F-Cooper implementation with point pillar backbone. + V2VAM implementation with a PointPillar backbone. """ + communication_adapter_class = SpatialFeatureCommunicationAdapter + def __init__(self, args): super(PointPillarintermediateV2VAM, self).__init__() @@ -67,6 +70,9 @@ def backbone_fix(self): p.requires_grad = False def forward(self, data_dict): + if "intermediate_features" in data_dict: + return self.fuse_agents(data_dict) + voxel_features = data_dict["processed_lidar"]["voxel_features"] voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] @@ -98,3 +104,59 @@ def forward(self, data_dict): output_dict = {"psm": psm, "rm": rm} return output_dict + + def encode_agent(self, data_dict): + """ + Encode one agent into the BEV feature exchanged by V2VAM. + + Parameters + ---------- + data_dict : dict + Sender-local preprocessed LiDAR input. + + Returns + ------- + dict + Learned spatial feature map before V2VAM fusion. + """ + processed_lidar = data_dict["processed_lidar"] + batch_dict = { + "voxel_features": processed_lidar["voxel_features"], + "voxel_coords": processed_lidar["voxel_coords"], + "voxel_num_points": processed_lidar["voxel_num_points"], + } + batch_dict = self.pillar_vfe(batch_dict) + batch_dict = self.scatter(batch_dict) + batch_dict = self.backbone(batch_dict) + spatial_features = batch_dict["spatial_features_2d"] + if self.shrink_flag: + spatial_features = self.shrink_conv(spatial_features) + if self.compression: + spatial_features = self.naive_compressor.encode(spatial_features) + return {"spatial_features": spatial_features} + + def fuse_agents(self, data_dict): + """ + Fuse delivered V2VAM features and run the detection heads. + + Parameters + ---------- + data_dict : dict + Receiver input containing learned features and ``record_len``. + + Returns + ------- + dict + Classification and regression maps. + """ + spatial_features = data_dict["intermediate_features"]["spatial_features"] + if self.compression: + spatial_features = self.naive_compressor.decode(spatial_features) + fused_feature = self.fusion_net( + spatial_features, + data_dict["record_len"], + ) + return { + "psm": self.cls_head(fused_feature), + "rm": self.reg_head(fused_feature), + } diff --git a/OpenCOOD/opencood/models/point_pillar_transformer.py b/OpenCOOD/opencood/models/point_pillar_transformer.py index 41f8e434..a3649181 100644 --- a/OpenCOOD/opencood/models/point_pillar_transformer.py +++ b/OpenCOOD/opencood/models/point_pillar_transformer.py @@ -1,6 +1,7 @@ import torch import torch.nn as nn +from opencood.models.communication_adapters.intermediate import SpatialFeatureCommunicationAdapter from opencood.models.sub_modules.pillar_vfe import PillarVFE from opencood.models.sub_modules.point_pillar_scatter import PointPillarScatter from opencood.models.sub_modules.base_bev_backbone import BaseBEVBackbone @@ -11,6 +12,8 @@ class PointPillarTransformer(nn.Module): + communication_adapter_class = SpatialFeatureCommunicationAdapter + def __init__(self, args): super(PointPillarTransformer, self).__init__() @@ -64,6 +67,9 @@ def backbone_fix(self): p.requires_grad = False def forward(self, data_dict): + if "intermediate_features" in data_dict: + return self.fuse_agents(data_dict) + voxel_features = data_dict["processed_lidar"]["voxel_features"] voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] @@ -106,3 +112,74 @@ def forward(self, data_dict): output_dict = {"psm": psm, "rm": rm} return output_dict + + def encode_agent(self, data_dict): + """ + Encode one agent into the feature map exchanged by V2X-ViT. + + Parameters + ---------- + data_dict : dict + Sender-local preprocessed LiDAR input. + + Returns + ------- + dict + Learned spatial feature map before transformer fusion. + """ + processed_lidar = data_dict["processed_lidar"] + batch_dict = { + "voxel_features": processed_lidar["voxel_features"], + "voxel_coords": processed_lidar["voxel_coords"], + "voxel_num_points": processed_lidar["voxel_num_points"], + } + batch_dict = self.pillar_vfe(batch_dict) + batch_dict = self.scatter(batch_dict) + batch_dict = self.backbone(batch_dict) + spatial_features = batch_dict["spatial_features_2d"] + if self.shrink_flag: + spatial_features = self.shrink_conv(spatial_features) + if self.compression: + spatial_features = self.naive_compressor.encode(spatial_features) + return {"spatial_features": spatial_features} + + def fuse_agents(self, data_dict): + """ + Fuse delivered V2X-ViT features and run the detection heads. + + Parameters + ---------- + data_dict : dict + Receiver input containing learned features and V2X-ViT metadata. + + Returns + ------- + dict + Classification and regression maps. + """ + spatial_features = data_dict["intermediate_features"]["spatial_features"] + if self.compression: + spatial_features = self.naive_compressor.decode(spatial_features) + regroup_feature, mask = regroup( + spatial_features, + data_dict["record_len"], + self.max_cav, + ) + prior_encoding = data_dict["prior_encoding"].unsqueeze(-1).unsqueeze(-1) + prior_encoding = prior_encoding.repeat( + 1, + 1, + 1, + regroup_feature.shape[3], + regroup_feature.shape[4], + ) + regroup_feature = torch.cat([regroup_feature, prior_encoding], dim=2) + fused_feature = self.fusion_net( + regroup_feature.permute(0, 1, 3, 4, 2), + mask, + data_dict["spatial_correction_matrix"], + ).permute(0, 3, 1, 2) + return { + "psm": self.cls_head(fused_feature), + "rm": self.reg_head(fused_feature), + } diff --git a/OpenCOOD/opencood/models/point_pillar_v2vnet.py b/OpenCOOD/opencood/models/point_pillar_v2vnet.py index f318e392..715bf8eb 100644 --- a/OpenCOOD/opencood/models/point_pillar_v2vnet.py +++ b/OpenCOOD/opencood/models/point_pillar_v2vnet.py @@ -1,5 +1,6 @@ import torch.nn as nn +from opencood.models.communication_adapters.intermediate import SpatialFeatureCommunicationAdapter from opencood.models.sub_modules.pillar_vfe import PillarVFE from opencood.models.sub_modules.point_pillar_scatter import PointPillarScatter from opencood.models.sub_modules.base_bev_backbone import BaseBEVBackbone @@ -9,6 +10,8 @@ class PointPillarV2VNet(nn.Module): + communication_adapter_class = SpatialFeatureCommunicationAdapter + def __init__(self, args): super(PointPillarV2VNet, self).__init__() @@ -62,6 +65,9 @@ def backbone_fix(self): p.requires_grad = False def forward(self, data_dict): + if "intermediate_features" in data_dict: + return self.fuse_agents(data_dict) + voxel_features = data_dict["processed_lidar"]["voxel_features"] voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] @@ -91,3 +97,60 @@ def forward(self, data_dict): output_dict = {"psm": psm, "rm": rm} return output_dict + + def encode_agent(self, data_dict): + """ + Encode one agent into the BEV feature exchanged by V2VNet. + + Parameters + ---------- + data_dict : dict + Sender-local preprocessed LiDAR input. + + Returns + ------- + dict + Learned spatial feature map before V2VNet fusion. + """ + processed_lidar = data_dict["processed_lidar"] + batch_dict = { + "voxel_features": processed_lidar["voxel_features"], + "voxel_coords": processed_lidar["voxel_coords"], + "voxel_num_points": processed_lidar["voxel_num_points"], + } + batch_dict = self.pillar_vfe(batch_dict) + batch_dict = self.scatter(batch_dict) + batch_dict = self.backbone(batch_dict) + spatial_features = batch_dict["spatial_features_2d"] + if self.shrink_flag: + spatial_features = self.shrink_conv(spatial_features) + if self.compression: + spatial_features = self.naive_compressor.encode(spatial_features) + return {"spatial_features": spatial_features} + + def fuse_agents(self, data_dict): + """ + Fuse delivered V2VNet features and run the detection heads. + + Parameters + ---------- + data_dict : dict + Receiver input containing features, poses, and ``record_len``. + + Returns + ------- + dict + Classification and regression maps. + """ + spatial_features = data_dict["intermediate_features"]["spatial_features"] + if self.compression: + spatial_features = self.naive_compressor.decode(spatial_features) + fused_feature = self.fusion_net( + spatial_features, + data_dict["record_len"], + data_dict["pairwise_t_matrix"], + ) + return { + "psm": self.cls_head(fused_feature), + "rm": self.reg_head(fused_feature), + } diff --git a/OpenCOOD/opencood/models/point_pillar_where2comm.py b/OpenCOOD/opencood/models/point_pillar_where2comm.py index 295bd814..2b0f6de5 100644 --- a/OpenCOOD/opencood/models/point_pillar_where2comm.py +++ b/OpenCOOD/opencood/models/point_pillar_where2comm.py @@ -1,5 +1,8 @@ +import torch import torch.nn as nn +import torch.nn.functional as F +from opencood.models.communication_adapters.where2comm import Where2CommCommunicationAdapter from opencood.models.sub_modules.base_bev_backbone import BaseBEVBackbone from opencood.models.fuse_modules.where2comm_fuse import Where2comm from opencood.models.sub_modules.downsample_conv import DownsampleConv @@ -9,6 +12,8 @@ class PointPillarWhere2comm(nn.Module): + communication_adapter_class = Where2CommCommunicationAdapter + def __init__(self, args): super(PointPillarWhere2comm, self).__init__() self.max_cav = args["max_cav"] @@ -66,6 +71,9 @@ def backbone_fix(self): p.requires_grad = False def forward(self, data_dict): + if "intermediate_features" in data_dict: + return self.fuse_agents(data_dict) + voxel_features = data_dict["processed_lidar"]["voxel_features"] voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] @@ -107,3 +115,113 @@ def forward(self, data_dict): output_dict = {"psm": psm, "rm": rm, "com": communication_rates} return output_dict + + def encode_agent(self, data_dict, *, apply_communication_mask): + """ + Encode one agent up to the Where2Comm transmission boundary. + + Parameters + ---------- + data_dict : dict + Sender-local preprocessed LiDAR input. + apply_communication_mask : bool + Whether to select only features intended for transmission. + + Returns + ------- + dict + Learned feature map and selected-location ratio. + """ + processed_lidar = data_dict["processed_lidar"] + batch_dict = { + "voxel_features": processed_lidar["voxel_features"], + "voxel_coords": processed_lidar["voxel_coords"], + "voxel_num_points": processed_lidar["voxel_num_points"], + "record_len": torch.ones(1, dtype=torch.long, device=processed_lidar["voxel_features"].device), + } + batch_dict = self.pillar_vfe(batch_dict) + batch_dict = self.scatter(batch_dict) + batch_dict = self.backbone(batch_dict) + + confidence_features = batch_dict["spatial_features_2d"] + if self.shrink_flag: + confidence_features = self.shrink_conv(confidence_features) + confidence_map = self.cls_head(confidence_features) + + if self.multi_scale: + spatial_features = self.backbone.blocks[0](batch_dict["spatial_features"]) + else: + spatial_features = confidence_features + if self.compression: + spatial_features = self.naive_compressor.encode(spatial_features) + + communication_rate = torch.ones((), device=spatial_features.device) + if not self.fusion_net.fully: + communication_mask, communication_rate = self.fusion_net.naive_communication.build_mask( + confidence_map, + force_first=False, + ) + if apply_communication_mask: + if spatial_features.shape[-2:] != communication_mask.shape[-2:]: + communication_mask = F.interpolate( + communication_mask, + size=spatial_features.shape[-2:], + mode="bilinear", + align_corners=False, + ) + spatial_features = spatial_features * communication_mask + + return { + "spatial_features": spatial_features, + "communication_rate": communication_rate, + } + + def fuse_agents(self, data_dict): + """ + Fuse successfully received learned features and run detection heads. + + Parameters + ---------- + data_dict : dict + Receiver input containing encoded features and ``record_len``. + + Returns + ------- + dict + Fused classification, regression, and communication outputs. + """ + intermediate_features = data_dict["intermediate_features"] + spatial_features = intermediate_features["spatial_features"] + record_len = data_dict["record_len"] + + if self.multi_scale: + batch_node_features = self.fusion_net.regroup(spatial_features, record_len) + fused_feature = torch.stack([self.fusion_net.fuse_modules[0](features) for features in batch_node_features]) + upsampled_features = [self.backbone.deblocks[0](fused_feature) if len(self.backbone.deblocks) > 0 else fused_feature] + + for level in range(1, self.fusion_net.num_levels): + spatial_features = self.backbone.blocks[level](spatial_features) + batch_node_features = self.fusion_net.regroup(spatial_features, record_len) + fused_feature = torch.stack([self.fusion_net.fuse_modules[level](features) for features in batch_node_features]) + if len(self.backbone.deblocks) > 0: + upsampled_features.append(self.backbone.deblocks[level](fused_feature)) + else: + upsampled_features.append(fused_feature) + + fused_feature = torch.cat(upsampled_features, dim=1) + if len(self.backbone.deblocks) > self.fusion_net.num_levels: + fused_feature = self.backbone.deblocks[-1](fused_feature) + if self.shrink_flag: + fused_feature = self.shrink_conv(fused_feature) + else: + if self.compression: + spatial_features = self.naive_compressor.decode(spatial_features) + batch_node_features = self.fusion_net.regroup(spatial_features, record_len) + fused_feature = torch.stack([self.fusion_net.fuse_modules(features) for features in batch_node_features]) + + communication_rates = intermediate_features["communication_rate"] + return { + "psm": self.cls_head(fused_feature), + "rm": self.reg_head(fused_feature), + "com": communication_rates.mean(), + } diff --git a/OpenCOOD/opencood/models/second.py b/OpenCOOD/opencood/models/second.py index 26f797d5..4f28467b 100644 --- a/OpenCOOD/opencood/models/second.py +++ b/OpenCOOD/opencood/models/second.py @@ -33,7 +33,7 @@ def forward(self, data_dict): "voxel_features": voxel_features, "voxel_coords": voxel_coords, "voxel_num_points": voxel_num_points, - "batch_size": self.batch_size, + "batch_size": data_dict.get("batch_size", self.batch_size), } batch_dict = self.mean_vfe(batch_dict) diff --git a/OpenCOOD/opencood/models/second_intermediate.py b/OpenCOOD/opencood/models/second_intermediate.py index 3f37d67b..aad6f284 100644 --- a/OpenCOOD/opencood/models/second_intermediate.py +++ b/OpenCOOD/opencood/models/second_intermediate.py @@ -1,6 +1,7 @@ import torch import torch.nn as nn +from opencood.models.communication_adapters.intermediate import MultiScaleFeatureCommunicationAdapter 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 @@ -8,6 +9,8 @@ class SecondIntermediate(nn.Module): + communication_adapter_class = MultiScaleFeatureCommunicationAdapter + def __init__(self, args): super(SecondIntermediate, self).__init__() @@ -26,6 +29,9 @@ def __init__(self, args): self.reg_head = nn.Conv2d(256 * 2, 7 * args["anchor_num"], kernel_size=1) def forward(self, data_dict): + if "intermediate_features" in data_dict: + return self.fuse_agents(data_dict) + voxel_features = data_dict["processed_lidar"]["voxel_features"] voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] @@ -52,3 +58,56 @@ def forward(self, data_dict): output_dict = {"psm": psm, "rm": rm} return output_dict + + def encode_agent(self, data_dict): + """ + Encode one agent into private multi-scale SECOND features. + + Parameters + ---------- + data_dict : dict + Sender-local preprocessed LiDAR input. + + Returns + ------- + dict + Ordered feature maps at the 2D attention boundaries. + """ + processed_lidar = data_dict["processed_lidar"] + batch_dict = { + "voxel_features": processed_lidar["voxel_features"], + "voxel_coords": processed_lidar["voxel_coords"], + "voxel_num_points": processed_lidar["voxel_num_points"], + "batch_size": data_dict.get("batch_size", 1), + } + batch_dict = self.mean_vfe(batch_dict) + batch_dict = self.backbone_3d(batch_dict) + batch_dict = self.height_compression(batch_dict) + return { + "feature_maps": self.backbone_2d.encode_agent(batch_dict["spatial_features"]), + } + + def fuse_agents(self, data_dict): + """ + Fuse received SECOND feature scales and run detection heads. + + Parameters + ---------- + data_dict : dict + Receiver input containing learned features and ``record_len``. + + Returns + ------- + dict + Classification and regression maps. + """ + intermediate_features = data_dict["intermediate_features"] + feature_maps = tuple(intermediate_features[f"feature_{index}"] for index in range(len(self.backbone_2d.fuse_modules))) + fused_features = self.backbone_2d.fuse_agents( + feature_maps, + data_dict["record_len"], + ) + return { + "psm": self.cls_head(fused_features), + "rm": self.reg_head(fused_features), + } diff --git a/OpenCOOD/opencood/models/sub_modules/att_bev_backbone.py b/OpenCOOD/opencood/models/sub_modules/att_bev_backbone.py index 51b188d5..b6ad7010 100644 --- a/OpenCOOD/opencood/models/sub_modules/att_bev_backbone.py +++ b/OpenCOOD/opencood/models/sub_modules/att_bev_backbone.py @@ -105,31 +105,92 @@ def forward(self, data_dict): spatial_features = data_dict["spatial_features"] record_len = data_dict["record_len"] - ups = [] - ret_dict = {} - x = spatial_features + upsampled_features = [] + features = spatial_features + for level, block in enumerate(self.blocks): + features = block(features) + if self.compress and level < len(self.compression_modules): + features = self.compression_modules[level](features) + fused_features = self.fuse_modules[level](features, record_len) + if len(self.deblocks) > 0: + upsampled_features.append(self.deblocks[level](fused_features)) + else: + upsampled_features.append(fused_features) - for i in range(len(self.blocks)): - x = self.blocks[i](x) - if self.compress and i < len(self.compression_modules): - x = self.compression_modules[i](x) - x_fuse = self.fuse_modules[i](x, record_len) + if len(upsampled_features) > 1: + features = torch.cat(upsampled_features, dim=1) + elif len(upsampled_features) == 1: + features = upsampled_features[0] + else: + raise ValueError("At least one backbone level is required") + + if len(self.deblocks) > len(self.blocks): + features = self.deblocks[-1](features) + data_dict["spatial_features_2d"] = features + return data_dict - stride = int(spatial_features.shape[2] / x.shape[2]) - ret_dict["spatial_features_%dx" % stride] = x + def encode_agent(self, spatial_features): + """ + Compute every private backbone scale before feature fusion. + + Parameters + ---------- + spatial_features : torch.Tensor + Unfused BEV features for one or more agents. + + Returns + ------- + tuple[torch.Tensor, ...] + Ordered private feature maps consumed by attention fusion. + """ + feature_maps = [] + features = spatial_features + for level, block in enumerate(self.blocks): + features = block(features) + if self.compress and level < len(self.compression_modules): + encoded_features = self.compression_modules[level].encode(features) + feature_maps.append(encoded_features) + features = self.compression_modules[level].decode(encoded_features) + else: + feature_maps.append(features) + return tuple(feature_maps) + + def fuse_agents(self, feature_maps, record_len): + """ + Fuse encoded feature scales and decode the final BEV representation. + + Parameters + ---------- + feature_maps : tuple[torch.Tensor, ...] + Concatenated per-agent feature maps at every scale. + record_len : torch.Tensor + Number of successfully available agents in each sample. + + Returns + ------- + torch.Tensor + Decoded fused BEV feature map. + """ + if len(feature_maps) != len(self.fuse_modules): + raise ValueError("Unexpected number of intermediate feature scales") + ups = [] + for level, features in enumerate(feature_maps): + if self.compress and level < len(self.compression_modules): + features = self.compression_modules[level].decode(features) + fused_features = self.fuse_modules[level](features, record_len) if len(self.deblocks) > 0: - ups.append(self.deblocks[i](x_fuse)) + ups.append(self.deblocks[level](fused_features)) else: - ups.append(x_fuse) + ups.append(fused_features) if len(ups) > 1: - x = torch.cat(ups, dim=1) + features = torch.cat(ups, dim=1) elif len(ups) == 1: - x = ups[0] + features = ups[0] + else: + raise ValueError("At least one intermediate feature scale is required") if len(self.deblocks) > len(self.blocks): - x = self.deblocks[-1](x) - - data_dict["spatial_features_2d"] = x - return data_dict + features = self.deblocks[-1](features) + return features diff --git a/OpenCOOD/opencood/models/sub_modules/auto_encoder.py b/OpenCOOD/opencood/models/sub_modules/auto_encoder.py index dcf51792..c5719d04 100644 --- a/OpenCOOD/opencood/models/sub_modules/auto_encoder.py +++ b/OpenCOOD/opencood/models/sub_modules/auto_encoder.py @@ -52,10 +52,40 @@ def __init__(self, feature_num, layer_num): feature_num //= 2 def forward(self, x): + return self.decode(self.encode(x)) + + def encode(self, x): + """ + Produce the compact spatial representation for transmission. + + Parameters + ---------- + x : torch.Tensor + Full-resolution feature map. + + Returns + ------- + torch.Tensor + Encoded feature map with reduced channels and resolution. + """ for i in range(len(self.encoder)): x = self.encoder[i](x) + return x + def decode(self, x): + """ + Restore a compact spatial representation for feature fusion. + + Parameters + ---------- + x : torch.Tensor + Encoded feature map received over V2X. + + Returns + ------- + torch.Tensor + Decoded feature map matching the fusion network input. + """ for i in range(len(self.decoder) - 1, -1, -1): x = self.decoder[i](x) - return x diff --git a/OpenCOOD/opencood/models/sub_modules/naive_compress.py b/OpenCOOD/opencood/models/sub_modules/naive_compress.py index f22788bf..8c2cfbd9 100644 --- a/OpenCOOD/opencood/models/sub_modules/naive_compress.py +++ b/OpenCOOD/opencood/models/sub_modules/naive_compress.py @@ -23,7 +23,36 @@ def __init__(self, input_dim, compress_raito): ) def forward(self, x): - x = self.encoder(x) - x = self.decoder(x) + return self.decode(self.encode(x)) - return x + def encode(self, x): + """ + Produce the channel-compressed representation for transmission. + + Parameters + ---------- + x : torch.Tensor + Full-channel spatial feature map. + + Returns + ------- + torch.Tensor + Channel-compressed feature map. + """ + return self.encoder(x) + + def decode(self, x): + """ + Restore a transmitted channel-compressed representation. + + Parameters + ---------- + x : torch.Tensor + Channel-compressed feature map. + + Returns + ------- + torch.Tensor + Feature map restored to the fusion network's channel count. + """ + return self.decoder(x) diff --git a/OpenCOOD/opencood/models/voxel_net.py b/OpenCOOD/opencood/models/voxel_net.py index 2b0b2e08..0d7799f3 100644 --- a/OpenCOOD/opencood/models/voxel_net.py +++ b/OpenCOOD/opencood/models/voxel_net.py @@ -1,10 +1,8 @@ import torch.nn as nn import torch.nn.functional as F import torch -from torch.autograd import Variable from opencood.models.sub_modules.pillar_vfe import PillarVFE -from opencood.utils.common_utils import torch_tensor_to_numpy # conv2d + bn + relu @@ -172,10 +170,35 @@ def __init__(self, args): self.T = args["T"] self.anchor_num = args["anchor_num"] - def voxel_indexing(self, sparse_features, coords): + def voxel_indexing(self, sparse_features, coords, batch_size): + """ + Scatter sparse voxel features into a dense model batch. + + Parameters + ---------- + sparse_features : torch.Tensor + Encoded features for non-empty voxels. + coords : torch.Tensor + Batched voxel coordinates. + batch_size : int + Number of point clouds represented by the coordinates. + + Returns + ------- + torch.Tensor + Dense voxel feature tensor. + """ dim = sparse_features.shape[-1] - - dense_feature = Variable(torch.zeros(dim, self.N, self.D, self.H, self.W).cuda()) + dense_feature = torch.zeros( + dim, + batch_size, + self.D, + self.H, + self.W, + device=sparse_features.device, + dtype=sparse_features.dtype, + ) + coords = coords.long() dense_feature[:, coords[:, 0], coords[:, 1], coords[:, 2], coords[:, 3]] = sparse_features.transpose(0, 1) @@ -185,14 +208,14 @@ def forward(self, data_dict): voxel_features = data_dict["processed_lidar"]["voxel_features"] voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] + batch_size = data_dict.get("batch_size", self.N) batch_dict = {"voxel_features": voxel_features, "voxel_coords": voxel_coords, "voxel_num_points": voxel_num_points} # feature learning network vwfs = self.svfe(batch_dict)["pillar_features"] - voxel_coords = torch_tensor_to_numpy(voxel_coords) - vwfs = self.voxel_indexing(vwfs, voxel_coords) + vwfs = self.voxel_indexing(vwfs, voxel_coords, batch_size) # convolutional middle network vwfs = self.cml(vwfs) @@ -201,7 +224,7 @@ def forward(self, data_dict): # merge the depth and feature dim into one, output probability score # map and regression map - psm, rm = self.rpn(vwfs.view(self.N, -1, self.H, self.W)) + psm, rm = self.rpn(vwfs.view(batch_size, -1, self.H, self.W)) output_dict = {"psm": psm, "rm": rm} diff --git a/OpenCOOD/opencood/models/voxel_net_intermediate.py b/OpenCOOD/opencood/models/voxel_net_intermediate.py index 4ed50a7d..97474b50 100644 --- a/OpenCOOD/opencood/models/voxel_net_intermediate.py +++ b/OpenCOOD/opencood/models/voxel_net_intermediate.py @@ -6,11 +6,10 @@ import torch.nn.functional as F import torch import numpy as np -from torch.autograd import Variable +from opencood.models.communication_adapters.intermediate import SpatialFeatureCommunicationAdapter from opencood.models.voxel_net import RPN, CML from opencood.models.sub_modules.pillar_vfe import PillarVFE -from opencood.utils.common_utils import torch_tensor_to_numpy from opencood.models.fuse_modules.self_attn import AttFusion from opencood.models.sub_modules.auto_encoder import AutoEncoder @@ -50,6 +49,8 @@ def forward(self, x): class VoxelNetIntermediate(nn.Module): + communication_adapter_class = SpatialFeatureCommunicationAdapter + def __init__(self, args): super(VoxelNetIntermediate, self).__init__() self.svfe = PillarVFE(args["pillar_vfe"], num_point_features=4, voxel_size=args["voxel_size"], point_cloud_range=args["lidar_range"]) @@ -69,10 +70,36 @@ def __init__(self, args): self.compression = True self.compression_layer = AutoEncoder(128, args["compression"]) - def voxel_indexing(self, sparse_features, coords): + def voxel_indexing(self, sparse_features, coords, batch_size): + """ + Scatter sparse VoxelNet features into a dense 3D tensor. + + Parameters + ---------- + sparse_features : torch.Tensor + Learned feature vector for every non-empty voxel. + coords : torch.Tensor + Batch, depth, height, and width index for every sparse feature. + batch_size : int + Number of agents represented by the sparse input. + + Returns + ------- + torch.Tensor + Dense tensor shaped ``(batch_size, channels, D, H, W)``. + """ dim = sparse_features.shape[-1] - dense_feature = Variable(torch.zeros(dim, self.N, self.D, self.H, self.W).cuda()) + dense_feature = torch.zeros( + dim, + batch_size, + self.D, + self.H, + self.W, + device=sparse_features.device, + dtype=sparse_features.dtype, + ) + coords = coords.long() dense_feature[:, coords[:, 0], coords[:, 1], coords[:, 2], coords[:, 3]] = sparse_features.transpose(0, 1) @@ -119,6 +146,9 @@ def regroup(self, dense_feature, record_len): return regroup_features def forward(self, data_dict): + if "intermediate_features" in data_dict: + return self.fuse_agents(data_dict) + voxel_features = data_dict["processed_lidar"]["voxel_features"] voxel_coords = data_dict["processed_lidar"]["voxel_coords"] voxel_num_points = data_dict["processed_lidar"]["voxel_num_points"] @@ -135,8 +165,7 @@ def forward(self, data_dict): # feature learning network vwfs = self.svfe(batch_dict)["pillar_features"] - voxel_coords = torch_tensor_to_numpy(voxel_coords) - vwfs = self.voxel_indexing(vwfs, voxel_coords) + vwfs = self.voxel_indexing(vwfs, voxel_coords, self.N) # convolutional middle network vwfs = self.cml(vwfs) @@ -158,3 +187,60 @@ def forward(self, data_dict): output_dict = {"psm": psm, "rm": rm} return output_dict + + def encode_agent(self, data_dict): + """ + Encode one agent through VoxelNet's convolutional middle layers. + + Parameters + ---------- + data_dict : dict + Sender-local preprocessed voxel input. + + Returns + ------- + dict + Learned spatial feature map before attention fusion. + """ + processed_lidar = data_dict["processed_lidar"] + batch_dict = { + "voxel_features": processed_lidar["voxel_features"], + "voxel_coords": processed_lidar["voxel_coords"], + "voxel_num_points": processed_lidar["voxel_num_points"], + } + voxel_features = self.svfe(batch_dict)["pillar_features"] + dense_features = self.voxel_indexing( + voxel_features, + processed_lidar["voxel_coords"], + data_dict.get("batch_size", 1), + ) + spatial_features = self.cml(dense_features).view( + data_dict.get("batch_size", 1), + -1, + self.H, + self.W, + ) + if self.compression: + spatial_features = self.compression_layer.encode(spatial_features) + return {"spatial_features": spatial_features} + + def fuse_agents(self, data_dict): + """ + Fuse received VoxelNet features and run its proposal network. + + Parameters + ---------- + data_dict : dict + Receiver input containing learned features and ``record_len``. + + Returns + ------- + dict + Classification and regression maps. + """ + spatial_features = data_dict["intermediate_features"]["spatial_features"] + if self.compression: + spatial_features = self.compression_layer.decode(spatial_features) + fused_features = self.fusion_net(spatial_features, data_dict["record_len"]) + psm, rm = self.rpn(fused_features) + return {"psm": psm, "rm": rm} diff --git a/OpenCOOD/opencood/tools/inference_utils.py b/OpenCOOD/opencood/tools/inference_utils.py index 86e7dcfc..e8902755 100644 --- a/OpenCOOD/opencood/tools/inference_utils.py +++ b/OpenCOOD/opencood/tools/inference_utils.py @@ -26,7 +26,8 @@ def inference_late_fusion(batch_data, model, dataset): output_dict = OrderedDict() for cav_id, cav_content in batch_data.items(): - output_dict[cav_id] = model(cav_content) + if "processed_lidar" in cav_content: + output_dict[cav_id] = model(cav_content) pred_box_tensor, pred_score, gt_box_tensor = dataset.post_process(batch_data, output_dict) diff --git a/OpenCOOD/test/test_communication_adapters.py b/OpenCOOD/test/test_communication_adapters.py index d07a43f3..0a6969e9 100644 --- a/OpenCOOD/test/test_communication_adapters.py +++ b/OpenCOOD/test/test_communication_adapters.py @@ -1,8 +1,9 @@ """Tests for model-owned communication adapters.""" +from collections import OrderedDict import sys from typing import Any -from unittest.mock import MagicMock +from unittest.mock import MagicMock, Mock import numpy as np import pytest @@ -17,16 +18,36 @@ torch = pytest.importorskip("torch") from torch import nn # noqa: E402 +from opencood.data_utils.datasets import early_fusion_dataset as early_dataset_module # noqa: E402 +from opencood.data_utils.datasets import late_fusion_dataset as late_dataset_module # noqa: E402 +from opencood.data_utils.datasets.early_fusion_dataset import EarlyFusionDataset # noqa: E402 +from opencood.data_utils.datasets.late_fusion_dataset import LateFusionDataset # noqa: E402 from opencood.models.communication_adapters import ( # noqa: E402 + BevInferenceInput, + DetectionInferenceInput, + EarlyFusionCommunicationAdapter, + EarlyFusionWirePayload, FpvrcnnAgentInferenceInput, FpvrcnnWirePayload, + IntermediateFeatureWirePayload, + LateFusionWirePayload, ModelCommunicationAdapter, - PendingModelCommunicationAdapter, + MultiScaleFeatureCommunicationAdapter, + MultiScaleFeatureInput, PoseFrameMetadata, + SpatialFeatureCommunicationAdapter, + SpatialFeatureInput, + Where2CommCommunicationAdapter, + Where2CommWirePayload, build_communication_adapter, ) +from opencood.models.communication_adapters import early as early_module # noqa: E402 from opencood.models.communication_adapters import fpvrcnn as fpvrcnn_module # noqa: E402 from opencood.models.communication_adapters.fpvrcnn import FpvrcnnCommunicationAdapter # noqa: E402 +from opencood.models.communication_adapters.late import LateFusionCommunicationAdapter # noqa: E402 +from opencood.models.sub_modules.auto_encoder import AutoEncoder # noqa: E402 +from opencood.models.sub_modules.naive_compress import NaiveCompressor # noqa: E402 +from opencood.tools.inference_utils import inference_late_fusion # noqa: E402 class DeclaredAdapter(ModelCommunicationAdapter): @@ -67,6 +88,20 @@ def generate_anchor_box(): return np.ones((2, 7), dtype=np.float32) +class FakeLatePostprocessor(FakePostprocessor): + @staticmethod + def decode_agent_predictions(model_input, output): + assert "processed_lidar" in model_input + assert output == {"head": "output"} + return ( + torch.tensor( + [[[1.0, 0.0, 2.0]] * 8], + dtype=torch.float32, + ), + torch.tensor([0.75], dtype=torch.float32), + ) + + class FakeFpvrcnnDataset: def __init__(self) -> None: self.pre_processor = FakePreprocessor() @@ -112,6 +147,104 @@ def encode_agent(self, batch): } +class FakeEarlyFusionDataset: + def __init__(self) -> None: + self.lidar_points = np.array( + [[1.0, 0.0, 2.0, 0.5]], + dtype=np.float32, + ) + self.metadata = PoseFrameMetadata( + lidar_pose=(10.0, 20.0, 0.0, 0.0, 0.0, 0.0), + capture_frame=8, + ) + self.published_payload = None + + def extract_data(self, idx, agent_payload_builder=None) -> None: + assert idx == 8 + assert agent_payload_builder is not None + self.published_payload = agent_payload_builder(self.lidar_points, self.metadata) + + +class FakeLateFusionDataset: + def __init__(self) -> None: + self.pre_processor = FakePreprocessor() + self.post_processor = FakeLatePostprocessor() + self.inference_input = BevInferenceInput(bev_input=np.ones((2, 3, 4), dtype=np.float32)) + self.metadata = PoseFrameMetadata( + lidar_pose=(10.0, 20.0, 0.0, 0.0, 0.0, 0.0), + capture_frame=8, + ) + self.published_payload = None + + def extract_data(self, idx, agent_payload_builder) -> None: + assert idx == 8 + self.published_payload = agent_payload_builder( + self.inference_input, + self.metadata, + ) + + +class FakeLateFusionModel(nn.Module): + def forward(self, model_input): + return {"head": "output"} + + +class FakeIntermediateDataset: + def __init__(self) -> None: + self.pre_processor = FakePreprocessor() + self.inference_input = BevInferenceInput(bev_input=np.ones((2, 3, 4), dtype=np.float32)) + self.metadata = PoseFrameMetadata( + lidar_pose=(10.0, 20.0, 0.0, 0.0, 0.0, 0.0), + capture_frame=8, + ) + self.published_payload = None + + def extract_data(self, idx, agent_payload_builder) -> None: + assert idx == 8 + self.published_payload = agent_payload_builder( + self.inference_input, + self.metadata, + ) + + +class FakeSpatialFeatureModel(nn.Module): + def encode_agent(self, model_input): + assert "processed_lidar" in model_input + return { + "spatial_features": torch.arange(8, dtype=torch.float32).reshape( + 1, + 2, + 2, + 2, + ) + } + + +class FakeMultiScaleFeatureModel(nn.Module): + def encode_agent(self, model_input): + assert "processed_lidar" in model_input + return { + "feature_maps": ( + torch.ones((1, 2, 2, 2), dtype=torch.float32), + torch.full((1, 4, 1, 1), 2.0, dtype=torch.float32), + ) + } + + +class FakeWhere2CommModel(nn.Module): + def encode_agent(self, model_input, *, apply_communication_mask): + assert "processed_lidar" in model_input + value = 0.0 if apply_communication_mask else 1.0 + return { + "spatial_features": torch.full( + (1, 2, 2, 2), + value, + dtype=torch.float32, + ), + "communication_rate": torch.tensor(0.25), + } + + def test_build_communication_adapter_uses_model_declaration(): model = DeclaredModel() @@ -121,10 +254,35 @@ def test_build_communication_adapter_uses_model_declaration(): assert adapter.model is model -def test_build_communication_adapter_uses_pending_fallback(): - adapter = build_communication_adapter(PlainModel(), torch.device("cpu")) +def test_build_communication_adapter_requires_model_declaration(): + with pytest.raises(TypeError, match="must declare communication_adapter_class"): + build_communication_adapter(PlainModel(), torch.device("cpu")) + + +def test_build_communication_adapter_selects_early_fusion_contract(): + model = PlainModel() + + adapter = build_communication_adapter( + model, + torch.device("cpu"), + "EarlyFusionDataset", + ) + + assert isinstance(adapter, EarlyFusionCommunicationAdapter) + assert adapter.model is model + + +def test_build_communication_adapter_selects_late_fusion_contract(): + model = PlainModel() + + adapter = build_communication_adapter( + model, + torch.device("cpu"), + "LateFusionDataset", + ) - assert isinstance(adapter, PendingModelCommunicationAdapter) + assert isinstance(adapter, LateFusionCommunicationAdapter) + assert adapter.model is model def test_build_communication_adapter_rejects_invalid_declaration(): @@ -132,13 +290,521 @@ def test_build_communication_adapter_rejects_invalid_declaration(): build_communication_adapter(InvalidAdapterModel(), torch.device("cpu")) -def test_pending_adapter_delegates_to_dataset_payload_path(): +def test_early_adapter_sends_local_lidar_with_pose_metadata(): + dataset = FakeEarlyFusionDataset() + adapter = EarlyFusionCommunicationAdapter(PlainModel(), torch.device("cpu")) + + adapter.prepare_transmission_payloads(dataset, idx=8) + + payload = dataset.published_payload + assert isinstance(payload, EarlyFusionWirePayload) + assert payload.metadata is dataset.metadata + np.testing.assert_array_equal(payload.lidar_points, dataset.lidar_points) + + +def test_early_dataset_gives_adapter_sender_local_points(monkeypatch): + monkeypatch.setattr( + early_dataset_module.np.random, + "permutation", + lambda size: np.arange(size), + ) + dataset = object.__new__(EarlyFusionDataset) + dataset.module_name = "OpenCOOD.EarlyFusionDataset" + dataset.payload_handler = MagicMock() + dataset.retrieve_base_data = Mock( + return_value={ + "cav-1": { + "lidar_np": np.array( + [ + [0.0, 0.0, 0.0, 1.0], + [4.0, 5.0, 6.0, 0.5], + ], + dtype=np.float32, + ), + "params": {"lidar_pose": [10.0, 20.0, 0.0, 0.0, 0.0, 0.0]}, + "time_delay": 2, + } + } + ) + payload = object() + payload_builder = Mock(return_value=payload) + + dataset.extract_data(10, payload_builder) + + lidar_points, metadata = payload_builder.call_args.args + np.testing.assert_array_equal( + lidar_points, + np.array([[4.0, 5.0, 6.0, 0.5]], dtype=np.float32), + ) + assert metadata == PoseFrameMetadata( + lidar_pose=[10.0, 20.0, 0.0, 0.0, 0.0, 0.0], + capture_frame=8, + ) + dataset.payload_handler.set_opencda_payload.assert_called_once_with( + "cav-1", + dataset.module_name, + payload, + ) + + +def test_early_adapter_projects_received_lidar_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(early_module, "x1_to_x2", lambda _sender, _receiver: transformation) + local_points = np.array([[1.0, 0.0, 2.0, 0.5]], dtype=np.float32) + payload = EarlyFusionWirePayload( + lidar_points=local_points, + metadata=PoseFrameMetadata( + lidar_pose=(1.0, 2.0, 0.0, 0.0, 0.0, 0.0), + capture_frame=5, + ), + ) + adapter = EarlyFusionCommunicationAdapter(PlainModel(), 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["projected_lidar"][0], [10.0, 21.0, 2.0, 0.5]) + np.testing.assert_array_equal(payload.lidar_points, local_points) + + +def test_early_adapter_rejects_payload_from_another_contract(): + adapter = EarlyFusionCommunicationAdapter(PlainModel(), torch.device("cpu")) + + with pytest.raises(TypeError, match="Expected EarlyFusionWirePayload"): + adapter.decode_received_payload( + object(), + receiver_pose=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ) + + +def test_early_dataset_ignores_looped_back_ego_payload(): + dataset = object.__new__(EarlyFusionDataset) + dataset.module_name = "OpenCOOD.EarlyFusionDataset" + dataset.payload_handler = MagicMock() + dataset.payload_handler.current_artery_payload = {"ego": object()} + remote_payload = object() + dataset.payload_handler.get_artery_payload.return_value = remote_payload + dataset.communication_adapter = MagicMock() + remote_lidar = np.full((1, 4), 2.0, dtype=np.float32) + dataset.communication_adapter.decode_received_payload.return_value = { + "projected_lidar": remote_lidar, + } + ego_lidar = np.ones((1, 4), dtype=np.float32) + dataset.get_item_single_car = MagicMock( + return_value={"projected_lidar": ego_lidar}, + ) + base_data = OrderedDict( + { + "ego": {"ego": True}, + "remote": {"ego": False}, + } + ) + ego_pose = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + + result = dataset._EarlyFusionDataset__process_with_messages( + "ego", + ego_pose, + base_data, + ) + + assert len(result["projected_lidar_stack"]) == 2 + np.testing.assert_array_equal(result["projected_lidar_stack"][0], ego_lidar) + np.testing.assert_array_equal(result["projected_lidar_stack"][1], remote_lidar) + dataset.payload_handler.get_artery_payload.assert_called_once_with( + "ego", + "remote", + dataset.module_name, + ) + + +def test_early_supervision_uses_complete_local_scene(): + dataset = object.__new__(EarlyFusionDataset) + dataset.post_processor = MagicMock() + boxes = np.ones((2, 7), dtype=np.float32) + mask = np.ones(2, dtype=np.float32) + object_ids = ["vehicle-1", "vehicle-2"] + dataset.post_processor.generate_object_center.return_value = ( + boxes, + mask, + object_ids, + ) + base_data = OrderedDict( + { + "ego": {"objects": ["vehicle-1"]}, + "remote": {"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_bbx_center"] is boxes + assert supervision["object_bbx_mask"] is mask + assert supervision["object_ids"] is object_ids + + +def test_late_adapter_sends_decoded_detection_candidates(): + dataset = FakeLateFusionDataset() + adapter = LateFusionCommunicationAdapter( + FakeLateFusionModel(), + torch.device("cpu"), + ) + + adapter.prepare_transmission_payloads(dataset, idx=8) + + payload = dataset.published_payload + assert isinstance(payload, LateFusionWirePayload) + assert payload.metadata is dataset.metadata + assert payload.detections.boxes.shape == (1, 8, 3) + np.testing.assert_array_equal( + payload.detections.scores, + np.array([0.75], dtype=np.float32), + ) + + +def test_late_adapter_projects_received_3d_detections(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( + "opencood.models.communication_adapters.late.x1_to_x2", + lambda _sender, _receiver: transformation, + ) + local_boxes = np.array([[[1.0, 0.0, 2.0]] * 8], dtype=np.float32) + payload = LateFusionWirePayload( + detections=DetectionInferenceInput( + boxes=local_boxes, + scores=np.array([0.9], dtype=np.float32), + ), + metadata=PoseFrameMetadata( + lidar_pose=(1.0, 2.0, 0.0, 0.0, 0.0, 0.0), + capture_frame=5, + ), + ) + adapter = LateFusionCommunicationAdapter(PlainModel(), 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["pred_box_tensor"][0, 0], [10.0, 21.0, 2.0]) + np.testing.assert_allclose(decoded["pred_score"], [0.9]) + np.testing.assert_array_equal(payload.detections.boxes, local_boxes) + + +def test_late_adapter_projects_received_2d_detections(monkeypatch): + transformation = np.array( + [ + [1.0, 0.0, 0.0, 10.0], + [0.0, 1.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( + "opencood.models.communication_adapters.late.x1_to_x2", + lambda _sender, _receiver: transformation, + ) + payload = LateFusionWirePayload( + detections=DetectionInferenceInput( + boxes=np.array( + [[[1.0, 2.0], [3.0, 2.0], [3.0, 4.0], [1.0, 4.0]]], + dtype=np.float32, + ), + scores=np.array([0.9], dtype=np.float32), + ), + metadata=PoseFrameMetadata( + lidar_pose=(1.0, 2.0, 0.0, 0.0, 0.0, 0.0), + capture_frame=5, + ), + ) + adapter = LateFusionCommunicationAdapter(PlainModel(), 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["pred_box_tensor"], + np.array( + [[[11.0, 22.0], [13.0, 22.0], [13.0, 24.0], [11.0, 24.0]]], + dtype=np.float32, + ), + ) + + +def test_late_inference_does_not_run_model_for_received_detections(): + local_input = {"processed_lidar": object()} + remote_detections = { + "pred_box_tensor": torch.ones((1, 8, 3)), + "pred_score": torch.tensor([0.8]), + } + batch = {"ego": local_input, "cav-2": remote_detections} + local_output = {"psm": object(), "rm": object()} + model = Mock(return_value=local_output) + expected = (object(), object(), object()) dataset = MagicMock() - adapter = PendingModelCommunicationAdapter(PlainModel(), torch.device("cpu")) + dataset.post_process.return_value = expected + + result = inference_late_fusion(batch, model, dataset) + + assert result == expected + model.assert_called_once_with(local_input) + dataset.post_process.assert_called_once() + post_process_batch, output_dict = dataset.post_process.call_args.args + assert post_process_batch is batch + assert output_dict == {"ego": local_output} + + +def test_late_dataset_ignores_looped_back_ego_payload(): + dataset = object.__new__(LateFusionDataset) + dataset.module_name = "OpenCOOD.LateFusionDataset" + dataset.payload_handler = MagicMock() + dataset.payload_handler.current_artery_payload = {"ego": object()} + remote_payload = object() + dataset.payload_handler.get_artery_payload.return_value = remote_payload + remote_detections = { + "pred_box_tensor": np.ones((1, 8, 3), dtype=np.float32), + "pred_score": np.ones(1, dtype=np.float32), + } + dataset.communication_adapter = MagicMock() + dataset.communication_adapter.decode_received_payload.return_value = remote_detections + local_input = object() + dataset._LateFusionDataset__build_model_input = MagicMock( + return_value={"inference_input": local_input}, + ) + base_data = OrderedDict( + { + "ego": {"ego": True}, + "remote": {"ego": False}, + } + ) + ego_pose = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + + result = dataset._LateFusionDataset__process_with_messages( + "ego", + ego_pose, + base_data, + ) + + assert result["ego"]["inference_input"] is local_input + assert result["remote"] is remote_detections + dataset.payload_handler.get_artery_payload.assert_called_once_with( + "ego", + "remote", + dataset.module_name, + ) + + +def test_late_dataset_merges_local_and_delivered_detections_before_global_nms( + monkeypatch, +): + monkeypatch.setattr( + late_dataset_module.box_utils, + "project_box3d", + lambda boxes, _transformation: boxes, + ) + dataset = object.__new__(LateFusionDataset) + local_boxes = torch.ones((1, 8, 3), dtype=torch.float32) + remote_boxes = torch.full((1, 8, 3), 2.0, dtype=torch.float32) + local_scores = torch.tensor([0.9], dtype=torch.float32) + remote_scores = torch.tensor([0.8], dtype=torch.float32) + expected_boxes = torch.full((1, 8, 3), 3.0, dtype=torch.float32) + expected_scores = torch.tensor([0.7], dtype=torch.float32) + expected_gt = object() + dataset.post_processor = MagicMock() + dataset.post_processor.decode_agent_predictions.return_value = ( + local_boxes, + local_scores, + ) + dataset.post_processor.post_process_detections.return_value = ( + expected_boxes, + expected_scores, + ) + dataset.post_processor.generate_gt_bbx.return_value = expected_gt + batch = OrderedDict( + { + "ego": { + "processed_lidar": object(), + "transformation_matrix": torch.eye(4), + }, + "remote": { + "pred_box_tensor": remote_boxes, + "pred_score": remote_scores, + }, + } + ) + + result = dataset.post_process(batch, {"ego": {"model": "output"}}) + + assert result[0] is expected_boxes + assert result[1] is expected_scores + assert result[2] is expected_gt + merged_boxes, merged_scores = dataset.post_processor.post_process_detections.call_args.args + torch.testing.assert_close( + merged_boxes, + torch.cat([local_boxes, remote_boxes], dim=0), + ) + torch.testing.assert_close( + merged_scores, + torch.cat([local_scores, remote_scores], dim=0), + ) + dataset.post_processor.decode_agent_predictions.assert_called_once_with( + batch["ego"], + {"model": "output"}, + ) + + +def test_late_supervision_and_labels_use_complete_local_scene(): + dataset = object.__new__(LateFusionDataset) + dataset.post_processor = MagicMock() + boxes = np.ones((2, 7), dtype=np.float32) + mask = np.ones(2, dtype=np.float32) + object_ids = ["vehicle-1", "vehicle-2"] + anchors = np.full((2, 7), 2.0, dtype=np.float32) + labels = {"targets": np.full((2, 7), 3.0, dtype=np.float32)} + dataset.post_processor.generate_object_center.return_value = ( + boxes, + mask, + object_ids, + ) + dataset.post_processor.generate_anchor_box.return_value = anchors + dataset.post_processor.generate_label.return_value = labels + base_data = OrderedDict( + { + "ego": {"objects": ["vehicle-1"]}, + "remote": {"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, + ) + dataset.post_processor.generate_label.assert_called_once_with( + gt_box_center=boxes, + anchors=anchors, + mask=mask, + ) + assert supervision["object_ids"] == object_ids + assert supervision["label_dict"] is labels + np.testing.assert_array_equal(supervision["anchor_box"], anchors) + + +def test_spatial_feature_adapter_sends_learned_feature_map(): + dataset = FakeIntermediateDataset() + adapter = SpatialFeatureCommunicationAdapter( + FakeSpatialFeatureModel(), + torch.device("cpu"), + ) + + adapter.prepare_transmission_payloads(dataset, idx=8) + + payload = dataset.published_payload + assert isinstance(payload, IntermediateFeatureWirePayload) + assert isinstance(payload.inference_input, SpatialFeatureInput) + assert payload.metadata is dataset.metadata + np.testing.assert_array_equal( + payload.inference_input.spatial_features, + np.arange(8, dtype=np.float32).reshape(1, 2, 2, 2), + ) + + +def test_multiscale_feature_adapter_preserves_ordered_scales(): + dataset = FakeIntermediateDataset() + adapter = MultiScaleFeatureCommunicationAdapter( + FakeMultiScaleFeatureModel(), + torch.device("cpu"), + ) + + adapter.prepare_transmission_payloads(dataset, idx=8) + payload = dataset.published_payload + decoded = adapter.decode_received_payload( + payload, + receiver_pose=[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ) + + assert isinstance(payload, IntermediateFeatureWirePayload) + assert isinstance(payload.inference_input, MultiScaleFeatureInput) + np.testing.assert_array_equal(decoded["feature_0"], np.ones((1, 2, 2, 2))) + np.testing.assert_array_equal(decoded["feature_1"], np.full((1, 4, 1, 1), 2.0)) + assert decoded["metadata"] is dataset.metadata + + +def test_where2comm_adapter_masks_only_transmitted_features(): + dataset = FakeIntermediateDataset() + adapter = Where2CommCommunicationAdapter( + FakeWhere2CommModel(), + torch.device("cpu"), + ) + + adapter.prepare_transmission_payloads(dataset, idx=8) + payload = dataset.published_payload + local_features = adapter.encode_local_intermediate_input( + dataset, + dataset.inference_input, + ) + + assert isinstance(payload, Where2CommWirePayload) + np.testing.assert_array_equal( + payload.inference_input.spatial_features, + np.zeros((1, 2, 2, 2)), + ) + np.testing.assert_array_equal( + local_features["spatial_features"], + np.ones((1, 2, 2, 2)), + ) + assert payload.inference_input.communication_rate == pytest.approx(0.25) + + +def test_channel_compressor_exposes_smaller_wire_representation(): + compressor = NaiveCompressor(input_dim=8, compress_raito=2).eval() + spatial_features = torch.ones((1, 8, 4, 4), dtype=torch.float32) + + encoded = compressor.encode(spatial_features) + decoded = compressor.decode(encoded) + + assert encoded.shape == (1, 4, 4, 4) + assert decoded.shape == spatial_features.shape + torch.testing.assert_close(compressor(spatial_features), decoded) + + +def test_autoencoder_exposes_smaller_wire_representation(): + compressor = AutoEncoder(feature_num=8, layer_num=1).eval() + spatial_features = torch.ones((1, 8, 4, 4), dtype=torch.float32) - adapter.prepare_transmission_payloads(dataset, idx=4) + encoded = compressor.encode(spatial_features) + decoded = compressor.decode(encoded) - dataset.extract_data.assert_called_once_with(4) + assert encoded.shape == (1, 4, 2, 2) + assert decoded.shape == spatial_features.shape + torch.testing.assert_close(compressor(spatial_features), decoded) def test_fpvrcnn_adapter_builds_payload_from_sender_model_outputs(): diff --git a/OpenCOOD/test/test_communication_contracts.py b/OpenCOOD/test/test_communication_contracts.py index d64b37d8..68718bab 100644 --- a/OpenCOOD/test/test_communication_contracts.py +++ b/OpenCOOD/test/test_communication_contracts.py @@ -18,14 +18,17 @@ from opencood.models.communication_adapters import ( # noqa: E402 BevInferenceInput, + DetectionInferenceInput, EarlyFusionWirePayload, FpvrcnnAgentInferenceInput, FpvrcnnWirePayload, - IntermediateFusionWirePayload, + IntermediateFeatureWirePayload, LateFusionWirePayload, PointCloudInferenceInput, PoseFrameMetadata, + SpatialFeatureInput, VoxelInferenceInput, + Where2CommWirePayload, build_inference_input, inference_input_to_dict, merge_inference_inputs, @@ -115,10 +118,11 @@ def test_wire_payloads_expose_only_network_contract_fields(): "targets", } expected_fields = { - EarlyFusionWirePayload: {"projected_lidar"}, - LateFusionWirePayload: {"inference_input", "metadata"}, - IntermediateFusionWirePayload: {"inference_input", "metadata"}, + EarlyFusionWirePayload: {"lidar_points", "metadata"}, + LateFusionWirePayload: {"detections", "metadata"}, + IntermediateFeatureWirePayload: {"inference_input", "metadata"}, FpvrcnnWirePayload: {"inference_input", "metadata"}, + Where2CommWirePayload: {"inference_input", "metadata"}, } for payload_type, expected in expected_fields.items(): @@ -143,9 +147,23 @@ def test_fpvrcnn_wire_payload_contains_stage_one_and_vsa_outputs(): assert payload.inference_input.point_features.shape == (3, 8) +def test_intermediate_wire_payload_contains_learned_features(): + payload = IntermediateFeatureWirePayload( + inference_input=SpatialFeatureInput( + spatial_features=np.ones((1, 64, 48, 176), dtype=np.float32), + ), + metadata=None, + ) + + assert payload.inference_input.spatial_features.shape == (1, 64, 48, 176) + + 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)), + detections=DetectionInferenceInput( + boxes=np.arange(24, dtype=np.float32).reshape(1, 8, 3), + scores=np.array([0.75], dtype=np.float32), + ), metadata=PoseFrameMetadata( lidar_pose=(1.0, 2.0, 3.0, 0.0, 0.0, 0.0), capture_frame=21, @@ -157,6 +175,6 @@ def test_typed_wire_payload_survives_v2x_pickle_round_trip(): assert isinstance(restored, LateFusionWirePayload) assert restored.metadata == payload.metadata np.testing.assert_array_equal( - restored.inference_input.bev_input, - payload.inference_input.bev_input, + restored.detections.boxes, + payload.detections.boxes, ) diff --git a/OpenCOOD/test/test_intermediate_fusion_contract.py b/OpenCOOD/test/test_intermediate_fusion_contract.py new file mode 100644 index 00000000..1a4e7a33 --- /dev/null +++ b/OpenCOOD/test/test_intermediate_fusion_contract.py @@ -0,0 +1,244 @@ +"""Tests for the learned-feature Intermediate fusion 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 import ( # noqa: E402 + IntermediateFusionDataset, +) +from opencood.models.communication_adapters import ( # noqa: E402 + PoseFrameMetadata, + V2XViTMetadata, +) + + +class RecordingPayloadHandler: + """Collect payloads published by a dataset extraction call.""" + + 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 test_extract_data_requires_adapter_to_build_learned_payloads(): + dataset = IntermediateFusionDataset.__new__(IntermediateFusionDataset) + dataset.cur_ego_pose_flag = True + dataset.module_name = "OpenCOOD.IntermediateFusionDataset" + dataset.model_name = "point_pillar_fcooper" + dataset.payload_handler = RecordingPayloadHandler() + dataset.retrieve_base_data = MagicMock( + return_value=OrderedDict( + { + "cav-1": { + "ego": True, + "params": {"lidar_pose": [1.0, 2.0, 0.0, 0.0, 0.0, 0.0]}, + "time_delay": 0, + }, + "cav-2": { + "ego": False, + "params": {"lidar_pose": [3.0, 4.0, 0.0, 0.0, 0.0, 0.0]}, + "time_delay": 1, + }, + } + ) + ) + inference_inputs = [object(), object()] + dataset.get_item_single_car = MagicMock( + side_effect=[ + {"inference_input": inference_inputs[0]}, + {"inference_input": inference_inputs[1]}, + ] + ) + payload_builder = MagicMock(side_effect=lambda model_input, metadata: (model_input, metadata)) + + dataset.extract_data(7, payload_builder) + + assert payload_builder.call_count == 2 + assert [entry[0] for entry in dataset.payload_handler.published] == [ + "cav-1", + "cav-2", + ] + assert dataset.payload_handler.published[0][2] == (inference_inputs[0], None) + assert dataset.payload_handler.published[1][2] == (inference_inputs[1], None) + + +def test_message_assembly_keeps_only_delivered_learned_features(): + dataset = IntermediateFusionDataset.__new__(IntermediateFusionDataset) + dataset.communication_adapter = MagicMock() + dataset.communication_adapter.encode_local_intermediate_input.return_value = {"spatial_features": np.ones((1, 2, 2, 2), dtype=np.float32)} + dataset.communication_adapter.decode_received_payload.return_value = { + "spatial_features": np.full((1, 2, 2, 2), 2.0, dtype=np.float32), + "metadata": None, + } + dataset.payload_handler = MagicMock() + dataset.payload_handler.current_artery_payload = {"cav-1": object()} + delivered_payload = object() + dataset.payload_handler.get_artery_payload.side_effect = lambda _receiver, sender, _module: delivered_payload if sender == "cav-2" else None + dataset.module_name = "OpenCOOD.IntermediateFusionDataset" + dataset.get_item_single_car = MagicMock(return_value={"inference_input": object()}) + dataset.build_model_metadata = MagicMock(return_value=None) + 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._IntermediateFusionDataset__process_with_messages( + "cav-1", + ego_pose, + base_data, + receive_frame=9, + ) + + assert len(result["intermediate_features"]) == 2 + assert result["metadata"] == [None, None] + dataset.communication_adapter.decode_received_payload.assert_called_once_with( + delivered_payload, + ego_pose, + ) + + +def test_inference_sample_concatenates_agent_feature_batches(): + dataset = IntermediateFusionDataset.__new__(IntermediateFusionDataset) + dataset.model_name = "point_pillar_fcooper" + dataset.visualize = False + local_supervision = { + "object_bbx_center": np.ones((2, 7), dtype=np.float32), + "object_bbx_mask": np.ones(2, dtype=np.float32), + "object_ids": ["vehicle-1"], + "anchor_box": np.ones((2, 7), dtype=np.float32), + "label_dict": {}, + } + + sample = dataset.assemble_inference_sample( + inference_input={ + "intermediate_features": [ + {"spatial_features": np.ones((1, 2, 2, 2), dtype=np.float32)}, + {"spatial_features": np.full((1, 2, 2, 2), 2.0, dtype=np.float32)}, + ], + "metadata": [None, None], + }, + local_supervision=local_supervision, + visualization_context={}, + receive_frame=3, + ) + + ego_sample = sample["ego"] + assert ego_sample["cav_num"] == 2 + assert "processed_lidar" not in ego_sample + np.testing.assert_array_equal( + ego_sample["intermediate_features"]["spatial_features"][:, 0, 0, 0], + np.array([1.0, 2.0], dtype=np.float32), + ) + + +def test_pairwise_transforms_are_derived_only_for_available_agents(monkeypatch): + dataset = IntermediateFusionDataset.__new__(IntermediateFusionDataset) + dataset.model_name = "point_pillar_v2vnet" + dataset.max_cav = 3 + dataset.proj_first = False + dataset.visualize = False + transform = MagicMock( + side_effect=lambda source, target: np.full( + (4, 4), + source[0] * 10 + target[0], + dtype=np.float64, + ) + ) + monkeypatch.setattr( + "opencood.data_utils.datasets.intermediate_fusion_dataset.x1_to_x2", + transform, + ) + metadata = [ + PoseFrameMetadata( + lidar_pose=(1.0, 0.0, 0.0, 0.0, 0.0, 0.0), + capture_frame=4, + ), + PoseFrameMetadata( + lidar_pose=(2.0, 0.0, 0.0, 0.0, 0.0, 0.0), + capture_frame=4, + ), + ] + + sample = dataset.assemble_inference_sample( + inference_input={ + "intermediate_features": [ + {"spatial_features": np.ones((1, 2, 2, 2), dtype=np.float32)}, + {"spatial_features": np.ones((1, 2, 2, 2), dtype=np.float32)}, + ], + "metadata": metadata, + }, + local_supervision={}, + visualization_context={}, + receive_frame=4, + ) + + pairwise = sample["ego"]["pairwise_t_matrix"] + assert pairwise.shape == (3, 3, 4, 4) + assert transform.call_count == 4 + np.testing.assert_array_equal(pairwise[0, 1], np.full((4, 4), 12.0)) + np.testing.assert_array_equal(pairwise[2], np.zeros((3, 4, 4))) + + +def test_v2xvit_metadata_is_padded_after_delivered_agents(): + dataset = IntermediateFusionDataset.__new__(IntermediateFusionDataset) + dataset.model_name = "point_pillar_transformer" + dataset.max_cav = 3 + dataset.visualize = False + first_correction = np.eye(4, dtype=np.float32) + second_correction = np.full((4, 4), 2.0, dtype=np.float32) + metadata = [ + V2XViTMetadata( + velocity=0.5, + agent_type=0, + capture_frame=8, + spatial_correction_matrix=first_correction, + ), + V2XViTMetadata( + velocity=0.25, + agent_type=1, + capture_frame=6, + spatial_correction_matrix=second_correction, + ), + ] + + sample = dataset.assemble_inference_sample( + inference_input={ + "intermediate_features": [ + {"spatial_features": np.ones((1, 2, 2, 2), dtype=np.float32)}, + {"spatial_features": np.ones((1, 2, 2, 2), dtype=np.float32)}, + ], + "metadata": metadata, + }, + local_supervision={}, + visualization_context={}, + receive_frame=10, + ) + + ego_sample = sample["ego"] + assert ego_sample["velocity"] == [0.5, 0.25, 0.0] + assert ego_sample["time_delay"] == [2.0, 4.0, 0.0] + assert ego_sample["infra"] == [0, 1, 0.0] + np.testing.assert_array_equal( + ego_sample["spatial_correction_matrix"], + np.stack([first_correction, second_correction, np.eye(4)]), + ) diff --git a/OpenCOOD/test/test_intermediate_model_split.py b/OpenCOOD/test/test_intermediate_model_split.py new file mode 100644 index 00000000..cf3af533 --- /dev/null +++ b/OpenCOOD/test/test_intermediate_model_split.py @@ -0,0 +1,94 @@ +"""Tests for distributed intermediate-model execution boundaries.""" + +import sys + +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 opencood.models.fuse_modules.where2comm_fuse import Communication # noqa: E402 +from opencood.models.sub_modules.att_bev_backbone import AttBEVBackbone # noqa: E402 + + +@pytest.mark.parametrize("compression", [0, 1]) +def test_multiscale_backbone_split_matches_centralized_path(compression): + model_config = { + "layer_nums": [0, 0], + "layer_strides": [1, 2], + "num_filters": [4, 8], + "upsample_strides": [1, 2], + "num_upsample_filter": [4, 4], + "compression": compression, + } + backbone = AttBEVBackbone(model_config, input_channels=4).eval() + spatial_features = ( + torch.arange( + 2 * 4 * 8 * 8, + dtype=torch.float32, + ).reshape(2, 4, 8, 8) + / 100 + ) + record_len = torch.tensor([2], dtype=torch.long) + + with torch.no_grad(): + centralized = backbone( + { + "spatial_features": spatial_features.clone(), + "record_len": record_len, + } + )["spatial_features_2d"] + encoded_features = backbone.encode_agent(spatial_features.clone()) + distributed = backbone.fuse_agents(encoded_features, record_len) + + torch.testing.assert_close(distributed, centralized) + if compression: + assert encoded_features[0].shape == (2, 2, 4, 4) + assert encoded_features[0].numel() < 2 * 4 * 8 * 8 + + +def test_where2comm_rate_is_measured_before_local_feature_is_forced(): + communication = Communication({"threshold": 0.5}).eval() + confidence_map = torch.tensor( + [ + [[[10.0, -10.0]]], + [[[-10.0, 10.0]]], + ], + dtype=torch.float32, + ) + + transmitted_mask, transmitted_rate = communication.build_mask( + confidence_map, + force_first=False, + ) + receiver_mask, receiver_rate = communication.build_mask( + confidence_map, + force_first=True, + ) + + torch.testing.assert_close( + transmitted_mask, + torch.tensor( + [ + [[[1.0, 0.0]]], + [[[0.0, 1.0]]], + ] + ), + ) + torch.testing.assert_close( + receiver_mask, + torch.tensor( + [ + [[[1.0, 1.0]]], + [[[0.0, 1.0]]], + ] + ), + ) + assert transmitted_rate.item() == pytest.approx(0.5) + assert receiver_rate.item() == pytest.approx(0.5) diff --git a/OpenCOOD/test/test_late_fusion_postprocessors.py b/OpenCOOD/test/test_late_fusion_postprocessors.py new file mode 100644 index 00000000..923bc982 --- /dev/null +++ b/OpenCOOD/test/test_late_fusion_postprocessors.py @@ -0,0 +1,123 @@ +"""Tests for sender-side late-fusion prediction decoding.""" + +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, + ) + +torch = pytest.importorskip("torch") + +from opencood.data_utils.post_processor import ciassd_postprocessor # noqa: E402 +from opencood.data_utils.post_processor import voxel_postprocessor # noqa: E402 +from opencood.data_utils.post_processor.bev_postprocessor import ( # noqa: E402 + BevPostprocessor, +) +from opencood.data_utils.post_processor.ciassd_postprocessor import ( # noqa: E402 + CiassdPostprocessor, +) +from opencood.data_utils.post_processor.voxel_postprocessor import ( # noqa: E402 + VoxelPostprocessor, +) + + +def test_voxel_sender_decodes_candidates_without_nms_or_projection(monkeypatch): + postprocessor = VoxelPostprocessor.__new__(VoxelPostprocessor) + postprocessor.params = { + "target_args": {"score_threshold": 0.5}, + "order": "hwl", + } + decoded_boxes = torch.tensor( + [[[1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.25]]], + dtype=torch.float32, + ) + expected_corners = torch.ones((1, 8, 3), dtype=torch.float32) + postprocessor.delta_to_boxes3d = MagicMock(return_value=decoded_boxes) + corners = MagicMock(return_value=expected_corners) + monkeypatch.setattr(voxel_postprocessor.box_utils, "boxes_to_corners_3d", corners) + cav_content = {"anchor_box": torch.ones((1, 7), dtype=torch.float32)} + output = { + "psm": torch.tensor([[[[10.0]]]], dtype=torch.float32), + "rm": torch.zeros((1, 7, 1, 1), dtype=torch.float32), + } + + boxes, scores = postprocessor.decode_agent_predictions(cav_content, output) + + assert boxes is expected_corners + assert scores.item() == pytest.approx(torch.sigmoid(torch.tensor(10.0)).item()) + postprocessor.delta_to_boxes3d.assert_called_once_with( + output["rm"], + cav_content["anchor_box"], + ) + corners.assert_called_once_with(decoded_boxes[0], order="hwl") + + +def test_bev_sender_decodes_only_above_threshold_candidates(): + postprocessor = BevPostprocessor.__new__(BevPostprocessor) + postprocessor.params = {"target_args": {"score_threshold": 0.5}} + postprocessor.denormalize_reg_map = MagicMock(side_effect=lambda value: value) + expected_corners = torch.ones((1, 4, 2), dtype=torch.float32) + postprocessor.reg_map_to_bbx_corners = MagicMock(return_value=expected_corners) + output = { + "cls": torch.tensor([[[[10.0, -10.0]]]], dtype=torch.float32), + "reg": torch.zeros((1, 6, 1, 2), dtype=torch.float32), + } + + boxes, scores = postprocessor.decode_agent_predictions({}, output) + + assert boxes is expected_corners + assert scores.shape == (1,) + assert scores.item() == pytest.approx(torch.sigmoid(torch.tensor(10.0)).item()) + mask = postprocessor.reg_map_to_bbx_corners.call_args.args[1] + torch.testing.assert_close(mask, torch.tensor([[True, False]])) + + +def test_ciassd_sender_applies_iou_and_direction_before_serialization( + monkeypatch, +): + postprocessor = CiassdPostprocessor.__new__(CiassdPostprocessor) + postprocessor.params = { + "target_args": {"score_threshold": 0.5}, + "order": "hwl", + } + decoded_boxes = torch.tensor( + [ + [ + [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 0.25], + [7.0, 8.0, 9.0, 4.0, 5.0, 6.0, -0.5], + ] + ], + dtype=torch.float32, + ) + expected_corners = torch.ones((1, 8, 3), dtype=torch.float32) + postprocessor.delta_to_boxes3d = MagicMock(return_value=decoded_boxes) + corners = MagicMock(return_value=expected_corners) + monkeypatch.setattr(ciassd_postprocessor.box_utils, "boxes_to_corners_3d", corners) + cav_content = {"anchor_box": torch.ones((1, 1, 2, 7), dtype=torch.float32)} + output = { + "preds_dict_stage1": { + "cls_preds": torch.tensor([[[[10.0]], [[-10.0]]]], dtype=torch.float32), + "box_preds": torch.zeros((1, 14, 1, 1), dtype=torch.float32), + "iou_preds": torch.tensor([[[[1.0]], [[-1.0]]]], dtype=torch.float32), + "dir_cls_preds": torch.tensor( + [[[[1.0]], [[0.0]], [[1.0]], [[0.0]]]], + dtype=torch.float32, + ), + } + } + + boxes, scores = postprocessor.decode_agent_predictions(cav_content, output) + + assert boxes is expected_corners + assert scores.shape == (1,) + assert scores.item() == pytest.approx(torch.sigmoid(torch.tensor(10.0)).item()) + selected_boxes = corners.call_args.args[0] + assert selected_boxes.shape == (1, 7) + assert selected_boxes[0, 6].item() == pytest.approx(0.25 + np.pi) diff --git a/opencda/core/common/coperception_model_manager.py b/opencda/core/common/coperception_model_manager.py index 2f688f4c..33d51d35 100644 --- a/opencda/core/common/coperception_model_manager.py +++ b/opencda/core/common/coperception_model_manager.py @@ -506,7 +506,11 @@ 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.communication_adapter = build_communication_adapter( + self.model, + self.device, + self.hypes.get("fusion", {}).get("core_method"), + ) self.opencood_dataset: DatasetOpenCOOD | None = None self.data_loader: DataLoader[Any] | None = None self.current_memory_data = None diff --git a/opencda/core/common/test/test_coperception_model_manager.py b/opencda/core/common/test/test_coperception_model_manager.py index 95512c15..f95d758c 100644 --- a/opencda/core/common/test/test_coperception_model_manager.py +++ b/opencda/core/common/test/test_coperception_model_manager.py @@ -153,7 +153,11 @@ def test_init_attaches_model_communication_adapter_to_dataset(self, manager_deps manager = CoperceptionModelManager(DummyOpt(), "2023_01_01") - manager_deps["build_communication_adapter"].assert_called_once_with(manager.model, manager.device) + manager_deps["build_communication_adapter"].assert_called_once_with( + manager.model, + manager.device, + "LateFusionDataset", + ) assert manager.communication_adapter is adapter assert dataset.communication_adapter is adapter From 687fa6758256a0fe28a506250d6b68aadca3df66 Mon Sep 17 00:00:00 2001 From: Grigoriy Khoroshilov Date: Fri, 4 Sep 2026 16:11:26 +0300 Subject: [PATCH 2/2] Small fix --- OpenCOOD/opencood/models/communication_adapters/early.py | 3 ++- OpenCOOD/opencood/models/communication_adapters/fpvrcnn.py | 3 ++- .../opencood/models/communication_adapters/intermediate.py | 3 ++- OpenCOOD/opencood/models/communication_adapters/late.py | 3 ++- OpenCOOD/opencood/models/communication_adapters/where2comm.py | 3 ++- OpenCOOD/test/test_communication_adapters.py | 2 +- 6 files changed, 11 insertions(+), 6 deletions(-) diff --git a/OpenCOOD/opencood/models/communication_adapters/early.py b/OpenCOOD/opencood/models/communication_adapters/early.py index c55c9faa..53c28606 100644 --- a/OpenCOOD/opencood/models/communication_adapters/early.py +++ b/OpenCOOD/opencood/models/communication_adapters/early.py @@ -21,7 +21,8 @@ def extract_data( self, idx: int, agent_payload_builder: Callable[[npt.NDArray[Any], PoseFrameMetadata], object], - ) -> None: ... + ) -> None: + raise NotImplementedError class EarlyFusionCommunicationAdapter(ModelCommunicationAdapter): diff --git a/OpenCOOD/opencood/models/communication_adapters/fpvrcnn.py b/OpenCOOD/opencood/models/communication_adapters/fpvrcnn.py index 823d2b7b..d45578ae 100644 --- a/OpenCOOD/opencood/models/communication_adapters/fpvrcnn.py +++ b/OpenCOOD/opencood/models/communication_adapters/fpvrcnn.py @@ -28,7 +28,8 @@ def extract_data( self, idx: int, agent_payload_builder: Callable[[dict[str, Any], PoseFrameMetadata], object] | None, - ) -> None: ... + ) -> None: + raise NotImplementedError class FpvrcnnCommunicationAdapter(ModelCommunicationAdapter): diff --git a/OpenCOOD/opencood/models/communication_adapters/intermediate.py b/OpenCOOD/opencood/models/communication_adapters/intermediate.py index ddbe86a9..e2a6adbb 100644 --- a/OpenCOOD/opencood/models/communication_adapters/intermediate.py +++ b/OpenCOOD/opencood/models/communication_adapters/intermediate.py @@ -30,7 +30,8 @@ def extract_data( [InferenceInput, IntermediateMetadata | None], object, ], - ) -> None: ... + ) -> None: + raise NotImplementedError class IntermediateFeatureCommunicationAdapter(ModelCommunicationAdapter): diff --git a/OpenCOOD/opencood/models/communication_adapters/late.py b/OpenCOOD/opencood/models/communication_adapters/late.py index 6d299094..e065b762 100644 --- a/OpenCOOD/opencood/models/communication_adapters/late.py +++ b/OpenCOOD/opencood/models/communication_adapters/late.py @@ -29,7 +29,8 @@ def extract_data( self, idx: int, agent_payload_builder: Callable[[InferenceInput, PoseFrameMetadata], object], - ) -> None: ... + ) -> None: + raise NotImplementedError class LateFusionCommunicationAdapter(ModelCommunicationAdapter): diff --git a/OpenCOOD/opencood/models/communication_adapters/where2comm.py b/OpenCOOD/opencood/models/communication_adapters/where2comm.py index 81ad0bdd..23c849ad 100644 --- a/OpenCOOD/opencood/models/communication_adapters/where2comm.py +++ b/OpenCOOD/opencood/models/communication_adapters/where2comm.py @@ -29,7 +29,8 @@ def extract_data( [InferenceInput, PoseFrameMetadata], object, ], - ) -> None: ... + ) -> None: + raise NotImplementedError class Where2CommCommunicationAdapter(ModelCommunicationAdapter): diff --git a/OpenCOOD/test/test_communication_adapters.py b/OpenCOOD/test/test_communication_adapters.py index 0a6969e9..1b8c3057 100644 --- a/OpenCOOD/test/test_communication_adapters.py +++ b/OpenCOOD/test/test_communication_adapters.py @@ -306,7 +306,7 @@ def test_early_dataset_gives_adapter_sender_local_points(monkeypatch): monkeypatch.setattr( early_dataset_module.np.random, "permutation", - lambda size: np.arange(size), + np.arange, ) dataset = object.__new__(EarlyFusionDataset) dataset.module_name = "OpenCOOD.EarlyFusionDataset"