[NeurIPS 2023] Robust Asynchronous Collaborative 3D Perception via Bird’s Eye View Flow
# Configuration file for the Sphinx documentation builder.
#
# This file only contains a selection of the most common options. For a full
# list see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Path setup --------------------------------------------------------------

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#
import os
import sys
# sys.path.insert(0, os.path.abspath('.'))
sys.path.insert(0, os.path.abspath('./../'))


# -- Project information -----------------------------------------------------

project = 'OpenCOOD'
copyright = '2021, UCLA Mobility Lab'
author = 'Runsheng Xu, Hao Xiang, Jiaqi Ma'

# The full version, including alpha/beta/rc tags
release = '0.1'


# -- General configuration ---------------------------------------------------

# Add any Sphinx extension module names here, as strings. They can be
# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
# ones.
extensions = [
    'sphinx.ext.autodoc',
    'sphinx.ext.napoleon',
    'sphinx.ext.doctest',
    'sphinx.ext.intersphinx',
    'sphinx.ext.todo',
    'sphinx.ext.coverage',
    'sphinx.ext.mathjax',
    'recommonmark',
]

# Add any paths that contain templates here, relative to this directory.
templates_path = ['_templates']

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
# This pattern also affects html_static_path and html_extra_path.
exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']

todo_include_todos = False

# -- Options for HTML output -------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
#
html_theme = 'sphinx_rtd_theme'

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ['_static']

docs/conf.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

from opencood.data_utils.datasets.late_fusion_dataset import LateFusionDataset
from opencood.data_utils.datasets.early_fusion_dataset import EarlyFusionDataset
from opencood.data_utils.datasets.early_fusion_dataset_v2x import EarlyFusionDatasetV2X
from opencood.data_utils.datasets.early_fusion_dataset_dair import EarlyFusionDatasetDAIR
from opencood.data_utils.datasets.intermediate_fusion_dataset import IntermediateFusionDataset
from opencood.data_utils.datasets.intermediate_fusion_dataset_opv2v_delay import IntermediateFusionDatasetAsync
from opencood.data_utils.datasets.intermediate_fusion_dataset_v2x import IntermediateFusionDatasetV2X
from opencood.data_utils.datasets.intermediate_fusion_dataset_dair import IntermediateFusionDatasetDAIR
from opencood.data_utils.datasets.late_fusion_dataset_v2x import LateFusionDatasetV2X
from opencood.data_utils.datasets.late_fusion_dataset_dair import LateFusionDatasetDAIR
from opencood.data_utils.datasets.intermediate_fusion_dataset_v2 import IntermediateFusionDatasetV2
from opencood.data_utils.datasets.intermediate_fusion_dataset_v2_v2x import IntermediateFusionDatasetV2V2X
from opencood.data_utils.datasets.intermediate_fusion_dataset_v2_dair import IntermediateFusionDatasetV2DAIR
from opencood.data_utils.datasets.intermediate_fusion_dataset_dair_delay import IntermediateFusionDatasetDAIRAsync
from opencood.data_utils.datasets.intermediate_fusion_dataset_dair_irregular import IntermediateFusionDatasetDAIRIrregular
from opencood.data_utils.datasets.intermediate_fusion_dataset_dair_multisweep import IntermediateFusionDatasetDAIRMultisweep
from opencood.data_utils.datasets.intermediate_fusion_dataset_opv2v_multisweep import IntermediateFusionDatasetMultisweep
from opencood.data_utils.datasets.intermediate_fusion_dataset_opv2v_irregular import IntermediateFusionDatasetIrregular
from opencood.data_utils.datasets.intermediate_fusion_dataset_opv2v_irregular_compensation import IntermediateFusionDatasetIrregularCompensation
from opencood.data_utils.datasets.late_fusion_dataset_opv2v_regular import LateFusionDatasetRegular
from opencood.data_utils.datasets.late_fusion_dataset_opv2v_irregular import LateFusionDatasetIrregular
from opencood.data_utils.datasets.late_fusion_dataset_opv2v_irregular_4flow import LateFusionDatasetIrregular4Flow
from opencood.data_utils.datasets.intermediate_fusion_dataset_opv2v_irregular_flow import IntermediateFusionDatasetIrregularFlow
from opencood.data_utils.datasets.intermediate_fusion_dataset_opv2v_irregular_flow_compensation import IntermediateFusionDatasetIrregularFlowCompensation
from opencood.data_utils.datasets.intermediate_fusion_dataset_opv2v_irregular_singleframe import IntermediateFusionDatasetIrregularSingleframe
from opencood.data_utils.datasets.intermediate_fusion_dataset_opv2v_irregular_flow_new import IntermediateFusionDatasetIrregularFlowNew
from opencood.data_utils.datasets.intermediate_fusion_dataset_opv2v_irregular_flow_variance import IntermediateFusionDatasetIrregularFlowVariance
from opencood.data_utils.datasets.intermediate_fusion_dataset_dair_irregular_multi import IntermediateFusionDatasetDAIRIrregularMulti
from opencood.data_utils.datasets.intermediate_fusion_dataset_opv2v_outage import IntermediateFusionDataset_outage
from opencood.data_utils.datasets.intermediate_fusion_dataset_v2xset_irregular_flow_new import IntermediateFusionDatasetV2XSETFlowNew
from opencood.data_utils.datasets.intermediate_fusion_dataset_v2xset_compensation import IntermediateFusionDatasetV2XSETCompensation
from opencood.data_utils.datasets.intermediate_fusion_dataset_v2xset import IntermediateFusionDatasetV2XSET
from opencood.data_utils.datasets.intermediate_fusion_dataset_v2xset_singleframe import IntermediateFusionDatasetV2XSETSingleframe

__all__ = {
    'LateFusionDataset': LateFusionDataset,
    'EarlyFusionDataset': EarlyFusionDataset,
    'EarlyFusionDatasetV2X': EarlyFusionDatasetV2X,
    'EarlyFusionDatasetDAIR': EarlyFusionDatasetDAIR,
    'IntermediateFusionDataset': IntermediateFusionDataset,
    'IntermediateFusionDatasetAsync': IntermediateFusionDatasetAsync, 
    'IntermediateFusionDatasetV2X': IntermediateFusionDatasetV2X,
    'IntermediateFusionDatasetDAIR': IntermediateFusionDatasetDAIR,
    'LateFusionDatasetV2X': LateFusionDatasetV2X,
    'LateFusionDatasetDAIR': LateFusionDatasetDAIR,
    'IntermediateFusionDatasetV2': IntermediateFusionDatasetV2,
    'IntermediateFusionDatasetV2V2X': IntermediateFusionDatasetV2V2X,
    'IntermediateFusionDatasetV2DAIR': IntermediateFusionDatasetV2DAIR,
    'IntermediateFusionDatasetDAIRAsync': IntermediateFusionDatasetDAIRAsync,
    'IntermediateFusionDatasetMultisweep': IntermediateFusionDatasetMultisweep,
    'IntermediateFusionDatasetIrregular': IntermediateFusionDatasetIrregular,
    'IntermediateFusionDatasetIrregularCompensation': IntermediateFusionDatasetIrregularCompensation,
    'LateFusionDatasetRegular': LateFusionDatasetRegular,
    'LateFusionDatasetIrregular': LateFusionDatasetIrregular, 
    'LateFusionDatasetIrregular4Flow': LateFusionDatasetIrregular4Flow,
    'IntermediateFusionDatasetIrregularFlow': IntermediateFusionDatasetIrregularFlow, 
    'IntermediateFusionDatasetIrregularFlowCompensation': IntermediateFusionDatasetIrregularFlowCompensation,
    'IntermediateFusionDatasetIrregularSingleframe': IntermediateFusionDatasetIrregularSingleframe,
    'IntermediateFusionDatasetIrregularFlowNew': IntermediateFusionDatasetIrregularFlowNew,
    'IntermediateFusionDatasetDAIRIrregular': IntermediateFusionDatasetDAIRIrregular,
    'IntermediateFusionDatasetDAIRIrregularMulti': IntermediateFusionDatasetDAIRIrregularMulti,
    'IntermediateFusionDatasetIrregularFlowVariance': IntermediateFusionDatasetIrregularFlowVariance,
    'IntermediateFusionDataset_outage': IntermediateFusionDataset_outage,
    'IntermediateFusionDatasetV2XSETFlowNew': IntermediateFusionDatasetV2XSETFlowNew,
    'IntermediateFusionDatasetV2XSETCompensation': IntermediateFusionDatasetV2XSETCompensation,
    'IntermediateFusionDatasetV2XSET': IntermediateFusionDatasetV2XSET,
    'IntermediateFusionDatasetV2XSETSingleframe': IntermediateFusionDatasetV2XSETSingleframe
}

# the final range for evaluation
GT_RANGE_OPV2V = [-140, -40, -3, 140, 40, 1] #[-140, -40, -3, 140, 40, 1], [-153.6, -51.2, -3, 153.6, 51.2, 1] [-140.8, -38.4, -3, 140.8, 38.4, 1]
GT_RANGE_V2XSIM = [-32, -32, -3, 32, 32, 1]
# The communication range for cavs
COM_RANGE = 70


def build_dataset(dataset_cfg, visualize=False, train=True):
    dataset_name = dataset_cfg['fusion']['core_method']
    error_message = f"{dataset_name} is not found. " \
                    f"Please add your processor file's name in opencood/" \
                    f"data_utils/datasets/init.py"

    dataset = __all__[dataset_name](
        params=dataset_cfg,
        visualize=visualize,
        train=train
    )

    return dataset

opencood/data_utils/datasets/__init__.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Basedataset class for all kinds of fusion.
"""

import os
from collections import OrderedDict
from pyparsing import Or

import json

import torch
import numpy as np
from torch.utils.data import Dataset
import pickle

import opencood.utils.pcd_utils as pcd_utils
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import downsample_lidar_minimum
from opencood.utils.transformation_utils import tfm_to_pose
from opencood.utils.keypoint_utils import bev_sample, get_keypoints


class BaseDataset(Dataset):
    """
    Base dataset for all kinds of fusion. Mainly used to initialize the
    database and associate the __get_item__ index with the correct timestamp
    and scenario.

    Parameters
    __________
    params : dict
        The dictionary contains all parameters for training/testing.

    visualize : false
        If set to true, the raw point cloud will be saved in the memory
        for visualization.

    Attributes
    ----------
    scenario_database : OrderedDict
        A structured dictionary contains all file information.

    len_record : list
        The list to record each scenario's data length. This is used to
        retrieve the correct index during training.

    pre_processor : opencood.pre_processor
        Used to preprocess the raw data.

    post_processor : opencood.post_processor
        Used to generate training labels and convert the model outputs to
        bbx formats.

    data_augmentor : opencood.data_augmentor
        Used to augment data.

    """

    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        # Structure: {scenario_id : {cav_1 : {timestamp1 : {yaml: path,
        # lidar: path, cameras:list of path}}}}
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # at least 1 cav should show up
            cav_list = sorted([x for x in os.listdir(scenario_folder)
                               if os.path.isdir(
                    os.path.join(scenario_folder, x))])
            assert len(cav_list) > 0

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                # save all yaml files to the dictionary
                cav_path = os.path.join(scenario_folder, cav_id)

                # use the frame number as key, the full path as the values
                yaml_files = \
                    sorted([os.path.join(cav_path, x)
                            for x in os.listdir(cav_path) if
                            x.endswith('.yaml')])
                timestamps = self.extract_timestamps(yaml_files)

                for timestamp in timestamps:
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        camera_files

                    if self.select_keypoint:
                        if (self.proj_first and j == 0) or (not self.proj_first):
                            # when init the dataset, it read over all pcd files.
                            # it maybe slow, but no need to perform keypoint sampling for each time.\
                            kp_file_path = f"{self.params['select_kp']['kp_path']}_{self.params['select_kp']['n_samples']}/{scenario_folders_name[i]}/{cav_id}/{timestamp}.npy"
                            kp_dir = kp_file_path.rsplit('/',1)[0]

                            if not os.path.exists(kp_dir):
                                os.makedirs(kp_dir)

                            if os.path.exists(kp_file_path):
                                keypoints = np.load(kp_file_path)
                            else:
                                pcd_np = pcd_utils.pcd_to_np(lidar_file)
                                bev_keypoints = bev_sample(pcd_np,
                                            self.params['select_kp']['lidar_range'],
                                            self.params['select_kp']['n_samples'],
                                            mode='uniform',
                                            all_samples=False)

                                np.save(kp_file_path, bev_keypoints)
                                print(f"saving {kp_file_path} in basedataset")
                                keypoints = bev_keypoints
                            # sampling again from bev keypoints
                            self.scenario_database[i][cav_id][timestamp]['lidar_keypoints_np'] = keypoints
                                    
                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    if not self.len_record:
                        self.len_record.append(len(timestamps))
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + len(timestamps))
                else:
                    self.scenario_database[i][cav_id]['ego'] = False

    def __len__(self):
        return self.len_record[-1]

    def __getitem__(self, idx):
        """
        Abstract method, needs to be define by the children class.
        """
        pass

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
        """
        # we loop the accumulated length list to see get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]

        # check the timestamp index
        timestamp_index = idx if scenario_index == 0 else \
            idx - self.len_record[scenario_index - 1]
        # retrieve the corresponding timestamp key
        timestamp_key = self.return_timestamp_key(scenario_database,
                                                  timestamp_index)
        data = OrderedDict()
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            data[cav_id] = OrderedDict()
            data[cav_id]['ego'] = cav_content['ego']

            # todo: load camera image in the future version
            # load the corresponding data into the dictionary
            
            # data[cav_id]['params'] = \
            #     load_yaml(cav_content[timestamp_key]['yaml'])
            json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['params'] = json.load(f)
            else:
                data[cav_id]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])

            # data[cav_id]['lidar_np'] = \
            #     pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

            npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
            if os.path.exists(npy_file):
                data[cav_id]['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])
            
            if self.select_keypoint and 'lidar_keypoints_np' in cav_content[timestamp_key]:
                data[cav_id]['lidar_keypoints_np'] = cav_content[timestamp_key]['lidar_keypoints_np']

        return data

    @staticmethod
    def extract_timestamps(yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]

            timestamp = res.replace('.yaml', '')
            timestamp = timestamp.replace('.json', '')
            timestamps.append(timestamp)

        return timestamps

    @staticmethod
    def return_timestamp_key(scenario_database, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # get all timestamp keys
        timestamp_keys = list(scenario_database.items())[0][1]
        # retrieve the correct index
        timestamp_key = list(timestamp_keys.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def load_camera_files(cav_path, timestamp):
        """
        Retrieve the paths to all camera files.

        Parameters
        ----------
        cav_path : str
            The full file path of current cav.

        timestamp : str
            Current timestamp

        Returns
        -------
        camera_files : list
            The list containing all camera png file paths.
        """
        camera0_file = os.path.join(cav_path,
                                    timestamp + '_camera0.png')
        camera1_file = os.path.join(cav_path,
                                    timestamp + '_camera1.png')
        camera2_file = os.path.join(cav_path,
                                    timestamp + '_camera2.png')
        camera3_file = os.path.join(cav_path,
                                    timestamp + '_camera3.png')
        return [camera0_file, camera1_file, camera2_file, camera3_file]

    def project_points_to_bev_map(self, points, ratio=0.1):
        """
        Project points to BEV occupancy map with default ratio=0.1.

        Parameters
        ----------
        points : np.ndarray
            (N, 3) / (N, 4)

        ratio : float
            Discretization parameters. Default is 0.1.

        Returns
        -------
        bev_map : np.ndarray
            BEV occupancy map including projected points
            with shape (img_row, img_col).

        """
        return self.pre_processor.project_points_to_bev_map(points, ratio)

    def augment(self, lidar_np, object_bbx_center, object_bbx_mask):
        """
        Given the raw point cloud, augment by flipping and rotation.

        Parameters
        ----------
        lidar_np : np.ndarray
            (n, 4) shape

        object_bbx_center : np.ndarray
            (n, 7) shape to represent bbx's x, y, z, h, w, l, yaw

        object_bbx_mask : np.ndarray
            Indicate which elements in object_bbx_center are padded.
        """
        tmp_dict = {'lidar_np': lidar_np,
                    'object_bbx_center': object_bbx_center,
                    'object_bbx_mask': object_bbx_mask}
        tmp_dict = self.data_augmentor.forward(tmp_dict)

        lidar_np = tmp_dict['lidar_np']
        object_bbx_center = tmp_dict['object_bbx_center']
        object_bbx_mask = tmp_dict['object_bbx_mask']

        return lidar_np, object_bbx_center, object_bbx_mask

    def collate_batch_train(self, batch):
        """
        Customized collate function for pytorch dataloader during training
        for early and late fusion dataset.

        Parameters
        ----------
        batch : dict

        Returns
        -------
        batch : dict
            Reformatted batch.
        """
        # during training, we only care about ego.
        output_dict = {'ego': {}}

        object_bbx_center = []
        object_bbx_mask = []
        processed_lidar_list = []
        label_dict_list = []

        if self.visualize:
            origin_lidar = []

        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            processed_lidar_list.append(ego_dict['processed_lidar'])
            label_dict_list.append(ego_dict['label_dict'])

            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))

        processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(processed_lidar_list)
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)
        output_dict['ego'].update({'object_bbx_center': object_bbx_center,
                                   'object_bbx_mask': object_bbx_mask,
                                   'processed_lidar': processed_lidar_torch_dict,
                                   'anchor_box': torch.from_numpy(ego_dict['anchor_box']),
                                   'label_dict': label_torch_dict})
        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})

        return output_dict

    def visualize_result(self, pred_box_tensor,
                         gt_tensor,
                         pcd,
                         show_vis,
                         save_path,
                         dataset=None):
        # visualize the model output
        self.post_processor.visualize(pred_box_tensor,
                                      gt_tensor,
                                      pcd,
                                      show_vis,
                                      save_path,
                                      dataset=dataset)

    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.
        The object_bbx_center is in ego coordinate.

        Notice: it is a wrap of postprocessor

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """
        return self.post_processor.generate_object_center(cav_contents,
                                                        reference_lidar_pose)



opencood/data_utils/datasets/basedataset.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for early fusion
"""
import random
import math
from collections import OrderedDict

import numpy as np
import torch
from torch.utils.data import DataLoader

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.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2


class EarlyFusionDataset(basedataset.BaseDataset):
    """
    This dataset is used for early fusion, where each CAV transmit the raw
    point cloud to the ego vehicle.
    """
    def __init__(self, params, visualize, train=True):
        super(EarlyFusionDataset, self).__init__(params, visualize, train)
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = build_postprocessor(params['postprocess'], train)

    def __getitem__(self, idx):
        base_data_dict = self.retrieve_base_data(idx)

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        ego_id = -1
        ego_lidar_pose = []

        # first find the ego vehicle's lidar pose
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['params']['lidar_pose']
                break

        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        projected_lidar_stack = []
        object_stack = []
        object_id_stack = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            distance = \
                math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
                           ego_lidar_pose[0]) ** 2 + (
                                  selected_cav_base['params'][
                                      'lidar_pose'][1] - ego_lidar_pose[
                                      1]) ** 2)
            if distance > self.params['comm_range']:
                continue

            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose)
            # all these lidar and object coordinates are projected to ego
            # already.
            projected_lidar_stack.append(
                selected_cav_processed['projected_lidar'])
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

        # exclude all repetitive objects
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        object_stack = np.vstack(object_stack)
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # convert list to numpy array, (N, 4)
        projected_lidar_stack = np.vstack(projected_lidar_stack)

        # data augmentation
        projected_lidar_stack, object_bbx_center, mask = \
            self.augment(projected_lidar_stack, object_bbx_center, mask)

        # we do lidar filtering in the stacked lidar
        projected_lidar_stack = mask_points_by_range(projected_lidar_stack,
                                                     self.params['preprocess'][
                                                         'cav_lidar_range'])
        # augmentation may remove some of the bbx out of range
        object_bbx_center_valid = object_bbx_center[mask == 1]
        object_bbx_center_valid, range_mask = \
            box_utils.mask_boxes_outside_range_numpy(object_bbx_center_valid,
                                                     self.params['preprocess'][
                                                         'cav_lidar_range'],
                                                     self.params['postprocess'][
                                                         'order'],
                                                     return_mask=True
                                                     )
        mask[object_bbx_center_valid.shape[0]:] = 0
        object_bbx_center[:object_bbx_center_valid.shape[0]] = \
            object_bbx_center_valid
        object_bbx_center[object_bbx_center_valid.shape[0]:] = 0
        unique_indices = list(np.array(unique_indices)[range_mask])

        # pre-process the lidar to voxel/bev/downsampled lidar
        lidar_dict = self.pre_processor.preprocess(projected_lidar_stack)

        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        processed_data_dict['ego'].update(
            {'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': lidar_dict,
             'label_dict': label_dict})

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                                                   projected_lidar_stack})

        return processed_data_dict

    def get_item_single_car(self, selected_cav_base, ego_pose):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information.
        ego_pose : list
            The ego vehicle lidar pose under world coordinate.

        Returns
        -------
        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)

        # retrieve objects under ego coordinates
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base],
                                                       ego_pose)

        # filter lidar
        lidar_np = selected_cav_base['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        # remove points that hit itself
        lidar_np = mask_ego_points(lidar_np)
        # project the lidar to ego space
        lidar_np[:, :3] = \
            box_utils.project_points_by_matrix_torch(lidar_np[:, :3],
                                                     transformation_matrix)

        selected_cav_processed.update(
            {'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'projected_lidar': lidar_np})

        return selected_cav_processed

    def collate_batch_test(self, batch):
        """
        Customized collate function for pytorch dataloader during testing
        for late fusion dataset.

        Parameters
        ----------
        batch : dict

        Returns
        -------
        batch : dict
            Reformatted batch.
        """
        # currently, we only support batch size of 1 during testing
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        batch = batch[0] # only ego

        output_dict = {}

        for cav_id, cav_content in batch.items():
            output_dict.update({cav_id: {}})
            # shape: (1, max_num, 7)
            object_bbx_center = \
                torch.from_numpy(np.array([cav_content['object_bbx_center']]))
            object_bbx_mask = \
                torch.from_numpy(np.array([cav_content['object_bbx_mask']]))
            object_ids = cav_content['object_ids']

            # the anchor box is the same for all bounding boxes usually, thus
            # we don't need the batch dimension.
            if cav_content['anchor_box'] is not None:
                output_dict[cav_id].update({'anchor_box':
                    torch.from_numpy(np.array(
                        cav_content[
                            'anchor_box']))})
            if self.visualize:
                origin_lidar = [cav_content['origin_lidar']]

            # processed lidar dictionary
            processed_lidar_torch_dict = \
                self.pre_processor.collate_batch(
                    [cav_content['processed_lidar']])
            # label dictionary
            label_torch_dict = \
                self.post_processor.collate_batch([cav_content['label_dict']])

            # save the transformation matrix (4, 4) to ego vehicle
            transformation_matrix_torch = \
                torch.from_numpy(np.identity(4)).float()
            transformation_matrix_clean_torch = \
                torch.from_numpy(np.identity(4)).float()

            output_dict[cav_id].update({'object_bbx_center': object_bbx_center,
                                        'object_bbx_mask': object_bbx_mask,
                                        'processed_lidar': processed_lidar_torch_dict,
                                        'label_dict': label_torch_dict,
                                        'object_ids': object_ids,
                                        'transformation_matrix': transformation_matrix_torch,
                                        'transformation_matrix_clean': transformation_matrix_clean_torch})

            if self.visualize:
                origin_lidar = \
                    np.array(
                        downsample_lidar_minimum(pcd_np_list=origin_lidar))
                origin_lidar = torch.from_numpy(origin_lidar)
                output_dict[cav_id].update({'origin_lidar': origin_lidar})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

opencood/data_utils/datasets/early_fusion_dataset.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Quanhao Li

"""
Dataset class for DAIR-V2X dataset early fusion
"""
import os
import random
import math
from collections import OrderedDict
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
import numpy as np
import torch
import json
import opencood.data_utils.datasets
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils import box_utils
from opencood.data_utils.post_processor import build_postprocessor
from opencood.data_utils.datasets import early_fusion_dataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2
from opencood.utils.transformation_utils import tfm_to_pose
from opencood.utils.transformation_utils import veh_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import inf_side_rot_and_trans_to_trasnformation_matrix

def load_json(path):
    with open(path, mode="r") as f:
        data = json.load(f)
    return data

class EarlyFusionDatasetDAIR(early_fusion_dataset.EarlyFusionDataset):
    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        self.max_cav = 2
        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if "kd_flag" in params.keys():
            self.kd_flag = params['kd_flag']
        else:
            self.kd_flag = False

        assert 'clip_pc' in params['fusion']['args']
        if params['fusion']['args']['clip_pc']:
            self.clip_pc = True
        else:
            self.clip_pc = False
        
        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None

        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = build_postprocessor(
            params['postprocess'],
            train)

        if self.train:
            split_dir = params['root_dir']
        else:
            split_dir = params['validate_dir']

        self.root_dir = params['data_dir']
        self.split_info = load_json(split_dir)
        co_datainfo = load_json(os.path.join(self.root_dir, 'cooperative/data_info.json'))
        self.co_data = OrderedDict()
        for frame_info in co_datainfo:
            veh_frame_id = frame_info['vehicle_image_path'].split("/")[-1].replace(".jpg", "")
            self.co_data[veh_frame_id] = frame_info

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.
        Parameters
        ----------
        idx : int
            Index given by dataloader.
        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
        """
        veh_frame_id = self.split_info[idx]
        frame_info = self.co_data[veh_frame_id]
        system_error_offset = frame_info["system_error_offset"]
        data = OrderedDict()
        data[0] = OrderedDict() # veh-side
        data[0]['ego'] = True
        data[1] = OrderedDict() # inf-side
        data[1]['ego'] = False
                
        data[0]['params'] = OrderedDict()
        data[0]['params']['vehicles'] = load_json(os.path.join(self.root_dir,frame_info['cooperative_label_path']))
        lidar_to_novatel_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/lidar_to_novatel/'+str(veh_frame_id)+'.json'))
        novatel_to_world_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/novatel_to_world/'+str(veh_frame_id)+'.json'))

        transformation_matrix = veh_side_rot_and_trans_to_trasnformation_matrix(lidar_to_novatel_json_file,novatel_to_world_json_file)

        data[0]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix)

        data[0]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["vehicle_pointcloud_path"]))
        if self.clip_pc:
            data[0]['lidar_np'] = data[0]['lidar_np'][data[0]['lidar_np'][:,0]>0]

        data[1]['params'] = OrderedDict()
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")
        data[1]['params']['vehicles'] = load_json(os.path.join(self.root_dir,frame_info['cooperative_label_path']))
        virtuallidar_to_world_json_file = load_json(os.path.join(self.root_dir,'infrastructure-side/calib/virtuallidar_to_world/'+str(inf_frame_id)+'.json'))
        transformation_matrix1 = inf_side_rot_and_trans_to_trasnformation_matrix(virtuallidar_to_world_json_file,system_error_offset)
        data[1]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix1)

        data[1]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["infrastructure_pointcloud_path"]))
        return data

    def __len__(self):
        return len(self.split_info)

    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_dairv2x(cav_contents,
                                                        reference_lidar_pose)
opencood/data_utils/datasets/early_fusion_dataset_dair.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu

"""
Dataset class for V2X-Sim dataset early fusion
"""
import random
import math
from collections import OrderedDict

import numpy as np
import torch
from torch.utils.data import DataLoader
import pickle

from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
import opencood.data_utils.post_processor as post_processor
from opencood.utils import box_utils
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.datasets import early_fusion_dataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2, x_to_world, tfm_to_pose

from icecream import ic


class EarlyFusionDatasetV2X(early_fusion_dataset.EarlyFusionDataset):
    """
        IntermediateFusionDatasetV2X is children of IntermediateFusionDataset
        It works almost the same as IntermediateFusionDataset,
        But it 
        rewrite __init__, __len__, retrieve_base_data , which is from BaseDataset
        rewrite generate_object_center, which is from IntermediateFusionDataset
    """

    ### rewrite __init__ ###
    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None
            
        # same as IntermediateFusionDataset
        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
            
        if "kd_flag" in params.keys():
            self.kd_flag = params['kd_flag']
        else:
            self.kd_flag = False

        if "box_align" in params.keys():
            self.box_align = True
        else:
            self.box_align = False
        
        print("Dataset dir:", root_dir)

        # v2x: read the pickle file
        with open(root_dir, 'rb') as f:
            dataset_infos = pickle.load(f)  # dataset_infos is a list 

        self.agent_start = eval(min([i[-1] for i in dataset_infos[0].keys() if i.startswith("lidar_pose")]))
        # self.agent_start = 1 when v2x-sim2.0,  = 0 when v2x-sim1.0


        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        self.keyframe_database = OrderedDict()
        self.len_record = len(dataset_infos)

        # loop over all keyframe.
        # data_info is one sample.
        # np.random.seed(303)
        for (i, data_info) in enumerate(dataset_infos):
            self.keyframe_database.update({i: OrderedDict()})

            # at least 1 cav should show up
            cav_num = data_info['agent_num']
            assert cav_num > 0
            select_ego = 1

            # in one keyframe, loop all agent
            for cav_id in range(self.agent_start, self.agent_start + cav_num):

                self.keyframe_database[i][cav_id] = OrderedDict()
                self.keyframe_database[i][cav_id]['lidar'] = data_info[f'lidar_path_{cav_id}']  # maybe add camera in the future
                self.keyframe_database[i][cav_id]['params'] = OrderedDict()
                self.keyframe_database[i][cav_id]['params']['lidar_pose'] = tfm_to_pose(data_info[f"lidar_pose_{cav_id}"]) # tfm in data_info, turn to [x,y,z,roll,yaw,pitch]
                self.keyframe_database[i][cav_id]['params']['vehicles'] = data_info[f'labels_{cav_id}']['gt_boxes_global']
                self.keyframe_database[i][cav_id]['params']['object_ids'] = data_info[f'labels_{cav_id}']['gt_object_ids'].tolist()

                # randomly select ego is better
                if cav_id == select_ego:
                    self.keyframe_database[i][cav_id]['ego'] = True
                    
                else:
                    self.keyframe_database[i][cav_id]['ego'] = False

    ### rewrite __len__ ###
    def __len__(self):
        return self.len_record

    ### rewrite retrieve_base_data ###
    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            lidar_np: (N, 4)
        """
        # we loop the accumulated length list to see get the scenario index
        keyframe = self.keyframe_database[idx]

        data = OrderedDict()
        ego_idx = -1
        # load files for all CAVs
        for cav_id, cav_content in keyframe.items():
            data[cav_id] = OrderedDict()
            data[cav_id]['ego'] = cav_content['ego']
            if cav_content['ego']:
                ego_idx = cav_id
            data[cav_id]['params'] = cav_content['params'] # lidar_pose, vehicles(gt_boxes), object_id(token)

            # load the corresponding data into the dictionary
            nbr_dims = 4 # x,y,z,intensity
            scan = np.fromfile(cav_content['lidar'], dtype='float32')
            points = scan.reshape((-1, 5))[:, :nbr_dims] 
            data[cav_id]['lidar_np'] = points

        # in fact we want to put ego in the first place
        # we will exchange two agent.
        if ego_idx != 0:
            if self.agent_start == 0:  # v2x-sim1.0
                data[ego_idx], data[0] = data[0], data[ego_idx]
            elif self.agent_start == 1:  # v2x-sim2.0
                data[0] = data[ego_idx]
                data.move_to_end(0, last=False) # move to first
                data.pop(ego_idx)

        return data

    ### rewrite generate_object_center ###
    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_v2x(cav_contents,
                                                        reference_lidar_pose)
opencood/data_utils/datasets/early_fusion_dataset_v2x.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Yifan Lu
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for early fusion
"""
from collections import OrderedDict

import numpy as np
import torch
import math
import copy
import time

import opencood.data_utils.post_processor as post_processor
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np


class IntermediateFusionDataset(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):
        super(IntermediateFusionDataset, self). \
            __init__(params, visualize, train)
        
        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        if "kd_flag" in params.keys():
            self.kd_flag = params['kd_flag']
        else:
            self.kd_flag = False
        '''
        if "box_align" in params.keys():
            self.box_align = True
            self.stage1_result_path = params['box_align']['train_result'] if train else params['box_align']['val_result']
            self.stage1_result = read_json(self.stage1_result_path)
            self.box_align_args = params['box_align']['args']
        
        else:
            self.box_align = False
        '''

        print("OPV2V dataset initialized! \
                {} samples totally!".format(self.len_record[-1]))

    def __getitem__(self, idx):

        base_data_dict = self.retrieve_base_data(idx)

        base_data_dict = add_noise_data_dict(base_data_dict,self.params['noise_setting'])

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        ego_id = -1
        ego_lidar_pose = []

        # first find the ego vehicle's lidar pose
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['params']['lidar_pose']
                ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
                break
            
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0


        processed_features = []
        object_stack = []
        object_id_stack = []
        too_far = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        projected_lidar_clean_list = []
        cav_id_list = []

        if self.visualize:
            projected_lidar_stack = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            distance = \
                math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
                           ego_lidar_pose[0]) ** 2 + (
                                  selected_cav_base['params'][
                                      'lidar_pose'][1] - ego_lidar_pose[
                                      1]) ** 2)

            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue


            lidar_pose_clean_list.append(selected_cav_base['params']['lidar_pose_clean'])
            lidar_pose_list.append(selected_cav_base['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)


        ########## Added by Yifan Lu 2022.8.14 ##############
        # box align to correct pose.
        '''
        if self.box_align and str(idx) in self.stage1_result.keys():
            stage1_content = self.stage1_result[str(idx)]
            if stage1_content is not None:
                cav_id_list_stage1 = stage1_content['cav_id_list']
                
                pred_corners_list = stage1_content['pred_corner3d_np_list']
                pred_corners_list = [np.array(corners, dtype=np.float64) for corners in pred_corners_list]
                uncertainty_list = stage1_content['uncertainty_np_list']
                uncertainty_list = [np.array(uncertainty, dtype=np.float64) for uncertain ty in uncertainty_list]
                stage1_lidar_pose_list = [base_data_dict[cav_id]['params']['lidar_pose'] for cav_id in cav_id_list_stage1]
                stage1_lidar_pose = np.array(stage1_lidar_pose_list)


                refined_pose = box_alignment_relative_sample_np(pred_corners_list,
                                                                stage1_lidar_pose, 
                                                                uncertainty_list=uncertainty_list, 
                                                                **self.box_align_args)
                stage1_lidar_pose[:,[0,1,4]] = refined_pose
                stage1_lidar_pose_refined_list = stage1_lidar_pose.tolist() # updated lidar_pose_list
                for cav_id, lidar_pose_refined in zip(cav_id_list_stage1, stage1_lidar_pose_refined_list):
                    if cav_id not in cav_id_list:
                        continue
                    idx_in_list = cav_id_list.index(cav_id)
                    lidar_pose_list[idx_in_list] = lidar_pose_refined
                    base_data_dict[cav_id]['params']['lidar_pose'] = lidar_pose_refined
        '''     
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]


            # obj: 
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose, 
                ego_lidar_pose_clean,
                idx)
                
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            processed_features.append(
                selected_cav_processed['processed_features'])
            if self.kd_flag:
                projected_lidar_clean_list.append(
                    selected_cav_processed['projected_lidar_clean'])

            if self.visualize:
                projected_lidar_stack.append(
                    selected_cav_processed['projected_lidar'])

        ########## Added by Yifan Lu 2022.4.5 ################
        # filter those out of communicate range
        # then we can calculate get_pairwise_transformation
        for cav_id in too_far:
            base_data_dict.pop(cav_id)
        
        pairwise_t_matrix = \
            self.get_pairwise_transformation(base_data_dict,
                                             self.max_cav)

        lidar_poses = np.array(lidar_pose_list).reshape(-1, 6)  # [N_cav, 6]
        lidar_poses_clean = np.array(lidar_pose_clean_list).reshape(-1, 6)  # [N_cav, 6]
        ######################################################

        ############ for disconet ###########
        if self.kd_flag:
            stack_lidar_np = np.vstack(projected_lidar_clean_list)
            stack_lidar_np = mask_points_by_range(stack_lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            stack_feature_processed = self.pre_processor.preprocess(stack_lidar_np)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        object_stack = np.vstack(object_stack)
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(processed_features)

        merged_feature_dict = self.merge_features_to_dict(processed_features)
        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)
        processed_data_dict['ego'].update(
            {'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'lidar_poses_clean': lidar_poses_clean,
             'lidar_poses': lidar_poses})

        if self.kd_flag:
            processed_data_dict['ego'].update({'teacher_processed_lidar':
                stack_feature_processed})

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(
                    projected_lidar_stack)})


        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, ego_pose_clean, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information.
        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.
        ego_pose_clean : list, length 6
            only used for gt box generation
        
        idx: int,
            debug use.

        Returns
        -------
        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) # T_ego_cav
        transformation_matrix_clean = \
            x1_to_x2(selected_cav_base['params']['lidar_pose_clean'],
                     ego_pose_clean)

        # retrieve objects under ego coordinates
        # this is used to generate accurate GT bounding box.
        object_bbx_center, object_bbx_mask, object_ids = self.generate_object_center([selected_cav_base],
                                                    ego_pose_clean)

        # filter lidar
        lidar_np = selected_cav_base['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        # remove points that hit itself
        lidar_np = mask_ego_points(lidar_np)
        # project the lidar to ego space
        # x,y,z in ego space
        projected_lidar = \
            box_utils.project_points_by_matrix_torch(lidar_np[:, :3],
                                                        transformation_matrix)
        if self.kd_flag:
            lidar_np_clean = copy.deepcopy(lidar_np)

        if self.proj_first:
            lidar_np[:, :3] = projected_lidar
            
        lidar_np = mask_points_by_range(lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
        

        processed_lidar = self.pre_processor.preprocess(lidar_np)

        selected_cav_processed.update(
            {'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'projected_lidar': projected_lidar,
             'processed_features': processed_lidar,
             'transformation_matrix': transformation_matrix,
             'transformation_matrix_clean': transformation_matrix_clean})

        if self.kd_flag:
            projected_lidar_clean = \
                box_utils.project_points_by_matrix_torch(lidar_np_clean[:, :3],
                                                    transformation_matrix_clean)
            lidar_np_clean[:, :3] = projected_lidar_clean
            lidar_np_clean = mask_points_by_range(lidar_np_clean,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            selected_cav_processed.update(
                {"projected_lidar_clean": lidar_np_clean}
            )

        return selected_cav_processed

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        if self.kd_flag:
            teacher_processed_lidar_list = []
        if self.visualize:
            origin_lidar = []

        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            lidar_pose_list.append(ego_dict['lidar_poses']) # ego_dict['lidar_pose'] is np.ndarray [N,6]
            lidar_pose_clean_list.append(ego_dict['lidar_poses_clean'])

            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])

            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])

            if self.kd_flag:
                teacher_processed_lidar_list.append(ego_dict['teacher_processed_lidar'])

            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))


        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)

        # [sum(record_len), C, H, W]
        processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(merged_feature_dict)  # coords 增加了 batch_id 维度，[x, y, z] -> [id, x, y, z]
        # [2, 3, 4, ..., M], M <= max_cav
        record_len = torch.from_numpy(np.array(record_len, dtype=int))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        lidar_pose = torch.from_numpy(np.concatenate(lidar_pose_list, axis=0))
        lidar_pose_clean = torch.from_numpy(np.concatenate(lidar_pose_clean_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'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],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'lidar_pose_clean': lidar_pose_clean,
                                   'lidar_pose': lidar_pose})


        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
        
        if self.kd_flag:
            teacher_processed_lidar_torch_dict = \
                self.pre_processor.collate_batch(teacher_processed_lidar_list)
            output_dict['ego'].update({'teacher_processed_lidar':teacher_processed_lidar_torch_dict})

        if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
            (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
            return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

opencood/data_utils/datasets/intermediate_fusion_dataset.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Quanhao Li

"""
Dataset class for early fusion
"""
import random
import math
from collections import OrderedDict

import os
import numpy as np
import torch
from torch.utils.data import DataLoader
import json
import opencood.data_utils.datasets
import opencood.data_utils.post_processor as post_processor
from opencood.utils import box_utils

from opencood.data_utils.datasets import intermediate_fusion_dataset
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.transformation_utils import tfm_to_pose
from opencood.utils.transformation_utils import veh_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import inf_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import x_to_world

def load_json(path):
    with open(path, mode="r") as f:
        data = json.load(f)
    return data

class IntermediateFusionDatasetDAIR(intermediate_fusion_dataset.IntermediateFusionDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        self.max_cav = 2
        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if "kd_flag" in params.keys():
            self.kd_flag = params['kd_flag']
        else:
            self.kd_flag = False

        assert 'clip_pc' in params['fusion']['args']
        if params['fusion']['args']['clip_pc']:
            self.clip_pc = True
        else:
            self.clip_pc = False
        
        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None

        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        if self.train:
            split_dir = params['root_dir']
        else:
            split_dir = params['validate_dir']

        self.root_dir = params['data_dir']
        self.split_info = load_json(split_dir)
        co_datainfo = load_json(os.path.join(self.root_dir, 'cooperative/data_info.json'))
        self.co_data = OrderedDict()
        for frame_info in co_datainfo:
            veh_frame_id = frame_info['vehicle_image_path'].split("/")[-1].replace(".jpg", "")
            self.co_data[veh_frame_id] = frame_info

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.
        Parameters
        ----------
        idx : int
            Index given by dataloader.
        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
        """
        
        veh_frame_id = self.split_info[idx]
        frame_info = self.co_data[veh_frame_id]
        system_error_offset = frame_info["system_error_offset"]
        data = OrderedDict()
        data[0] = OrderedDict() # veh-side
        data[0]['ego'] = True
        data[1] = OrderedDict() # inf-side
        data[1]['ego'] = False

        data[0]['params'] = OrderedDict()
        data[0]['params']['vehicles'] = load_json(os.path.join(self.root_dir,frame_info['cooperative_label_path']))
        # print(data[0]['params']['vehicles'])
        lidar_to_novatel_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/lidar_to_novatel/'+str(veh_frame_id)+'.json'))
        novatel_to_world_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/novatel_to_world/'+str(veh_frame_id)+'.json'))

        transformation_matrix = veh_side_rot_and_trans_to_trasnformation_matrix(lidar_to_novatel_json_file,novatel_to_world_json_file)

        data[0]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix)

        data[0]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["vehicle_pointcloud_path"]))
        if self.clip_pc:
            data[0]['lidar_np'] = data[0]['lidar_np'][data[0]['lidar_np'][:,0]>0]

        data[1]['params'] = OrderedDict()
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")

        data[1]['params']['vehicles'] = [] # we only load cooperative label in vehicle side

        virtuallidar_to_world_json_file = load_json(os.path.join(self.root_dir,'infrastructure-side/calib/virtuallidar_to_world/'+str(inf_frame_id)+'.json'))

        transformation_matrix1 = inf_side_rot_and_trans_to_trasnformation_matrix(virtuallidar_to_world_json_file,system_error_offset)
        data[1]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix1)

        data[1]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["infrastructure_pointcloud_path"]))
        return data

    def __len__(self):
        return len(self.split_info)

    ### rewrite generate_object_center ###
    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_dairv2x(cav_contents,
                                                        reference_lidar_pose)

opencood/data_utils/datasets/intermediate_fusion_dataset_dair.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Modified: Sizhe Wei <sizhewei@sjtu.edu.cn>

"""
Dataset class for intermediate fusion with time delay k
"""
import random
import math
from collections import OrderedDict

import os
import os.path as osp
import numpy as np
import torch
from torch.utils.data import DataLoader
import json
import opencood.data_utils.datasets
import opencood.data_utils.post_processor as post_processor
from opencood.utils import box_utils

from opencood.data_utils.datasets import intermediate_fusion_dataset
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.transformation_utils import tfm_to_pose
from opencood.utils.transformation_utils import veh_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import inf_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import x_to_world
from opencood.utils.pose_utils import add_noise_data_dict

def load_json(path):
    with open(path, mode="r") as f:
        data = json.load(f)
    return data

def build_idx_to_info(data):
    idx2info = {}
    for elem in data:
        if elem["pointcloud_path"] == "":
            continue
        idx = elem["pointcloud_path"].split("/")[-1].replace(".pcd", "")
        idx2info[idx] = elem
    return idx2info

def build_idx_to_co_info(data):
    idx2info = {}
    for elem in data:
        if elem["vehicle_pointcloud_path"] == "":
            continue
        idx = elem["vehicle_pointcloud_path"].split("/")[-1].replace(".pcd", "")
        idx2info[idx] = elem
    return idx2info

def id_to_str(id, digits=6):
    result = ""
    for i in range(digits):
        result = str(id % 10) + result
        id //= 10
    return result

class IntermediateFusionDatasetDAIRAsync(intermediate_fusion_dataset.IntermediateFusionDataset):
    """
    Written by sizhewei @ 2022/09/28
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):
        #注意yaml文件应该有sensor_type：lidar/camera
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        self.max_cav = 2
        self.k = params['time_delay']
        
        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if "kd_flag" in params.keys():
            self.kd_flag = params['kd_flag']
        else:
            self.kd_flag = False

        if "box_align" in params.keys():
            self.box_align = True
            self.stage1_result_path = params['box_align']['train_result'] if train else params['box_align']['val_result']
            self.stage1_result = load_json(self.stage1_result_path)
            self.box_align_args = params['box_align']['args']
        
        else:
            self.box_align = False

        assert 'clip_pc' in params['fusion']['args']
        if params['fusion']['args']['clip_pc']:
            self.clip_pc = True
        else:
            self.clip_pc = False
            
        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None


        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        #这里root_dir是一个json文件！--> 代表一个split
        if self.train:
            split_dir = params['root_dir']
        else:
            split_dir = params['validate_dir']

        self.root_dir = '/remote-home/share/my_dair_v2x/v2x_c/cooperative-vehicle-infrastructure' # TODO:
        self.inf_idx2info = build_idx_to_info(
            load_json(osp.join(self.root_dir, "infrastructure-side/data_info.json"))
        )
        self.co_idx2info = build_idx_to_co_info(
            load_json(osp.join(self.root_dir, "cooperative/data_info.json"))
        )

        self.data_split = load_json(split_dir)
        self.data = []
        for veh_idx in self.data_split:
            if self.is_valid_id(veh_idx):
                self.data.append(veh_idx)

        print("ASync dataset with {} time delay initialized! {} samples totally!".format(self.k, len(self.data)))
    
    def is_valid_id(self, veh_frame_id):
        """
        Written by sizhewei @ 2022/10/05
        Given veh_frame_id, determine whether there is a corresponding inf_frame that meets the k delay requirement.

        Parameters
        ----------
        veh_frame_id : 05d
            Vehicle frame id

        Returns
        -------
        bool valud
            True means there is a corresponding road-side frame.
        """
        # print('veh_frame_id: ',veh_frame_id,'\n')
        frame_info = {}
        
        frame_info = self.co_idx2info[veh_frame_id]
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")
        cur_inf_info = self.inf_idx2info[inf_frame_id]
        if (
            int(inf_frame_id) - self.k < int(cur_inf_info["batch_start_id"])
            or id_to_str(int(inf_frame_id) - self.k) not in self.inf_idx2info
        ):
            return False

        return True

    def retrieve_base_data(self, idx):
        """
        Modified by sizhewei @ 2022/09/28
        Given the index, return the corresponding async data (time delay: self.k).

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
        """
        veh_frame_id = self.data[idx]
        # print('veh_frame_id: ',veh_frame_id,'\n')
        frame_info = {}
        system_error_offset = {}
        
        frame_info = self.co_idx2info[veh_frame_id]
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")

        inf_frame_id = id_to_str(int(inf_frame_id) - self.k)

        system_error_offset = frame_info["system_error_offset"]
        data = OrderedDict()
        #cav_id=0是车端，1是路边单元
        data[0] = OrderedDict()
        data[0]['ego'] = True
        data[0]['params'] = OrderedDict()
        data[0]['params']['vehicles'] = load_json(os.path.join(self.root_dir,frame_info['cooperative_label_path']))
        # print(data[0]['params']['vehicles'])
        lidar_to_novatel_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/lidar_to_novatel/'+str(veh_frame_id)+'.json'))
        novatel_to_world_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/novatel_to_world/'+str(veh_frame_id)+'.json'))

        transformation_matrix = veh_side_rot_and_trans_to_trasnformation_matrix(lidar_to_novatel_json_file,novatel_to_world_json_file)

        data[0]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix)

        data[0]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["vehicle_pointcloud_path"]))
        if self.clip_pc:
            data[0]['lidar_np'] = data[0]['lidar_np'][data[0]['lidar_np'][:,0]>0]
            
        data[1] = OrderedDict()
        data[1]['ego'] = False

        data[1]['params'] = OrderedDict()

        # data[1]['params']['vehicles'] = load_json(os.path.join(self.root_dir,frame_info['cooperative_label_path']))
        data[1]['params']['vehicles'] = [] # we only load cooperative label in vehicle side

        virtuallidar_to_world_json_file = load_json(os.path.join(self.root_dir,'infrastructure-side/calib/virtuallidar_to_world/'+str(inf_frame_id)+'.json'))

        transformation_matrix1 = inf_side_rot_and_trans_to_trasnformation_matrix(virtuallidar_to_world_json_file,system_error_offset)
        data[1]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix1)

        cur_inf_info = self.inf_idx2info[inf_frame_id]
        data[1]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir, \
            'infrastructure-side', cur_inf_info['pointcloud_path']))
        data[0]['veh_frame_id'] = veh_frame_id
        data[1]['veh_frame_id'] = inf_frame_id
        return data

    def __getitem__(self, idx):
        # base_data_dict = self.retrieve_base_data(idx)
        base_data_dict = self.retrieve_base_data(idx)

        base_data_dict = add_noise_data_dict(base_data_dict,self.params['noise_setting'])

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        ego_id = -1
        ego_lidar_pose = []

        # first find the ego vehicle's lidar pose
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['params']['lidar_pose']
                ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
                break
            
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0


        processed_features = []
        object_stack = []
        object_id_stack = []
        too_far = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        projected_lidar_clean_list = []
        cav_id_list = []

        if self.visualize:
            projected_lidar_stack = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            distance = \
                math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
                           ego_lidar_pose[0]) ** 2 + (
                                  selected_cav_base['params'][
                                      'lidar_pose'][1] - ego_lidar_pose[
                                      1]) ** 2)

            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue


            lidar_pose_clean_list.append(selected_cav_base['params']['lidar_pose_clean'])
            lidar_pose_list.append(selected_cav_base['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)


        ########## Added by Yifan Lu 2022.8.14 ##############
        # box align to correct pose.
        '''
        if self.box_align and str(idx) in self.stage1_result.keys():
            stage1_content = self.stage1_result[str(idx)]
            if stage1_content is not None:
                cav_id_list_stage1 = stage1_content['cav_id_list']
                
                pred_corners_list = stage1_content['pred_corner3d_np_list']
                pred_corners_list = [np.array(corners, dtype=np.float64) for corners in pred_corners_list]
                uncertainty_list = stage1_content['uncertainty_np_list']
                uncertainty_list = [np.array(uncertainty, dtype=np.float64) for uncertainty in uncertainty_list]
                stage1_lidar_pose_list = [base_data_dict[cav_id]['params']['lidar_pose'] for cav_id in cav_id_list_stage1]
                stage1_lidar_pose = np.array(stage1_lidar_pose_list)


                refined_pose = box_alignment_relative_sample_np(pred_corners_list,
                                                                stage1_lidar_pose, 
                                                                uncertainty_list=uncertainty_list, 
                                                                **self.box_align_args)
                stage1_lidar_pose[:,[0,1,4]] = refined_pose
                stage1_lidar_pose_refined_list = stage1_lidar_pose.tolist() # updated lidar_pose_list
                for cav_id, lidar_pose_refined in zip(cav_id_list_stage1, stage1_lidar_pose_refined_list):
                    if cav_id not in cav_id_list:
                        continue
                    idx_in_list = cav_id_list.index(cav_id)
                    lidar_pose_list[idx_in_list] = lidar_pose_refined
                    base_data_dict[cav_id]['params']['lidar_pose'] = lidar_pose_refined
        '''     


        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]



            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose, 
                ego_lidar_pose_clean,
                idx)
                
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            processed_features.append(
                selected_cav_processed['processed_features'])
            if self.kd_flag:
                projected_lidar_clean_list.append(
                    selected_cav_processed['projected_lidar_clean'])

            if self.visualize:
                projected_lidar_stack.append(
                    selected_cav_processed['projected_lidar'])





        ########## Added by Yifan Lu 2022.4.5 ################
        # filter those out of communicate range
        # then we can calculate get_pairwise_transformation
        for cav_id in too_far:
            base_data_dict.pop(cav_id)
        
        pairwise_t_matrix = \
            self.get_pairwise_transformation(base_data_dict,
                                             self.max_cav)

        lidar_poses = np.array(lidar_pose_list).reshape(-1, 6)  # [N_cav, 6]
        lidar_poses_clean = np.array(lidar_pose_clean_list).reshape(-1, 6)  # [N_cav, 6]
        ######################################################

        ############ for disconet ###########
        if self.kd_flag:
            stack_lidar_np = np.vstack(projected_lidar_clean_list)
            stack_lidar_np = mask_points_by_range(stack_lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            stack_feature_processed = self.pre_processor.preprocess(stack_lidar_np)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        object_stack = np.vstack(object_stack)
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(processed_features)

        merged_feature_dict = self.merge_features_to_dict(processed_features)

        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        processed_data_dict['ego'].update(
            {'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'lidar_poses_clean': lidar_poses_clean,
             'lidar_poses': lidar_poses})

        if self.kd_flag:
            processed_data_dict['ego'].update({'teacher_processed_lidar':
                stack_feature_processed})

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(
                    projected_lidar_stack)})


        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})

        # processed_data_dict['ego'].update({'veh_frame_id': base_data_dict[0]['veh_frame_id']})

        return processed_data_dict

    def __len__(self):
        # 符合条件的 frame 的数量
        return len(self.data)

    ### rewrite generate_object_center ###
    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_dairv2x(cav_contents,
                                                        reference_lidar_pose)

opencood/data_utils/datasets/intermediate_fusion_dataset_dair_delay.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Modified: Sizhe Wei <sizhewei@sjtu.edu.cn>

"""
Dataset class for intermediate fusion with time delay k
"""
import random
import math
from collections import OrderedDict

import os
import os.path as osp
import numpy as np
import torch
from torch.utils.data import DataLoader
import json
import opencood.data_utils.datasets
import opencood.data_utils.post_processor as post_processor
from opencood.utils import box_utils
from scipy import stats

from opencood.data_utils.datasets import intermediate_fusion_dataset
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.transformation_utils import tfm_to_pose
from opencood.utils.transformation_utils import veh_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import inf_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import x_to_world
from opencood.utils.pose_utils import add_noise_data_dict

def load_json(path):
    with open(path, mode="r") as f:
        data = json.load(f)
    return data

def build_idx_to_info(data):
    idx2info = {}
    for elem in data:
        if elem["pointcloud_path"] == "":
            continue
        idx = elem["pointcloud_path"].split("/")[-1].replace(".pcd", "")
        idx2info[idx] = elem
    return idx2info

def build_idx_to_co_info(data):
    idx2info = {}
    for elem in data:
        if elem["vehicle_pointcloud_path"] == "":
            continue
        idx = elem["vehicle_pointcloud_path"].split("/")[-1].replace(".pcd", "")
        idx2info[idx] = elem
    return idx2info

def build_inf_fid_to_veh_fid(data):
    inf_fid2veh_fid = {}
    for elem in data:
        veh_fid = elem["vehicle_pointcloud_path"].split("/")[-1].rstrip('.pcd')
        inf_fid = elem["infrastructure_pointcloud_path"].split("/")[-1].rstrip('.pcd')
        inf_fid2veh_fid[inf_fid] = veh_fid
    return inf_fid2veh_fid

def id_to_str(id, digits=6):
    result = ""
    for i in range(digits):
        result = str(id % 10) + result
        id //= 10
    return result

class IntermediateFusionDatasetDAIRIrregular(intermediate_fusion_dataset.IntermediateFusionDataset):
    """
    Written by sizhewei @ 2022/09/28
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):
        #注意yaml文件应该有sensor_type：lidar/camera
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        self.max_cav = 2
        
        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 1

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 10

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 0
        
        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if "kd_flag" in params.keys():
            self.kd_flag = params['kd_flag']
        else:
            self.kd_flag = False

        if "box_align" in params.keys():
            self.box_align = True
            self.stage1_result_path = params['box_align']['train_result'] if train else params['box_align']['val_result']
            self.stage1_result = load_json(self.stage1_result_path)
            self.box_align_args = params['box_align']['args']
        
        else:
            self.box_align = False

        assert 'clip_pc' in params['fusion']['args']
        if params['fusion']['args']['clip_pc']:
            self.clip_pc = True
        else:
            self.clip_pc = False
            
        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None


        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)
        self.anchor_box = self.post_processor.generate_anchor_box()

        #这里root_dir是一个json文件！--> 代表一个split
        if self.train:
            split_dir = params['root_dir']
        else:
            split_dir = params['validate_dir']

        self.root_dir = '/remote-home/share/my_dair_v2x/v2x_c/cooperative-vehicle-infrastructure' # TODO:
        self.inf_idx2info = build_idx_to_info(
            load_json(osp.join(self.root_dir, "infrastructure-side/data_info.json"))
        )
        self.co_idx2info = build_idx_to_co_info(
            load_json(osp.join(self.root_dir, "cooperative/data_info.json"))
        )
        self.inf_fid2veh_fid = build_inf_fid_to_veh_fid(load_json(osp.join(self.root_dir, "cooperative/data_info.json"))
        )

        self.data_split = load_json(split_dir)
        self.data = []
        for veh_idx in self.data_split:
            if self.is_valid_id(veh_idx):
                self.data.append(veh_idx)

        print("Irregular async dataset with past %d frames and expectation time delay = %d initialized! %d samples totally!" % (self.k, int(self.binomial_n*self.binomial_p), len(self.data)))

    def is_valid_id(self, veh_frame_id):
        """
        Written by sizhewei @ 2022/10/05
        Given veh_frame_id, determine whether there is a corresponding inf_frame that meets the k delay requirement.

        Parameters
        ----------
        veh_frame_id : 05d
            Vehicle frame id

        Returns
        -------
        bool valud
            True means there is a corresponding road-side frame.
        """
        # print('veh_frame_id: ',veh_frame_id,'\n')
        frame_info = {}
        
        frame_info = self.co_idx2info[veh_frame_id]
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")
        cur_inf_info = self.inf_idx2info[inf_frame_id]
        if (int(inf_frame_id) - self.binomial_n*self.k < int(cur_inf_info["batch_start_id"])):
            return False
        for i in range(3*self.binomial_n * self.k):# TODO: delete 3
            delay_id = id_to_str(int(inf_frame_id) - i) 
            if delay_id not in self.inf_fid2veh_fid.keys():
                return False

        return True

    def retrieve_base_data(self, idx):
        """
        Modified by sizhewei @ 2022/09/28
        Given the index, return the corresponding async data (time delay expection = sum(B(n, p)) ).

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            {
                [0]  / [1]: {
                    'ego' : True / False,
                    'params' : {
                        'vehicles':
                        'lidar_pose':
                    },
                    'lidar_np' : 
                    'veh_frame_id' :
                    'avg_time_delay' : ([1] only)
                }
            }
        """
        veh_frame_id = self.data[idx]
        # print('veh_frame_id: ',veh_frame_id,'\n')
        frame_info = {}
        system_error_offset = {}
        
        frame_info = self.co_idx2info[veh_frame_id]
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")

        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 
        # B(n, p)
        trails = bernoulliDist.rvs(self.binomial_n)
        sample_interval = sum(trails)
        
        curr_inf_frame_id = inf_frame_id
        inf_frame_id = id_to_str(int(inf_frame_id) - sample_interval)

        system_error_offset = frame_info["system_error_offset"]
        data = OrderedDict()
        #cav_id=0是车端，1是路边单元
        data[0] = OrderedDict()
        data[0]['ego'] = True
        data[0]['params'] = OrderedDict()
        data[0]['params']['vehicles'] = load_json(os.path.join(self.root_dir,frame_info['cooperative_label_path']))
        # print(data[0]['params']['vehicles'])
        lidar_to_novatel_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/lidar_to_novatel/'+str(veh_frame_id)+'.json'))
        novatel_to_world_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/novatel_to_world/'+str(veh_frame_id)+'.json'))

        transformation_matrix = veh_side_rot_and_trans_to_trasnformation_matrix(lidar_to_novatel_json_file,novatel_to_world_json_file)

        data[0]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix)

        data[0]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["vehicle_pointcloud_path"]))
        if self.clip_pc:
            data[0]['lidar_np'] = data[0]['lidar_np'][data[0]['lidar_np'][:,0]>0]

        data[0]['lidar_np_curr'] = data[0]['lidar_np'] # vehicle side lidar is no delay
            
        data[1] = OrderedDict()
        data[1]['ego'] = False

        data[1]['params'] = OrderedDict()

        # data[1]['params']['vehicles'] = load_json(os.path.join(self.root_dir,frame_info['cooperative_label_path']))
        data[1]['params']['vehicles'] = [] # we only load cooperative label in vehicle side

        virtuallidar_to_world_json_file = load_json(os.path.join(self.root_dir,'infrastructure-side/calib/virtuallidar_to_world/'+str(inf_frame_id)+'.json'))

        transformation_matrix1 = inf_side_rot_and_trans_to_trasnformation_matrix(virtuallidar_to_world_json_file,system_error_offset)
        data[1]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix1)


        # single label 
        data[0]['params']['vehicles_single'] = load_json(os.path.join(self.root_dir, \
                                'vehicle-side/label/lidar/{}.json'.format(veh_frame_id)))
        data[1]['params']['vehicles_single'] = load_json(os.path.join(self.root_dir, \
                                'infrastructure-side/label/virtuallidar/{}.json'.format(inf_frame_id)))

        cur_inf_info = self.inf_idx2info[inf_frame_id]
        data[1]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir, \
            'infrastructure-side', cur_inf_info['pointcloud_path']))

        curr_inf_info = self.inf_idx2info[curr_inf_frame_id]
        data[1]['lidar_np_curr'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir, \
            'infrastructure-side', curr_inf_info['pointcloud_path']))

        data[0]['veh_frame_id'] = veh_frame_id
        data[1]['veh_frame_id'] = inf_frame_id
        data[1]['avg_time_delay'] = sample_interval
        return data


    def get_item_single_car(self, selected_cav_base, ego_pose, ego_pose_clean, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information.
        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.
        ego_pose_clean : list, length 6
            only used for gt box generation
        
        idx: int,
            debug use.

        Returns
        -------
        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) # T_ego_cav
        transformation_matrix_clean = \
            x1_to_x2(selected_cav_base['params']['lidar_pose_clean'],
                     ego_pose_clean)

        # retrieve objects under ego coordinates
        # this is used to generate accurate GT bounding box.
        object_bbx_center, object_bbx_mask, object_ids = self.generate_object_center([selected_cav_base],
                                                    ego_pose_clean)

        # filter lidar
        lidar_np = selected_cav_base['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        # remove points that hit itself
        lidar_np = mask_ego_points(lidar_np)

        # current lidar np, w.o. delay
        # filter lidar
        curr_lidar_np = selected_cav_base['lidar_np_curr']
        curr_lidar_np = shuffle_points(curr_lidar_np)
        # remove points that hit itself
        curr_lidar_np = mask_ego_points(curr_lidar_np)

        # project the lidar to ego space
        # x,y,z in ego space
        projected_lidar = \
            box_utils.project_points_by_matrix_torch(curr_lidar_np[:, :3],
                                                        transformation_matrix)
        if self.kd_flag:
            lidar_np_clean = copy.deepcopy(lidar_np)

        if self.proj_first:
            lidar_np[:, :3] = projected_lidar
            
        lidar_np = mask_points_by_range(lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
        

        processed_lidar = self.pre_processor.preprocess(lidar_np)

        selected_cav_processed.update(
            {'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'projected_lidar': projected_lidar,
             'processed_features': processed_lidar,
             'transformation_matrix': transformation_matrix,
             'transformation_matrix_clean': transformation_matrix_clean})

        if self.kd_flag:
            projected_lidar_clean = \
                box_utils.project_points_by_matrix_torch(lidar_np_clean[:, :3],
                                                    transformation_matrix_clean)
            lidar_np_clean[:, :3] = projected_lidar_clean
            lidar_np_clean = mask_points_by_range(lidar_np_clean,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            selected_cav_processed.update(
                {"projected_lidar_clean": lidar_np_clean}
            )

        # generate targets label single GT, note the reference pose is itself.
        object_bbx_center, object_bbx_mask, object_ids = self.generate_object_center_single(
            [selected_cav_base], selected_cav_base['params']['lidar_pose']
        )
        
        label_dict = self.post_processor.generate_label(
            gt_box_center=object_bbx_center, anchors=self.anchor_box, mask=object_bbx_mask
        )
        selected_cav_processed.update({
                            "single_label_dict": label_dict,
                            "single_object_bbx_center": object_bbx_center,
                            "single_object_bbx_mask": object_bbx_mask})

        return selected_cav_processed

    def __getitem__(self, idx):
        # base_data_dict = self.retrieve_base_data(idx)
        base_data_dict = self.retrieve_base_data(idx)

        base_data_dict = add_noise_data_dict(base_data_dict,self.params['noise_setting'])

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        ego_id = -1
        ego_lidar_pose = []

        # first find the ego vehicle's lidar pose
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['params']['lidar_pose']
                ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
                break
            
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0


        processed_features = []
        object_stack = []
        object_id_stack = []
        too_far = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        projected_lidar_clean_list = []
        cav_id_list = []
        single_label_list = []
        single_object_bbx_center_list = []
        single_object_bbx_mask_list = []

        if self.visualize:
            projected_lidar_stack = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            distance = \
                math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
                           ego_lidar_pose[0]) ** 2 + (
                                  selected_cav_base['params'][
                                      'lidar_pose'][1] - ego_lidar_pose[
                                      1]) ** 2)

            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue


            lidar_pose_clean_list.append(selected_cav_base['params']['lidar_pose_clean'])
            lidar_pose_list.append(selected_cav_base['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)


        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]

            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose, 
                ego_lidar_pose_clean,
                idx)
                
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            processed_features.append(
                selected_cav_processed['processed_features'])
            if self.kd_flag:
                projected_lidar_clean_list.append(
                    selected_cav_processed['projected_lidar_clean'])

            if self.visualize:
                projected_lidar_stack.append(
                    selected_cav_processed['projected_lidar'])

            single_label_list.append(selected_cav_processed['single_label_dict'])
            single_object_bbx_center_list.append(selected_cav_processed['single_object_bbx_center'])
            single_object_bbx_mask_list.append(selected_cav_processed['single_object_bbx_mask'])

        ########## Added by Yifan Lu 2022.4.5 ################
        # filter those out of communicate range
        # then we can calculate get_pairwise_transformation
        for cav_id in too_far:
            base_data_dict.pop(cav_id)
        
        pairwise_t_matrix = \
            self.get_pairwise_transformation(base_data_dict,
                                             self.max_cav)

        lidar_poses = np.array(lidar_pose_list).reshape(-1, 6)  # [N_cav, 6]
        lidar_poses_clean = np.array(lidar_pose_clean_list).reshape(-1, 6)  # [N_cav, 6]
        ######################################################

        # single label
        single_label_dicts = self.post_processor.collate_batch(single_label_list)
        single_object_bbx_center = torch.from_numpy(np.array(single_object_bbx_center_list))
        single_object_bbx_mask = torch.from_numpy(np.array(single_object_bbx_mask_list))
        processed_data_dict['ego'].update({
            "single_label_dict_torch": single_label_dicts,
            "single_object_bbx_center_torch": single_object_bbx_center,
            "single_object_bbx_mask_torch": single_object_bbx_mask,
            })


        ############ for disconet ###########
        if self.kd_flag:
            stack_lidar_np = np.vstack(projected_lidar_clean_list)
            stack_lidar_np = mask_points_by_range(stack_lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            stack_feature_processed = self.pre_processor.preprocess(stack_lidar_np)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        object_stack = np.vstack(object_stack)
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(processed_features)

        merged_feature_dict = self.merge_features_to_dict(processed_features)

        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        processed_data_dict['ego'].update(
            {'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'lidar_poses_clean': lidar_poses_clean,
             'lidar_poses': lidar_poses})

        # for irregular time delay use:
        if len(too_far) == 0:
            avg_time_delay = base_data_dict[1]['avg_time_delay']
            cp_rate = 1
        else:
            avg_time_delay = self.binomial_n
            cp_rate = 0
        processed_data_dict['ego'].update(
            {'avg_time_delay': avg_time_delay,
                'cp_rate': cp_rate})

        if self.kd_flag:
            processed_data_dict['ego'].update({'teacher_processed_lidar':
                stack_feature_processed})

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(
                    projected_lidar_stack)})


        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})

        # processed_data_dict['ego'].update({'veh_frame_id': base_data_dict[0]['veh_frame_id']})

        return processed_data_dict

    def __len__(self):
        # 符合条件的 frame 的数量
        return len(self.data)

    ### rewrite generate_object_center ###
    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_dairv2x(cav_contents,
                                                        reference_lidar_pose)

    ### Add new func for single side
    def generate_object_center_single(self,
                               cav_contents,
                               reference_lidar_pose,
                               **kwargs):
        """
        veh or inf 's coordinate
        """
        suffix = "_single"

        return self.post_processor.generate_object_center_dairv2x_single(cav_contents, suffix)

    def collate_batch_train(self, batch):
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []
        object_bbx_center_single = []
        object_bbx_mask_single = []


        # average time delay
        avg_time_delay = []
        # cp rate
        cp_rate = []
        
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        if self.kd_flag:
            teacher_processed_lidar_list = []
        if self.visualize:
            origin_lidar = []

        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            lidar_pose_list.append(ego_dict['lidar_poses']) # ego_dict['lidar_pose'] is np.ndarray [N,6]
            lidar_pose_clean_list.append(ego_dict['lidar_poses_clean'])

            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])

            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])

            avg_time_delay.append(ego_dict['avg_time_delay'])
            cp_rate.append(ego_dict['cp_rate'])

            if self.kd_flag:
                teacher_processed_lidar_list.append(ego_dict['teacher_processed_lidar'])

            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])

            pos_equal_one_single.append(ego_dict['single_label_dict_torch']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_label_dict_torch']['neg_equal_one'])
            targets_single.append(ego_dict['single_label_dict_torch']['targets'])
            object_bbx_center_single.append(ego_dict['single_object_bbx_center_torch'])
            object_bbx_mask_single.append(ego_dict['single_object_bbx_mask_torch'])

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))


        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)

        # [sum(record_len), C, H, W]
        processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(merged_feature_dict)  # coords 增加了 batch_id 维度，[x, y, z] -> [id, x, y, z]
        # [2, 3, 4, ..., M], M <= max_cav
        record_len = torch.from_numpy(np.array(record_len, dtype=int))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        lidar_pose = torch.from_numpy(np.concatenate(lidar_pose_list, axis=0))
        lidar_pose_clean = torch.from_numpy(np.concatenate(lidar_pose_clean_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'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],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'lidar_pose_clean': lidar_pose_clean,
                                   'lidar_pose': lidar_pose})

        avg_time_delay = float(np.mean(np.array(avg_time_delay)))
        cp_rate = float(np.mean(np.array(cp_rate)))
        output_dict['ego'].update({'avg_time_delay': avg_time_delay,
                                      'cp_rate': cp_rate})  

        output_dict['ego'].update({
            "single_object_label":{ # the same name as opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_irregular_flow_new.py 
                    "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                    "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                    "targets": torch.cat(targets_single, dim=0),
                    # for centerpoint
                    "object_bbx_center_single": torch.cat(object_bbx_center_single, dim=0),
                    "object_bbx_mask_single": torch.cat(object_bbx_mask_single, dim=0)
                },
            "object_bbx_center_single": torch.cat(object_bbx_center_single, dim=0),
            "object_bbx_mask_single": torch.cat(object_bbx_mask_single, dim=0)
        })


        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
        
        if self.kd_flag:
            teacher_processed_lidar_torch_dict = \
                self.pre_processor.collate_batch(teacher_processed_lidar_list)
            output_dict['ego'].update({'teacher_processed_lidar':teacher_processed_lidar_torch_dict})

        if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
            (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
            return None

        return output_dict

opencood/data_utils/datasets/intermediate_fusion_dataset_dair_irregular.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Modified: Sizhe Wei <sizhewei@sjtu.edu.cn>

"""
Dataset class for intermediate fusion with time delay k
"""
import random
import math
from collections import OrderedDict

import os
import os.path as osp
import numpy as np
import torch
from torch.utils.data import DataLoader
import json
import opencood.data_utils.datasets
import opencood.data_utils.post_processor as post_processor
from opencood.utils import box_utils
from scipy import stats

from opencood.data_utils.datasets import intermediate_fusion_dataset
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.transformation_utils import tfm_to_pose
from opencood.utils.transformation_utils import veh_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import inf_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import x_to_world
from opencood.utils.pose_utils import add_noise_data_dict

def load_json(path):
    with open(path, mode="r") as f:
        data = json.load(f)
    return data

def build_idx_to_info(data):
    idx2info = {}
    for elem in data:
        if elem["pointcloud_path"] == "":
            continue
        idx = elem["pointcloud_path"].split("/")[-1].replace(".pcd", "")
        idx2info[idx] = elem
    return idx2info

def build_idx_to_co_info(data):
    idx2info = {}
    for elem in data:
        if elem["vehicle_pointcloud_path"] == "":
            continue
        idx = elem["vehicle_pointcloud_path"].split("/")[-1].replace(".pcd", "")
        idx2info[idx] = elem
    return idx2info

def build_inf_fid_to_veh_fid(data):
    inf_fid2veh_fid = {}
    for elem in data:
        veh_fid = elem["vehicle_pointcloud_path"].split("/")[-1].rstrip('.pcd')
        inf_fid = elem["infrastructure_pointcloud_path"].split("/")[-1].rstrip('.pcd')
        inf_fid2veh_fid[inf_fid] = veh_fid
    return inf_fid2veh_fid

def id_to_str(id, digits=6):
    result = ""
    for i in range(digits):
        result = str(id % 10) + result
        id //= 10
    return result

class IntermediateFusionDatasetDAIRIrregular(intermediate_fusion_dataset.IntermediateFusionDataset):
    """
    Written by sizhewei @ 2022/09/28
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):
        #注意yaml文件应该有sensor_type：lidar/camera
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        self.max_cav = 2
        
        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 1

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 10

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 0
        
        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if "kd_flag" in params.keys():
            self.kd_flag = params['kd_flag']
        else:
            self.kd_flag = False

        if "box_align" in params.keys():
            self.box_align = True
            self.stage1_result_path = params['box_align']['train_result'] if train else params['box_align']['val_result']
            self.stage1_result = load_json(self.stage1_result_path)
            self.box_align_args = params['box_align']['args']
        
        else:
            self.box_align = False

        assert 'clip_pc' in params['fusion']['args']
        if params['fusion']['args']['clip_pc']:
            self.clip_pc = True
        else:
            self.clip_pc = False
            
        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None


        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        #这里root_dir是一个json文件！--> 代表一个split
        if self.train:
            split_dir = params['root_dir']
        else:
            split_dir = params['validate_dir']

        self.root_dir = '/remote-home/share/my_dair_v2x/v2x_c/cooperative-vehicle-infrastructure' # TODO:
        self.inf_idx2info = build_idx_to_info(
            load_json(osp.join(self.root_dir, "infrastructure-side/data_info.json"))
        )
        self.co_idx2info = build_idx_to_co_info(
            load_json(osp.join(self.root_dir, "cooperative/data_info.json"))
        )
        self.inf_fid2veh_fid = build_inf_fid_to_veh_fid(load_json(osp.join(self.root_dir, "cooperative/data_info.json"))
        )

        self.data_split = load_json(split_dir)
        self.data = []
        for veh_idx in self.data_split:
            if self.is_valid_id(veh_idx):
                self.data.append(veh_idx)

        print("Irregular async dataset with past %d frames and expectation time delay = %d initialized! %d samples totally!" % (self.k, int(self.binomial_n*self.binomial_p), len(self.data)))

    def is_valid_id(self, veh_frame_id):
        """
        Written by sizhewei @ 2022/10/05
        Given veh_frame_id, determine whether there is a corresponding inf_frame that meets the k delay requirement.

        Parameters
        ----------
        veh_frame_id : 05d
            Vehicle frame id

        Returns
        -------
        bool valud
            True means there is a corresponding road-side frame.
        """
        # print('veh_frame_id: ',veh_frame_id,'\n')
        frame_info = {}
        
        frame_info = self.co_idx2info[veh_frame_id]
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")
        cur_inf_info = self.inf_idx2info[inf_frame_id]
        if (int(inf_frame_id) - self.binomial_n*self.k < int(cur_inf_info["batch_start_id"])):
            return False
        for i in range(3*self.binomial_n * self.k):# TODO: delete 3
            delay_id = id_to_str(int(inf_frame_id) - i) 
            if delay_id not in self.inf_fid2veh_fid.keys():
                return False

        return True

    def retrieve_base_data(self, idx):
        """
        Modified by sizhewei @ 2022/09/28
        Given the index, return the corresponding async data (time delay expection = sum(B(n, p)) ).

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            {
                [0]  / [1]: {
                    'ego' : True / False,
                    'params' : {
                        'vehicles':
                        'lidar_pose':
                    },
                    'lidar_np' : 
                    'veh_frame_id' :
                    'avg_time_delay' : ([1] only)
                }
            }
        """
        veh_frame_id = self.data[idx]
        # print('veh_frame_id: ',veh_frame_id,'\n')
        frame_info = {}
        system_error_offset = {}
        
        frame_info = self.co_idx2info[veh_frame_id]
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")

        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 
        # B(n, p)
        trails = bernoulliDist.rvs(self.binomial_n)
        sample_interval = sum(trails)

        inf_frame_id = id_to_str(int(inf_frame_id) - sample_interval)

        system_error_offset = frame_info["system_error_offset"]
        data = OrderedDict()
        #cav_id=0是车端，1是路边单元
        data[0] = OrderedDict()
        data[0]['ego'] = True
        data[0]['params'] = OrderedDict()
        data[0]['params']['vehicles'] = load_json(os.path.join(self.root_dir,frame_info['cooperative_label_path']))
        # print(data[0]['params']['vehicles'])
        lidar_to_novatel_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/lidar_to_novatel/'+str(veh_frame_id)+'.json'))
        novatel_to_world_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/novatel_to_world/'+str(veh_frame_id)+'.json'))

        transformation_matrix = veh_side_rot_and_trans_to_trasnformation_matrix(lidar_to_novatel_json_file,novatel_to_world_json_file)

        data[0]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix)

        data[0]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["vehicle_pointcloud_path"]))
        if self.clip_pc:
            data[0]['lidar_np'] = data[0]['lidar_np'][data[0]['lidar_np'][:,0]>0]
            
        data[1] = OrderedDict()
        data[1]['ego'] = False

        data[1]['params'] = OrderedDict()

        # data[1]['params']['vehicles'] = load_json(os.path.join(self.root_dir,frame_info['cooperative_label_path']))
        data[1]['params']['vehicles'] = [] # we only load cooperative label in vehicle side

        virtuallidar_to_world_json_file = load_json(os.path.join(self.root_dir,'infrastructure-side/calib/virtuallidar_to_world/'+str(inf_frame_id)+'.json'))

        transformation_matrix1 = inf_side_rot_and_trans_to_trasnformation_matrix(virtuallidar_to_world_json_file,system_error_offset)
        data[1]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix1)

        cur_inf_info = self.inf_idx2info[inf_frame_id]
        data[1]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir, \
            'infrastructure-side', cur_inf_info['pointcloud_path']))
        data[0]['veh_frame_id'] = veh_frame_id
        data[1]['veh_frame_id'] = inf_frame_id
        data[1]['avg_time_delay'] = sample_interval
        return data

    def __getitem__(self, idx):
        # base_data_dict = self.retrieve_base_data(idx)
        base_data_dict = self.retrieve_base_data(idx)

        base_data_dict = add_noise_data_dict(base_data_dict,self.params['noise_setting'])

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        ego_id = -1
        ego_lidar_pose = []

        # first find the ego vehicle's lidar pose
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['params']['lidar_pose']
                ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
                break
            
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0


        processed_features = []
        object_stack = []
        object_id_stack = []
        too_far = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        projected_lidar_clean_list = []
        cav_id_list = []

        if self.visualize:
            projected_lidar_stack = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            distance = \
                math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
                           ego_lidar_pose[0]) ** 2 + (
                                  selected_cav_base['params'][
                                      'lidar_pose'][1] - ego_lidar_pose[
                                      1]) ** 2)

            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue


            lidar_pose_clean_list.append(selected_cav_base['params']['lidar_pose_clean'])
            lidar_pose_list.append(selected_cav_base['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)


        ########## Added by Yifan Lu 2022.8.14 ##############
        # box align to correct pose.
        '''
        if self.box_align and str(idx) in self.stage1_result.keys():
            stage1_content = self.stage1_result[str(idx)]
            if stage1_content is not None:
                cav_id_list_stage1 = stage1_content['cav_id_list']
                
                pred_corners_list = stage1_content['pred_corner3d_np_list']
                pred_corners_list = [np.array(corners, dtype=np.float64) for corners in pred_corners_list]
                uncertainty_list = stage1_content['uncertainty_np_list']
                uncertainty_list = [np.array(uncertainty, dtype=np.float64) for uncertainty in uncertainty_list]
                stage1_lidar_pose_list = [base_data_dict[cav_id]['params']['lidar_pose'] for cav_id in cav_id_list_stage1]
                stage1_lidar_pose = np.array(stage1_lidar_pose_list)


                refined_pose = box_alignment_relative_sample_np(pred_corners_list,
                                                                stage1_lidar_pose, 
                                                                uncertainty_list=uncertainty_list, 
                                                                **self.box_align_args)
                stage1_lidar_pose[:,[0,1,4]] = refined_pose
                stage1_lidar_pose_refined_list = stage1_lidar_pose.tolist() # updated lidar_pose_list
                for cav_id, lidar_pose_refined in zip(cav_id_list_stage1, stage1_lidar_pose_refined_list):
                    if cav_id not in cav_id_list:
                        continue
                    idx_in_list = cav_id_list.index(cav_id)
                    lidar_pose_list[idx_in_list] = lidar_pose_refined
                    base_data_dict[cav_id]['params']['lidar_pose'] = lidar_pose_refined
        '''     


        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]



            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose, 
                ego_lidar_pose_clean,
                idx)
                
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            processed_features.append(
                selected_cav_processed['processed_features'])
            if self.kd_flag:
                projected_lidar_clean_list.append(
                    selected_cav_processed['projected_lidar_clean'])

            if self.visualize:
                projected_lidar_stack.append(
                    selected_cav_processed['projected_lidar'])





        ########## Added by Yifan Lu 2022.4.5 ################
        # filter those out of communicate range
        # then we can calculate get_pairwise_transformation
        for cav_id in too_far:
            base_data_dict.pop(cav_id)
        
        pairwise_t_matrix = \
            self.get_pairwise_transformation(base_data_dict,
                                             self.max_cav)

        lidar_poses = np.array(lidar_pose_list).reshape(-1, 6)  # [N_cav, 6]
        lidar_poses_clean = np.array(lidar_pose_clean_list).reshape(-1, 6)  # [N_cav, 6]
        ######################################################

        ############ for disconet ###########
        if self.kd_flag:
            stack_lidar_np = np.vstack(projected_lidar_clean_list)
            stack_lidar_np = mask_points_by_range(stack_lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            stack_feature_processed = self.pre_processor.preprocess(stack_lidar_np)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        object_stack = np.vstack(object_stack)
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(processed_features)

        merged_feature_dict = self.merge_features_to_dict(processed_features)

        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        processed_data_dict['ego'].update(
            {'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'lidar_poses_clean': lidar_poses_clean,
             'lidar_poses': lidar_poses})

        # for irregular time delay use:
        if len(too_far) == 0:
            avg_time_delay = base_data_dict[1]['avg_time_delay']
            cp_rate = 1
        else:
            avg_time_delay = self.binomial_n
            cp_rate = 0
        processed_data_dict['ego'].update(
            {'avg_time_delay': avg_time_delay,
                'cp_rate': cp_rate})

        if self.kd_flag:
            processed_data_dict['ego'].update({'teacher_processed_lidar':
                stack_feature_processed})

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(
                    projected_lidar_stack)})


        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})

        # processed_data_dict['ego'].update({'veh_frame_id': base_data_dict[0]['veh_frame_id']})

        return processed_data_dict

    def __len__(self):
        # 符合条件的 frame 的数量
        return len(self.data)

    ### rewrite generate_object_center ###
    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_dairv2x(cav_contents,
                                                        reference_lidar_pose)

    def collate_batch_train(self, batch):
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        lidar_pose_list = []
        lidar_pose_clean_list = []

        # average time delay
        avg_time_delay = []
        # cp rate
        cp_rate = []
        
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        if self.kd_flag:
            teacher_processed_lidar_list = []
        if self.visualize:
            origin_lidar = []

        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            lidar_pose_list.append(ego_dict['lidar_poses']) # ego_dict['lidar_pose'] is np.ndarray [N,6]
            lidar_pose_clean_list.append(ego_dict['lidar_poses_clean'])

            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])

            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])

            avg_time_delay.append(ego_dict['avg_time_delay'])
            cp_rate.append(ego_dict['cp_rate'])

            if self.kd_flag:
                teacher_processed_lidar_list.append(ego_dict['teacher_processed_lidar'])

            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))


        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)

        # [sum(record_len), C, H, W]
        processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(merged_feature_dict)  # coords 增加了 batch_id 维度，[x, y, z] -> [id, x, y, z]
        # [2, 3, 4, ..., M], M <= max_cav
        record_len = torch.from_numpy(np.array(record_len, dtype=int))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        lidar_pose = torch.from_numpy(np.concatenate(lidar_pose_list, axis=0))
        lidar_pose_clean = torch.from_numpy(np.concatenate(lidar_pose_clean_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'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],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'lidar_pose_clean': lidar_pose_clean,
                                   'lidar_pose': lidar_pose})

        avg_time_delay = float(np.mean(np.array(avg_time_delay)))
        cp_rate = float(np.mean(np.array(cp_rate)))
        output_dict['ego'].update({'avg_time_delay': avg_time_delay,
                                      'cp_rate': cp_rate})  


        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
        
        if self.kd_flag:
            teacher_processed_lidar_torch_dict = \
                self.pre_processor.collate_batch(teacher_processed_lidar_list)
            output_dict['ego'].update({'teacher_processed_lidar':teacher_processed_lidar_torch_dict})

        if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
            (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
            return None

        return output_dict

opencood/data_utils/datasets/intermediate_fusion_dataset_dair_irregular_backup.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Modified: Sizhe Wei <sizhewei@sjtu.edu.cn>

"""
Dataset class for intermediate fusion with time delay k
"""
import random
import math
from collections import OrderedDict

import os
import os.path as osp
import numpy as np
import torch
from torch.utils.data import DataLoader
import json
import opencood.data_utils.datasets
import opencood.data_utils.post_processor as post_processor
from opencood.utils import box_utils
from scipy import stats

from opencood.data_utils.datasets import intermediate_fusion_dataset_opv2v_irregular_flow_new
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.transformation_utils import tfm_to_pose
from opencood.utils.transformation_utils import veh_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import inf_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import x_to_world
from opencood.utils.pose_utils import add_noise_data_dict
from opencood.utils.flow_utils import generate_flow_map, generate_flow_map_szwei



def load_json(path):
    with open(path, mode="r") as f:
        data = json.load(f)
    return data

def build_idx_to_info(data):
    idx2info = {}
    for elem in data:
        if elem["pointcloud_path"] == "":
            continue
        idx = elem["pointcloud_path"].split("/")[-1].replace(".pcd", "")
        idx2info[idx] = elem
    return idx2info

def build_idx_to_co_info(data):
    idx2info = {}
    for elem in data:
        if elem["vehicle_pointcloud_path"] == "":
            continue
        idx = elem["vehicle_pointcloud_path"].split("/")[-1].replace(".pcd", "")
        idx2info[idx] = elem
    return idx2info

def build_inf_fid_to_veh_fid(data):
    inf_fid2veh_fid = {}
    for elem in data:
        veh_fid = elem["vehicle_pointcloud_path"].split("/")[-1].rstrip('.pcd')
        inf_fid = elem["infrastructure_pointcloud_path"].split("/")[-1].rstrip('.pcd')
        inf_fid2veh_fid[inf_fid] = veh_fid
    return inf_fid2veh_fid

def id_to_str(id, digits=6):
    result = ""
    for i in range(digits):
        result = str(id % 10) + result
        id //= 10
    return result

class IntermediateFusionDatasetDAIRIrregularMulti(intermediate_fusion_dataset_opv2v_irregular_flow_new.IntermediateFusionDatasetIrregularFlowNew):
    """
    Written by sizhewei @ 2022/09/28
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):
        #注意yaml文件应该有sensor_type：lidar/camera
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        self.max_cav = 2
        
        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 1

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 10

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 0

        self.num_roi_thres = -1
        if 'num_roi_thres' in params:
            self.num_roi_thres = params['num_roi_thres']

        # 控制是否需要生成GT flow
        self.is_generate_gt_flow = False
        if 'is_generate_gt_flow' in params and params['is_generate_gt_flow']:
            self.is_generate_gt_flow = True

        self.viz_bbx_flag = False
        
        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False


        assert 'clip_pc' in params['fusion']['args']
        if params['fusion']['args']['clip_pc']:
            self.clip_pc = True
        else:
            self.clip_pc = False
            

        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        #这里root_dir是一个json文件！--> 代表一个split
        if self.train:
            split_dir = params['root_dir']
        else:
            split_dir = params['validate_dir']

        self.root_dir = params['data_dir']

        self.inf_idx2info = build_idx_to_info(
            load_json(osp.join(self.root_dir, "infrastructure-side/data_info.json"))
        )
        self.co_idx2info = build_idx_to_co_info(
            load_json(osp.join(self.root_dir, "cooperative/data_info.json"))
        )
        self.inf_fid2veh_fid = build_inf_fid_to_veh_fid(load_json(osp.join(self.root_dir, "cooperative/data_info.json"))
        )

        self.data_split = load_json(split_dir)
        self.data = []
        for veh_idx in self.data_split:
            if self.is_valid_id(veh_idx):
                self.data.append(veh_idx)
        
        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead. 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("Irregular async dataset with past %d frames and expectation time delay = %d initialized! %d samples totally!" % (self.k, int(self.binomial_n*self.binomial_p), len(self.data)))

    def get_vehicle_trans(self, veh_frame_id):

        lidar_to_novatel_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/lidar_to_novatel/'+str(veh_frame_id)+'.json'))
        novatel_to_world_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/novatel_to_world/'+str(veh_frame_id)+'.json'))

        transformation_matrix = veh_side_rot_and_trans_to_trasnformation_matrix(lidar_to_novatel_json_file,novatel_to_world_json_file)
        trans = tfm_to_pose(transformation_matrix)

        return trans
    
    def get_inf_trans(self, inf_frame_id, system_error_offset):
        virtuallidar_to_world_json_file = load_json(os.path.join(self.root_dir,'infrastructure-side/calib/virtuallidar_to_world/'+str(inf_frame_id)+'.json'))
        transformation_matrix1 = inf_side_rot_and_trans_to_trasnformation_matrix(virtuallidar_to_world_json_file,system_error_offset)
        trans = tfm_to_pose(transformation_matrix1)

        return trans

    def is_valid_id(self, veh_frame_id):
        """
        Written by sizhewei @ 2022/10/05
        Given veh_frame_id, determine whether there is a corresponding inf_frame that meets the k delay requirement.

        Parameters
        ----------
        veh_frame_id : 05d
            Vehicle frame id

        Returns
        -------
        bool valud
            True means there is a corresponding road-side frame.
        """
        # print('veh_frame_id: ',veh_frame_id,'\n')
        frame_info = {}
        
        frame_info = self.co_idx2info[veh_frame_id]
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")
        cur_inf_info = self.inf_idx2info[inf_frame_id]
        if (int(inf_frame_id) - self.binomial_n*self.k < int(cur_inf_info["batch_start_id"])):
            return False
        for i in range(self.binomial_n * self.k): 
            delay_id = id_to_str(int(inf_frame_id) - i) 
            if delay_id not in self.inf_fid2veh_fid.keys():
                return False

        return True

    def retrieve_base_data(self, idx):
        """
        Modified by sizhewei @ 2022/09/28
        Given the index, return the corresponding async data (time delay expection = sum(B(n, p)) ).

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            {
                [0]  / [1]: {
                    'ego' : True / False,
                    'params' : {
                        'vehicles':
                        'lidar_pose':
                    },
                    'lidar_np' : 
                    'veh_frame_id' :
                    'avg_time_delay' : ([1] only)
                }
            }
        """
        final_data = OrderedDict()
        veh_frame_id = self.data[idx]
        
        # print('veh_frame_id: ',veh_frame_id,'\n')
        frame_info = {}
        system_error_offset = {}
        
        frame_info = self.co_idx2info[veh_frame_id]
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        '''
        final_data : {
            [0] / [1] : {
                'ego' : True / False,
                'curr' : {
                    'frame_id' :
                }
                'past_k' : {
                    [0] / [1] / ... / [k-1] : {
                        'frame_id'
                    }
                }
            }

        }
        '''

        final_data = OrderedDict()
        curr_veh_frame_id = self.data[idx]
        frame_info = self.co_idx2info[veh_frame_id]
        system_error_offset = frame_info['system_error_offset']
        curr_inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")
        for i in range(2):
            final_data[i] = OrderedDict()
            final_data[i]['ego'] = True if i == 0 else False # first one is veh, second is inf
            final_data[i]['curr'] = {}
            final_data[i]['curr']['frame_id'] = curr_veh_frame_id if i == 0 else curr_inf_frame_id
            final_data[i]['curr']['timestamp'] = curr_veh_frame_id if i == 0 else curr_inf_frame_id
            final_data[i]['curr']['time_diff'] = 0
            final_data[i]['curr']['sample_interval'] = 0

            final_data[i]['curr']['lidar_np'] = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["vehicle_pointcloud_path"]))[0] if i==0 else \
                pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["infrastructure_pointcloud_path"]))[0]
            final_data[i]['curr']['params'] = OrderedDict()
            final_data[i]['curr']['params']['vehicles'] = \
                load_json(osp.join(self.root_dir, frame_info['cooperative_label_path'])) if i == 0 else []
            final_data[i]['curr']['params']['vehicles_single'] = \
                load_json(os.path.join(self.root_dir, 'vehicle-side/label/lidar/{}.json'.format(curr_veh_frame_id))) if i==0 else \
                load_json(os.path.join(self.root_dir, 'infrastructure-side/label/virtuallidar/{}.json'.format(curr_inf_frame_id)))
            
            final_data[i]['curr']['params']['lidar_pose'] = self.get_vehicle_trans(curr_veh_frame_id) if i==0 else \
                self.get_inf_trans(curr_inf_frame_id, system_error_offset)
                          
            final_data[i]['past_k'] = OrderedDict()
            if i==0:
                for j in range(self.k):
                    final_data[i]['past_k'][j] = final_data[0]['curr']
            else:
                latest_frame_id = curr_inf_frame_id
                for j in range(self.k):
                    # B(n, p)
                    trails = bernoulliDist.rvs(self.binomial_n)
                    sample_interval = sum(trails)
                    # sample_interval = 5
                    # if j == 0:
                    #     sample_interval = 3
                    # elif j == 1:
                    #     sample_interval = 1
                    # else:
                    #     sample_interval = 1
                    latest_frame_id = id_to_str(int(latest_frame_id) - sample_interval)
                    try:
                        veh_frame_id_of_inf = self.inf_fid2veh_fid[latest_frame_id]
                    except KeyError:
                        print("=====Error====")
                    frame_info = self.co_idx2info[veh_frame_id_of_inf]

                    system_offset = self.co_idx2info[veh_frame_id_of_inf]['system_error_offset']
                    final_data[i]['past_k'][j] = {}
                    final_data[i]['past_k'][j]['frame_id'] = latest_frame_id
                    final_data[i]['past_k'][j]['timestamp'] = latest_frame_id
                    final_data[i]['past_k'][j]['time_diff'] = int(curr_inf_frame_id) - int(latest_frame_id)
                    final_data[i]['past_k'][j]['sample_interval'] = - sample_interval
                    final_data[i]['past_k'][j]['lidar_np'] = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["infrastructure_pointcloud_path"]))[0]
                    final_data[i]['past_k'][j]['params'] = OrderedDict()
                    final_data[i]['past_k'][j]['params']['vehicles'] = []
                    final_data[i]['past_k'][j]['params']['lidar_pose'] = self.get_inf_trans(latest_frame_id, system_offset)
        
        """
        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            {
                [0] / [1]: {
                    'ego' : True / False,
                    'curr': {
                        'frame_id' :
                        'params': vehicle & pose
                    }
                    'past_k' : {
                        [0] / [1] / ... / [k-1] : {
                            'frame_id' :
                            'params': vehicle & pose
                        }
                    }
                    'avg_time_delay' : ([1] only)
                }
            }
        """



        return final_data



    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                "single_object_bbx_center": single_object_bbx_center,       # 用于viz single view
                "single_object_ids": single_object_ids,                # 用于viz single view
                'flow_gt':              # single view flow
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'pastk_2_past0_tr_mats':    # list of past k frames' transformation matrix to past 0 frame
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        if self.viz_bbx_flag:
            selected_cav_processed.update({'single_lidar': lidar_np[:, :3]})
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past_k to past_0 tansfomation matrix
        pastk_2_past0_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)
            # past_k trans past_0 matrix
            pastk_2_past0 = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], selected_cav_base['past_k'][0]['params']['lidar_pose'])
            pastk_2_past0_tr_mats.append(pastk_2_past0)
            
            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            # # 5. single view label
            # # past_i label at past_i single view
            # # opencood/data_utils/post_processor/base_postprocessor.py
            # object_bbx_center, object_bbx_mask, object_ids = \
            #     self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # # generate the anchor boxes
            # # opencood/data_utils/post_processor/voxel_postprocessor.py
            # anchor_box = self.anchor_box
            # single_view_label_dict = self.post_processor.generate_label(
            #         gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            #     )
            # past_k_label_dicts.append(single_view_label_dict)
        
        

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)
        pastk_2_past0_tr_mats = np.stack(pastk_2_past0_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        single_object_bbx_center, single_object_bbx_mask, single_object_ids = \
            self.generate_object_center_dair_single([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=single_object_bbx_center, anchors=anchor_box, mask=single_object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
        # print("selected_cav_base['ego]", selected_cav_base['ego'])
        # print("object_ids:", object_ids)
        # print()
            
        selected_cav_processed.update({
            "single_label_dict": label_dict,
            "single_object_bbx_center": single_object_bbx_center[single_object_bbx_mask == 1],
            "single_object_ids": single_object_ids,
            "curr_feature": curr_feature,
            'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
            'object_ids': object_ids,
            'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
            'past_k_tr_mats': past_k_tr_mats,
            'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats,
            'past_k_poses': past_k_poses,
            'past_k_features': past_k_features,
            'past_k_time_diffs': past_k_time_diffs,
            'past_k_sample_interval': past_k_sample_interval,
        #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
        #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
        #  'past_k_label_dicts': past_k_label_dicts,
            'if_no_point': if_no_point
            })

        if self.is_generate_gt_flow:
            # generate flow, from past_0 and curr
            prev_object_id_stack = {}
            prev_object_stack = {}
            for t_i in range(2):
                split_part = selected_cav_base['past_k'][0] if t_i == 0 else selected_cav_base['curr'] # TODO: 这里面的 prev 和 curr 可能反了
                object_bbx_center, object_bbx_mask, object_ids = \
                    self.generate_object_center([split_part], selected_cav_base['past_k'][0]['params']['lidar_pose'])
                prev_object_id_stack[t_i] = object_ids
                prev_object_stack[t_i] = object_bbx_center
            
            for t_i in range(2):
                unique_object_ids = list(set(prev_object_id_stack[t_i]))
                unique_indices = \
                    [prev_object_id_stack[t_i].index(x) for x in unique_object_ids]
                prev_object_stack[t_i] = np.vstack(prev_object_stack[t_i])
                prev_object_stack[t_i] = prev_object_stack[t_i][unique_indices]
                prev_object_id_stack[t_i] = unique_object_ids
            
            # TODO: generate_flow_map: yhu, generate_flow_map_szwei: szwei
            flow_map, warp_mask = generate_flow_map_szwei(prev_object_stack,
                                            prev_object_id_stack,
                                            self.params['preprocess']['cav_lidar_range'],
                                            self.params['preprocess']['args']['voxel_size'],
                                            past_k=1)

            selected_cav_processed.update({'flow_gt': flow_map})
            selected_cav_processed.update({'warp_mask': warp_mask})

        return selected_cav_processed

    def __getitem__(self, idx):
        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            },
            cav_id_2 : { ... }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []
        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = math.sqrt( \
                (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  

        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)

        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        pastk_2_past0_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        if self.is_generate_gt_flow:
            flow_gt = []
            warp_mask = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []

        if self.viz_bbx_flag:
            single_lidar_stack = []
            single_object_stack = []
            single_object_id_stack = []
            single_mask_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换, 用于单体监督            cav view + curr 的label
                'single_object_bbx_center'
                'single_object_ids'
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
                'flow_gt':               # [2, H, W] flow ground truth
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            if self.viz_bbx_flag:
                single_lidar_stack.append(selected_cav_processed['single_lidar'])
                single_object_stack.append(selected_cav_processed['single_object_bbx_center'])
                single_object_id_stack.append(selected_cav_processed['single_object_ids'])
                # mask = np.zeros(self.params['postprocess']['max_num'])
                # mask[:single_object_stack.shape[0]] = 1
                # single_mask_stack.append(mask)
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: N, k
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k frames to cav past 0 frame trans matrix, list of len=N, pastk_2_past0_tr_mats[i]: ndarray(k, 4, 4)
            pastk_2_past0_tr_mats.append(selected_cav_processed['pastk_2_past0_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        
            if self.is_generate_gt_flow:
                # for flow
                flow_gt.append(selected_cav_processed['flow_gt'])
                warp_mask.append(selected_cav_processed['warp_mask'])
        
        pastk_2_past0_tr_mats = np.stack(pastk_2_past0_tr_mats, axis=0) # N, k, 4, 4

        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) TODO: 这里面没有搞懂为什么不用 past_k_tr_mats

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack, 
             'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats})
            #  'times': self.times})

        if self.is_generate_gt_flow:
            flow_gt = np.vstack(flow_gt) # (N, H, W, 2)
            processed_data_dict['ego'].update({'flow_gt': flow_gt})
            warp_mask = np.vstack(warp_mask) # (N, C, H, W)
            processed_data_dict['ego'].update({'warp_mask': warp_mask})
        
        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            tmp = past_k_time_diffs_stack[self.k:].reshape(-1, self.k) # (N, k)
            tmp_past_k_time_diffs = np.concatenate((tmp[:, :1] , (tmp[:, 1:] - tmp[:, :-1])), axis=1) # (N, k)
            avg_time_diff = sum(tmp_past_k_time_diffs.reshape(-1)) / tmp_past_k_time_diffs.reshape(-1).shape[0]
            processed_data_dict['ego'].update({'avg_sample_interval':\
                sum(past_k_sample_interval_stack[self.k:]) / len(past_k_sample_interval_stack[self.k:])})
            processed_data_dict['ego'].update({'avg_time_delay':\
                avg_time_diff})
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        if self.viz_bbx_flag:
            for id, cav in enumerate(cav_id_list):
                processed_data_dict[id] = {}
                processed_data_dict[id].update({
                    'single_lidar': single_lidar_stack[id],
                    'single_object_bbx_center': single_object_stack[id],
                    # 'single_object_bbx_mask': single_mask_stack[i],
                    'single_object_ids': single_object_id_stack[id]
                })

        return processed_data_dict

    def __len__(self):
        # 符合条件的 frame 的数量
        return len(self.data)

    ### rewrite generate_object_center ###
    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_dairv2x(cav_contents,
                                                        reference_lidar_pose)
    
    ### Add new func for single side
    def generate_object_center_dair_single(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        veh or inf 's coordinate. 

        reference_lidar_pose is of no use.
        """
        suffix = "_single"
        return self.post_processor.generate_object_center_dairv2x_single(cav_contents, suffix)


opencood/data_utils/datasets/intermediate_fusion_dataset_dair_irregular_multi.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/4/15
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
two stage flow update framework

"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
from opencood.utils.flow_utils import generate_flow_map, generate_flow_map_szwei

from opencood.utils.box_utils import boxes_to_corners2d # for debug use
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetDAIRMultisweep(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):

        self.times = []

        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
            print("!!! Absolutely Regular !!!")

        # 控制是否需要生成GT flow
        self.is_generate_gt_flow = False
        if 'is_generate_gt_flow' in params and params['is_generate_gt_flow']:
            self.is_generate_gt_flow = True
        
        # 只有在绘制每个sample的匹配框时用到 sizhewei
        self.viz_bbx_flag = False
        if 'viz_bbx_flag' in params and params['viz_bbx_flag']:
            self.viz_bbx_flag = True
        
        self.sample_interval_exp = int(self.binomial_n * self.binomial_p)

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        '''
        scenario_database Structure: 
        {
            scenario_id : {
                cav_1 : {
                    'ego' : true / false , 
                    timestamp1 : {
                        yaml: path,
                        lidar: path, 
                        cameras: list of path
                    },
                    ...
                },
                ...
            }
        }
        '''
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # copy timestamps npy file, store all irregular timestamps of each non-ego vehicle in this scenario
            timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
            time_annotations = np.load(timestamps_file)

            # at least 1 cav should show up
            cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
            assert len(cav_list) > 0

            # use the frame number as key, the full path as the values, store all json or yaml files in this scenario
            yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
            if len(yaml_files)==0:
                yaml_files = sorted([x 
                            for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                            x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

            start_timestamp = int(float(self.extract_timestamps(yaml_files)[0]))
            while(1):
                time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
                time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
                if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
                    start_timestamp += 1
                else:
                    break

            end_timestamp = int(float(self.extract_timestamps(yaml_files)[-1]))
            if start_timestamp%2 == 0:
                # even
                end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
            else:
                end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
            num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)
            regular_timestamps = [start_timestamp+2*i for i in range(num_timestamps)]

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                cav_path = os.path.join(scenario_folder, cav_id)

                if j==0: # ego
                    timestamps = regular_timestamps
                else:
                    if self.is_ab_regular:
                        timestamps = regular_timestamps
                    else:
                        timestamps = list(time_annotations[j-1, :])

                for timestamp in timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files
                
                # regular的timestamps 用于做 curr 真实时刻的ground truth
                self.scenario_database[i][cav_id]['regular'] = OrderedDict()
                for timestamp in regular_timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id]['regular'][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id]['regular'][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id]['regular'][timestamp]['lidar'] = \
                        lidar_file

                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    # num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    num_ego_timestamps = len(timestamps) - self.binomial_n * self.k
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)
                else:
                    self.scenario_database[i][cav_id]['ego'] = False
                    

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead. 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("=== OPV2V-Irregular Multi-sweep dataset with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def extract_timestamps(self, yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]
            if res.endswith('.yaml'):
                timestamp = res.replace('.yaml', '')
            elif res.endswith('.json'):
                timestamp = res.replace('.json', '')
            else:
                print("Woops! There is no processing method for file {}".format(res))
                sys.exit(1)
            timestamps.append(timestamp)

        return timestamps

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)
                        ...,		
                        [k-1] : {} (id-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k-1]:{}  (id-\tau-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * self.k
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # TODO: debug use, uncomment: to use regular version yaml GT
            # time_new = str( '%06d' % int(float(timestamp_key)))
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # TODO: debug use, uncomment: to use regular version lidar input
            # time_new = str( '%06d' % int(float(timestamp_key))) 
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                # sample_interval
                if data[cav_id]['ego']:             # ego sample_interval = E(B(n, p))
                    if i == 0: # ego-past-0 与 ego-curr 是一样的
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    sample_interval = self.sample_interval_exp
                    if sample_interval == 0:
                        sample_interval = 1
                else:                               # non-ego sample_interval ~ B(n, p)
                    # if i == 0:
                    #     sample_interval = 5
                    # elif i ==1:
                    #     sample_interval = 2
                    # else:
                    #     sample_interval = 1
                    if self.sample_interval_exp==0 \
                        and self.is_no_shift \
                            and i == 0:
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==0: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>0: # 过去的几帧不要重复
                            sample_interval = 1                

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': # current 在各自view 下 feature
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float
        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            },
            cav_id_2 : { ... }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = math.sqrt( \
                (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        pastk_2_past0_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        if self.is_generate_gt_flow:
            flow_gt = []
            warp_mask = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []

        if self.viz_bbx_flag:
            single_lidar_stack = []
            single_object_stack = []
            single_object_id_stack = []
            single_mask_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换, 用于单体监督            cav view + curr 的label
                'single_object_bbx_center'
                'single_object_ids'
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
                'flow_gt':               # [2, H, W] flow ground truth
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            if self.viz_bbx_flag:
                single_lidar_stack.append(selected_cav_processed['single_lidar'])
                single_object_stack.append(selected_cav_processed['single_object_bbx_center'])
                single_object_id_stack.append(selected_cav_processed['single_object_ids'])
                # mask = np.zeros(self.params['postprocess']['max_num'])
                # mask[:single_object_stack.shape[0]] = 1
                # single_mask_stack.append(mask)
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: N, k
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k frames to cav past 0 frame trans matrix, list of len=N, pastk_2_past0_tr_mats[i]: ndarray(k, 4, 4)
            pastk_2_past0_tr_mats.append(selected_cav_processed['pastk_2_past0_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        
            if self.is_generate_gt_flow:
                # for flow
                flow_gt.append(selected_cav_processed['flow_gt'])
                warp_mask.append(selected_cav_processed['warp_mask'])
        
        pastk_2_past0_tr_mats = np.stack(pastk_2_past0_tr_mats, axis=0) # N, k, 4, 4

        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) TODO: 这里面没有搞懂为什么不用 past_k_tr_mats

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack, 
             'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats})
            #  'times': self.times})

        if self.is_generate_gt_flow:
            flow_gt = np.vstack(flow_gt) # (N, H, W, 2)
            processed_data_dict['ego'].update({'flow_gt': flow_gt})
            warp_mask = np.vstack(warp_mask) # (N, C, H, W)
            processed_data_dict['ego'].update({'warp_mask': warp_mask})
        
        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            tmp = past_k_time_diffs_stack[self.k:].reshape(-1, self.k) # (N, k)
            tmp_past_k_time_diffs = np.concatenate((tmp[:, :1] , (tmp[:, 1:] - tmp[:, :-1])), axis=1) # (N, k)
            avg_time_diff = sum(tmp_past_k_time_diffs.reshape(-1)) / tmp_past_k_time_diffs.reshape(-1).shape[0]
            processed_data_dict['ego'].update({'avg_sample_interval':\
                sum(past_k_sample_interval_stack[self.k:]) / len(past_k_sample_interval_stack[self.k:])})
            processed_data_dict['ego'].update({'avg_time_delay':\
                avg_time_diff})
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        if self.viz_bbx_flag:
            for id, cav in enumerate(cav_id_list):
                processed_data_dict[id] = {}
                processed_data_dict[id].update({
                    'single_lidar': single_lidar_stack[id],
                    'single_object_bbx_center': single_object_stack[id],
                    # 'single_object_bbx_mask': single_mask_stack[i],
                    'single_object_ids': single_object_id_stack[id]
                })

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                "single_object_bbx_center": single_object_bbx_center,       # 用于viz single view
                "single_object_ids": single_object_ids,                # 用于viz single view
                'flow_gt':              # single view flow
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'pastk_2_past0_tr_mats':    # list of past k frames' transformation matrix to past 0 frame
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        if self.viz_bbx_flag:
            selected_cav_processed.update({'single_lidar': lidar_np[:, :3]})
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past_k to past_0 tansfomation matrix
        pastk_2_past0_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)
            # past_k trans past_0 matrix
            pastk_2_past0 = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], selected_cav_base['past_k'][0]['params']['lidar_pose'])
            pastk_2_past0_tr_mats.append(pastk_2_past0)
            
            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            # # 5. single view label
            # # past_i label at past_i single view
            # # opencood/data_utils/post_processor/base_postprocessor.py
            # object_bbx_center, object_bbx_mask, object_ids = \
            #     self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # # generate the anchor boxes
            # # opencood/data_utils/post_processor/voxel_postprocessor.py
            # anchor_box = self.anchor_box
            # single_view_label_dict = self.post_processor.generate_label(
            #         gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            #     )
            # past_k_label_dicts.append(single_view_label_dict)
        
        

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)
        pastk_2_past0_tr_mats = np.stack(pastk_2_past0_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        single_object_bbx_center, single_object_bbx_mask, single_object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=single_object_bbx_center, anchors=anchor_box, mask=single_object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update({
            "single_label_dict": label_dict,
            "single_object_bbx_center": single_object_bbx_center[single_object_bbx_mask == 1],
            "single_object_ids": single_object_ids,
            "curr_feature": curr_feature,
            'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
            'object_ids': object_ids,
            'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
            'past_k_tr_mats': past_k_tr_mats,
            'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats,
            'past_k_poses': past_k_poses,
            'past_k_features': past_k_features,
            'past_k_time_diffs': past_k_time_diffs,
            'past_k_sample_interval': past_k_sample_interval,
        #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
        #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
        #  'past_k_label_dicts': past_k_label_dicts,
            'if_no_point': if_no_point
            })

        if self.is_generate_gt_flow:
            # generate flow, from past_0 and curr
            prev_object_id_stack = {}
            prev_object_stack = {}
            for t_i in range(2):
                split_part = selected_cav_base['past_k'][0] if t_i == 0 else selected_cav_base['curr'] # TODO: 这里面的 prev 和 curr 可能反了
                object_bbx_center, object_bbx_mask, object_ids = \
                    self.generate_object_center([split_part], selected_cav_base['past_k'][0]['params']['lidar_pose'])
                prev_object_id_stack[t_i] = object_ids
                prev_object_stack[t_i] = object_bbx_center
            
            for t_i in range(2):
                unique_object_ids = list(set(prev_object_id_stack[t_i]))
                unique_indices = \
                    [prev_object_id_stack[t_i].index(x) for x in unique_object_ids]
                prev_object_stack[t_i] = np.vstack(prev_object_stack[t_i])
                prev_object_stack[t_i] = prev_object_stack[t_i][unique_indices]
                prev_object_id_stack[t_i] = unique_object_ids
            
            # TODO: generate_flow_map: yhu, generate_flow_map_szwei: szwei
            flow_map, warp_mask = generate_flow_map_szwei(prev_object_stack,
                                            prev_object_id_stack,
                                            self.params['preprocess']['cav_lidar_range'],
                                            self.params['preprocess']['args']['voxel_size'],
                                            past_k=1)

            selected_cav_processed.update({'flow_gt': flow_map})
            selected_cav_processed.update({'warp_mask': warp_mask})

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def dist_time(ts1, ts2, i = -1):
        """caculate the time interval between two timestamps

        Args:
            ts1 (string): time stamp at some time
            ts2 (string): current time stamp
            i (int, optional): past frame id, for debug use. Defaults to -1.
        
        Returns:
            time_diff (float): time interval (ts1 - ts2)
        """
        if not i==-1:
            return -i
        else:
            return (float(ts1) - float(ts2))

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'flow_gt': [N, C, H, W]
        }

        Returns:
        ----------
        batch['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack, (Optional) or just use current
            'curr_processed_lidar': merged_curr_feature_dict, (Optional)
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices], 是不是不需要了？
            'processed_lidar': merged_feature_dict, 模型输入, 带有延迟的 多帧
            'record_len': record_len, (B, ) 每个batch有几辆车 模型输入
            'anchor_box': anchor_box,
            'label_dict': label_dict, 在current ego 坐标系下
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix, (B, L, k, 4, 4) 所有车辆、所有帧 转换到 ego current
            'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats, 过去k帧到过去0帧的变换矩阵 [BxN, 3, 4, 4]
            'curr_lidar_pose': curr_lidar_pose, [BxN, 6]
            'past_lidar_pose': past_k_lidar_pose, 模型输入 [BxN, k, 6] 
            'past_k_time_interval': past_k_time_diff, 模型输入 [BxNxk] 距离当前帧的时间间隔
            'past_k_sample_interval': past_k_sample_interval, int [BxNxk] 两两之间的采样间隔
            'avg_sample_interval': avg_sample_interval,
            'avg_time_delay': avg_time_delay
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        pastk_2_past0_tr_mats = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []
        
        if self.is_generate_gt_flow:
            # flow gt
            flow_gt_list = []
            warp_mask_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            pastk_2_past0_tr_mats.append(ego_dict['pastk_2_past0_tr_mats'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            if self.is_generate_gt_flow:
                flow_gt_list.append(ego_dict['flow_gt'])
                warp_mask_list.append(ego_dict['warp_mask'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        pastk_2_past0_tr_mats = torch.from_numpy(np.vstack(pastk_2_past0_tr_mats))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        if self.is_generate_gt_flow:
            # for flow
            flow_gt = torch.from_numpy(np.vstack(flow_gt_list))
            label_torch_dict.update({'flow_gt': flow_gt})
            warp_mask = torch.from_numpy(np.vstack(warp_mask_list))
            label_torch_dict.update({'warp_mask': warp_mask})

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay})
                                #    'times': time_consume})
        output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(self.anchor_box))})
        
        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})

        if self.viz_bbx_flag:
            single_lidar = []
            single_object_bbx_center = []
            single_object_mask = []
            single_object_ids = []
            for i in range(len(batch[0].keys())-1):
                single_lidar.append(torch.from_numpy(batch[0][i]['single_lidar']))
                single_object_bbx_center.append(torch.from_numpy(batch[0][i]['single_object_bbx_center']))
                single_object_ids.append(batch[0][i]['single_object_ids'])
            output_dict['ego'].update({
                'single_lidar_list': single_lidar,
                'single_object_bbx_center': single_object_bbx_center,
                'single_object_ids': single_object_ids
            })
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

    def generate_pred_bbx_frames(self, m_single, trans_mat_pastk_2_past0, past_time_diff, anchor_box):
        '''
        m_single : {
            psm_single, rm_single, (dm_single)
        }
        box_result : {
            'past_k_time_diff' : 
            [0] : {
                pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. (n, 8, 3)
                pred_box_center_tensor : (n, 7)
                scores: (n, )
            },
            ...
            [k-1] : { ... }
        }
        '''
        box_results = self.post_processor.single_post_process(m_single, trans_mat_pastk_2_past0, past_time_diff, anchor_box, self.k)
        return box_results

# if __name__ == '__main__':   



# # -*- coding: utf-8 -*-
# # Modified: Sizhe Wei <sizhewei@sjtu.edu.cn>

# """
# Dataset class for intermediate fusion with past k frames
# """
# import random
# import math
# from collections import OrderedDict

# import os
# import os.path as osp
# import numpy as np
# import torch
# from torch.utils.data import DataLoader
# import json
# import opencood.data_utils.datasets
# import opencood.data_utils.post_processor as post_processor
# from opencood.utils import box_utils

# from opencood.data_utils.datasets import intermediate_fusion_dataset
# from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
# from opencood.data_utils.pre_processor import build_preprocessor
# from opencood.hypes_yaml.yaml_utils import load_yaml
# from opencood.utils.pcd_utils import \
#     mask_points_by_range, mask_ego_points, shuffle_points, \
#     downsample_lidar_minimum
# from opencood.utils.transformation_utils import x1_to_x2
# import opencood.utils.pcd_utils as pcd_utils
# from opencood.utils.transformation_utils import tfm_to_pose
# from opencood.utils.transformation_utils import veh_side_rot_and_trans_to_trasnformation_matrix
# from opencood.utils.transformation_utils import inf_side_rot_and_trans_to_trasnformation_matrix
# from opencood.utils.transformation_utils import x_to_world
# from opencood.utils.pose_utils import add_noise_data_dict

# def load_json(path):
#     with open(path, mode="r") as f:
#         data = json.load(f)
#     return data

# def build_idx_to_info(data):
#     idx2info = {}
#     for elem in data:
#         if elem["pointcloud_path"] == "":
#             continue
#         idx = elem["pointcloud_path"].split("/")[-1].replace(".pcd", "")
#         idx2info[idx] = elem
#     return idx2info

# def build_idx_to_co_info(data):
#     idx2info = {}
#     for elem in data:
#         if elem["vehicle_pointcloud_path"] == "":
#             continue
#         idx = elem["vehicle_pointcloud_path"].split("/")[-1].replace(".pcd", "")
#         idx2info[idx] = elem
#     return idx2info

# def id_to_str(id, digits=6):
#     result = ""
#     for i in range(digits):
#         result = str(id % 10) + result
#         id //= 10
#     return result

# class IntermediateFusionDatasetDAIRMultisweep(intermediate_fusion_dataset.IntermediateFusionDataset):
#     """
#     Written by sizhewei @ 2022/09/28
#     This class is for intermediate fusion where each vehicle transmit the
#     deep features to ego.
#     """
#     def __init__(self, params, visualize, train=True):
#         #注意yaml文件应该有sensor_type：lidar/camera
#         self.params = params
#         self.visualize = visualize
#         self.train = train
#         self.data_augmentor = DataAugmentor(params['data_augment'],
#                                             train)
#         self.max_cav = 2
#         self.k = params['time_delay']
        
#         # if project first, cav's lidar will first be projected to
#         # the ego's coordinate frame. otherwise, the feature will be
#         # projected instead.
#         assert 'proj_first' in params['fusion']['args']
#         if params['fusion']['args']['proj_first']:
#             self.proj_first = True
#         else:
#             self.proj_first = False

#         if "kd_flag" in params.keys():
#             self.kd_flag = params['kd_flag']
#         else:
#             self.kd_flag = False

#         if "box_align" in params.keys():
#             self.box_align = True
#             self.stage1_result_path = params['box_align']['train_result'] if train else params['box_align']['val_result']
#             self.stage1_result = load_json(self.stage1_result_path)
#             self.box_align_args = params['box_align']['args']
        
#         else:
#             self.box_align = False

#         assert 'clip_pc' in params['fusion']['args']
#         if params['fusion']['args']['clip_pc']:
#             self.clip_pc = True
#         else:
#             self.clip_pc = False
            
#         if 'select_kp' in params:
#             self.select_keypoint = params['select_kp']
#         else:
#             self.select_keypoint = None


#         self.pre_processor = build_preprocessor(params['preprocess'],
#                                                 train)
#         self.post_processor = post_processor.build_postprocessor(
#             params['postprocess'],
#             train)

#         #这里root_dir是一个json文件！--> 代表一个split
#         if self.train:
#             split_dir = params['root_dir']
#         else:
#             split_dir = params['validate_dir']

#         self.root_dir = '/GPFS/rhome/quanhaoli/workspace/dataset/my_dair_v2x/v2x_c/cooperative-vehicle-infrastructure'
#         self.inf_idx2info = build_idx_to_info(
#             load_json(osp.join(self.root_dir, "infrastructure-side/data_info.json"))
#         )
#         self.co_idx2info = build_idx_to_co_info(
#             load_json(osp.join(self.root_dir, "cooperative/data_info_deleted.json"))
#         )

#         self.data_split = load_json(split_dir)
#         self.data = []
#         for veh_idx in self.data_split:
#             if self.is_valid_id(veh_idx):
#                 self.data.append(veh_idx)

#         print("ASync dataset with {} time delay initialized! {} samples totally!".format(self.k, len(self.data)))
#     def is_valid_id(self, veh_frame_id):
#         """
#         Written by sizhewei @ 2022/10/05
#         Given veh_frame_id, determine whether there is a corresponding inf_frame that meets the k delay requirement.

#         Parameters
#         ----------
#         veh_frame_id : 05d
#             Vehicle frame id

#         Returns
#         -------
#         bool valud
#             True means there is a corresponding road-side frame.
#         """
#         # print('veh_frame_id: ',veh_frame_id,'\n')
#         frame_info = {}
        
#         frame_info = self.co_idx2info[veh_frame_id]
#         inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")
#         cur_inf_info = self.inf_idx2info[inf_frame_id]
#         if (
#             int(inf_frame_id) - self.k < int(cur_inf_info["batch_start_id"])
#             or id_to_str(int(inf_frame_id) - self.k) not in self.inf_idx2info
#         ):
#             return False

#         return True

#     def retrieve_base_data(self, idx):
#         """
#         Modified by sizhewei @ 2022/09/28
#         Given the index, return the corresponding async data (time delay: self.k).

#         Parameters
#         ----------
#         idx : int
#             Index given by dataloader.

#         Returns
#         -------
#         data : dict
#             The dictionary contains loaded yaml params and lidar data for
#             each cav.
#         """
#         veh_frame_id = self.data[idx]
#         # print('veh_frame_id: ',veh_frame_id,'\n')
#         frame_info = {}
#         system_error_offset = {}
        
#         frame_info = self.co_idx2info[veh_frame_id]
#         inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")

#         inf_frame_id = id_to_str(int(inf_frame_id) - self.k)

#         system_error_offset = frame_info["system_error_offset"]
#         data = OrderedDict()
#         #cav_id=0是车端，1是路边单元
#         data[0] = OrderedDict()
#         data[0]['ego'] = True
#         data[0]['params'] = OrderedDict()
#         data[0]['params']['vehicles'] = load_json(os.path.join(self.root_dir,frame_info['cooperative_label_path']))
#         # print(data[0]['params']['vehicles'])
#         lidar_to_novatel_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/lidar_to_novatel/'+str(veh_frame_id)+'.json'))
#         novatel_to_world_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/novatel_to_world/'+str(veh_frame_id)+'.json'))

#         transformation_matrix = veh_side_rot_and_trans_to_trasnformation_matrix(lidar_to_novatel_json_file,novatel_to_world_json_file)

#         data[0]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix)

#         data[0]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["vehicle_pointcloud_path"]))
#         if self.clip_pc:
#             data[0]['lidar_np'] = data[0]['lidar_np'][data[0]['lidar_np'][:,0]>0]
            
#         data[1] = OrderedDict()
#         data[1]['ego'] = False

#         data[1]['params'] = OrderedDict()

#         # data[1]['params']['vehicles'] = load_json(os.path.join(self.root_dir,frame_info['cooperative_label_path']))
#         data[1]['params']['vehicles'] = [] # we only load cooperative label in vehicle side

#         virtuallidar_to_world_json_file = load_json(os.path.join(self.root_dir,'infrastructure-side/calib/virtuallidar_to_world/'+str(inf_frame_id)+'.json'))

#         transformation_matrix1 = inf_side_rot_and_trans_to_trasnformation_matrix(virtuallidar_to_world_json_file,system_error_offset)
#         data[1]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix1)

#         cur_inf_info = self.inf_idx2info[inf_frame_id]
#         data[1]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir, \
#             'infrastructure-side', cur_inf_info['pointcloud_path']))
#         data[0]['veh_frame_id'] = veh_frame_id
#         data[1]['veh_frame_id'] = inf_frame_id
#         return data

#     def __getitem__(self, idx):
#         # base_data_dict = self.retrieve_base_data(idx)
#         base_data_dict = self.retrieve_base_data(idx)

#         base_data_dict = add_noise_data_dict(base_data_dict,self.params['noise_setting'])

#         processed_data_dict = OrderedDict()
#         processed_data_dict['ego'] = {}

#         ego_id = -1
#         ego_lidar_pose = []

#         # first find the ego vehicle's lidar pose
#         for cav_id, cav_content in base_data_dict.items():
#             if cav_content['ego']:
#                 ego_id = cav_id
#                 ego_lidar_pose = cav_content['params']['lidar_pose']
#                 ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
#                 break
            
#         assert cav_id == list(base_data_dict.keys())[
#             0], "The first element in the OrderedDict must be ego"
#         assert ego_id != -1
#         assert len(ego_lidar_pose) > 0


#         processed_features = []
#         object_stack = []
#         object_id_stack = []
#         too_far = []
#         lidar_pose_list = []
#         lidar_pose_clean_list = []
#         projected_lidar_clean_list = []
#         cav_id_list = []

#         if self.visualize:
#             projected_lidar_stack = []

#         # loop over all CAVs to process information
#         for cav_id, selected_cav_base in base_data_dict.items():
#             # check if the cav is within the communication range with ego
#             distance = \
#                 math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
#                            ego_lidar_pose[0]) ** 2 + (
#                                   selected_cav_base['params'][
#                                       'lidar_pose'][1] - ego_lidar_pose[
#                                       1]) ** 2)

#             # if distance is too far, we will just skip this agent
#             if distance > self.params['comm_range']:
#                 too_far.append(cav_id)
#                 continue


#             lidar_pose_clean_list.append(selected_cav_base['params']['lidar_pose_clean'])
#             lidar_pose_list.append(selected_cav_base['params']['lidar_pose']) # 6dof pose
#             cav_id_list.append(cav_id)


#         ########## Added by Yifan Lu 2022.8.14 ##############
#         # box align to correct pose.
#         '''
#         if self.box_align and str(idx) in self.stage1_result.keys():
#             stage1_content = self.stage1_result[str(idx)]
#             if stage1_content is not None:
#                 cav_id_list_stage1 = stage1_content['cav_id_list']
                
#                 pred_corners_list = stage1_content['pred_corner3d_np_list']
#                 pred_corners_list = [np.array(corners, dtype=np.float64) for corners in pred_corners_list]
#                 uncertainty_list = stage1_content['uncertainty_np_list']
#                 uncertainty_list = [np.array(uncertainty, dtype=np.float64) for uncertainty in uncertainty_list]
#                 stage1_lidar_pose_list = [base_data_dict[cav_id]['params']['lidar_pose'] for cav_id in cav_id_list_stage1]
#                 stage1_lidar_pose = np.array(stage1_lidar_pose_list)


#                 refined_pose = box_alignment_relative_sample_np(pred_corners_list,
#                                                                 stage1_lidar_pose, 
#                                                                 uncertainty_list=uncertainty_list, 
#                                                                 **self.box_align_args)
#                 stage1_lidar_pose[:,[0,1,4]] = refined_pose
#                 stage1_lidar_pose_refined_list = stage1_lidar_pose.tolist() # updated lidar_pose_list
#                 for cav_id, lidar_pose_refined in zip(cav_id_list_stage1, stage1_lidar_pose_refined_list):
#                     if cav_id not in cav_id_list:
#                         continue
#                     idx_in_list = cav_id_list.index(cav_id)
#                     lidar_pose_list[idx_in_list] = lidar_pose_refined
#                     base_data_dict[cav_id]['params']['lidar_pose'] = lidar_pose_refined
#         '''     


#         for cav_id in cav_id_list:
#             selected_cav_base = base_data_dict[cav_id]



#             selected_cav_processed = self.get_item_single_car(
#                 selected_cav_base,
#                 ego_lidar_pose, 
#                 ego_lidar_pose_clean,
#                 idx)
                
#             object_stack.append(selected_cav_processed['object_bbx_center'])
#             object_id_stack += selected_cav_processed['object_ids']

#             processed_features.append(
#                 selected_cav_processed['processed_features'])
#             if self.kd_flag:
#                 projected_lidar_clean_list.append(
#                     selected_cav_processed['projected_lidar_clean'])

#             if self.visualize:
#                 projected_lidar_stack.append(
#                     selected_cav_processed['projected_lidar'])





#         ########## Added by Yifan Lu 2022.4.5 ################
#         # filter those out of communicate range
#         # then we can calculate get_pairwise_transformation
#         for cav_id in too_far:
#             base_data_dict.pop(cav_id)
        
#         pairwise_t_matrix = \
#             self.get_pairwise_transformation(base_data_dict,
#                                              self.max_cav)

#         lidar_poses = np.array(lidar_pose_list).reshape(-1, 6)  # [N_cav, 6]
#         lidar_poses_clean = np.array(lidar_pose_clean_list).reshape(-1, 6)  # [N_cav, 6]
#         ######################################################

#         ############ for disconet ###########
#         if self.kd_flag:
#             stack_lidar_np = np.vstack(projected_lidar_clean_list)
#             stack_lidar_np = mask_points_by_range(stack_lidar_np,
#                                         self.params['preprocess'][
#                                             'cav_lidar_range'])
#             stack_feature_processed = self.pre_processor.preprocess(stack_lidar_np)

#         # exclude all repetitive objects    
#         unique_indices = \
#             [object_id_stack.index(x) for x in set(object_id_stack)]
#         object_stack = np.vstack(object_stack)
#         object_stack = object_stack[unique_indices]

#         # make sure bounding boxes across all frames have the same number
#         object_bbx_center = \
#             np.zeros((self.params['postprocess']['max_num'], 7))
#         mask = np.zeros(self.params['postprocess']['max_num'])
#         object_bbx_center[:object_stack.shape[0], :] = object_stack
#         mask[:object_stack.shape[0]] = 1

#         # merge preprocessed features from different cavs into the same dict
#         cav_num = len(processed_features)

#         merged_feature_dict = self.merge_features_to_dict(processed_features)

#         # generate the anchor boxes
#         anchor_box = self.post_processor.generate_anchor_box()

#         # generate targets label
#         label_dict = \
#             self.post_processor.generate_label(
#                 gt_box_center=object_bbx_center,
#                 anchors=anchor_box,
#                 mask=mask)

#         processed_data_dict['ego'].update(
#             {'object_bbx_center': object_bbx_center,
#              'object_bbx_mask': mask,
#              'object_ids': [object_id_stack[i] for i in unique_indices],
#              'anchor_box': anchor_box,
#              'processed_lidar': merged_feature_dict,
#              'label_dict': label_dict,
#              'cav_num': cav_num,
#              'pairwise_t_matrix': pairwise_t_matrix,
#              'lidar_poses_clean': lidar_poses_clean,
#              'lidar_poses': lidar_poses})

#         if self.kd_flag:
#             processed_data_dict['ego'].update({'teacher_processed_lidar':
#                 stack_feature_processed})

#         if self.visualize:
#             processed_data_dict['ego'].update({'origin_lidar':
#                 np.vstack(
#                     projected_lidar_stack)})


#         processed_data_dict['ego'].update({'sample_idx': idx,
#                                             'cav_id_list': cav_id_list})

#         # processed_data_dict['ego'].update({'veh_frame_id': base_data_dict[0]['veh_frame_id']})

#         return processed_data_dict

#     def __len__(self):
#         # 符合条件的 frame 的数量
#         return len(self.data)

#     ### rewrite generate_object_center ###
#     def generate_object_center(self,
#                                cav_contents,
#                                reference_lidar_pose):
#         """
#         Retrieve all objects in a format of (n, 7), where 7 represents
#         x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

#         Notice: it is a wrap of postprocessor function

#         Parameters
#         ----------
#         cav_contents : list
#             List of dictionary, save all cavs' information.
#             in fact it is used in get_item_single_car, so the list length is 1

#         reference_lidar_pose : list
#             The final target lidar pose with length 6.

#         Returns
#         -------
#         object_np : np.ndarray
#             Shape is (max_num, 7).
#         mask : np.ndarray
#             Shape is (max_num,).
#         object_ids : list
#             Length is number of bbx in current sample.
#         """

#         return self.post_processor.generate_object_center_dairv2x(cav_contents,
#                                                         reference_lidar_pose)

opencood/data_utils/datasets/intermediate_fusion_dataset_dair_multisweep.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Yifan Lu
# Modified: sizhewei @ 2022/10/12
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for early fusion
"""
from collections import OrderedDict
import os
import numpy as np
import torch
import math
import copy

import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np


class IntermediateFusionDatasetAsync(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'time_delay' in params: # number of frames-delay
            self.k = params['time_delay'] 
        else:
            self.k = 0

        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        # Structure: {scenario_id : {cav_1 : {timestamp1 : {yaml: path,
        # lidar: path, cameras:list of path}}}}
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # at least 1 cav should show up
            cav_list = sorted([x for x in os.listdir(scenario_folder)
                               if os.path.isdir(
                    os.path.join(scenario_folder, x))])
            assert len(cav_list) > 0

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                # save all yaml files to the dictionary
                cav_path = os.path.join(scenario_folder, cav_id)

                # use the frame number as key, the full path as the values
                yaml_files = \
                    sorted([os.path.join(cav_path, x)
                            for x in os.listdir(cav_path) if
                            x.endswith('.yaml')])
                timestamps = self.extract_timestamps(yaml_files)

                if self.k > 0:
                    if j == 0: # ego
                        timestamps = timestamps[self.k:]        # 从第k个往后
                    else:
                        timestamps = timestamps[:-self.k]       # 丢掉最后k个 

                for timestamp in timestamps:
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        camera_files

                    if self.select_keypoint:
                        if (self.proj_first and j == 0) or (not self.proj_first):
                            # when init the dataset, it read over all pcd files.
                            # it maybe slow, but no need to perform keypoint sampling for each time.\
                            kp_file_path = f"{self.params['select_kp']['kp_path']}_{self.params['select_kp']['n_samples']}/{scenario_folders_name[i]}/{cav_id}/{timestamp}.npy"
                            kp_dir = kp_file_path.rsplit('/',1)[0]

                            if not os.path.exists(kp_dir):
                                os.makedirs(kp_dir)

                            if os.path.exists(kp_file_path):
                                keypoints = np.load(kp_file_path)
                            else:
                                pcd_np = pcd_utils.pcd_to_np(lidar_file)
                                bev_keypoints = bev_sample(pcd_np,
                                            self.params['select_kp']['lidar_range'],
                                            self.params['select_kp']['n_samples'],
                                            mode='uniform',
                                            all_samples=False)

                                np.save(kp_file_path, bev_keypoints)
                                print(f"saving {kp_file_path} in basedataset")
                                keypoints = bev_keypoints
                            # sampling again from bev keypoints
                            self.scenario_database[i][cav_id][timestamp]['lidar_keypoints_np'] = keypoints
                                    
                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    if not self.len_record:
                        self.len_record.append(len(timestamps))
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + len(timestamps))
                else:
                    self.scenario_database[i][cav_id]['ego'] = False

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        if "kd_flag" in params.keys():
            self.kd_flag = params['kd_flag']
        else:
            self.kd_flag = False

        print("OPV2V ASync dataset with {} time delay initialized! {} samples totally!".format(self.k, self.len_record[-1]))

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
        """
        # we loop the accumulated length list to see get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]

        data = OrderedDict()
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            # check the timestamp index
            timestamp_index = idx if scenario_index == 0 else \
                idx - self.len_record[scenario_index - 1]
            # retrieve the corresponding timestamp key
            # timestamp_key = self.return_timestamp_key_async(scenario_database, cav_id, 
            #                                         timestamp_index)
            timestamp_key = list(cav_content.items())[timestamp_index][0]
            data[cav_id] = OrderedDict()
            data[cav_id]['ego'] = cav_content['ego']

            # todo: load camera image in the future version
            # load the corresponding data into the dictionary
            data[cav_id]['params'] = \
                load_yaml(cav_content[timestamp_key]['yaml'])
            data[cav_id]['lidar_np'] = \
                pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])
            
            if self.select_keypoint and 'lidar_keypoints_np' in cav_content[timestamp_key]:
                data[cav_id]['lidar_keypoints_np'] = cav_content[timestamp_key]['lidar_keypoints_np']

        return data


    def __getitem__(self, idx):
        base_data_dict = self.retrieve_base_data(idx)

        base_data_dict = add_noise_data_dict(base_data_dict,self.params['noise_setting'])

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        ego_id = -1
        ego_lidar_pose = []

        # first find the ego vehicle's lidar pose
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['params']['lidar_pose']
                ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
                break
            
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0


        processed_features = []
        object_stack = []
        object_id_stack = []
        too_far = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        projected_lidar_clean_list = []
        cav_id_list = []

        if self.visualize:
            projected_lidar_stack = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            distance = \
                math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
                           ego_lidar_pose[0]) ** 2 + (
                                  selected_cav_base['params'][
                                      'lidar_pose'][1] - ego_lidar_pose[
                                      1]) ** 2)

            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue


            lidar_pose_clean_list.append(selected_cav_base['params']['lidar_pose_clean'])
            lidar_pose_list.append(selected_cav_base['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)


        ########## Added by Yifan Lu 2022.8.14 ##############
        # box align to correct pose.
        '''
        if self.box_align and str(idx) in self.stage1_result.keys():
            stage1_content = self.stage1_result[str(idx)]
            if stage1_content is not None:
                cav_id_list_stage1 = stage1_content['cav_id_list']
                
                pred_corners_list = stage1_content['pred_corner3d_np_list']
                pred_corners_list = [np.array(corners, dtype=np.float64) for corners in pred_corners_list]
                uncertainty_list = stage1_content['uncertainty_np_list']
                uncertainty_list = [np.array(uncertainty, dtype=np.float64) for uncertainty in uncertainty_list]
                stage1_lidar_pose_list = [base_data_dict[cav_id]['params']['lidar_pose'] for cav_id in cav_id_list_stage1]
                stage1_lidar_pose = np.array(stage1_lidar_pose_list)


                refined_pose = box_alignment_relative_sample_np(pred_corners_list,
                                                                stage1_lidar_pose, 
                                                                uncertainty_list=uncertainty_list, 
                                                                **self.box_align_args)
                stage1_lidar_pose[:,[0,1,4]] = refined_pose
                stage1_lidar_pose_refined_list = stage1_lidar_pose.tolist() # updated lidar_pose_list
                for cav_id, lidar_pose_refined in zip(cav_id_list_stage1, stage1_lidar_pose_refined_list):
                    if cav_id not in cav_id_list:
                        continue
                    idx_in_list = cav_id_list.index(cav_id)
                    lidar_pose_list[idx_in_list] = lidar_pose_refined
                    base_data_dict[cav_id]['params']['lidar_pose'] = lidar_pose_refined
        '''     

        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]


            # obj: 
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose, 
                ego_lidar_pose_clean,
                idx)
                
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            processed_features.append(
                selected_cav_processed['processed_features'])
            if self.kd_flag:
                projected_lidar_clean_list.append(
                    selected_cav_processed['projected_lidar_clean'])

            if self.visualize:
                projected_lidar_stack.append(
                    selected_cav_processed['projected_lidar'])


        ########## Added by Yifan Lu 2022.4.5 ################
        # filter those out of communicate range
        # then we can calculate get_pairwise_transformation
        for cav_id in too_far:
            base_data_dict.pop(cav_id)
        
        pairwise_t_matrix = \
            self.get_pairwise_transformation(base_data_dict,
                                             self.max_cav)

        lidar_poses = np.array(lidar_pose_list).reshape(-1, 6)  # [N_cav, 6]
        lidar_poses_clean = np.array(lidar_pose_clean_list).reshape(-1, 6)  # [N_cav, 6]
        ######################################################

        ############ for disconet ###########
        if self.kd_flag:
            stack_lidar_np = np.vstack(projected_lidar_clean_list)
            stack_lidar_np = mask_points_by_range(stack_lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            stack_feature_processed = self.pre_processor.preprocess(stack_lidar_np)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        object_stack = np.vstack(object_stack)
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(processed_features)

        merged_feature_dict = self.merge_features_to_dict(processed_features)

        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        processed_data_dict['ego'].update(
            {'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'lidar_poses_clean': lidar_poses_clean,
             'lidar_poses': lidar_poses})

        if self.kd_flag:
            processed_data_dict['ego'].update({'teacher_processed_lidar':
                stack_feature_processed})

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(
                    projected_lidar_stack)})


        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, ego_pose_clean, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information.
        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.
        ego_pose_clean : list, length 6
            only used for gt box generation
        
        idx: int,
            debug use.

        Returns
        -------
        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) # T_ego_cav
        transformation_matrix_clean = \
            x1_to_x2(selected_cav_base['params']['lidar_pose_clean'],
                     ego_pose_clean)

        # retrieve objects under ego coordinates
        # this is used to generate accurate GT bounding box.
        object_bbx_center, object_bbx_mask, object_ids = self.generate_object_center([selected_cav_base],
                                                    ego_pose_clean)

        # filter lidar
        lidar_np = selected_cav_base['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        # remove points that hit itself
        lidar_np = mask_ego_points(lidar_np)
        # project the lidar to ego space
        # x,y,z in ego space
        projected_lidar = \
            box_utils.project_points_by_matrix_torch(lidar_np[:, :3],
                                                        transformation_matrix)
        if self.kd_flag:
            lidar_np_clean = copy.deepcopy(lidar_np)

        if self.proj_first:
            lidar_np[:, :3] = projected_lidar
            
        lidar_np = mask_points_by_range(lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
        

        processed_lidar = self.pre_processor.preprocess(lidar_np)

        selected_cav_processed.update(
            {'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'projected_lidar': projected_lidar,
             'processed_features': processed_lidar,
             'transformation_matrix': transformation_matrix,
             'transformation_matrix_clean': transformation_matrix_clean})

        if self.kd_flag:
            projected_lidar_clean = \
                box_utils.project_points_by_matrix_torch(lidar_np_clean[:, :3],
                                                    transformation_matrix_clean)
            lidar_np_clean[:, :3] = projected_lidar_clean
            lidar_np_clean = mask_points_by_range(lidar_np_clean,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            selected_cav_processed.update(
                {"projected_lidar_clean": lidar_np_clean}
            )

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        if self.kd_flag:
            teacher_processed_lidar_list = []
        if self.visualize:
            origin_lidar = []

        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            lidar_pose_list.append(ego_dict['lidar_poses']) # ego_dict['lidar_pose'] is np.ndarray [N,6]
            lidar_pose_clean_list.append(ego_dict['lidar_poses_clean'])

            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])

            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])

            if self.kd_flag:
                teacher_processed_lidar_list.append(ego_dict['teacher_processed_lidar'])

            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))


        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)

        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        lidar_pose = torch.from_numpy(np.concatenate(lidar_pose_list, axis=0))
        lidar_pose_clean = torch.from_numpy(np.concatenate(lidar_pose_clean_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'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],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'lidar_pose_clean': lidar_pose_clean,
                                   'lidar_pose': lidar_pose})


        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
        
        if self.kd_flag:
            teacher_processed_lidar_torch_dict = \
                self.pre_processor.collate_batch(teacher_processed_lidar_list)
            output_dict['ego'].update({'teacher_processed_lidar':teacher_processed_lidar_torch_dict})

        if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
            (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
            return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_delay.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Yifan Lu
# Modified: sizhewei @ 2022/10/12
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for early fusion
"""
from collections import OrderedDict
import os
import numpy as np
import torch
import math
import copy

import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np


class IntermediateFusionDatasetAsync(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
            print(f'======= is_ab_regular? {self.is_ab_regular} =======')

        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        # Structure: {scenario_id : {cav_1 : {timestamp1 : {yaml: path,
        # lidar: path, cameras:list of path}}}}
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # at least 1 cav should show up
            cav_list = sorted([x for x in os.listdir(scenario_folder)
                               if os.path.isdir(
                    os.path.join(scenario_folder, x))])
            assert len(cav_list) > 0

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                # save all yaml files to the dictionary
                cav_path = os.path.join(scenario_folder, cav_id)

                # use the frame number as key, the full path as the values
                yaml_files = \
                    sorted([os.path.join(cav_path, x)
                            for x in os.listdir(cav_path) if
                            x.endswith('.yaml')])
                timestamps = self.extract_timestamps(yaml_files)

                if self.k > 0:
                    if j == 0: # ego
                        timestamps = timestamps[self.k:]        # 从第k个往后
                    else:
                        timestamps = timestamps[:-self.k]       # 丢掉最后k个 

                for timestamp in timestamps:
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        camera_files

                    if self.select_keypoint:
                        if (self.proj_first and j == 0) or (not self.proj_first):
                            # when init the dataset, it read over all pcd files.
                            # it maybe slow, but no need to perform keypoint sampling for each time.\
                            kp_file_path = f"{self.params['select_kp']['kp_path']}_{self.params['select_kp']['n_samples']}/{scenario_folders_name[i]}/{cav_id}/{timestamp}.npy"
                            kp_dir = kp_file_path.rsplit('/',1)[0]

                            if not os.path.exists(kp_dir):
                                os.makedirs(kp_dir)

                            if os.path.exists(kp_file_path):
                                keypoints = np.load(kp_file_path)
                            else:
                                pcd_np = pcd_utils.pcd_to_np(lidar_file)
                                bev_keypoints = bev_sample(pcd_np,
                                            self.params['select_kp']['lidar_range'],
                                            self.params['select_kp']['n_samples'],
                                            mode='uniform',
                                            all_samples=False)

                                np.save(kp_file_path, bev_keypoints)
                                print(f"saving {kp_file_path} in basedataset")
                                keypoints = bev_keypoints
                            # sampling again from bev keypoints
                            self.scenario_database[i][cav_id][timestamp]['lidar_keypoints_np'] = keypoints
                                    
                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    if not self.len_record:
                        self.len_record.append(len(timestamps))
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + len(timestamps))
                else:
                    self.scenario_database[i][cav_id]['ego'] = False

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        if "kd_flag" in params.keys():
            self.kd_flag = params['kd_flag']
        else:
            self.kd_flag = False

        print("OPV2V ASync dataset with {} time delay initialized! {} samples totally!".format(self.k, self.len_record[-1]))

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
        """
        # we loop the accumulated length list to see get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]

        data = OrderedDict()
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            # check the timestamp index
            timestamp_index = idx if scenario_index == 0 else \
                idx - self.len_record[scenario_index - 1]
            # retrieve the corresponding timestamp key
            # timestamp_key = self.return_timestamp_key_async(scenario_database, cav_id, 
            #                                         timestamp_index)
            timestamp_key = list(cav_content.items())[timestamp_index][0]
            data[cav_id] = OrderedDict()
            data[cav_id]['ego'] = cav_content['ego']

            # todo: load camera image in the future version
            # load the corresponding data into the dictionary
            data[cav_id]['params'] = \
                load_yaml(cav_content[timestamp_key]['yaml'])
            data[cav_id]['lidar_np'] = \
                pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])
            
            if self.select_keypoint and 'lidar_keypoints_np' in cav_content[timestamp_key]:
                data[cav_id]['lidar_keypoints_np'] = cav_content[timestamp_key]['lidar_keypoints_np']

        return data


    def __getitem__(self, idx):
        base_data_dict = self.retrieve_base_data(idx)

        base_data_dict = add_noise_data_dict(base_data_dict,self.params['noise_setting'])

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        ego_id = -1
        ego_lidar_pose = []

        # first find the ego vehicle's lidar pose
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['params']['lidar_pose']
                ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
                break
            
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0


        processed_features = []
        object_stack = []
        object_id_stack = []
        too_far = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        projected_lidar_clean_list = []
        cav_id_list = []

        if self.visualize:
            projected_lidar_stack = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            distance = \
                math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
                           ego_lidar_pose[0]) ** 2 + (
                                  selected_cav_base['params'][
                                      'lidar_pose'][1] - ego_lidar_pose[
                                      1]) ** 2)

            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue


            lidar_pose_clean_list.append(selected_cav_base['params']['lidar_pose_clean'])
            lidar_pose_list.append(selected_cav_base['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)


        ########## Added by Yifan Lu 2022.8.14 ##############
        # box align to correct pose.
        '''
        if self.box_align and str(idx) in self.stage1_result.keys():
            stage1_content = self.stage1_result[str(idx)]
            if stage1_content is not None:
                cav_id_list_stage1 = stage1_content['cav_id_list']
                
                pred_corners_list = stage1_content['pred_corner3d_np_list']
                pred_corners_list = [np.array(corners, dtype=np.float64) for corners in pred_corners_list]
                uncertainty_list = stage1_content['uncertainty_np_list']
                uncertainty_list = [np.array(uncertainty, dtype=np.float64) for uncertainty in uncertainty_list]
                stage1_lidar_pose_list = [base_data_dict[cav_id]['params']['lidar_pose'] for cav_id in cav_id_list_stage1]
                stage1_lidar_pose = np.array(stage1_lidar_pose_list)


                refined_pose = box_alignment_relative_sample_np(pred_corners_list,
                                                                stage1_lidar_pose, 
                                                                uncertainty_list=uncertainty_list, 
                                                                **self.box_align_args)
                stage1_lidar_pose[:,[0,1,4]] = refined_pose
                stage1_lidar_pose_refined_list = stage1_lidar_pose.tolist() # updated lidar_pose_list
                for cav_id, lidar_pose_refined in zip(cav_id_list_stage1, stage1_lidar_pose_refined_list):
                    if cav_id not in cav_id_list:
                        continue
                    idx_in_list = cav_id_list.index(cav_id)
                    lidar_pose_list[idx_in_list] = lidar_pose_refined
                    base_data_dict[cav_id]['params']['lidar_pose'] = lidar_pose_refined
        '''     

        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]


            # obj: 
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose, 
                ego_lidar_pose_clean,
                idx)
                
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            processed_features.append(
                selected_cav_processed['processed_features'])
            if self.kd_flag:
                projected_lidar_clean_list.append(
                    selected_cav_processed['projected_lidar_clean'])

            if self.visualize:
                projected_lidar_stack.append(
                    selected_cav_processed['projected_lidar'])


        ########## Added by Yifan Lu 2022.4.5 ################
        # filter those out of communicate range
        # then we can calculate get_pairwise_transformation
        for cav_id in too_far:
            base_data_dict.pop(cav_id)
        
        pairwise_t_matrix = \
            self.get_pairwise_transformation(base_data_dict,
                                             self.max_cav)

        lidar_poses = np.array(lidar_pose_list).reshape(-1, 6)  # [N_cav, 6]
        lidar_poses_clean = np.array(lidar_pose_clean_list).reshape(-1, 6)  # [N_cav, 6]
        ######################################################

        ############ for disconet ###########
        if self.kd_flag:
            stack_lidar_np = np.vstack(projected_lidar_clean_list)
            stack_lidar_np = mask_points_by_range(stack_lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            stack_feature_processed = self.pre_processor.preprocess(stack_lidar_np)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        object_stack = np.vstack(object_stack)
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(processed_features)

        merged_feature_dict = self.merge_features_to_dict(processed_features)

        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        processed_data_dict['ego'].update(
            {'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'lidar_poses_clean': lidar_poses_clean,
             'lidar_poses': lidar_poses})

        if self.kd_flag:
            processed_data_dict['ego'].update({'teacher_processed_lidar':
                stack_feature_processed})

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(
                    projected_lidar_stack)})


        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, ego_pose_clean, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information.
        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.
        ego_pose_clean : list, length 6
            only used for gt box generation
        
        idx: int,
            debug use.

        Returns
        -------
        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) # T_ego_cav
        transformation_matrix_clean = \
            x1_to_x2(selected_cav_base['params']['lidar_pose_clean'],
                     ego_pose_clean)

        # retrieve objects under ego coordinates
        # this is used to generate accurate GT bounding box.
        object_bbx_center, object_bbx_mask, object_ids = self.generate_object_center([selected_cav_base],
                                                    ego_pose_clean)

        # filter lidar
        lidar_np = selected_cav_base['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        # remove points that hit itself
        lidar_np = mask_ego_points(lidar_np)
        # project the lidar to ego space
        # x,y,z in ego space
        projected_lidar = \
            box_utils.project_points_by_matrix_torch(lidar_np[:, :3],
                                                        transformation_matrix)
        if self.kd_flag:
            lidar_np_clean = copy.deepcopy(lidar_np)

        if self.proj_first:
            lidar_np[:, :3] = projected_lidar
            
        lidar_np = mask_points_by_range(lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
        

        processed_lidar = self.pre_processor.preprocess(lidar_np)

        selected_cav_processed.update(
            {'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'projected_lidar': projected_lidar,
             'processed_features': processed_lidar,
             'transformation_matrix': transformation_matrix,
             'transformation_matrix_clean': transformation_matrix_clean})

        if self.kd_flag:
            projected_lidar_clean = \
                box_utils.project_points_by_matrix_torch(lidar_np_clean[:, :3],
                                                    transformation_matrix_clean)
            lidar_np_clean[:, :3] = projected_lidar_clean
            lidar_np_clean = mask_points_by_range(lidar_np_clean,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            selected_cav_processed.update(
                {"projected_lidar_clean": lidar_np_clean}
            )

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        if self.kd_flag:
            teacher_processed_lidar_list = []
        if self.visualize:
            origin_lidar = []

        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            lidar_pose_list.append(ego_dict['lidar_poses']) # ego_dict['lidar_pose'] is np.ndarray [N,6]
            lidar_pose_clean_list.append(ego_dict['lidar_poses_clean'])

            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])

            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])

            if self.kd_flag:
                teacher_processed_lidar_list.append(ego_dict['teacher_processed_lidar'])

            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))


        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)

        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        lidar_pose = torch.from_numpy(np.concatenate(lidar_pose_list, axis=0))
        lidar_pose_clean = torch.from_numpy(np.concatenate(lidar_pose_clean_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'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],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'lidar_pose_clean': lidar_pose_clean,
                                   'lidar_pose': lidar_pose})


        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
        
        if self.kd_flag:
            teacher_processed_lidar_torch_dict = \
                self.pre_processor.collate_batch(teacher_processed_lidar_list)
            output_dict['ego'].update({'teacher_processed_lidar':teacher_processed_lidar_torch_dict})

        if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
            (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
            return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_delay_new.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/1/27
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for intermediate fusion with past k frames on irregular OPV2V
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetIrregular(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):

        self.times = []

        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
            print(f'======= is_ab_regular? {self.is_ab_regular} =======')
        
        self.sample_interval_exp = int(self.binomial_n * self.binomial_p)

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        '''
        scenario_database Structure: 
        {
            scenario_id : {
                cav_1 : {
                    'ego' : true / false , 
                    timestamp1 : {
                        yaml: path,
                        lidar: path, 
                        cameras: list of path
                    },
                    ...
                },
                ...
            }
        }
        '''
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # copy timestamps npy file
            timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
            time_annotations = np.load(timestamps_file)

            # at least 1 cav should show up
            cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
            assert len(cav_list) > 0

            # use the frame number as key, the full path as the values
            yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
            if len(yaml_files)==0:
                yaml_files = sorted([x 
                            for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                            x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

            start_timestamp = int(float(self.extract_timestamps(yaml_files)[0]))
            while(1):
                time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
                time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
                if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
                    start_timestamp += 1
                else:
                    break

            end_timestamp = int(float(self.extract_timestamps(yaml_files)[-1]))
            if start_timestamp%2 == 0:
                # even
                end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
            else:
                end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
            num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)
            regular_timestamps = [start_timestamp+2*i for i in range(num_timestamps)]

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                cav_path = os.path.join(scenario_folder, cav_id)

                if j==0: # ego
                    timestamps = regular_timestamps
                else:
                    if self.is_ab_regular:
                        timestamps = regular_timestamps
                    else:
                        timestamps = list(time_annotations[j-1, :])

                for timestamp in timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files
                
                # regular的timestamps 用于做 curr 真实时刻的ground truth
                self.scenario_database[i][cav_id]['regular'] = OrderedDict()
                for timestamp in regular_timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id]['regular'][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id]['regular'][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id]['regular'][timestamp]['lidar'] = \
                        lidar_file

                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    # num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    num_ego_timestamps = len(timestamps) - self.binomial_n * self.k
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)
                else:
                    self.scenario_database[i][cav_id]['ego'] = False
                    

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead. 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("=== OPV2V-Irregular Multi-sweep dataset with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def extract_timestamps(self, yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]
            if res.endswith('.yaml'):
                timestamp = res.replace('.yaml', '')
            elif res.endswith('.json'):
                timestamp = res.replace('.json', '')
            else:
                print("Woops! There is no processing method for file {}".format(res))
                sys.exit(1)
            timestamps.append(timestamp)

        return timestamps

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)
                        ...,		
                        [k-1] : {} (id-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k-1]:{}  (id-\tau-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * self.k
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # TODO: debug use, uncomment: to use regular version yaml GT
            # time_new = str( '%06d' % int(float(timestamp_key)))
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # TODO: debug use, uncomment: to use regular version lidar input
            # time_new = str( '%06d' % int(float(timestamp_key))) 
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                # sample_interval
                if data[cav_id]['ego']:             # ego sample_interval = E(B(n, p))
                    if i == 0: # ego-past-0 与 ego-curr 是一样的
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    sample_interval = self.sample_interval_exp
                    if sample_interval == 0:
                        sample_interval = 1
                else:                               # non-ego sample_interval ~ B(n, p)
                    if self.sample_interval_exp==0 \
                        and self.is_no_shift \
                            and i == 0:
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==0: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>0: # 过去的几帧不要重复
                            sample_interval = 1                

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float

        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            },
            cav_id_2 : { ... }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = math.sqrt( \
                (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换, 用于单体监督            cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: N, k
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) TODO: 这里面没有搞懂为什么不用 past_k_tr_mats

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack})
            #  'times': self.times})

        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            processed_data_dict['ego'].update({'avg_sample_interval':\
                sum(past_k_sample_interval_stack[self.k:]) / len(past_k_sample_interval_stack[self.k:])})
            processed_data_dict['ego'].update({'avg_time_delay':\
                sum(past_k_time_diffs_stack[self.k:]) / len(past_k_time_diffs_stack[self.k:])})
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            '''
            # 5. single view label
            # past_i label at past_i single view
            # opencood/data_utils/post_processor/base_postprocessor.py
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # generate the anchor boxes
            # opencood/data_utils/post_processor/voxel_postprocessor.py
            anchor_box = self.anchor_box
            single_view_label_dict = self.post_processor.generate_label(
                    gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
                )
            past_k_label_dicts.append(single_view_label_dict)
            '''

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
             'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
            #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
            #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
            #  'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def dist_time(ts1, ts2, i = -1):
        """caculate the time interval between two timestamps

        Args:
            ts1 (string): time stamp at some time
            ts2 (string): current time stamp
            i (int, optional): past frame id, for debug use. Defaults to -1.
        
        Returns:
            time_diff (float): time interval (ts1 - ts2)
        """
        if not i==-1:
            return -i
        else:
            return (float(ts1) - float(ts2))

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay})
                                #    'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_irregular.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/2/13

"""
Dataset class for intermediate fusion with 
1) past k frames as input 
2) current frame as supervision
on irregular OPV2V
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset, intermediate_fusion_dataset_opv2v_irregular
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetIrregularCompensation(intermediate_fusion_dataset_opv2v_irregular.IntermediateFusionDatasetIrregular):
    """
    This class is for intermediate fusion on OPV2V_irregular, with past k frames and current frame as input. 
    Note: current frame is used for supervison, not for feature fusion!
    """
    def __init__(self, params, visualize, train=True):
        super(IntermediateFusionDatasetIrregularCompensation, self).__init__(params, visualize, train)
        self.k = self.k + 1 # 用来在 past_k 中添加 current frame
        print("=== OPV2V-Irregular Multi-sweep dataloader (Compensation Ver.) for intermediate fusion with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k-1, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{                       # current
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)          # past 1
                        ...,		
                        [k] : {} (id-(k-1))         # past k
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k] : {}
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * (self.k-1)
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            try:
                timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            except IndexError:
                print('IndexError!!!')
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # time_new = str( '%06d' % int(float(timestamp_key))) # TODO: debug用 使用regular的GT
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0]) # TODO:
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # time_new = str( '%06d' % int(float(timestamp_key))) # TODO: debug用 使用regular的lidar
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                if i == 0: # ego-past-0 与 ego-curr 是一样的, 用于feature监督
                    data[cav_id]['past_k'][i] = data[cav_id]['curr']
                    continue
                # sample_interval
                if data[cav_id]['ego']:             # ego: do not need past frame
                    break
                    # sample_interval = self.sample_interval_exp
                else:                               # non-ego sample_interval ~ B(n, p)
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==0: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>0: # 过去的几帧不要重复
                            sample_interval = 1  

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                # uncomment code, use updated gt
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float

        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (1) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    }
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
                
            }, 
            cav_id_2 : {
                'ego' : false,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)          # past 1
                    ...,		                # ...
                    [k] : {} (id-(k-1))         # past k
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            if cav_id != ego_id:
                distance = math.sqrt( \
                    (selected_cav_base['past_k'][1]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                        (selected_cav_base['past_k'][1]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
                # if distance is too far, we will just skip this agent
                if distance > self.params['comm_range']:
                    too_far.append(cav_id)
                    continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true / false , 
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (1 / k) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)          # past 1 
                    ...,		                # ...
                    [k] : {} (id-(k-1))         # past k
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k+1, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k+1, 6
            past_k_pose_stack += selected_cav_processed['past_k_poses']
            # time differences: Nx(k+1)
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: Nx(k+1)
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k+1, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) 

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, 6)  # (Nx(k-1)+1, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack})
            #  'times': self.times})

        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            tmp = past_k_time_diffs_stack[1:].reshape(-1, self.k)
            tmp_past_k_time_diffs = tmp[:, 1:] - tmp[:, :-1]
            avg_time_diff = sum(tmp_past_k_time_diffs.reshape(-1)) / tmp_past_k_time_diffs.reshape(-1).shape[0]
            processed_data_dict['ego'].update({'avg_time_delay':\
                avg_time_diff})
            
            tmp_past_k_sample_interval = past_k_sample_interval_stack[1:].reshape(-1, self.k)[:, 1:]
            avg_past_k_sample_interval = sum(tmp_past_k_sample_interval.reshape(-1)) / tmp_past_k_sample_interval.reshape(-1).shape[0]
            processed_data_dict['ego'].update({'avg_sample_interval':\
                avg_past_k_sample_interval})
            
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict

    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        if selected_cav_base['ego']:
            k = 1
        else:
            k = self.k
        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            '''
            # 5. single view label
            # past_i label at past_i single view
            # opencood/data_utils/post_processor/base_postprocessor.py
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # generate the anchor boxes
            # opencood/data_utils/post_processor/voxel_postprocessor.py
            anchor_box = self.anchor_box
            single_view_label_dict = self.post_processor.generate_label(
                    gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
                )
            past_k_label_dicts.append(single_view_label_dict)
            '''

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
             'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
            #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
            #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
            #  'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        return selected_cav_processed

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(len(processed_feature_list[cav_id])):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict


    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_label':          single_label_dict_stack,
            'curr_processed_lidar':         merged_curr_feature_dict,
            'object_bbx_center':            object_bbx_center,
            'object_bbx_mask':              mask,
            # 'anchor_box':                   anchor_box,
            'processed_lidar':              merged_feature_dict,
            'label_dict':                   label_dict,
            'object_ids':                   [object_id_stack[i] for i in unique_indices],
            'single_past_dict':             
            'pairwise_t_matrix':            pairwise_t_matrix,
            'cav_num':                      cav_num,
            'curr_lidar_poses':             curr_lidar_poses,
            'past_lidar_poses':             past_k_lidar_poses,
            'past_k_time_interval':         
            'past_k_sample_interval':   
            'avg_sample_interval': , 
            'avg_time_delay':   
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [(N-1)xk+1,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)
        # TODO: 计算方式需要重新改一下
        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay})
                                #    'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    if cav_content['ego']:
                        past_k_poses.append(x_to_world(cav_content['past_k'][0]['params']['lidar_pose']))
                        continue
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_irregular_compensation.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/1/27
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for intermediate fusion with past k frames on irregular OPV2V
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
from opencood.utils.flow_utils import generate_flow_map

from opencood.utils.box_utils import boxes_to_corners2d # for debug use
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetIrregularFlow(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):

        self.times = []

        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
        
        self.sample_interval_exp = int(self.binomial_n * self.binomial_p)

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        '''
        scenario_database Structure: 
        {
            scenario_id : {
                cav_1 : {
                    'ego' : true / false , 
                    timestamp1 : {
                        yaml: path,
                        lidar: path, 
                        cameras: list of path
                    },
                    ...
                },
                ...
            }
        }
        '''
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # copy timestamps npy file
            timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
            time_annotations = np.load(timestamps_file)

            # at least 1 cav should show up
            cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
            assert len(cav_list) > 0

            # use the frame number as key, the full path as the values
            yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
            if len(yaml_files)==0:
                yaml_files = sorted([x 
                            for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                            x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

            start_timestamp = int(float(self.extract_timestamps(yaml_files)[0]))
            while(1):
                time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
                time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
                if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
                    start_timestamp += 1
                else:
                    break

            end_timestamp = int(float(self.extract_timestamps(yaml_files)[-1]))
            if start_timestamp%2 == 0:
                # even
                end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
            else:
                end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
            num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)
            regular_timestamps = [start_timestamp+2*i for i in range(num_timestamps)]

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                cav_path = os.path.join(scenario_folder, cav_id)

                if j==0: # ego
                    timestamps = regular_timestamps
                else:
                    timestamps = list(time_annotations[j-1, :])

                for timestamp in timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files
                
                # regular的timestamps 用于做 curr 真实时刻的ground truth
                self.scenario_database[i][cav_id]['regular'] = OrderedDict()
                for timestamp in regular_timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id]['regular'][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id]['regular'][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id]['regular'][timestamp]['lidar'] = \
                        lidar_file

                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    # num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    num_ego_timestamps = len(timestamps) - self.binomial_n * self.k
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)
                else:
                    self.scenario_database[i][cav_id]['ego'] = False
                    

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead. 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("=== OPV2V-Irregular Multi-sweep dataset with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def extract_timestamps(self, yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]
            if res.endswith('.yaml'):
                timestamp = res.replace('.yaml', '')
            elif res.endswith('.json'):
                timestamp = res.replace('.json', '')
            else:
                print("Woops! There is no processing method for file {}".format(res))
                sys.exit(1)
            timestamps.append(timestamp)

        return timestamps

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)
                        ...,		
                        [k-1] : {} (id-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k-1]:{}  (id-\tau-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * self.k
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # TODO: debug use, uncomment: to use regular version yaml GT
            # time_new = str( '%06d' % int(float(timestamp_key)))
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # TODO: debug use, uncomment: to use regular version lidar input
            # time_new = str( '%06d' % int(float(timestamp_key))) 
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                # sample_interval
                if data[cav_id]['ego']:             # ego sample_interval = E(B(n, p))
                    if i == 0: # ego-past-0 与 ego-curr 是一样的
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    sample_interval = self.sample_interval_exp
                    if sample_interval == 0:
                        sample_interval = 1
                else:                               # non-ego sample_interval ~ B(n, p)
                    if self.sample_interval_exp==0 \
                        and self.is_no_shift \
                            and i == 0:
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==0: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>0: # 过去的几帧不要重复
                            sample_interval = 1                

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float

        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            },
            cav_id_2 : { ... }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = math.sqrt( \
                (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        flow_gt = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换, 用于单体监督            cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
                'flow_gt':               # [2, H, W] flow ground truth
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: N, k
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        
            # for flow
            flow_gt.append(selected_cav_processed['flow_gt'])
        
        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) TODO: 这里面没有搞懂为什么不用 past_k_tr_mats

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack})
            #  'times': self.times})

        flow_gt = np.vstack(flow_gt) # (N, 2, H, W)
        processed_data_dict['ego'].update({'flow_gt': flow_gt})
        
        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            processed_data_dict['ego'].update({'avg_sample_interval':\
                sum(past_k_sample_interval_stack[self.k:]) / len(past_k_sample_interval_stack[self.k:])})
            processed_data_dict['ego'].update({'avg_time_delay':\
                sum(past_k_time_diffs_stack[self.k:]) / len(past_k_time_diffs_stack[self.k:])})
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'flow_gt':              # single view flow
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            '''
            # 5. single view label
            # past_i label at past_i single view
            # opencood/data_utils/post_processor/base_postprocessor.py
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # generate the anchor boxes
            # opencood/data_utils/post_processor/voxel_postprocessor.py
            anchor_box = self.anchor_box
            single_view_label_dict = self.post_processor.generate_label(
                    gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
                )
            past_k_label_dicts.append(single_view_label_dict)
            '''
        
        

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
             'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
            #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
            #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
            #  'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        # generate flow, from past_0 and curr
        prev_object_id_stack = {}
        prev_object_stack = {}
        for t_i in range(2):
            split_part = selected_cav_base['past_k'][0] if t_i == 0 else selected_cav_base['curr']
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([split_part], selected_cav_base['past_k'][0]['params']['lidar_pose'])
            prev_object_id_stack[t_i] = object_ids
            prev_object_stack[t_i] = object_bbx_center
        
        for t_i in range(2):
            unique_object_ids = list(set(prev_object_id_stack[t_i]))
            unique_indices = \
                [prev_object_id_stack[t_i].index(x) for x in unique_object_ids]
            prev_object_stack[t_i] = np.vstack(prev_object_stack[t_i])
            prev_object_stack[t_i] = prev_object_stack[t_i][unique_indices]
            prev_object_id_stack[t_i] = unique_object_ids

        # TODO: for debug use
        # previous_cood = prev_object_stack[0][0]
        # # test_prev_box_2d = boxes_to_corners2d(np.expand_dims(previous_cood, axis=0), order='hwl')
        # # print('test_prev_box_2d', test_prev_box_2d)
        # #   [[[52.839737   5.977605  -1.9081762] 
        # #   [52.84148    7.9091973 -1.9081762]
        # #   [48.236977   7.913352  -1.9081762]
        # #   [48.235233   5.9817595 -1.9081762]]]
        # current_cood = previous_cood.copy()
        # current_cood[0] += 100
        # current_cood[1] += 30
        # object_id = prev_object_id_stack[0][0]
        # prev_object_id_stack = {}
        # prev_object_stack = {}
        # for i in range(2):
        #     prev_object_id_stack[i] = [object_id]
        #     prev_object_stack[i] = [previous_cood] if i == 0 else [current_cood]
        ########################################
        
        flow_map = generate_flow_map(prev_object_stack,
                                        prev_object_id_stack,
                                        self.params['preprocess']['cav_lidar_range'],
                                        self.params['preprocess']['args']['voxel_size'],
                                        past_k=1)

        # # TODO: debug use
        # debug_path = '/remote-home/share/sizhewei/logs/where2comm_flow_debug/viz_flow/'
        # torch.save(flow_map, debug_path + 'flow_map_true.pt')
        # torch.save(prev_object_stack, debug_path + 'bbx_0.pt')
        # ########################################

        selected_cav_processed.update({'flow_gt': flow_map})

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def dist_time(ts1, ts2, i = -1):
        """caculate the time interval between two timestamps

        Args:
            ts1 (string): time stamp at some time
            ts2 (string): current time stamp
            i (int, optional): past frame id, for debug use. Defaults to -1.
        
        Returns:
            time_diff (float): time interval (ts1 - ts2)
        """
        if not i==-1:
            return -i
        else:
            return (float(ts1) - float(ts2))

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'flow_gt': [N, C, H, W]
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []
        # flow gt
        flow_gt_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            flow_gt_list.append(ego_dict['flow_gt'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for flow
        flow_gt = torch.from_numpy(np.vstack(flow_gt_list))
        label_torch_dict.update({'flow_gt': flow_gt})

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay})
                                #    'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_irregular_flow.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/2/13

"""
Dataset class for intermediate fusion with 
1) past k frames as input 
2) current frame as supervision
3) gt flow
on irregular OPV2V
This dataloader is just for debug, to viz :
    past_0 feature + flow map == curr feature ?
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset, intermediate_fusion_dataset_opv2v_irregular
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
from opencood.utils.flow_utils import generate_flow_map

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetIrregularFlowCompensation(intermediate_fusion_dataset_opv2v_irregular.IntermediateFusionDatasetIrregular):
    """
    This class is for intermediate fusion on OPV2V_irregular, with past k frames and current frame as input. 
    Note: current frame is used for supervison, not for feature fusion!
    """
    def __init__(self, params, visualize, train=True):
        super(IntermediateFusionDatasetIrregularFlowCompensation, self).__init__(params, visualize, train)
        self.k = self.k + 1 # 用来在 past_k 中添加 current frame
        print("=== OPV2V-Irregular Multi-sweep dataloader (Flow Compensation Ver.) for intermediate fusion with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k-1, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{                       # current
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)          # past 1
                        ...,		
                        [k] : {} (id-(k-1))         # past k
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k] : {}
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * (self.k-1)
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            try:
                timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            except IndexError:
                print('IndexError!!!')
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # time_new = str( '%06d' % int(float(timestamp_key))) # TODO: debug用 使用regular的GT
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0]) # TODO:
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # time_new = str( '%06d' % int(float(timestamp_key))) # TODO: debug用 使用regular的lidar
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                # ego-past-0 与 ego-curr 是一样的, 用于feature监督
                if i == 0: 
                    data[cav_id]['past_k'][i] = data[cav_id]['curr']
                    continue
                
                # sample_interval
                if data[cav_id]['ego']:             # ego: do not need past frame
                    if i == 1: # ego-past-1 与 ego-curr 是一样的
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    sample_interval = self.sample_interval_exp
                    if sample_interval == 0:
                        sample_interval = 1
                else:                               # non-ego sample_interval ~ B(n, p)
                    if self.sample_interval_exp==0 \
                        and self.is_no_shift \
                            and i == 1:
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==1: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>1: # 过去的几帧不要重复
                            sample_interval = 1  

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                # uncomment code, use updated gt
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float

        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (1) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    }
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
                
            }, 
            cav_id_2 : {
                'ego' : false,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)          # past 1
                    ...,		                # ...
                    [k] : {} (id-(k-1))         # past k
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            if cav_id != ego_id:
                distance = math.sqrt( \
                    (selected_cav_base['past_k'][1]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                        (selected_cav_base['past_k'][1]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
                # if distance is too far, we will just skip this agent
                if distance > self.params['comm_range']:
                    too_far.append(cav_id)
                    continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        flow_gt = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true / false , 
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)          # past 1 
                    ...,		                # ...
                    [k] : {} (id-(k-1))         # past k
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k+1, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k+1, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: Nx(k+1)
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: Nx(k+1)
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k+1, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])

            # for flow
            flow_gt.append(selected_cav_processed['flow_gt'])
        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) 

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (Nx(k-1)+1, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack})
            #  'times': self.times})

        flow_gt = np.vstack(flow_gt) # (N, 2, H, W)
        processed_data_dict['ego'].update({'flow_gt': flow_gt})

        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            processed_data_dict['ego'].update({'avg_sample_interval':\
                sum(past_k_sample_interval_stack[self.k:]) / len(past_k_sample_interval_stack[self.k:])})
            processed_data_dict['ego'].update({'avg_time_delay':\
                sum(past_k_time_diffs_stack[self.k:]) / len(past_k_time_diffs_stack[self.k:])})
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict

    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        k = self.k
        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            '''
            # 5. single view label
            # past_i label at past_i single view
            # opencood/data_utils/post_processor/base_postprocessor.py
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # generate the anchor boxes
            # opencood/data_utils/post_processor/voxel_postprocessor.py
            anchor_box = self.anchor_box
            single_view_label_dict = self.post_processor.generate_label(
                    gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
                )
            past_k_label_dicts.append(single_view_label_dict)
            '''

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
             'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
            #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
            #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
            #  'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        # generate flow, from past_1 and curr
        prev_object_id_stack = {}
        prev_object_stack = {}
        for t_i in range(2):
            split_part = selected_cav_base['past_k'][1] if t_i == 0 else selected_cav_base['curr']
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([split_part], selected_cav_base['past_k'][1]['params']['lidar_pose'])
            prev_object_id_stack[t_i] = object_ids
            prev_object_stack[t_i] = object_bbx_center
        
        for t_i in range(2):
            unique_object_ids = list(set(prev_object_id_stack[t_i]))
            unique_indices = \
                [prev_object_id_stack[t_i].index(x) for x in unique_object_ids]
            prev_object_stack[t_i] = np.vstack(prev_object_stack[t_i])
            prev_object_stack[t_i] = prev_object_stack[t_i][unique_indices]
            prev_object_id_stack[t_i] = unique_object_ids

        flow_map = generate_flow_map(prev_object_stack,
                                        prev_object_id_stack,
                                        self.params['preprocess']['cav_lidar_range'],
                                        self.params['preprocess']['args']['voxel_size'],
                                        past_k=1)

        selected_cav_processed.update({'flow_gt': flow_map})

        return selected_cav_processed

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(len(processed_feature_list[cav_id])):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict


    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_label':          single_label_dict_stack,
            'curr_processed_lidar':         merged_curr_feature_dict,
            'object_bbx_center':            object_bbx_center,
            'object_bbx_mask':              mask,
            # 'anchor_box':                   anchor_box,
            'processed_lidar':              merged_feature_dict,
            'label_dict':                   label_dict,
            'object_ids':                   [object_id_stack[i] for i in unique_indices],
            'single_past_dict':             
            'pairwise_t_matrix':            pairwise_t_matrix,
            'cav_num':                      cav_num,
            'curr_lidar_poses':             curr_lidar_poses,
            'past_lidar_poses':             past_k_lidar_poses,
            'past_k_time_interval':         
            'past_k_sample_interval':   
            'avg_sample_interval': , 
            'avg_time_delay':   
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []
        # flow gt
        flow_gt_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [(N-1)xk+1,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            flow_gt_list.append(ego_dict['flow_gt'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)
        # TODO: 计算方式需要重新改一下
        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for flow
        flow_gt = torch.from_numpy(np.vstack(flow_gt_list))
        label_torch_dict.update({'flow_gt': flow_gt})

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay})
                                #    'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_irregular_flow_compensation.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/4/15
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
two stage flow update framework

"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
from opencood.utils.flow_utils import generate_flow_map, generate_flow_map_szwei

from opencood.utils.box_utils import boxes_to_corners2d # for debug use
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetIrregularFlowNew(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):

        self.times = []

        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
            print("!!! Absolutely Regular !!!")

        # 控制是否需要生成GT flow
        self.is_generate_gt_flow = False
        if 'is_generate_gt_flow' in params and params['is_generate_gt_flow']:
            self.is_generate_gt_flow = True
        
        # 只有在绘制每个sample的匹配框时用到 sizhewei
        self.viz_bbx_flag = False
        if 'viz_bbx_flag' in params and params['viz_bbx_flag']:
            self.viz_bbx_flag = True

        self.num_roi_thres = -1
        if 'num_roi_thres' in params:
            self.num_roi_thres = params['num_roi_thres']
        
        self.sample_interval_exp = int(self.binomial_n * self.binomial_p)

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        '''
        scenario_database Structure: 
        {
            scenario_id : {
                cav_1 : {
                    'ego' : true / false , 
                    timestamp1 : {
                        yaml: path,
                        lidar: path, 
                        cameras: list of path
                    },
                    ...
                },
                ...
            }
        }
        '''
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # copy timestamps npy file, store all irregular timestamps of each non-ego vehicle in this scenario
            timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
            time_annotations = np.load(timestamps_file)

            # at least 1 cav should show up
            cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
            assert len(cav_list) > 0

            # use the frame number as key, the full path as the values, store all json or yaml files in this scenario
            yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
            if len(yaml_files)==0:
                yaml_files = sorted([x 
                            for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                            x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

            start_timestamp = int(float(self.extract_timestamps(yaml_files)[0]))
            while(1):
                time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
                time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
                if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
                    start_timestamp += 1
                else:
                    break

            end_timestamp = int(float(self.extract_timestamps(yaml_files)[-1]))
            if start_timestamp%2 == 0:
                # even
                end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
            else:
                end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
            num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)
            regular_timestamps = [start_timestamp+2*i for i in range(num_timestamps)]

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                cav_path = os.path.join(scenario_folder, cav_id)

                if j==0: # ego
                    timestamps = regular_timestamps
                else:
                    if self.is_ab_regular:
                        timestamps = regular_timestamps
                    else:
                        timestamps = list(time_annotations[j-1, :])

                for timestamp in timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files
                
                # regular的timestamps 用于做 curr 真实时刻的ground truth
                self.scenario_database[i][cav_id]['regular'] = OrderedDict()
                for timestamp in regular_timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id]['regular'][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id]['regular'][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id]['regular'][timestamp]['lidar'] = \
                        lidar_file

                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    # num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    num_ego_timestamps = len(timestamps) - self.binomial_n * self.k #* 3 # TODO: 
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)
                else:
                    self.scenario_database[i][cav_id]['ego'] = False
                    

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead. 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("=== OPV2V-Irregular Multi-sweep dataset with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def extract_timestamps(self, yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]
            if res.endswith('.yaml'):
                timestamp = res.replace('.yaml', '')
            elif res.endswith('.json'):
                timestamp = res.replace('.json', '')
            else:
                print("Woops! There is no processing method for file {}".format(res))
                sys.exit(1)
            timestamps.append(timestamp)

        return timestamps

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)
                        ...,		
                        [k-1] : {} (id-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k-1]:{}  (id-\tau-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * self.k #*3 # TODO:
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # TODO: debug use, uncomment: to use regular version yaml GT
            # time_new = str( '%06d' % int(float(timestamp_key)))
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # TODO: debug use, uncomment: to use regular version lidar input
            # time_new = str( '%06d' % int(float(timestamp_key))) 
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                # sample_interval
                if data[cav_id]['ego']:             # ego sample_interval = E(B(n, p))
                    if i == 0: # ego-past-0 与 ego-curr 是一样的
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    sample_interval = self.sample_interval_exp
                    if sample_interval == 0:
                        sample_interval = 1
                else:                               # non-ego sample_interval ~ B(n, p)
                    # delay_debug = 6
                    if i == 0:
                        sample_interval = 3 #delay_debug
                    elif i ==1:
                        sample_interval = 3
                        # sample_set = [2,3]
                        # import random
                        # sample_interval = random.sample(sample_set, 1)[0] #3 #10 - delay_debug
                        # trails = bernoulliDist.rvs(self.binomial_n)
                        # sample_interval = sum(trails)
                    else:
                        sample_interval = 3
                        # trails = bernoulliDist.rvs(self.binomial_n)
                        # sample_interval = sum(trails)
                    # if self.sample_interval_exp==0 \
                    #     and self.is_no_shift \
                    #         and i == 0:
                    #     data[cav_id]['past_k'][i] = data[cav_id]['curr']
                    #     continue
                    # if self.is_same_sample_interval:
                    #     sample_interval = self.sample_interval_exp
                    # else:
                    #     # B(n, p)
                    #     trails = bernoulliDist.rvs(self.binomial_n)
                    #     sample_interval = sum(trails)
                    # if sample_interval==0:
                    #     if i==0: # 检查past 0 的实际时间是否在curr 的后面
                    #         tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                    #         if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                    #             sample_interval = 1
                    #     if i>0: # 过去的几帧不要重复
                    #         sample_interval = 1                

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float
        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            },
            cav_id_2 : { ... }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = math.sqrt( \
                (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        pastk_2_past0_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        if self.is_generate_gt_flow:
            flow_gt = []
            warp_mask = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []

        if self.viz_bbx_flag:
            single_lidar_stack = []
            single_object_stack = []
            single_object_id_stack = []
            single_mask_stack = []
            single_past_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换, 用于单体监督            cav view + curr 的label
                'single_object_bbx_center'
                'single_object_ids'
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
                'flow_gt':               # [2, H, W] flow ground truth
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            if self.viz_bbx_flag:
                single_lidar_stack.append(selected_cav_processed['single_lidar'])
                single_object_stack.append(selected_cav_processed['single_object_bbx_center'])
                single_object_id_stack.append(selected_cav_processed['single_object_ids'])
                single_past_lidar_stack.append(selected_cav_processed['single_past_lidar'])
                # mask = np.zeros(self.params['postprocess']['max_num'])
                # mask[:single_object_stack.shape[0]] = 1
                # single_mask_stack.append(mask)
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: N, k
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k frames to cav past 0 frame trans matrix, list of len=N, pastk_2_past0_tr_mats[i]: ndarray(k, 4, 4)
            pastk_2_past0_tr_mats.append(selected_cav_processed['pastk_2_past0_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        
            if self.is_generate_gt_flow:
                # for flow
                flow_gt.append(selected_cav_processed['flow_gt'])
                warp_mask.append(selected_cav_processed['warp_mask'])
        
        pastk_2_past0_tr_mats = np.stack(pastk_2_past0_tr_mats, axis=0) # N, k, 4, 4

        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) TODO: 这里面没有搞懂为什么不用 past_k_tr_mats

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack, 
             'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats})
            #  'times': self.times})

        if self.is_generate_gt_flow:
            flow_gt = np.vstack(flow_gt) # (N, H, W, 2)
            processed_data_dict['ego'].update({'flow_gt': flow_gt})
            warp_mask = np.vstack(warp_mask) # (N, C, H, W)
            processed_data_dict['ego'].update({'warp_mask': warp_mask})
        
        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            tmp = past_k_time_diffs_stack[self.k:].reshape(-1, self.k) # (N, k)
            
            # tmp_past_k_time_diffs = np.concatenate((tmp[:, :1] , (tmp[:, 1:] - tmp[:, :-1])), axis=1) # (N, k) # TODO: 不同的方法这个计算方式不一样
            tmp_past_k_time_diffs = tmp[:, :1] # (N, 1) irregular setting 下，只用最近的一个时间间隔
            
            avg_time_diff = sum(tmp_past_k_time_diffs.reshape(-1)) / tmp_past_k_time_diffs.reshape(-1).shape[0]
            processed_data_dict['ego'].update({'avg_time_delay':\
                avg_time_diff})

            tmp = past_k_sample_interval_stack[self.k:].reshape(-1, self.k) # (N, k)
            # avg_sample_interval = sum(tmp.reshape(-1)) / len(tmp.reshape(-1)) # TODO: 不同的方法这个计算方式不一样
            avg_sample_interval = sum(tmp[:, :1].reshape(-1)) / len(tmp[:, :1].reshape(-1)) # (N, 1) irregular setting 下，只用最近的一个时间间隔
            processed_data_dict['ego'].update({'avg_sample_interval':\
                avg_sample_interval})

            tmp_var = np.var(tmp, axis=1) # (N, )
            avg_var = np.mean(tmp_var)
            processed_data_dict['ego'].update({'avg_var':\
                avg_var})
            
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        if self.viz_bbx_flag:
            for id, cav in enumerate(cav_id_list):
                processed_data_dict[id] = {}
                processed_data_dict[id].update({
                    'single_lidar': single_lidar_stack[id],
                    'single_object_bbx_center': single_object_stack[id],
                    # 'single_object_bbx_mask': single_mask_stack[i],
                    'single_object_ids': single_object_id_stack[id],
                    'single_past_lidar': single_past_lidar_stack[id]
                })

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                "single_object_bbx_center": single_object_bbx_center,       # 用于viz single view
                "single_object_ids": single_object_ids,                # 用于viz single view
                'flow_gt':              # single view flow
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'pastk_2_past0_tr_mats':    # list of past k frames' transformation matrix to past 0 frame
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        if self.viz_bbx_flag:
            selected_cav_processed.update({'single_lidar': lidar_np[:, :3]})
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past_k to past_0 tansfomation matrix
        pastk_2_past0_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)
            # past_k trans past_0 matrix
            pastk_2_past0 = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], selected_cav_base['past_k'][0]['params']['lidar_pose'])
            pastk_2_past0_tr_mats.append(pastk_2_past0)
            
            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            if self.viz_bbx_flag:
                selected_cav_processed.update({'single_past_lidar': lidar_np[:, :3]})
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            # # 5. single view label
            # # past_i label at past_i single view
            # # opencood/data_utils/post_processor/base_postprocessor.py
            # object_bbx_center, object_bbx_mask, object_ids = \
            #     self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # # generate the anchor boxes
            # # opencood/data_utils/post_processor/voxel_postprocessor.py
            # anchor_box = self.anchor_box
            # single_view_label_dict = self.post_processor.generate_label(
            #         gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            #     )
            # past_k_label_dicts.append(single_view_label_dict)
        
        

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)
        pastk_2_past0_tr_mats = np.stack(pastk_2_past0_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        single_object_bbx_center, single_object_bbx_mask, single_object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=single_object_bbx_center, anchors=anchor_box, mask=single_object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update({
            "single_label_dict": label_dict,
            "single_object_bbx_center": single_object_bbx_center[single_object_bbx_mask == 1],
            "single_object_ids": single_object_ids,
            "curr_feature": curr_feature,
            'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
            'object_ids': object_ids,
            'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
            'past_k_tr_mats': past_k_tr_mats,
            'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats,
            'past_k_poses': past_k_poses,
            'past_k_features': past_k_features,
            'past_k_time_diffs': past_k_time_diffs,
            'past_k_sample_interval': past_k_sample_interval,
        #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
        #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
        #  'past_k_label_dicts': past_k_label_dicts,
            'if_no_point': if_no_point
            })

        if self.is_generate_gt_flow:
            # generate flow, from past_0 and curr
            prev_object_id_stack = {}
            prev_object_stack = {}
            for t_i in range(2):
                split_part = selected_cav_base['past_k'][0] if t_i == 0 else selected_cav_base['curr'] # TODO: 这里面的 prev 和 curr 可能反了
                object_bbx_center, object_bbx_mask, object_ids = \
                    self.generate_object_center([split_part], selected_cav_base['past_k'][0]['params']['lidar_pose'])
                prev_object_id_stack[t_i] = object_ids
                prev_object_stack[t_i] = object_bbx_center
            
            for t_i in range(2):
                unique_object_ids = list(set(prev_object_id_stack[t_i]))
                unique_indices = \
                    [prev_object_id_stack[t_i].index(x) for x in unique_object_ids]
                prev_object_stack[t_i] = np.vstack(prev_object_stack[t_i])
                prev_object_stack[t_i] = prev_object_stack[t_i][unique_indices]
                prev_object_id_stack[t_i] = unique_object_ids
            
            # TODO: generate_flow_map: yhu, generate_flow_map_szwei: szwei
            flow_map, warp_mask = generate_flow_map_szwei(prev_object_stack,
                                            prev_object_id_stack,
                                            self.params['preprocess']['cav_lidar_range'],
                                            self.params['preprocess']['args']['voxel_size'],
                                            past_k=1)

            selected_cav_processed.update({'flow_gt': flow_map})
            selected_cav_processed.update({'warp_mask': warp_mask})

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def dist_time(ts1, ts2, i = -1):
        """caculate the time interval between two timestamps

        Args:
            ts1 (string): time stamp at some time
            ts2 (string): current time stamp
            i (int, optional): past frame id, for debug use. Defaults to -1.
        
        Returns:
            time_diff (float): time interval (ts1 - ts2)
        """
        if not i==-1:
            return -i
        else:
            return (float(ts1) - float(ts2))

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'flow_gt': [N, C, H, W]
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        past_k_avg_time_var = []
        pastk_2_past0_tr_mats = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []
        
        if self.is_generate_gt_flow:
            # flow gt
            flow_gt_list = []
            warp_mask_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            try:
                past_k_avg_time_var.append(ego_dict['avg_var'])
            except KeyError:
                past_k_avg_time_var.append(-1.0)
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            pastk_2_past0_tr_mats.append(ego_dict['pastk_2_past0_tr_mats'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            if self.is_generate_gt_flow:
                flow_gt_list.append(ego_dict['flow_gt'])
                warp_mask_list.append(ego_dict['warp_mask'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        past_k_avg_time_var = np.array(past_k_avg_time_var)
        avg_time_var = float(sum(past_k_avg_time_var) / len(past_k_avg_time_var))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        pastk_2_past0_tr_mats = torch.from_numpy(np.vstack(pastk_2_past0_tr_mats))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        if self.is_generate_gt_flow:
            # for flow
            flow_gt = torch.from_numpy(np.vstack(flow_gt_list))
            label_torch_dict.update({'flow_gt': flow_gt})
            warp_mask = torch.from_numpy(np.vstack(warp_mask_list))
            label_torch_dict.update({'warp_mask': warp_mask})

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay,
                                   'avg_time_var': avg_time_var})
                                #    'times': time_consume})
        output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(self.anchor_box))})
        
        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})

        if self.viz_bbx_flag:
            single_lidar = []
            single_object_bbx_center = []
            single_object_mask = []
            single_object_ids = []
            single_past_lidar = []
            for i in range(len(batch[0].keys())-1):
                single_lidar.append(torch.from_numpy(batch[0][i]['single_lidar']))
                single_past_lidar.append(torch.from_numpy(batch[0][i]['single_past_lidar']))
                single_object_bbx_center.append(torch.from_numpy(batch[0][i]['single_object_bbx_center']))
                single_object_ids.append(batch[0][i]['single_object_ids'])
            output_dict['ego'].update({
                'single_lidar_list': single_lidar,
                'single_object_bbx_center': single_object_bbx_center,
                'single_object_ids': single_object_ids,
                'single_past_lidar_list': single_past_lidar
            })
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

    def generate_pred_bbx_frames(self, m_single, trans_mat_pastk_2_past0, past_time_diff, anchor_box):
        '''
        m_single : {
            psm_single, rm_single, (dm_single)
        }
        box_result : {
            'past_k_time_diff' : 
            [0] : {
                pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. (n, 8, 3)
                pred_box_center_tensor : (n, 7)
                scores: (n, )
            },
            ...
            [k-1] : { ... }
        }
        '''
        box_results = self.post_processor.single_post_process(m_single, trans_mat_pastk_2_past0, past_time_diff, anchor_box, self.k, self.num_roi_thres)
        return box_results

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_irregular_flow_new.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/4/30
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
two stage flow update framework
带有variance的irregular dataset
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset, intermediate_fusion_dataset_opv2v_irregular_flow_new
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
from opencood.utils.flow_utils import generate_flow_map, generate_flow_map_szwei

from opencood.utils.box_utils import boxes_to_corners2d # for debug use
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np


# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetIrregularFlowVariance(intermediate_fusion_dataset_opv2v_irregular_flow_new.IntermediateFusionDatasetIrregularFlowNew):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):
        super(IntermediateFusionDatasetIrregularFlowVariance, self).__init__(params, visualize, train)

        # random sample range 
        if 'ir_range' in params:
            self.sample_range = params['ir_range']
        else:
            self.sample_range = 0

        self.sample_anchor = list(np.array([-0, -4, -8]) - int(self.sample_interval_exp))

        print("=== Random Sampling Range Scale = {}. ===".format(self.sample_range))

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)
                        ...,		
                        [k-1] : {} (id-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k-1]:{}  (id-\tau-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * self.k
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # TODO: debug use, uncomment: to use regular version yaml GT
            # time_new = str( '%06d' % int(float(timestamp_key)))
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # TODO: debug use, uncomment: to use regular version lidar input
            # time_new = str( '%06d' % int(float(timestamp_key))) 
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            sample_interval = 0
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                # sample_interval
                if data[cav_id]['ego']:             # ego sample_interval = E(B(n, p))
                    if i == 0: # ego-past-0 与 ego-curr 是一样的
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    this_sample_stamp_idx = curr_timestamp_idx + self.sample_anchor[i]
                    sample_interval = latest_sample_stamp_idx - this_sample_stamp_idx
                    latest_sample_stamp_idx = this_sample_stamp_idx
                else:                               # non-ego sample_interval ~ B(n, p)
                    if i==0:  # if non-ego-past-0, must late than current frame # TODO:
                        unit_offset = random.randint(-self.sample_range, 1)
                    else:
                        unit_offset = random.randint(-self.sample_range, self.sample_range)
                    this_sample_stamp_idx = curr_timestamp_idx + self.sample_anchor[i] + unit_offset
                    sample_interval = latest_sample_stamp_idx - this_sample_stamp_idx
                    latest_sample_stamp_idx = this_sample_stamp_idx

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float

        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            },
            cav_id_2 : { ... }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = math.sqrt( \
                (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        pastk_2_past0_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        if self.is_generate_gt_flow:
            flow_gt = []
            warp_mask = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换, 用于单体监督            cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
                'flow_gt':               # [2, H, W] flow ground truth
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: N, k
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k frames to cav past 0 frame trans matrix, list of len=N, pastk_2_past0_tr_mats[i]: ndarray(k, 4, 4)
            pastk_2_past0_tr_mats.append(selected_cav_processed['pastk_2_past0_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        
            if self.is_generate_gt_flow:
                # for flow
                flow_gt.append(selected_cav_processed['flow_gt'])
                warp_mask.append(selected_cav_processed['warp_mask'])
        
        pastk_2_past0_tr_mats = np.stack(pastk_2_past0_tr_mats, axis=0) # N, k, 4, 4

        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) TODO: 这里面没有搞懂为什么不用 past_k_tr_mats

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack, 
             'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats})
            #  'times': self.times})

        if self.is_generate_gt_flow:
            flow_gt = np.vstack(flow_gt) # (N, H, W, 2)
            processed_data_dict['ego'].update({'flow_gt': flow_gt})
            warp_mask = np.vstack(warp_mask) # (N, C, H, W)
            processed_data_dict['ego'].update({'warp_mask': warp_mask})
        
        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            tmp = past_k_time_diffs_stack[self.k:].reshape(-1, self.k) # (N, k)
            tmp_past_k_time_diffs = np.concatenate((tmp[:, :1] , (tmp[:, 1:] - tmp[:, :-1])), axis=1) # (N, k)
            avg_time_diff = sum(tmp_past_k_time_diffs.reshape(-1)) / tmp_past_k_time_diffs.reshape(-1).shape[0]
            processed_data_dict['ego'].update({'avg_sample_interval':\
                sum(past_k_sample_interval_stack[self.k:]) / len(past_k_sample_interval_stack[self.k:])})
            processed_data_dict['ego'].update({'avg_time_delay':\
                avg_time_diff})
            tmp_var = np.var(tmp_past_k_time_diffs, axis=1) # (N, )
            avg_var = np.mean(tmp_var)
            processed_data_dict['ego'].update({'avg_var':\
                avg_var})
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict

    # def get_item_single_car(self, selected_cav_base, ego_pose, idx):
    #     """
    #     Project the lidar and bbx to ego space first, and then do clipping.

    #     Parameters
    #     ----------
    #     selected_cav_base : dict
    #         The dictionary contains a single CAV's raw information, 
    #         structure: {
    #             'ego' : true,
    #             'curr' : {
    #                 'params': (yaml),
    #                 'lidar_np': (numpy),
    #                 'timestamp': string
    #             },
    #             'past_k' : {		           # (k) totally
    #                 [0]:{
    #                     'params': (yaml),
    #                     'lidar_np': (numpy),
    #                     'timestamp': string,
    #                     'time_diff': float,
    #                     'sample_interval': int
    #                 },
    #                 [1] : {},
    #                 ...,		
    #                 [k-1] : {}
    #             },
    #             'debug' : {                     # debug use
    #                 scene_name : string         
    #                 cav_id : string
    #             }
    #         }

    #     ego_pose : list, length 6
    #         The ego vehicle lidar pose under world coordinate.

    #     idx: int,
    #         debug use.

    #     Returns
    #     -------
    #     selected_cav_processed : dict
    #         The dictionary contains the cav's processed information.
    #         {
    #             'projected_lidar':      # lidar in ego space, 用于viz
    #             'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
    #             'flow_gt':              # single view flow
    #             'curr_feature':         # current feature, lidar预处理得到的feature
    #             'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
    #             'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
    #             'curr_pose':			# current pose, list, len = 6
    #             'past_k_poses': 		    # list of past k frames' poses
    #             'past_k_features': 		    # list of past k frames' lidar
    #             'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
    #             'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
    #             'pastk_2_past0_tr_mats':    # list of past k frames' transformation matrix to past 0 frame
    #             'past_k_sample_interval':   # list of past k frames' sample interval with later frame
    #             'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
    #             'avg_past_k_sample_interval': # avg_past_k_sample_interval,
    #             'if_no_point':              # bool, 用于判断是否合法
    #         }
    #     """
    #     selected_cav_processed = {}

    #     # curr lidar feature
    #     lidar_np = selected_cav_base['curr']['lidar_np']
    #     lidar_np = shuffle_points(lidar_np)
    #     lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

    #     if self.visualize:
    #         # trans matrix
    #         transformation_matrix = \
    #             x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
    #         projected_lidar = \
    #             box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
    #         selected_cav_processed.update({'projected_lidar': projected_lidar})

    #     lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
    #     curr_feature = self.pre_processor.preprocess(lidar_np)
        
    #     # past k transfomation matrix
    #     past_k_tr_mats = []
    #     # past_k to past_0 tansfomation matrix
    #     pastk_2_past0_tr_mats = []
    #     # past k lidars
    #     past_k_features = []
    #     # past k poses
    #     past_k_poses = []
    #     # past k timestamps
    #     past_k_time_diffs = []
    #     # past k sample intervals
    #     past_k_sample_interval = []

    #     # for debug use
    #     # avg_past_k_time_diff = 0
    #     # avg_past_k_sample_interval = 0
        
    #     # past k label 
    #     # past_k_label_dicts = [] # todo 这个部分可以删掉

    #     # 判断点的数量是否合法
    #     if_no_point = False

    #     # past k frames [trans matrix], [lidar feature], [pose], [time interval]
    #     for i in range(self.k):
    #         # 1. trans matrix
    #         transformation_matrix = \
    #             x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
    #         past_k_tr_mats.append(transformation_matrix)
    #         # past_k trans past_0 matrix
    #         pastk_2_past0 = \
    #             x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], selected_cav_base['past_k'][0]['params']['lidar_pose'])
    #         pastk_2_past0_tr_mats.append(pastk_2_past0)
            
    #         # 2. lidar feature
    #         lidar_np = selected_cav_base['past_k'][i]['lidar_np']
    #         lidar_np = shuffle_points(lidar_np)
    #         lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
    #         lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
    #         processed_features = self.pre_processor.preprocess(lidar_np)
    #         past_k_features.append(processed_features)

    #         if lidar_np.shape[0] == 0: # 没有点留下
    #             if_no_point = True

    #         # 3. pose
    #         past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

    #         # 4. time interval and sample interval
    #         past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
    #         past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

    #         ################################################################
    #         # sizhewei
    #         # for past k frames' single view label
    #         ################################################################
    #         '''
    #         # 5. single view label
    #         # past_i label at past_i single view
    #         # opencood/data_utils/post_processor/base_postprocessor.py
    #         object_bbx_center, object_bbx_mask, object_ids = \
    #             self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
    #         # generate the anchor boxes
    #         # opencood/data_utils/post_processor/voxel_postprocessor.py
    #         anchor_box = self.anchor_box
    #         single_view_label_dict = self.post_processor.generate_label(
    #                 gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
    #             )
    #         past_k_label_dicts.append(single_view_label_dict)
    #         '''
        
        

    #     '''
    #     # past k merge
    #     past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
    #     '''

    #     past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)
    #     pastk_2_past0_tr_mats = np.stack(pastk_2_past0_tr_mats, axis=0) # (k, 4, 4)

    #     # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
    #     # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

    #     # curr label at single view
    #     # opencood/data_utils/post_processor/base_postprocessor.py
    #     object_bbx_center, object_bbx_mask, object_ids = \
    #         self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
    #     # generate the anchor boxes
    #     # opencood/data_utils/post_processor/voxel_postprocessor.py
    #     anchor_box = self.anchor_box
    #     label_dict = self.post_processor.generate_label(
    #             gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
    #         )
        
    #     # curr label at ego view
    #     object_bbx_center, object_bbx_mask, object_ids = \
    #         self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
    #     selected_cav_processed.update(
    #         {"single_label_dict": label_dict,
    #          "curr_feature": curr_feature,
    #          'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
    #          'object_ids': object_ids,
    #          'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
    #          'past_k_tr_mats': past_k_tr_mats,
    #          'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats,
    #          'past_k_poses': past_k_poses,
    #          'past_k_features': past_k_features,
    #          'past_k_time_diffs': past_k_time_diffs,
    #          'past_k_sample_interval': past_k_sample_interval,
    #         #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
    #         #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
    #         #  'past_k_label_dicts': past_k_label_dicts,
    #          'if_no_point': if_no_point
    #          })

    #     if self.is_generate_gt_flow:
    #         # generate flow, from past_0 and curr
    #         prev_object_id_stack = {}
    #         prev_object_stack = {}
    #         for t_i in range(2):
    #             split_part = selected_cav_base['past_k'][0] if t_i == 0 else selected_cav_base['curr'] # TODO: 这里面的 prev 和 curr 可能反了
    #             object_bbx_center, object_bbx_mask, object_ids = \
    #                 self.generate_object_center([split_part], selected_cav_base['past_k'][0]['params']['lidar_pose'])
    #             prev_object_id_stack[t_i] = object_ids
    #             prev_object_stack[t_i] = object_bbx_center
            
    #         for t_i in range(2):
    #             unique_object_ids = list(set(prev_object_id_stack[t_i]))
    #             unique_indices = \
    #                 [prev_object_id_stack[t_i].index(x) for x in unique_object_ids]
    #             prev_object_stack[t_i] = np.vstack(prev_object_stack[t_i])
    #             prev_object_stack[t_i] = prev_object_stack[t_i][unique_indices]
    #             prev_object_id_stack[t_i] = unique_object_ids
            
    #         # TODO: generate_flow_map: yhu, generate_flow_map_szwei: szwei
    #         flow_map, warp_mask = generate_flow_map_szwei(prev_object_stack,
    #                                         prev_object_id_stack,
    #                                         self.params['preprocess']['cav_lidar_range'],
    #                                         self.params['preprocess']['args']['voxel_size'],
    #                                         past_k=1)

    #         selected_cav_processed.update({'flow_gt': flow_map})
    #         selected_cav_processed.update({'warp_mask': warp_mask})

    #     return selected_cav_processed

    # def merge_past_k_features_to_dict(self, processed_feature_list):
    #     """
    #     Merge the preprocessed features from different cavs to the same
    #     dictionary.

    #     Parameters
    #     ----------
    #     processed_feature_list : list
    #         A list of dictionary containing all processed features from
    #         different cavs.

    #     Returns
    #     -------
    #     merged_feature_dict: dict
    #         key: feature names, value: list of features.
    #     """

    #     merged_feature_dict = OrderedDict()

    #     for cav_id in range(len(processed_feature_list)):
    #         for time_id in range(self.k):
    #             for feature_name, feature in processed_feature_list[cav_id][time_id].items():
    #                 if feature_name not in merged_feature_dict:
    #                     merged_feature_dict[feature_name] = []
    #                 if isinstance(feature, list):
    #                     merged_feature_dict[feature_name] += feature
    #                 else:
    #                     merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
    #     return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'flow_gt': [N, C, H, W]
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        past_k_avg_time_var = []
        pastk_2_past0_tr_mats = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []
        
        if self.is_generate_gt_flow:
            # flow gt
            flow_gt_list = []
            warp_mask_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_var.append(ego_dict['avg_var'])
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            pastk_2_past0_tr_mats.append(ego_dict['pastk_2_past0_tr_mats'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            if self.is_generate_gt_flow:
                flow_gt_list.append(ego_dict['flow_gt'])
                warp_mask_list.append(ego_dict['warp_mask'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        past_k_avg_time_var = np.array(past_k_avg_time_var)
        avg_time_var = float(sum(past_k_avg_time_var) / len(past_k_avg_time_var))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        pastk_2_past0_tr_mats = torch.from_numpy(np.vstack(pastk_2_past0_tr_mats))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        if self.is_generate_gt_flow:
            # for flow
            flow_gt = torch.from_numpy(np.vstack(flow_gt_list))
            label_torch_dict.update({'flow_gt': flow_gt})
            warp_mask = torch.from_numpy(np.vstack(warp_mask_list))
            label_torch_dict.update({'warp_mask': warp_mask})

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay,
                                   'avg_time_var': avg_time_var})
                                #    'times': time_consume})
        output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(self.anchor_box))})
        
        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict


opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_irregular_flow_variance.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/1/27
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for intermediate fusion with past k frames on irregular OPV2V
对比 intermediate_fusion_dataset_opv2v_irregular_flow , 区别:
该类含有 with_regular 参数, 用于控制是否使用绝对regular的sample.
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
from opencood.utils.flow_utils import generate_flow_map

from opencood.utils.box_utils import boxes_to_corners2d # for debug use
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetIrregularFlowWRegular(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):

        self.times = []

        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
        
        self.sample_interval_exp = int(self.binomial_n * self.binomial_p)

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        '''
        scenario_database Structure: 
        {
            scenario_id : {
                cav_1 : {
                    'ego' : true / false , 
                    timestamp1 : {
                        yaml: path,
                        lidar: path, 
                        cameras: list of path
                    },
                    ...
                },
                ...
            }
        }
        '''
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # copy timestamps npy file
            timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
            time_annotations = np.load(timestamps_file)

            # at least 1 cav should show up
            cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
            assert len(cav_list) > 0

            # use the frame number as key, the full path as the values
            yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
            if len(yaml_files)==0:
                yaml_files = sorted([x 
                            for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                            x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

            start_timestamp = int(float(self.extract_timestamps(yaml_files)[0]))
            while(1):
                time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
                time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
                if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
                    start_timestamp += 1
                else:
                    break

            end_timestamp = int(float(self.extract_timestamps(yaml_files)[-1]))
            if start_timestamp%2 == 0:
                # even
                end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
            else:
                end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
            num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)
            regular_timestamps = [start_timestamp+2*i for i in range(num_timestamps)]

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                cav_path = os.path.join(scenario_folder, cav_id)

                if j==0: # ego
                    timestamps = regular_timestamps
                else:
                    timestamps = list(time_annotations[j-1, :])

                for timestamp in timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files
                
                # regular的timestamps 用于做 curr 真实时刻的ground truth
                self.scenario_database[i][cav_id]['regular'] = OrderedDict()
                for timestamp in regular_timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id]['regular'][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id]['regular'][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id]['regular'][timestamp]['lidar'] = \
                        lidar_file

                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    # num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    num_ego_timestamps = len(timestamps) - self.binomial_n * self.k
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)
                else:
                    self.scenario_database[i][cav_id]['ego'] = False
                    

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead. 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("=== OPV2V-Irregular Multi-sweep dataset with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def extract_timestamps(self, yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]
            if res.endswith('.yaml'):
                timestamp = res.replace('.yaml', '')
            elif res.endswith('.json'):
                timestamp = res.replace('.json', '')
            else:
                print("Woops! There is no processing method for file {}".format(res))
                sys.exit(1)
            timestamps.append(timestamp)

        return timestamps

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)
                        ...,		
                        [k-1] : {} (id-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k-1]:{}  (id-\tau-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * self.k
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # TODO: debug use, uncomment: to use regular version yaml GT
            # time_new = str( '%06d' % int(float(timestamp_key)))
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # TODO: debug use, uncomment: to use regular version lidar input
            # time_new = str( '%06d' % int(float(timestamp_key))) 
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                # sample_interval
                if data[cav_id]['ego']:             # ego sample_interval = E(B(n, p))
                    if i == 0: # ego-past-0 与 ego-curr 是一样的
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    sample_interval = self.sample_interval_exp
                    if sample_interval == 0:
                        sample_interval = 1
                else:                               # non-ego sample_interval ~ B(n, p)
                    if self.sample_interval_exp==0 \
                        and self.is_no_shift \
                            and i == 0:
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==0: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>0: # 过去的几帧不要重复
                            sample_interval = 1                

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float

        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            },
            cav_id_2 : { ... }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = math.sqrt( \
                (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        flow_gt = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换, 用于单体监督            cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
                'flow_gt':               # [2, H, W] flow ground truth
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: N, k
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        
            # for flow
            flow_gt.append(selected_cav_processed['flow_gt'])
        
        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) TODO: 这里面没有搞懂为什么不用 past_k_tr_mats

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack})
            #  'times': self.times})

        flow_gt = np.vstack(flow_gt) # (N, 2, H, W)
        processed_data_dict['ego'].update({'flow_gt': flow_gt})
        
        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            processed_data_dict['ego'].update({'avg_sample_interval':\
                sum(past_k_sample_interval_stack[self.k:]) / len(past_k_sample_interval_stack[self.k:])})
            processed_data_dict['ego'].update({'avg_time_delay':\
                sum(past_k_time_diffs_stack[self.k:]) / len(past_k_time_diffs_stack[self.k:])})
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'flow_gt':              # single view flow
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            '''
            # 5. single view label
            # past_i label at past_i single view
            # opencood/data_utils/post_processor/base_postprocessor.py
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # generate the anchor boxes
            # opencood/data_utils/post_processor/voxel_postprocessor.py
            anchor_box = self.anchor_box
            single_view_label_dict = self.post_processor.generate_label(
                    gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
                )
            past_k_label_dicts.append(single_view_label_dict)
            '''
        
        

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
             'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
            #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
            #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
            #  'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        # generate flow, from past_0 and curr
        prev_object_id_stack = {}
        prev_object_stack = {}
        for t_i in range(2):
            split_part = selected_cav_base['past_k'][0] if t_i == 0 else selected_cav_base['curr']
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([split_part], selected_cav_base['past_k'][0]['params']['lidar_pose'])
            prev_object_id_stack[t_i] = object_ids
            prev_object_stack[t_i] = object_bbx_center
        
        for t_i in range(2):
            unique_object_ids = list(set(prev_object_id_stack[t_i]))
            unique_indices = \
                [prev_object_id_stack[t_i].index(x) for x in unique_object_ids]
            prev_object_stack[t_i] = np.vstack(prev_object_stack[t_i])
            prev_object_stack[t_i] = prev_object_stack[t_i][unique_indices]
            prev_object_id_stack[t_i] = unique_object_ids

        # TODO: for debug use
        # previous_cood = prev_object_stack[0][0]
        # # test_prev_box_2d = boxes_to_corners2d(np.expand_dims(previous_cood, axis=0), order='hwl')
        # # print('test_prev_box_2d', test_prev_box_2d)
        # #   [[[52.839737   5.977605  -1.9081762] 
        # #   [52.84148    7.9091973 -1.9081762]
        # #   [48.236977   7.913352  -1.9081762]
        # #   [48.235233   5.9817595 -1.9081762]]]
        # current_cood = previous_cood.copy()
        # current_cood[0] += 100
        # current_cood[1] += 30
        # object_id = prev_object_id_stack[0][0]
        # prev_object_id_stack = {}
        # prev_object_stack = {}
        # for i in range(2):
        #     prev_object_id_stack[i] = [object_id]
        #     prev_object_stack[i] = [previous_cood] if i == 0 else [current_cood]
        ########################################
        
        flow_map = generate_flow_map(prev_object_stack,
                                        prev_object_id_stack,
                                        self.params['preprocess']['cav_lidar_range'],
                                        self.params['preprocess']['args']['voxel_size'],
                                        past_k=1)

        # # TODO: debug use
        # debug_path = '/remote-home/share/sizhewei/logs/where2comm_flow_debug/viz_flow/'
        # torch.save(flow_map, debug_path + 'flow_map_true.pt')
        # torch.save(prev_object_stack, debug_path + 'bbx_0.pt')
        # ########################################

        selected_cav_processed.update({'flow_gt': flow_map})

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def dist_time(ts1, ts2, i = -1):
        """caculate the time interval between two timestamps

        Args:
            ts1 (string): time stamp at some time
            ts2 (string): current time stamp
            i (int, optional): past frame id, for debug use. Defaults to -1.
        
        Returns:
            time_diff (float): time interval (ts1 - ts2)
        """
        if not i==-1:
            return -i
        else:
            return (float(ts1) - float(ts2))

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'flow_gt': [N, C, H, W]
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []
        # flow gt
        flow_gt_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            flow_gt_list.append(ego_dict['flow_gt'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for flow
        flow_gt = torch.from_numpy(np.vstack(flow_gt_list))
        label_torch_dict.update({'flow_gt': flow_gt})

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay})
                                #    'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_irregular_flow_w_regular.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/1/27
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
为了给v2vnet 等baseline方法训练 用的数据集
有一些baseline方法需要使用多个车之间的pairwise_t_matrix 在多帧的数据集里面忽略了这点 该类将多车之间的坐标系变换补充好
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetIrregularSingleframe(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):

        self.times = []

        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
        
        self.sample_interval_exp = int(self.binomial_n * self.binomial_p)

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        '''
        scenario_database Structure: 
        {
            scenario_id : {
                cav_1 : {
                    'ego' : true / false , 
                    timestamp1 : {
                        yaml: path,
                        lidar: path, 
                        cameras: list of path
                    },
                    ...
                },
                ...
            }
        }
        '''
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # copy timestamps npy file
            timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
            time_annotations = np.load(timestamps_file)

            # at least 1 cav should show up
            cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
            assert len(cav_list) > 0

            # use the frame number as key, the full path as the values
            yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
            if len(yaml_files)==0:
                yaml_files = sorted([x 
                            for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                            x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

            start_timestamp = int(float(self.extract_timestamps(yaml_files)[0]))
            while(1):
                time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
                time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
                if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
                    start_timestamp += 1
                else:
                    break

            end_timestamp = int(float(self.extract_timestamps(yaml_files)[-1]))
            if start_timestamp%2 == 0:
                # even
                end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
            else:
                end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
            num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)
            regular_timestamps = [start_timestamp+2*i for i in range(num_timestamps)]

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                cav_path = os.path.join(scenario_folder, cav_id)

                if j==0: # ego
                    timestamps = regular_timestamps
                else:
                    timestamps = list(time_annotations[j-1, :])

                for timestamp in timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files
                
                # regular的timestamps 用于做 curr 真实时刻的ground truth
                self.scenario_database[i][cav_id]['regular'] = OrderedDict()
                for timestamp in regular_timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id]['regular'][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id]['regular'][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id]['regular'][timestamp]['lidar'] = \
                        lidar_file

                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    # num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    num_ego_timestamps = len(timestamps) - self.binomial_n * self.k * 3 # TODO:
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)
                else:
                    self.scenario_database[i][cav_id]['ego'] = False
                    

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead. 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("=== OPV2V-Irregular Multi-sweep dataset with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def extract_timestamps(self, yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]
            if res.endswith('.yaml'):
                timestamp = res.replace('.yaml', '')
            elif res.endswith('.json'):
                timestamp = res.replace('.json', '')
            else:
                print("Woops! There is no processing method for file {}".format(res))
                sys.exit(1)
            timestamps.append(timestamp)

        return timestamps

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)
                        ...,		
                        [k-1] : {} (id-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k-1]:{}  (id-\tau-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * self.k * 3 # TODO:
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # TODO: debug use, uncomment: to use regular version yaml GT
            # time_new = str( '%06d' % int(float(timestamp_key)))
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # TODO: debug use, uncomment: to use regular version lidar input
            # time_new = str( '%06d' % int(float(timestamp_key))) 
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                # sample_interval
                if data[cav_id]['ego']:             # ego sample_interval = E(B(n, p))
                    if i == 0: # ego-past-0 与 ego-curr 是一样的
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    sample_interval = self.sample_interval_exp
                    if sample_interval == 0:
                        sample_interval = 1
                else:                               # non-ego sample_interval ~ B(n, p)
                    if self.sample_interval_exp==0 \
                        and self.is_no_shift \
                            and i == 0:
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==0: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>0: # 过去的几帧不要重复
                            sample_interval = 1                

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float

        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            },
            cav_id_2 : { ... }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = math.sqrt( \
                (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换, 用于单体监督            cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: N, k
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_pairwise_transformation(base_data_dict, self.max_cav)
            # self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            # ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) TODO: 这里面没有搞懂为什么不用 past_k_tr_mats

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack})
            #  'times': self.times})

        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            processed_data_dict['ego'].update({'avg_sample_interval':\
                sum(past_k_sample_interval_stack[self.k:]) / len(past_k_sample_interval_stack[self.k:])})
            processed_data_dict['ego'].update({'avg_time_delay':\
                sum(past_k_time_diffs_stack[self.k:]) / len(past_k_time_diffs_stack[self.k:])})
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            '''
            # 5. single view label
            # past_i label at past_i single view
            # opencood/data_utils/post_processor/base_postprocessor.py
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # generate the anchor boxes
            # opencood/data_utils/post_processor/voxel_postprocessor.py
            anchor_box = self.anchor_box
            single_view_label_dict = self.post_processor.generate_label(
                    gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
                )
            past_k_label_dicts.append(single_view_label_dict)
            '''

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
             'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
            #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
            #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
            #  'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def dist_time(ts1, ts2, i = -1):
        """caculate the time interval between two timestamps

        Args:
            ts1 (string): time stamp at some time
            ts2 (string): current time stamp
            i (int, optional): past frame id, for debug use. Defaults to -1.
        
        Returns:
            time_diff (float): time interval (ts1 - ts2)
        """
        if not i==-1:
            return -i
        else:
            return (float(ts1) - float(ts2))

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay})
                                #    'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                # lidar_pose = cav_content['curr']['params']['lidar_pose']
                lidar_pose = cav_content['past_k'][0]['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_irregular_singleframe.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Yifan Lu
# Modified: sizhewei @ 2022/10/12
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for early fusion
"""
from collections import OrderedDict
import os
import numpy as np
import torch
import math
import copy

import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np


class IntermediateFusionDatasetAsync(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        self.k = 1 # only one frame we need

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
        
        self.sample_interval_exp = int(self.binomial_n * self.binomial_p)

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        # Structure: {scenario_id : {cav_1 : {timestamp1 : {yaml: path,
        # lidar: path, cameras:list of path}}}}
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # copy timestamps npy file
            timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
            time_annotations = np.load(timestamps_file)

            # at least 1 cav should show up
            cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
            assert len(cav_list) > 0

            # use the frame number as key, the full path as the values
            yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
            if len(yaml_files)==0:
                yaml_files = sorted([x 
                            for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                            x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

            start_timestamp = int(float(self.extract_timestamps(yaml_files)[0]))
            while(1):
                time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
                time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
                if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
                    start_timestamp += 1
                else:
                    break

            end_timestamp = int(float(self.extract_timestamps(yaml_files)[-1]))
            if start_timestamp%2 == 0:
                # even
                end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
            else:
                end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
            num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)
            regular_timestamps = [start_timestamp+2*i for i in range(num_timestamps)]

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                # save all yaml files to the dictionary
                cav_path = os.path.join(scenario_folder, cav_id)

                if j==0: # ego
                    timestamps = regular_timestamps
                else:
                    timestamps = list(time_annotations[j-1, :])

                for timestamp in timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files
                # regular的timestamps 用于做 curr 真实时刻的ground truth
                self.scenario_database[i][cav_id]['regular'] = OrderedDict()
                for timestamp in regular_timestamps:
                    timestamp = "%.3f" % float(timestamp)
                    self.scenario_database[i][cav_id]['regular'][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id]['regular'][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id]['regular'][timestamp]['lidar'] = \
                        lidar_file

                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    # num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    num_ego_timestamps = len(timestamps) - self.binomial_n * self.k
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)
                else:
                    self.scenario_database[i][cav_id]['ego'] = False


        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("=== OPV2V-Irregular Single-sweep dataset with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def extract_timestamps(self, yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]
            if res.endswith('.yaml'):
                timestamp = res.replace('.yaml', '')
            elif res.endswith('.json'):
                timestamp = res.replace('.json', '')
            else:
                print("Woops! There is no processing method for file {}".format(res))
                sys.exit(1)
            timestamps.append(timestamp)

        return timestamps


    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
        """
        # we loop the accumulated length list to see get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]

        data = OrderedDict()
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            # check the timestamp index
            timestamp_index = idx if scenario_index == 0 else \
                idx - self.len_record[scenario_index - 1]
            # retrieve the corresponding timestamp key
            # timestamp_key = self.return_timestamp_key_async(scenario_database, cav_id, 
            #                                         timestamp_index)
            timestamp_key = list(cav_content.items())[timestamp_index][0]
            data[cav_id] = OrderedDict()
            data[cav_id]['ego'] = cav_content['ego']

            # todo: load camera image in the future version
            # load the corresponding data into the dictionary
            data[cav_id]['params'] = \
                load_yaml(cav_content[timestamp_key]['yaml'])
            data[cav_id]['lidar_np'] = \
                pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])
            
            if self.select_keypoint and 'lidar_keypoints_np' in cav_content[timestamp_key]:
                data[cav_id]['lidar_keypoints_np'] = cav_content[timestamp_key]['lidar_keypoints_np']

        return data


    def __getitem__(self, idx):
        base_data_dict = self.retrieve_base_data(idx)

        base_data_dict = add_noise_data_dict(base_data_dict,self.params['noise_setting'])

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        ego_id = -1
        ego_lidar_pose = []

        # first find the ego vehicle's lidar pose
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['params']['lidar_pose']
                ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
                break
            
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0


        processed_features = []
        object_stack = []
        object_id_stack = []
        too_far = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        projected_lidar_clean_list = []
        cav_id_list = []

        if self.visualize:
            projected_lidar_stack = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            distance = \
                math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
                           ego_lidar_pose[0]) ** 2 + (
                                  selected_cav_base['params'][
                                      'lidar_pose'][1] - ego_lidar_pose[
                                      1]) ** 2)

            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue


            lidar_pose_clean_list.append(selected_cav_base['params']['lidar_pose_clean'])
            lidar_pose_list.append(selected_cav_base['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)


        ########## Added by Yifan Lu 2022.8.14 ##############
        # box align to correct pose.
        '''
        if self.box_align and str(idx) in self.stage1_result.keys():
            stage1_content = self.stage1_result[str(idx)]
            if stage1_content is not None:
                cav_id_list_stage1 = stage1_content['cav_id_list']
                
                pred_corners_list = stage1_content['pred_corner3d_np_list']
                pred_corners_list = [np.array(corners, dtype=np.float64) for corners in pred_corners_list]
                uncertainty_list = stage1_content['uncertainty_np_list']
                uncertainty_list = [np.array(uncertainty, dtype=np.float64) for uncertainty in uncertainty_list]
                stage1_lidar_pose_list = [base_data_dict[cav_id]['params']['lidar_pose'] for cav_id in cav_id_list_stage1]
                stage1_lidar_pose = np.array(stage1_lidar_pose_list)


                refined_pose = box_alignment_relative_sample_np(pred_corners_list,
                                                                stage1_lidar_pose, 
                                                                uncertainty_list=uncertainty_list, 
                                                                **self.box_align_args)
                stage1_lidar_pose[:,[0,1,4]] = refined_pose
                stage1_lidar_pose_refined_list = stage1_lidar_pose.tolist() # updated lidar_pose_list
                for cav_id, lidar_pose_refined in zip(cav_id_list_stage1, stage1_lidar_pose_refined_list):
                    if cav_id not in cav_id_list:
                        continue
                    idx_in_list = cav_id_list.index(cav_id)
                    lidar_pose_list[idx_in_list] = lidar_pose_refined
                    base_data_dict[cav_id]['params']['lidar_pose'] = lidar_pose_refined
        '''     

        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]


            # obj: 
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose, 
                ego_lidar_pose_clean,
                idx)
                
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            processed_features.append(
                selected_cav_processed['processed_features'])
            if self.kd_flag:
                projected_lidar_clean_list.append(
                    selected_cav_processed['projected_lidar_clean'])

            if self.visualize:
                projected_lidar_stack.append(
                    selected_cav_processed['projected_lidar'])


        ########## Added by Yifan Lu 2022.4.5 ################
        # filter those out of communicate range
        # then we can calculate get_pairwise_transformation
        for cav_id in too_far:
            base_data_dict.pop(cav_id)
        
        pairwise_t_matrix = \
            self.get_pairwise_transformation(base_data_dict,
                                             self.max_cav)

        lidar_poses = np.array(lidar_pose_list).reshape(-1, 6)  # [N_cav, 6]
        lidar_poses_clean = np.array(lidar_pose_clean_list).reshape(-1, 6)  # [N_cav, 6]
        ######################################################

        ############ for disconet ###########
        if self.kd_flag:
            stack_lidar_np = np.vstack(projected_lidar_clean_list)
            stack_lidar_np = mask_points_by_range(stack_lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            stack_feature_processed = self.pre_processor.preprocess(stack_lidar_np)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        object_stack = np.vstack(object_stack)
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(processed_features)

        merged_feature_dict = self.merge_features_to_dict(processed_features)

        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        processed_data_dict['ego'].update(
            {'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'lidar_poses_clean': lidar_poses_clean,
             'lidar_poses': lidar_poses})

        if self.kd_flag:
            processed_data_dict['ego'].update({'teacher_processed_lidar':
                stack_feature_processed})

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(
                    projected_lidar_stack)})


        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, ego_pose_clean, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information.
        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.
        ego_pose_clean : list, length 6
            only used for gt box generation
        
        idx: int,
            debug use.

        Returns
        -------
        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) # T_ego_cav
        transformation_matrix_clean = \
            x1_to_x2(selected_cav_base['params']['lidar_pose_clean'],
                     ego_pose_clean)

        # retrieve objects under ego coordinates
        # this is used to generate accurate GT bounding box.
        object_bbx_center, object_bbx_mask, object_ids = self.generate_object_center([selected_cav_base],
                                                    ego_pose_clean)

        # filter lidar
        lidar_np = selected_cav_base['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        # remove points that hit itself
        lidar_np = mask_ego_points(lidar_np)
        # project the lidar to ego space
        # x,y,z in ego space
        projected_lidar = \
            box_utils.project_points_by_matrix_torch(lidar_np[:, :3],
                                                        transformation_matrix)
        if self.kd_flag:
            lidar_np_clean = copy.deepcopy(lidar_np)

        if self.proj_first:
            lidar_np[:, :3] = projected_lidar
            
        lidar_np = mask_points_by_range(lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
        

        processed_lidar = self.pre_processor.preprocess(lidar_np)

        selected_cav_processed.update(
            {'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'projected_lidar': projected_lidar,
             'processed_features': processed_lidar,
             'transformation_matrix': transformation_matrix,
             'transformation_matrix_clean': transformation_matrix_clean})

        if self.kd_flag:
            projected_lidar_clean = \
                box_utils.project_points_by_matrix_torch(lidar_np_clean[:, :3],
                                                    transformation_matrix_clean)
            lidar_np_clean[:, :3] = projected_lidar_clean
            lidar_np_clean = mask_points_by_range(lidar_np_clean,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
            selected_cav_processed.update(
                {"projected_lidar_clean": lidar_np_clean}
            )

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        if self.kd_flag:
            teacher_processed_lidar_list = []
        if self.visualize:
            origin_lidar = []

        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            lidar_pose_list.append(ego_dict['lidar_poses']) # ego_dict['lidar_pose'] is np.ndarray [N,6]
            lidar_pose_clean_list.append(ego_dict['lidar_poses_clean'])

            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])

            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])

            if self.kd_flag:
                teacher_processed_lidar_list.append(ego_dict['teacher_processed_lidar'])

            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))


        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)

        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        lidar_pose = torch.from_numpy(np.concatenate(lidar_pose_list, axis=0))
        lidar_pose_clean = torch.from_numpy(np.concatenate(lidar_pose_clean_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'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],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'lidar_pose_clean': lidar_pose_clean,
                                   'lidar_pose': lidar_pose})


        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
        
        if self.kd_flag:
            teacher_processed_lidar_torch_dict = \
                self.pre_processor.collate_batch(teacher_processed_lidar_list)
            output_dict['ego'].update({'teacher_processed_lidar':teacher_processed_lidar_torch_dict})

        if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
            (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
            return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_irregular_singleframe_fromdelay.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Yifan Lu
# License: TDG-Attribution-NonCommercial-NoDistrib
# Modified: sizhewei @ 2022/11/05

"""
Dataset class for intermediate fusion with past k frames
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import copy

import time
import json
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np


class IntermediateFusionDatasetMultisweep(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):

        self.times = []

        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        '''
        scenario_database Structure: 
        {
            scenario_id : {
                cav_1 : {
                    'ego' : true / false , 
                    timestamp1 : {
                        yaml: path,
                        lidar: path, 
                        cameras: list of path
                    },
                    ...
                },
                ...
            }
        }
        '''
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # at least 1 cav should show up
            cav_list = sorted([x 
                               for x in os.listdir(scenario_folder) if 
                               os.path.isdir(os.path.join(scenario_folder, x))])
            assert len(cav_list) > 0

            # roadside unit data's id is always negative, so here we want to
            # make sure they will be in the end of the list as they shouldn't
            # be ego vehicle.
            if int(cav_list[0]) < 0:
                cav_list = cav_list[1:] + [cav_list[0]]
                
            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                # save all yaml files to the dictionary
                cav_path = os.path.join(scenario_folder, cav_id)

                # use the frame number as key, the full path as the values
                yaml_files = \
                    sorted([os.path.join(cav_path, x)
                            for x in os.listdir(cav_path) if
                            x.endswith('.yaml')])
                if len(yaml_files)==0:
                    yaml_files = \
                        sorted([os.path.join(cav_path, x)
                                for x in os.listdir(cav_path) if
                                x.endswith('.json')])
                timestamps = self.extract_timestamps(yaml_files)	

                for timestamp in timestamps:
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files

                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)

                else:
                    self.scenario_database[i][cav_id]['ego'] = False
                    

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        # TODO: check pre_process & post_processor 是否符合新的数据结构 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("OPV2V Multi-sweep dataset with non-ego cavs' {} time delay and past {} frames collected initialized! \
                {} samples totally!".format(self.tau, self.k, self.len_record[-1]))

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {		(id)			# pose | lidar | label
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		# (k) totally
                        [0]:{		(id)			# pose | lidar | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                        },
                        [1] : {},	(id-1)			# pose | lidar | label
                        ...,						# pose | lidar | label
                        [k-1] : {} (id-(k-1))		# pose | lidar | label
                    }
                    
                }, 
                cav_id_2 : {
                    'ego': false, 
                    curr : 	{		(id)			# pose | lidar | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {		# (k) totally
                        [0] : {		(id - \tau)		# pose | lidar | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                        }			
                        ..., 						# pose | lidar | label
                        [k-1]:{}  (id-\tau-(k-1))	# pose | lidar | label
                    },
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]

        debug_times = np.zeros((3,))
        start_time = time.time()

        data = OrderedDict()
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                'ego' : true / false , 
                timestamp1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
            },
            '''
            data[cav_id] = OrderedDict()
            data[cav_id]['ego'] = cav_content['ego']
            data[cav_id]['past_k'] = OrderedDict()
            
            # current frame, for co-perception lable use
            data[cav_id]['curr'] = {}
            timestamp_index = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
            timestamp_index = timestamp_index + self.tau + self.k - 1
            timestamp_key = list(cav_content.items())[timestamp_index][0]
            
            time_s = time.time()
            # load param file: json is faster than yaml
            json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
            time_e = time.time()
            debug_times[0] += (time_e - time_s)

            time_s = time.time()
            # load lidar file: npy is faster than pcd
            npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
            if os.path.exists(npy_file):
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])
            time_e = time.time()
            debug_times[1] += (time_e - time_s)

            data[cav_id]['curr']['timestamp'] = \
                    timestamp_key
            data[cav_id]['curr']['time_diff'] = 0

            # past k frames, pose | lidar | label(for single view confidence map generator use)
            if data[cav_id]['ego']:
                temp = self.tau
            else:
                temp = 0
            for i in range(self.k):
                # check the timestamp index
                data[cav_id]['past_k'][i] = OrderedDict()
                timestamp_index = idx if scenario_index == 0 else \
                    idx - self.len_record[scenario_index - 1] # TODO: 这里面原来错加了一个i，记得改正
                timestamp_index = timestamp_index + self.k - 1 - i + temp
                timestamp_key = list(cav_content.items())[timestamp_index][0]
                # load the corresponding data into the dictionary
                time_s = time.time()
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                time_e = time.time()
                debug_times[0] += (time_e - time_s)

                time_s = time.time()
                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])
                time_e = time.time()
                debug_times[1] += (time_e - time_s)

                data[cav_id]['past_k'][i]['timestamp'] = \
                    timestamp_key
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])

        end_time = time.time()

        debug_times[2] += (end_time - start_time)

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
        }
        '''
        # base_data_dict
        base_data_dict = self.retrieve_base_data(idx)

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                # print(ego_lidar_pose)
                # ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        if self.visualize:
            projected_lidar_stack = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = \
                math.sqrt( \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)

            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue

            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  

        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_time_diffs_stack = []
        past_k_tr_mats = []
        past_k_label_dicts_stack = []

        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            '''
            selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'single_label_dict':	# single view label
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7)
                'object_ids':			# ego view label index. list. length is (max_num, 7)
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		# list of past k frames' poses
                'past_k_features': 		# list of past k frames' lidar
                'past_k_time_diffs': 	# list of past k frames' time diff with current frame
                'past_k_tr_mats': 		# list of past k frames' transformation matrix to current ego coordinate
                'past_k_label_dicts'    # [], [TBD] list of past k frames' label dict 
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose, 
                idx
            )

            if self.visualize:
                projected_lidar_stack.append(
                    selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack)

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)

        ########## Added by Yifan Lu 2022.4.5 ################
        # filter those out of communicate range
        # then we can calculate get_pairwise_transformation
        for cav_id in too_far:
            base_data_dict.pop(cav_id)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4)

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        object_stack = np.vstack(object_stack)
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'past_k_time_diffs': past_k_time_diffs_stack})

        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true / false,
                curr : {		(id)							#      |       | label
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		# (k) totally
                    [0]:{		(id) / (id-\tau-0)				# pose | lidar | label
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    [1] : {},	(id-1) / (id-\tau-1))			# pose | lidar | label
                    ...,										# pose | lidar | label
                    [k-1] : {} 	(id-(k-1)) / (id-\tau-(k-1))	# pose | lidar | label
                }	
            }, 
        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.
        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                single_label_dict:		# single view label
                object_bbx_center:		# ego view label. np.ndarray. Shape is (max_num, 7)
                object_ids:				# ego view label index. list. length is (max_num, 7)
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		# list of past k frames' poses
                'past_k_features': 		# list of past k frames' lidar
                'past_k_time_diffs': 	# list of past k frames' time diff with current frame
                'past_k_tr_mats': 		# list of past k frames' transformation matrix to current ego coordinate
                'past_k_label_dicts'    # [], [TBD] list of past k frames' label dict 
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update(
                {
                    'projected_lidar': projected_lidar
                }
            )

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        
        # past k label 
        past_k_label_dicts = [] # TODO: 这个部分可以删掉

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval
            if selected_cav_base['ego']:
                past_k_time_diffs.append(-i)
            else:
                past_k_time_diffs.append(-(self.tau+i))

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            '''
            # 5. single view label
            # past_i label at past_i single view
            # opencood/data_utils/post_processor/base_postprocessor.py
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # generate the anchor boxes
            # opencood/data_utils/post_processor/voxel_postprocessor.py
            anchor_box = self.anchor_box
            single_view_label_dict = self.post_processor.generate_label(
                    gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
                )
            past_k_label_dicts.append(single_view_label_dict)
            '''

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        selected_cav_processed.update({})
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
             'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_label_dicts': past_k_label_dicts
             })

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def dist_time(ts1, ts2, i = -1):
        """caculate the time interval between two timestamps

        Args:
            ts1 (string): time stamp at some time
            ts2 (string): current time stamp
            i (int, optional): past frame id, for debug use. Defaults to -1.
        
        Returns:
            time_diff (float): time interval (ts1 - ts2)
        """
        if not i==-1:
            return -i
        else:
            return (int(ts1) - int(ts2))

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
        }
        '''
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_interval.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                "targets": torch.cat(targets_single, dim=0)}

        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_interval = np.hstack(past_k_time_interval)
        past_k_time_interval = torch.from_numpy(past_k_time_interval)

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({
            'single_object_label': single_object_label,
            'curr_processed_lidar': curr_processed_lidar_torch_dict,
            '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,
            'single_past_dict': past_k_single_label_torch_dict,
            'object_ids': object_ids[0],
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_pose': curr_lidar_pose,
            'past_lidar_pose': past_k_lidar_pose,
            'past_k_time_interval': past_k_time_interval})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({
            'transformation_matrix': transformation_matrix_torch,
            'transformation_matrix_clean': transformation_matrix_clean_torch})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_multisweep.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/2/13

"""
Dataset class for intermediate fusion with 
1) past k frames as input 
2) current frame as supervision
on irregular OPV2V
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset, intermediate_fusion_dataset_opv2v_irregular
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetRegularCompensation(basedataset.BaseDataset):
    """
    This class is for intermediate fusion on OPV2V_irregular, with past k frames and current frame as input. 
    Note: current frame is used for supervison, not for feature fusion!
    """
    def __init__(self, params, visualize, train=True):
        super(IntermediateFusionDatasetRegularCompensation, self).__init__(params, visualize, train)

        self.times = []

        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
            print(f'======= is_ab_regular? {self.is_ab_regular} =======')
        
        self.sample_interval_exp = int(self.binomial_n * self.binomial_p)

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        '''
        scenario_database Structure: 
        {
            scenario_id : {
                cav_1 : {
                    'ego' : true / false , 
                    timestamp1 : {
                        yaml: path,
                        lidar: path, 
                        cameras: list of path
                    },
                    ...
                },
                ...
            }
        }
        '''
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # # copy timestamps npy file
            # timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
            # time_annotations = np.load(timestamps_file)

            # at least 1 cav should show up
            cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
            assert len(cav_list) > 0

            # use the frame number as key, the full path as the values
            yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
            if len(yaml_files)==0:
                yaml_files = sorted([x 
                            for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                            x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

            start_timestamp = int(float(self.extract_timestamps(yaml_files)[0]))
            # while(1):
            #     time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
            #     time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
            #     if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
            #         start_timestamp += 1
            #     else:
            #         break

            end_timestamp = int(float(self.extract_timestamps(yaml_files)[-1]))
            if start_timestamp%2 == 0:
                # even
                end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
            else:
                end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
            num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)
            regular_timestamps = [start_timestamp+2*i for i in range(num_timestamps)]

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                cav_path = os.path.join(scenario_folder, cav_id)

                # if j==0: # ego
                timestamps = regular_timestamps
                # else:
                #     if self.is_ab_regular:
                #         timestamps = regular_timestamps
                #     else:
                #         timestamps = list(time_annotations[j-1, :])

                for timestamp in timestamps:
                    timestamp = "%06d" % int(timestamp)
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    # print(yaml_file)
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files
                
                # regular的timestamps 用于做 curr 真实时刻的ground truth
                self.scenario_database[i][cav_id]['regular'] = OrderedDict()
                for timestamp in regular_timestamps:
                    timestamp = "%06d" % int(timestamp)
                    self.scenario_database[i][cav_id]['regular'][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id]['regular'][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id]['regular'][timestamp]['lidar'] = \
                        lidar_file

                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    # num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    num_ego_timestamps = len(timestamps) - self.binomial_n * self.k
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)
                else:
                    self.scenario_database[i][cav_id]['ego'] = False
                    

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead. 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("=== OPV2V-Irregular Multi-sweep dataset with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

        self.k = self.k + 1 # 用来在 past_k 中添加 current frame
        print("=== OPV2V-Irregular Multi-sweep dataloader (Compensation Ver.) for intermediate fusion with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k-1, self.binomial_n * self.binomial_p, self.len_record[-1]))
    
    def extract_timestamps(self, yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]
            if res.endswith('.yaml'):
                timestamp = res.replace('.yaml', '')
            elif res.endswith('.json'):
                timestamp = res.replace('.json', '')
            else:
                print("Woops! There is no processing method for file {}".format(res))
                sys.exit(1)
            timestamps.append(timestamp)

        return timestamps

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{                       # current
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)          # past 1
                        ...,		
                        [k] : {} (id-(k-1))         # past k
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k] : {}
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * (self.k-1)
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            try:
                timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            except IndexError:
                print('IndexError!!!')
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # time_new = str( '%06d' % int(float(timestamp_key))) # TODO: debug用 使用regular的GT
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0]) # TODO:
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # time_new = str( '%06d' % int(float(timestamp_key))) # TODO: debug用 使用regular的lidar
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                if i == 0: # ego-past-0 与 ego-curr 是一样的, 用于feature监督
                    data[cav_id]['past_k'][i] = data[cav_id]['curr']
                    continue
                # sample_interval
                if data[cav_id]['ego']:             # ego: do not need past frame
                    break
                    # sample_interval = self.sample_interval_exp
                else:                               # non-ego sample_interval ~ B(n, p)
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==0: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>0: # 过去的几帧不要重复
                            sample_interval = 1  

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                # uncomment code, use updated gt
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float

        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (1) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    }
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
                
            }, 
            cav_id_2 : {
                'ego' : false,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)          # past 1
                    ...,		                # ...
                    [k] : {} (id-(k-1))         # past k
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            if cav_id != ego_id:
                distance = math.sqrt( \
                    (selected_cav_base['past_k'][1]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                        (selected_cav_base['past_k'][1]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
                # if distance is too far, we will just skip this agent
                if distance > self.params['comm_range']:
                    too_far.append(cav_id)
                    continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true / false , 
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (1 / k) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)          # past 1 
                    ...,		                # ...
                    [k] : {} (id-(k-1))         # past k
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k+1, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k+1, 6
            past_k_pose_stack += selected_cav_processed['past_k_poses']
            # time differences: Nx(k+1)
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: Nx(k+1)
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k+1, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) 

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, 6)  # (Nx(k-1)+1, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack})
            #  'times': self.times})

        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            tmp = past_k_time_diffs_stack[1:].reshape(-1, self.k)
            tmp_past_k_time_diffs = tmp[:, 1:] - tmp[:, :-1]
            avg_time_diff = sum(tmp_past_k_time_diffs.reshape(-1)) / tmp_past_k_time_diffs.reshape(-1).shape[0]
            processed_data_dict['ego'].update({'avg_time_delay':\
                avg_time_diff})
            
            tmp_past_k_sample_interval = past_k_sample_interval_stack[1:].reshape(-1, self.k)[:, 1:]
            avg_past_k_sample_interval = sum(tmp_past_k_sample_interval.reshape(-1)) / tmp_past_k_sample_interval.reshape(-1).shape[0]
            processed_data_dict['ego'].update({'avg_sample_interval':\
                avg_past_k_sample_interval})
            
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict
    
    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def dist_time(ts1, ts2, i = -1):
        """caculate the time interval between two timestamps

        Args:
            ts1 (string): time stamp at some time
            ts2 (string): current time stamp
            i (int, optional): past frame id, for debug use. Defaults to -1.
        
        Returns:
            time_diff (float): time interval (ts1 - ts2)
        """
        if not i==-1:
            return -i
        else:
            return (float(ts1) - float(ts2))


    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        if selected_cav_base['ego']:
            k = 1
        else:
            k = self.k
        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            '''
            # 5. single view label
            # past_i label at past_i single view
            # opencood/data_utils/post_processor/base_postprocessor.py
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # generate the anchor boxes
            # opencood/data_utils/post_processor/voxel_postprocessor.py
            anchor_box = self.anchor_box
            single_view_label_dict = self.post_processor.generate_label(
                    gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
                )
            past_k_label_dicts.append(single_view_label_dict)
            '''

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
             'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
            #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
            #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
            #  'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        return selected_cav_processed

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(len(processed_feature_list[cav_id])):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict


    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_label':          single_label_dict_stack,
            'curr_processed_lidar':         merged_curr_feature_dict,
            'object_bbx_center':            object_bbx_center,
            'object_bbx_mask':              mask,
            # 'anchor_box':                   anchor_box,
            'processed_lidar':              merged_feature_dict,
            'label_dict':                   label_dict,
            'object_ids':                   [object_id_stack[i] for i in unique_indices],
            'single_past_dict':             
            'pairwise_t_matrix':            pairwise_t_matrix,
            'cav_num':                      cav_num,
            'curr_lidar_poses':             curr_lidar_poses,
            'past_lidar_poses':             past_k_lidar_poses,
            'past_k_time_interval':         
            'past_k_sample_interval':   
            'avg_sample_interval': , 
            'avg_time_delay':   
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [(N-1)xk+1,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)
        # TODO: 计算方式需要重新改一下
        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay})
                                #    'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    if cav_content['ego']:
                        past_k_poses.append(x_to_world(cav_content['past_k'][0]['params']['lidar_pose']))
                        continue
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_opv2v_regular_compensation.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yunshuang Yuan <yunshuang.yuan@ikg.uni-hannover.de>
# Modified by: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for 2-stage backbone intermediate fusion
Only used in FPV-RCNN
"""
import math
from collections import OrderedDict

import numpy as np
import torch
import copy

import opencood.data_utils.datasets
import opencood.data_utils.post_processor as post_processor
from opencood.utils import box_utils
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world

class IntermediateFusionDatasetV2(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """

    def __init__(self, params, visualize, train=True):
        super(IntermediateFusionDatasetV2, self). \
            __init__(params, visualize, train)
        self.pre_processor = \
            build_preprocessor(params['preprocess'], train)
        self.post_processor = \
            post_processor.build_postprocessor(params['postprocess'], train)


    def __getitem__(self, idx):
        # put here to avoid initialization error
        from opencood.pcdet_utils.roiaware_pool3d.roiaware_pool3d_utils \
            import points_in_boxes_cpu
        base_data_dict = self.retrieve_base_data(idx)
        base_data_dict = add_noise_data_dict(base_data_dict,self.params['noise_setting'])

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        ego_id = -1
        ego_lidar_pose = []

        # first find the ego vehicle's lidar pose
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['params']['lidar_pose']
                ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
                break
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        processed_features = []
        object_stack = []
        object_id_stack = []
        cav_object_stack = []
        lidar_pose_stack = []
        lidar_pose_clean_stack = []


        projected_lidar_stack = []
        no_projected_lidar_stack = []


        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            distance = \
                math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
                           ego_lidar_pose[0]) ** 2 + (
                                  selected_cav_base['params'][
                                      'lidar_pose'][1] - ego_lidar_pose[
                                      1]) ** 2)
            if distance > self.params['comm_range']:
                continue

            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                ego_lidar_pose_clean)
            if len(selected_cav_processed['projected_lidar']) > 10:
                object_stack.append(
                    selected_cav_processed['object_bbx_center'])
                object_id_stack += selected_cav_processed['object_ids']
                
                cav_object_stack.append(selected_cav_processed['cav_object_bbx_center'])

                processed_features.append(
                    selected_cav_processed['processed_features'])


                projected_lidar_stack.append(
                    selected_cav_processed['projected_lidar'])

                no_projected_lidar_stack.append(
                    selected_cav_processed['no_projected_lidar'])

                lidar_pose_stack.append(
                        selected_cav_base['params']['lidar_pose'])

                lidar_pose_clean_stack.append(
                    selected_cav_base['params']['lidar_pose_clean'])
                
        if self.visualize:
            projected_lidar_stack = projected_lidar_stack

        # exclude all repetitive objects
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        object_stack_all = np.vstack(object_stack)
        object_stack_all = object_stack_all[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack_all.shape[0], :] = object_stack_all
        mask[:object_stack_all.shape[0]] = 1

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(processed_features)
        merged_feature_dict = self.merge_features_to_dict(processed_features)

        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center, # hwl
                anchors=anchor_box,
                mask=mask)

        # Filter empty boxes
        object_stack_filtered = []
        label_dict_no_coop = []
        cav_label_dict_no_coop = []
        idx_cnt = 0
        for boxes, points in zip(object_stack, projected_lidar_stack):
            point_indices = points_in_boxes_cpu(points[:, :3], boxes[:,
                                                               [0, 1, 2, 5, 4,
                                                                3, 6]])
            cur_mask = point_indices.sum(axis=1) > 0

            # added by yifanlu
            cav_itself_object_center = np.array(cav_object_stack[idx_cnt])
            bbx_center = \
                np.zeros((self.params['postprocess']['max_num'], 7))
            bbx_mask = np.zeros(self.params['postprocess']['max_num'])
            bbx_center[:cav_itself_object_center.shape[0], :] = cav_itself_object_center
            bbx_mask[:cav_itself_object_center.shape[0]] = 1

            cav_label_dict_no_coop.append(
                self.post_processor.generate_label(
                    gt_box_center=bbx_center, # hwl
                    anchors=anchor_box,
                    mask=bbx_mask)
            )
            idx_cnt += 1
            ###

            if cur_mask.sum() == 0:
                label_dict_no_coop.append({
                    'pos_equal_one': np.zeros((*anchor_box.shape[:2],
                                               self.post_processor.anchor_num)),
                    'neg_equal_one': np.ones((*anchor_box.shape[:2],
                                              self.post_processor.anchor_num)),
                    'targets': np.zeros((*anchor_box.shape[:2],
                                         self.post_processor.anchor_num * 7))
                })
                continue
            object_stack_filtered.append(boxes[cur_mask])
            bbx_center = \
                np.zeros((self.params['postprocess']['max_num'], 7))
            bbx_mask = np.zeros(self.params['postprocess']['max_num'])
            bbx_center[:boxes[cur_mask].shape[0], :] = boxes[cur_mask]
            bbx_mask[:boxes[cur_mask].shape[0]] = 1
            label_dict_no_coop.append(
                self.post_processor.generate_label(
                    gt_box_center=bbx_center, # hwl
                    anchors=anchor_box,
                    mask=bbx_mask)
            )

 

        if not self.proj_first:
            label_dict_no_coop = cav_label_dict_no_coop

        label_dict = {
            'stage1': label_dict_no_coop,
            'stage2': label_dict
        }
        processed_data_dict['ego'].update(
            {'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'lidar_poses_clean': np.stack(lidar_pose_clean_stack),
             'lidar_poses': np.stack(lidar_pose_stack),
             'pairwise_t_matrix':self.get_pairwise_transformation(lidar_pose_stack, self.max_cav)})

        if self.proj_first:
            processed_data_dict['ego'].update({'origin_lidar':
                                                projected_lidar_stack})
        else:
            processed_data_dict['ego'].update({'origin_lidar':
                                                no_projected_lidar_stack})
        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar_vis':
                                                np.vstack(
                                                    projected_lidar_stack)})
        return processed_data_dict

    def get_item_single_car(self, selected_cav_base, ego_pose, ego_pose_clean):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        WHYYY YOU CAN PROJECT the LIDAR FIRST!!

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information.
        ego_pose : list
            The ego vehicle lidar pose under world coordinate.
        ego_pose_clean: list
            The clean ego vehicle lidar pose under world coordinate

        Returns
        -------
        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) 

        transformation_matrix_clean = \
            x1_to_x2(selected_cav_base['params']['lidar_pose_clean'],
                    ego_pose_clean)

        # retrieve objects under ego coordinates
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base],
                                                    ego_pose_clean)
        cav_object_bbx_center, cav_object_bbx_mask, cav_object_ids = \
            self.generate_object_center([selected_cav_base],
                                                    selected_cav_base['params']['lidar_pose_clean'])

        # filter lidar
        lidar_np = selected_cav_base['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        # remove points that hit itself
        lidar_np = mask_ego_points(lidar_np)
        # project the lidar to ego space
        projected_lidar = \
            box_utils.project_points_by_matrix_torch(lidar_np[:, :3],
                                                     transformation_matrix)
        
        no_projected_lidar = copy.deepcopy(lidar_np)

        lidar_np[:, :3] = projected_lidar
        
        projected_lidar = mask_points_by_range(lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])

        no_projected_lidar = mask_points_by_range(no_projected_lidar,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])

        if self.proj_first:
            lidar_np = projected_lidar
        else:
            lidar_np = no_projected_lidar
        processed_lidar = self.pre_processor.preprocess(lidar_np)

        selected_cav_processed.update(
            {'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'cav_object_bbx_center': cav_object_bbx_center[cav_object_bbx_mask == 1],
             'projected_lidar': projected_lidar,
             'no_projected_lidar':  no_projected_lidar,
             'processed_features': processed_lidar,
             'transformation_matrix': transformation_matrix,
             'transformation_matrix_clean': transformation_matrix_clean})

        return selected_cav_processed

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature)

        return merged_feature_dict

    def collate_batch_train(self, batch):
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        label_dict_no_coop_list = []
        lidar_pose_list = []
        lidar_pose_clean_list = []
        pairwise_t_matrix_list = []

        origin_lidar = []
        if self.visualize:
            origin_lidar_vis = []

        # added by yys, fpvrcnn needs anchors for
        # first stage proposal generation
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            lidar_pose_list.append(ego_dict['lidar_poses']) # ego_dict['lidar_pose'] is np.ndarray [N,6]
            lidar_pose_clean_list.append(ego_dict['lidar_poses_clean'])

            processed_lidar_list.append(ego_dict['processed_lidar'])
            record_len.append(ego_dict['cav_num'])
            label_dict_no_coop_list.append(ego_dict['label_dict']['stage1'])
            label_dict_list.append(ego_dict['label_dict']['stage2'])
            origin_lidar.append(ego_dict['origin_lidar'])
            if self.visualize:
                origin_lidar_vis.append(ego_dict['origin_lidar_vis'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))

        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        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 <= 5
        record_len = torch.from_numpy(np.array(record_len, dtype=int))
        lidar_pose = torch.from_numpy(np.concatenate(lidar_pose_list, axis=0))
        lidar_pose_clean = torch.from_numpy(np.concatenate(lidar_pose_clean_list, axis=0))
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))
        
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)
        label_dict_no_coop_list_ = [label_dict for label_list in
                                    label_dict_no_coop_list for label_dict in
                                    label_list]
        for i in range(len(label_dict_no_coop_list_)):
            if isinstance(label_dict_no_coop_list_[i], list):
                print('debug')
        label_no_coop_torch_dict = \
            self.post_processor.collate_batch(label_dict_no_coop_list_)
        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'object_bbx_center': object_bbx_center,
                                   'object_bbx_mask': object_bbx_mask,
                                   'processed_lidar': processed_lidar_torch_dict,
                                   'record_len': record_len,
                                   'label_dict': {
                                       'stage1': label_no_coop_torch_dict,
                                       'stage2': label_torch_dict},
                                   'object_ids': object_ids[0],
                                   'lidar_pose_clean': lidar_pose_clean,
                                   'lidar_pose': lidar_pose,
                                   'pairwise_t_matrix': pairwise_t_matrix,})

        coords = []
        idx = 0
        for b in range(len(batch)):
            for points in origin_lidar[b]:
                assert len(points) != 0
                coor_pad = np.pad(points, ((0, 0), (1, 0)),
                                  mode="constant", constant_values=idx)
                coords.append(coor_pad)
                idx += 1
        origin_lidar_for_vsa = np.concatenate(coords, axis=0)

        origin_lidar_for_vsa = torch.from_numpy(origin_lidar_for_vsa)
        output_dict['ego'].update({'origin_lidar_for_vsa': origin_lidar_for_vsa})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar_vis))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        return output_dict

    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.
        The object_bbx_center is in ego coordinate.

        Notice: it is a wrap of postprocessor

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """
        return self.post_processor.generate_object_center(cav_contents,
                                                        reference_lidar_pose)

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def visualize_result(self, pred_box_tensor,
                         gt_tensor,
                         pcd,
                         show_vis,
                         save_path,
                         dataset=None):
        # we need to convert the pcd from [n, 5] -> [n, 4]
        pcd = pcd[:, 1:]
        # visualize the model output
        self.post_processor.visualize(pred_box_tensor,
                                      gt_tensor,
                                      pcd,
                                      show_vis,
                                      save_path,
                                      dataset=dataset)
    
    def get_pairwise_transformation(self, lidar_pose_list, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        lidar_pose_list

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for lidar_pose in lidar_pose_list:
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
opencood/data_utils/datasets/intermediate_fusion_dataset_v2.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Quanhao Li, Yifan Lu
"""
Dataset class for DAIR-V2X interemdiate fusion
Only used by FPV-RCNN 
"""
import random
import math
from collections import OrderedDict

import os
import numpy as np
import torch
from torch.utils.data import DataLoader
import json
import opencood.data_utils.datasets
import opencood.data_utils.post_processor as post_processor
from opencood.utils import box_utils, common_utils

from opencood.data_utils.datasets import intermediate_fusion_dataset_v2
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.transformation_utils import tfm_to_pose
from opencood.utils.transformation_utils import veh_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import inf_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import x_to_world
from opencood.utils.pose_utils import add_noise_data_dict

def load_json(path):
    with open(path, mode="r") as f:
        data = json.load(f)
    return data

class IntermediateFusionDatasetV2DAIR(intermediate_fusion_dataset_v2.IntermediateFusionDatasetV2):
    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        self.max_cav = 2
        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False


        assert 'clip_pc' in params['fusion']['args']
        if params['fusion']['args']['clip_pc']:
            self.clip_pc = True
        else:
            self.clip_pc = False
        
        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None

        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        if self.train:
            split_dir = params['root_dir']
        else:
            split_dir = params['validate_dir']

        self.root_dir = params['data_dir']

        self.split_info = load_json(split_dir)
        co_datainfo = load_json(os.path.join(self.root_dir, 'cooperative/data_info.json'))
        self.co_data = OrderedDict()
        for frame_info in co_datainfo:
            veh_frame_id = frame_info['vehicle_image_path'].split("/")[-1].replace(".jpg", "")
            self.co_data[veh_frame_id] = frame_info

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.
        NOTICE!
        It is different from Intermediate Fusion and Early Fusion
        Parameters
        ----------
        idx : int
            Index given by dataloader.
        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
        """
        veh_frame_id = self.split_info[idx]
        frame_info = self.co_data[veh_frame_id]
        system_error_offset = frame_info["system_error_offset"]
        data = OrderedDict()

        data[0] = OrderedDict()
        data[0]['ego'] = True
        data[1] = OrderedDict()
        data[1]['ego'] = False
                
        data[0]['params'] = OrderedDict()
        data[0]['params']['vehicles'] = load_json(os.path.join(self.root_dir, 'vehicle-side/label/lidar/',veh_frame_id + '.json'))
        lidar_to_novatel_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/lidar_to_novatel/'+str(veh_frame_id)+'.json'))
        novatel_to_world_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/novatel_to_world/'+str(veh_frame_id)+'.json'))

        transformation_matrix = veh_side_rot_and_trans_to_trasnformation_matrix(lidar_to_novatel_json_file,novatel_to_world_json_file)

        data[0]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix)

        data[0]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["vehicle_pointcloud_path"]))

        data[1]['params'] = OrderedDict()
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")

        data[1]['params']['vehicles'] = load_json(os.path.join(self.root_dir, 'infrastructure-side/label/virtuallidar/',inf_frame_id + '.json'))
        virtuallidar_to_world_json_file = load_json(os.path.join(self.root_dir,'infrastructure-side/calib/virtuallidar_to_world/'+str(inf_frame_id)+'.json'))
        transformation_matrix1 = inf_side_rot_and_trans_to_trasnformation_matrix(virtuallidar_to_world_json_file,system_error_offset)
        data[1]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix1)

        data[1]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["infrastructure_pointcloud_path"]))
        return data

    def __len__(self):
        return len(self.split_info)

    ### rewrite generate_object_center ###
    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_dairv2x_late(cav_contents,
                                                        reference_lidar_pose)


    ### rewrite post_process ###
    """
    We have to rewrite post_process for LateFusionDatasetDAIR
    because the object id can not used for identifying the same object
    
    here we will to use the IoU to determine it.
    """
    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx_by_iou(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor



    ### rewrite __getitem___
    ### we should use iou to filter repetitive boxes, instead of object id.
    def __getitem__(self, idx):
        # put here to avoid initialization error
        from opencood.pcdet_utils.roiaware_pool3d.roiaware_pool3d_utils \
            import points_in_boxes_cpu
        base_data_dict = self.retrieve_base_data(idx)
        base_data_dict = add_noise_data_dict(base_data_dict,self.params['noise_setting'])

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        ego_id = -1
        ego_lidar_pose = []

        # first find the ego vehicle's lidar pose
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['params']['lidar_pose']
                ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
                break
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        processed_features = []
        object_stack = []
        object_id_stack = []
        cav_object_stack = []
        lidar_pose_stack = []
        lidar_pose_clean_stack = []


        projected_lidar_stack = []
        no_projected_lidar_stack = []


        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            distance = \
                math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
                           ego_lidar_pose[0]) ** 2 + (
                                  selected_cav_base['params'][
                                      'lidar_pose'][1] - ego_lidar_pose[
                                      1]) ** 2)
            if distance > self.params['comm_range']:
                continue

            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                ego_lidar_pose_clean)
            if len(selected_cav_processed['projected_lidar']) > 10:
                object_stack.append(
                    selected_cav_processed['object_bbx_center'])
                object_id_stack += selected_cav_processed['object_ids']
                
                cav_object_stack.append(selected_cav_processed['cav_object_bbx_center'])

                processed_features.append(
                    selected_cav_processed['processed_features'])


                projected_lidar_stack.append(
                    selected_cav_processed['projected_lidar'])

                no_projected_lidar_stack.append(
                    selected_cav_processed['no_projected_lidar'])

                lidar_pose_stack.append(
                        selected_cav_base['params']['lidar_pose'])

                lidar_pose_clean_stack.append(
                    selected_cav_base['params']['lidar_pose_clean'])
                
        if self.visualize:
            projected_lidar_stack = projected_lidar_stack

        # exclude all repetitive objects for cooperative label
        # DAIR-V2X only has two agent. And cav_object_stack
        # Modified by Yifan Lu 2022.09.24
        if len(object_stack) == 1:
            object_stack_all = object_stack[0]
        else:
            veh_corners_np = box_utils.boxes_to_corners_3d(object_stack[0], self.order)
            inf_corners_np = box_utils.boxes_to_corners_3d(object_stack[1], self.order)
            # inf corners are in inf-coord, should transform to veh-coord
            inf_to_veh = x1_to_x2(lidar_pose_clean_stack[1], lidar_pose_clean_stack[0])
            inf_corners_np = box_utils.project_box3d(inf_corners_np, inf_to_veh)

            inf_polygon_list = list(common_utils.convert_format(inf_corners_np))
            veh_polygon_list = list(common_utils.convert_format(veh_corners_np))
            iou_thresh = 0.05 

            gt_from_inf = []
            for i in range(len(inf_polygon_list)):
                inf_polygon = inf_polygon_list[i]
                ious = common_utils.compute_iou(inf_polygon, veh_polygon_list)
                if (ious > iou_thresh).any():
                    continue
                gt_from_inf.append(inf_corners_np[i])
            
            if len(gt_from_inf):
                gt_from_inf = np.stack(gt_from_inf)
                co_corners_np = np.vstack([veh_corners_np, gt_from_inf])
            else:
                co_corners_np = veh_corners_np
            object_stack_all = box_utils.corner_to_center(co_corners_np, order=self.order)


        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack_all.shape[0], :] = object_stack_all
        mask[:object_stack_all.shape[0]] = 1

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(processed_features)
        merged_feature_dict = self.merge_features_to_dict(processed_features)

        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center, # hwl
                anchors=anchor_box,
                mask=mask)

        # Filter empty boxes
        object_stack_filtered = []
        label_dict_no_coop = []
        cav_label_dict_no_coop = []
        idx_cnt = 0
        for boxes, points in zip(object_stack, projected_lidar_stack):
            point_indices = points_in_boxes_cpu(points[:, :3], boxes[:,
                                                               [0, 1, 2, 5, 4,
                                                                3, 6]])
            cur_mask = point_indices.sum(axis=1) > 0

            # added by yifanlu
            cav_itself_object_center = np.array(cav_object_stack[idx_cnt])
            bbx_center = \
                np.zeros((self.params['postprocess']['max_num'], 7))
            bbx_mask = np.zeros(self.params['postprocess']['max_num'])
            bbx_center[:cav_itself_object_center.shape[0], :] = cav_itself_object_center
            bbx_mask[:cav_itself_object_center.shape[0]] = 1

            cav_label_dict_no_coop.append(
                self.post_processor.generate_label(
                    gt_box_center=bbx_center, # hwl
                    anchors=anchor_box,
                    mask=bbx_mask)
            )
            idx_cnt += 1
            ###

            if cur_mask.sum() == 0:
                label_dict_no_coop.append({
                    'pos_equal_one': np.zeros((*anchor_box.shape[:2],
                                               self.post_processor.anchor_num)),
                    'neg_equal_one': np.ones((*anchor_box.shape[:2],
                                              self.post_processor.anchor_num)),
                    'targets': np.zeros((*anchor_box.shape[:2],
                                         self.post_processor.anchor_num * 7))
                })
                continue
            object_stack_filtered.append(boxes[cur_mask])
            bbx_center = \
                np.zeros((self.params['postprocess']['max_num'], 7))
            bbx_mask = np.zeros(self.params['postprocess']['max_num'])
            bbx_center[:boxes[cur_mask].shape[0], :] = boxes[cur_mask]
            bbx_mask[:boxes[cur_mask].shape[0]] = 1
            label_dict_no_coop.append(
                self.post_processor.generate_label(
                    gt_box_center=bbx_center, # hwl
                    anchors=anchor_box,
                    mask=bbx_mask)
            )

        if not self.proj_first:
            label_dict_no_coop = cav_label_dict_no_coop

        label_dict = {
            'stage1': label_dict_no_coop,
            'stage2': label_dict
        }
        processed_data_dict['ego'].update(
            {'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': list(range(object_stack_all.shape[0])), # meaningless
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'lidar_poses_clean': np.stack(lidar_pose_clean_stack),
             'lidar_poses': np.stack(lidar_pose_stack),
             'pairwise_t_matrix':self.get_pairwise_transformation(lidar_pose_stack, self.max_cav)})

        if self.proj_first:
            processed_data_dict['ego'].update({'origin_lidar':
                                                projected_lidar_stack})
        else:
            processed_data_dict['ego'].update({'origin_lidar':
                                                no_projected_lidar_stack})
        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar_vis':
                                                np.vstack(
                                                    projected_lidar_stack)})
        return processed_data_dict

opencood/data_utils/datasets/intermediate_fusion_dataset_v2_dair.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu

"""
Dataset class for V2X-Sim 2.0 intermediate fusion
Only used in FPV-RCNN
"""
import random
import math
from collections import OrderedDict

import numpy as np
import torch
from torch.utils.data import DataLoader
import pickle

from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
import opencood.data_utils.post_processor as post_processor
from opencood.utils import box_utils
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.datasets import intermediate_fusion_dataset_v2
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2, x_to_world, tfm_to_pose


class IntermediateFusionDatasetV2V2X(intermediate_fusion_dataset_v2.IntermediateFusionDatasetV2):
    """
        IntermediateFusionDatasetV2V2X is children of IntermediateFusionDatasetV2
        It works almost the same as IntermediateFusionDataset,
        But it 
        rewrite __init__, __len__, retrieve_base_data , which is from BaseDataset
        rewrite generate_object_center, which is from IntermediateFusionDataset
    """

    ### rewrite __init__ ###
    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None
            
        # same as IntermediateFusionDataset
        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
            
        if "kd_flag" in params.keys():
            self.kd_flag = params['kd_flag']
        else:
            self.kd_flag = False

        if "box_align" in params.keys():
            self.box_align = True
        else:
            self.box_align = False
        
        print("Dataset dir:", root_dir)

        # v2x: read the pickle file
        with open(root_dir, 'rb') as f:
            dataset_infos = pickle.load(f)  # dataset_infos is a list 

        self.agent_start = eval(min([i[-1] for i in dataset_infos[0].keys() if i.startswith("lidar_pose")]))
        # self.agent_start = 1 when v2x-sim2.0,  = 0 when v2x-sim1.0


        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        self.keyframe_database = OrderedDict()
        self.len_record = len(dataset_infos)

        # loop over all keyframe.
        # data_info is one sample.
        # np.random.seed(303)
        for (i, data_info) in enumerate(dataset_infos):
            self.keyframe_database.update({i: OrderedDict()})

            # at least 1 cav should show up
            cav_num = data_info['agent_num']
            assert cav_num > 0
            select_ego = 1

            # in one keyframe, loop all agent
            for cav_id in range(self.agent_start, self.agent_start + cav_num):

                self.keyframe_database[i][cav_id] = OrderedDict()
                self.keyframe_database[i][cav_id]['lidar'] = data_info[f'lidar_path_{cav_id}']  # maybe add camera in the future
                self.keyframe_database[i][cav_id]['params'] = OrderedDict()
                self.keyframe_database[i][cav_id]['params']['lidar_pose'] = tfm_to_pose(data_info[f"lidar_pose_{cav_id}"]) # tfm in data_info, turn to [x,y,z,roll,yaw,pitch]
                self.keyframe_database[i][cav_id]['params']['vehicles'] = data_info[f'labels_{cav_id}']['gt_boxes_global']
                self.keyframe_database[i][cav_id]['params']['object_ids'] = data_info[f'labels_{cav_id}']['gt_object_ids'].tolist()

                # randomly select ego is better
                if cav_id == select_ego:
                    self.keyframe_database[i][cav_id]['ego'] = True
                    
                else:
                    self.keyframe_database[i][cav_id]['ego'] = False

    ### rewrite __len__ ###
    def __len__(self):
        return self.len_record

    ### rewrite retrieve_base_data ###
    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            lidar_np: (N, 4)
        """
        # we loop the accumulated length list to see get the scenario index
        keyframe = self.keyframe_database[idx]

        data = OrderedDict()
        ego_idx = -1
        # load files for all CAVs
        for cav_id, cav_content in keyframe.items():
            data[cav_id] = OrderedDict()
            data[cav_id]['ego'] = cav_content['ego']
            if cav_content['ego']:
                ego_idx = cav_id
            data[cav_id]['params'] = cav_content['params'] # lidar_pose, vehicles(gt_boxes), object_id(token)

            # load the corresponding data into the dictionary
            nbr_dims = 4 # x,y,z,intensity
            scan = np.fromfile(cav_content['lidar'], dtype='float32')
            points = scan.reshape((-1, 5))[:, :nbr_dims] 
            data[cav_id]['lidar_np'] = points

        # in fact we want to put ego in the first place
        # we will exchange two agent.
        if ego_idx != 0:
            if self.agent_start == 0:  # v2x-sim1.0
                data[ego_idx], data[0] = data[0], data[ego_idx]
            elif self.agent_start == 1:  # v2x-sim2.0
                data[0] = data[ego_idx]
                data.move_to_end(0, last=False) # move to first
                data.pop(ego_idx)

        return data

    ### rewrite generate_object_center ###
    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_v2x(cav_contents,
                                                        reference_lidar_pose)
opencood/data_utils/datasets/intermediate_fusion_dataset_v2_v2x.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu

"""
Dataset class for V2X-Sim 2.0 dataset intermediate fusion
"""
import random
import math
from collections import OrderedDict

import numpy as np
import torch
from torch.utils.data import DataLoader
import pickle

from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
import opencood.data_utils.post_processor as post_processor
from opencood.utils import box_utils
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.datasets import intermediate_fusion_dataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2, x_to_world, tfm_to_pose
from opencood.utils.common_utils import read_json
from icecream import ic


class IntermediateFusionDatasetV2X(intermediate_fusion_dataset.IntermediateFusionDataset):
    """
        IntermediateFusionDatasetV2X is children of IntermediateFusionDataset
        It works almost the same as IntermediateFusionDataset,
        But it 
        rewrite __init__, __len__, retrieve_base_data , which is from BaseDataset
        rewrite generate_object_center, which is from IntermediateFusionDataset
    """

    ### rewrite __init__ ###
    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None
            
        # same as IntermediateFusionDataset
        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
            
        if "kd_flag" in params.keys():
            self.kd_flag = params['kd_flag']
        else:
            self.kd_flag = False

        if "box_align" in params.keys():
            self.box_align = True
            self.stage1_result_path = params['box_align']['train_result'] if train else params['box_align']['val_result']
            self.stage1_result = read_json(self.stage1_result_path)
            self.box_align_args = params['box_align']['args']
        
        else:
            self.box_align = False
        
        print("Dataset dir:", root_dir)

        # v2x: read the pickle file
        with open(root_dir, 'rb') as f:
            dataset_infos = pickle.load(f)  # dataset_infos is a list 

        self.agent_start = eval(min([i[-1] for i in dataset_infos[0].keys() if i.startswith("lidar_pose")]))
        # self.agent_start = 1 when v2x-sim2.0,  = 0 when v2x-sim1.0


        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        self.keyframe_database = OrderedDict()
        self.len_record = len(dataset_infos)

        # loop over all keyframe.
        # data_info is one sample.
        # np.random.seed(303)
        for (i, data_info) in enumerate(dataset_infos):
            self.keyframe_database.update({i: OrderedDict()})

            # at least 1 cav should show up
            cav_num = data_info['agent_num']
            assert cav_num > 0
            select_ego = 1

            # in one keyframe, loop all agent
            for cav_id in range(self.agent_start, self.agent_start + cav_num):

                self.keyframe_database[i][cav_id] = OrderedDict()
                self.keyframe_database[i][cav_id]['lidar'] = data_info[f'lidar_path_{cav_id}']  # maybe add camera in the future
                self.keyframe_database[i][cav_id]['params'] = OrderedDict()
                self.keyframe_database[i][cav_id]['params']['lidar_pose'] = tfm_to_pose(data_info[f"lidar_pose_{cav_id}"]) # tfm in data_info, turn to [x,y,z,roll,yaw,pitch]
                self.keyframe_database[i][cav_id]['params']['vehicles'] = data_info[f'labels_{cav_id}']['gt_boxes_global']
                self.keyframe_database[i][cav_id]['params']['object_ids'] = data_info[f'labels_{cav_id}']['gt_object_ids'].tolist()

                if cav_id == select_ego:
                    self.keyframe_database[i][cav_id]['ego'] = True
                    
                else:
                    self.keyframe_database[i][cav_id]['ego'] = False

    ### rewrite __len__ ###
    def __len__(self):
        return self.len_record

    ### rewrite retrieve_base_data ###
    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            lidar_np: (N, 4)
        """
        # we loop the accumulated length list to see get the scenario index
        keyframe = self.keyframe_database[idx]

        data = OrderedDict()
        ego_idx = -1
        # load files for all CAVs
        for cav_id, cav_content in keyframe.items():
            data[cav_id] = OrderedDict()
            data[cav_id]['ego'] = cav_content['ego']
            if cav_content['ego']:
                ego_idx = cav_id
            data[cav_id]['params'] = cav_content['params'] # lidar_pose, vehicles(gt_boxes), object_id(token)

            # load the corresponding data into the dictionary
            nbr_dims = 4 # x,y,z,intensity
            scan = np.fromfile(cav_content['lidar'], dtype='float32')
            points = scan.reshape((-1, 5))[:, :nbr_dims] 
            data[cav_id]['lidar_np'] = points

        # in fact we want to put ego in the first place
        # we will exchange two agent.
        if ego_idx != 0:
            if self.agent_start == 0:  # v2x-sim1.0
                data[ego_idx], data[0] = data[0], data[ego_idx]
            elif self.agent_start == 1:  # v2x-sim2.0
                data[0] = data[ego_idx]
                data.move_to_end(0, last=False) # move to first
                data.pop(ego_idx)

        return data

    ### rewrite generate_object_center ###
    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_v2x(cav_contents,
                                                        reference_lidar_pose)
opencood/data_utils/datasets/intermediate_fusion_dataset_v2x.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/1/27
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for intermediate fusion with past k frames on irregular OPV2V
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetV2XSET(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):

        self.times = []

        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
            print(f'======= is_ab_regular? {self.is_ab_regular} =======')
        
        self.sample_interval_exp = int(self.binomial_n * self.binomial_p)

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        '''
        scenario_database Structure: 
        {
            scenario_id : {
                cav_1 : {
                    'ego' : true / false , 
                    timestamp1 : {
                        yaml: path,
                        lidar: path, 
                        cameras: list of path
                    },
                    ...
                },
                ...
            }
        }
        '''
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # # copy timestamps npy file
            # timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
            # time_annotations = np.load(timestamps_file)

            # at least 1 cav should show up
            cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
            assert len(cav_list) > 0

            # roadside unit data's id is always negative, so here we want to
            # make sure they will be in the end of the list as they shouldn't
            # be ego vehicle.
            if int(cav_list[0]) < 0:
                cav_list = cav_list[1:] + [cav_list[0]]

            # use the frame number as key, the full path as the values
            yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
            if len(yaml_files)==0:
                yaml_files = sorted([x 
                            for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                            x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

            regular_timestamps = self.extract_timestamps(yaml_files)

            # start_timestamp = int(float(self.extract_timestamps(yaml_files)[0]))
            # while(1):
            #     time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
            #     time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
            #     if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
            #         start_timestamp += 1
            #     else:
            #         break

            # end_timestamp = int(float(self.extract_timestamps(yaml_files)[-1]))
            # if start_timestamp%2 == 0:
            #     # even
            #     end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
            # else:
            #     end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
            # num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)
            # regular_timestamps = [start_timestamp+2*i for i in range(num_timestamps)]

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                cav_path = os.path.join(scenario_folder, cav_id)

                timestamps = regular_timestamps

                for timestamp in timestamps:
                    timestamp = "%06d" % float(timestamp)
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files
                
                self.scenario_database[i][cav_id]['regular'] = self.scenario_database[i][cav_id]
                
                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    # num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    num_ego_timestamps = len(timestamps) - self.binomial_n * self.k
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)
                else:
                    self.scenario_database[i][cav_id]['ego'] = False
                    

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead. 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("=== OPV2V-Irregular Multi-sweep dataset with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def extract_timestamps(self, yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]
            if res.endswith('.yaml'):
                timestamp = res.replace('.yaml', '')
            elif res.endswith('.json'):
                timestamp = res.replace('.json', '')
            else:
                print("Woops! There is no processing method for file {}".format(res))
                sys.exit(1)
            timestamps.append(timestamp)

        return timestamps

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)
                        ...,		
                        [k-1] : {} (id-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k-1]:{}  (id-\tau-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * self.k
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # TODO: debug use, uncomment: to use regular version yaml GT
            # time_new = str( '%06d' % int(float(timestamp_key)))
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # TODO: debug use, uncomment: to use regular version lidar input
            # time_new = str( '%06d' % int(float(timestamp_key))) 
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                # sample_interval
                if data[cav_id]['ego']:             # ego sample_interval = E(B(n, p))
                    if i == 0: # ego-past-0 与 ego-curr 是一样的
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    sample_interval = self.sample_interval_exp
                    if sample_interval == 0:
                        sample_interval = 1
                else:                               # non-ego sample_interval ~ B(n, p)
                    if self.sample_interval_exp==0 \
                        and self.is_no_shift \
                            and i == 0:
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==0: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>0: # 过去的几帧不要重复
                            sample_interval = 1                

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float

        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            },
            cav_id_2 : { ... }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = math.sqrt( \
                (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换, 用于单体监督            cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: N, k
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) TODO: 这里面没有搞懂为什么不用 past_k_tr_mats

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack})
            #  'times': self.times})

        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            processed_data_dict['ego'].update({'avg_sample_interval':\
                sum(past_k_sample_interval_stack[self.k:]) / len(past_k_sample_interval_stack[self.k:])})
            processed_data_dict['ego'].update({'avg_time_delay':\
                sum(past_k_time_diffs_stack[self.k:]) / len(past_k_time_diffs_stack[self.k:])})
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            '''
            # 5. single view label
            # past_i label at past_i single view
            # opencood/data_utils/post_processor/base_postprocessor.py
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # generate the anchor boxes
            # opencood/data_utils/post_processor/voxel_postprocessor.py
            anchor_box = self.anchor_box
            single_view_label_dict = self.post_processor.generate_label(
                    gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
                )
            past_k_label_dicts.append(single_view_label_dict)
            '''

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
             'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
            #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
            #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
            #  'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def dist_time(ts1, ts2, i = -1):
        """caculate the time interval between two timestamps

        Args:
            ts1 (string): time stamp at some time
            ts2 (string): current time stamp
            i (int, optional): past frame id, for debug use. Defaults to -1.
        
        Returns:
            time_diff (float): time interval (ts1 - ts2)
        """
        if not i==-1:
            return -i
        else:
            return (float(ts1) - float(ts2))

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay})
                                #    'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_v2xset.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/2/13

"""
Dataset class for intermediate fusion with 
1) past k frames as input 
2) current frame as supervision
on irregular OPV2V
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset, intermediate_fusion_dataset_v2xset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetV2XSETCompensation(intermediate_fusion_dataset_v2xset.IntermediateFusionDatasetV2XSET):
    """
    This class is for intermediate fusion on OPV2V_irregular, with past k frames and current frame as input. 
    Note: current frame is used for supervison, not for feature fusion!
    """
    def __init__(self, params, visualize, train=True):
        super(IntermediateFusionDatasetV2XSETCompensation, self).__init__(params, visualize, train)
        self.k = self.k + 1 # 用来在 past_k 中添加 current frame
        print("=== OPV2V-Irregular Multi-sweep dataloader (Compensation Ver.) for intermediate fusion with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k-1, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{                       # current
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)          # past 1
                        ...,		
                        [k] : {} (id-(k-1))         # past k
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k] : {}
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * (self.k-1)
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            try:
                timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            except IndexError:
                print('IndexError!!!')
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # time_new = str( '%06d' % int(float(timestamp_key))) # TODO: debug用 使用regular的GT
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0]) # TODO:
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # time_new = str( '%06d' % int(float(timestamp_key))) # TODO: debug用 使用regular的lidar
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                if i == 0: # ego-past-0 与 ego-curr 是一样的, 用于feature监督
                    data[cav_id]['past_k'][i] = data[cav_id]['curr']
                    continue
                # sample_interval
                if data[cav_id]['ego']:             # ego: do not need past frame
                    break
                    # sample_interval = self.sample_interval_exp
                else:                               # non-ego sample_interval ~ B(n, p)
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==0: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>0: # 过去的几帧不要重复
                            sample_interval = 1  

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                # uncomment code, use updated gt
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float

        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (1) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    }
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
                
            }, 
            cav_id_2 : {
                'ego' : false,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)          # past 1
                    ...,		                # ...
                    [k] : {} (id-(k-1))         # past k
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            if cav_id != ego_id:
                distance = math.sqrt( \
                    (selected_cav_base['past_k'][1]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                        (selected_cav_base['past_k'][1]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
                # if distance is too far, we will just skip this agent
                if distance > self.params['comm_range']:
                    too_far.append(cav_id)
                    continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true / false , 
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (1 / k) totally
                    [0]:{                       # current
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)          # past 1 
                    ...,		                # ...
                    [k] : {} (id-(k-1))         # past k
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k+1, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k+1, 6
            past_k_pose_stack += selected_cav_processed['past_k_poses']
            # time differences: Nx(k+1)
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: Nx(k+1)
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k+1, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) 

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, 6)  # (Nx(k-1)+1, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack})
            #  'times': self.times})

        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            tmp = past_k_time_diffs_stack[1:].reshape(-1, self.k)
            tmp_past_k_time_diffs = tmp[:, 1:] - tmp[:, :-1]
            avg_time_diff = sum(tmp_past_k_time_diffs.reshape(-1)) / tmp_past_k_time_diffs.reshape(-1).shape[0]
            processed_data_dict['ego'].update({'avg_time_delay':\
                avg_time_diff})
            
            tmp_past_k_sample_interval = past_k_sample_interval_stack[1:].reshape(-1, self.k)[:, 1:]
            avg_past_k_sample_interval = sum(tmp_past_k_sample_interval.reshape(-1)) / tmp_past_k_sample_interval.reshape(-1).shape[0]
            processed_data_dict['ego'].update({'avg_sample_interval':\
                avg_past_k_sample_interval})
            
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict

    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        if selected_cav_base['ego']:
            k = 1
        else:
            k = self.k
        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            '''
            # 5. single view label
            # past_i label at past_i single view
            # opencood/data_utils/post_processor/base_postprocessor.py
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # generate the anchor boxes
            # opencood/data_utils/post_processor/voxel_postprocessor.py
            anchor_box = self.anchor_box
            single_view_label_dict = self.post_processor.generate_label(
                    gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
                )
            past_k_label_dicts.append(single_view_label_dict)
            '''

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
             'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
            #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
            #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
            #  'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        return selected_cav_processed

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(len(processed_feature_list[cav_id])):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict


    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_label':          single_label_dict_stack,
            'curr_processed_lidar':         merged_curr_feature_dict,
            'object_bbx_center':            object_bbx_center,
            'object_bbx_mask':              mask,
            # 'anchor_box':                   anchor_box,
            'processed_lidar':              merged_feature_dict,
            'label_dict':                   label_dict,
            'object_ids':                   [object_id_stack[i] for i in unique_indices],
            'single_past_dict':             
            'pairwise_t_matrix':            pairwise_t_matrix,
            'cav_num':                      cav_num,
            'curr_lidar_poses':             curr_lidar_poses,
            'past_lidar_poses':             past_k_lidar_poses,
            'past_k_time_interval':         
            'past_k_sample_interval':   
            'avg_sample_interval': , 
            'avg_time_delay':   
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [(N-1)xk+1,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)
        # TODO: 计算方式需要重新改一下
        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay})
                                #    'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    if cav_content['ego']:
                        past_k_poses.append(x_to_world(cav_content['past_k'][0]['params']['lidar_pose']))
                        continue
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_v2xset_compensation.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/4/15
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
two stage flow update framework

"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
from opencood.utils.flow_utils import generate_flow_map, generate_flow_map_szwei

from opencood.utils.box_utils import boxes_to_corners2d # for debug use
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetV2XSETFlowNew(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):

        self.times = []

        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
            print("!!! Absolutely Regular !!!")

        # 控制是否需要生成GT flow
        self.is_generate_gt_flow = False
        if 'is_generate_gt_flow' in params and params['is_generate_gt_flow']:
            self.is_generate_gt_flow = True
        
        # 只有在绘制每个sample的匹配框时用到 sizhewei
        self.viz_bbx_flag = False
        if 'viz_bbx_flag' in params and params['viz_bbx_flag']:
            self.viz_bbx_flag = True

        self.num_roi_thres = -1
        if 'num_roi_thres' in params:
            self.num_roi_thres = params['num_roi_thres']
        
        self.sample_interval_exp = int(self.binomial_n * self.binomial_p)

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        '''
        scenario_database Structure: 
        {
            scenario_id : {
                cav_1 : {
                    'ego' : true / false , 
                    timestamp1 : {
                        yaml: path,
                        lidar: path, 
                        cameras: list of path
                    },
                    ...
                },
                ...
            }
        }
        '''
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # # copy timestamps npy file, store all irregular timestamps of each non-ego vehicle in this scenario
            # timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
            # time_annotations = np.load(timestamps_file)

            # at least 1 cav should show up
            cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
            assert len(cav_list) > 0

            # roadside unit data's id is always negative, so here we want to
            # make sure they will be in the end of the list as they shouldn't
            # be ego vehicle.
            if int(cav_list[0]) < 0:
                cav_list = cav_list[1:] + [cav_list[0]]

            # use the frame number as key, the full path as the values, store all json or yaml files in this scenario
            yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
            if len(yaml_files)==0:
                yaml_files = sorted([x 
                            for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                            x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

            regular_timestamps = self.extract_timestamps(yaml_files)

            # start_timestamp = int(float(self.extract_timestamps(yaml_files)[0]))
            # while(1):
            #     time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
            #     time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
            #     if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
            #         start_timestamp += 1
            #     else:
            #         break

            # end_timestamp = int(float(self.extract_timestamps(yaml_files)[-1]))
            # if start_timestamp%2 == 0:
            #     # even
            #     end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
            # else:
            #     end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
            # num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)
            # regular_timestamps = [start_timestamp+2*i for i in range(num_timestamps)]

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                cav_path = os.path.join(scenario_folder, cav_id)

                timestamps = regular_timestamps
                
                for timestamp in timestamps:
                    timestamp = "%06d" % int(timestamp)
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files
                
                # regular的timestamps 用于做 curr 真实时刻的ground truth
                self.scenario_database[i][cav_id]['regular'] = self.scenario_database[i][cav_id]

                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    # num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    num_ego_timestamps = len(timestamps) - self.binomial_n * self.k # * 3 # TODO: 
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)
                else:
                    self.scenario_database[i][cav_id]['ego'] = False
                    

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead. 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("=== OPV2V-Irregular Multi-sweep dataset with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def extract_timestamps(self, yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]
            if res.endswith('.yaml'):
                timestamp = res.replace('.yaml', '')
            elif res.endswith('.json'):
                timestamp = res.replace('.json', '')
            else:
                print("Woops! There is no processing method for file {}".format(res))
                sys.exit(1)
            timestamps.append(timestamp)

        return timestamps

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)
                        ...,		
                        [k-1] : {} (id-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k-1]:{}  (id-\tau-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * self.k # * 3 # TODO:
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # TODO: debug use, uncomment: to use regular version yaml GT
            # time_new = str( '%06d' % int(float(timestamp_key)))
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # TODO: debug use, uncomment: to use regular version lidar input
            # time_new = str( '%06d' % int(float(timestamp_key))) 
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                # sample_interval
                if data[cav_id]['ego']:             # ego sample_interval = E(B(n, p))
                    if i == 0: # ego-past-0 与 ego-curr 是一样的
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    sample_interval = self.sample_interval_exp
                    if sample_interval == 0:
                        sample_interval = 1
                else:                               # non-ego sample_interval ~ B(n, p)
                    # delay_debug = 6
                    # if i == 0:
                    #     sample_interval = 3 #delay_debug
                    # elif i ==1:
                    #     sample_interval = 3
                    #     # sample_set = [2,3]
                    #     # import random
                    #     # sample_interval = random.sample(sample_set, 1)[0] #3 #10 - delay_debug
                    #     # trails = bernoulliDist.rvs(self.binomial_n)
                    #     # sample_interval = sum(trails)
                    # else:
                    #     sample_interval = 3
                    #     # trails = bernoulliDist.rvs(self.binomial_n)
                    #     # sample_interval = sum(trails)
                    if self.sample_interval_exp==0 \
                        and self.is_no_shift \
                            and i == 0:
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==0: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>0: # 过去的几帧不要重复
                            sample_interval = 1                

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float
        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            },
            cav_id_2 : { ... }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = math.sqrt( \
                (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        pastk_2_past0_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        if self.is_generate_gt_flow:
            flow_gt = []
            warp_mask = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []

        if self.viz_bbx_flag:
            single_lidar_stack = []
            single_object_stack = []
            single_object_id_stack = []
            single_mask_stack = []
            single_past_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换, 用于单体监督            cav view + curr 的label
                'single_object_bbx_center'
                'single_object_ids'
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
                'flow_gt':               # [2, H, W] flow ground truth
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            if self.viz_bbx_flag:
                single_lidar_stack.append(selected_cav_processed['single_lidar'])
                single_object_stack.append(selected_cav_processed['single_object_bbx_center'])
                single_object_id_stack.append(selected_cav_processed['single_object_ids'])
                single_past_lidar_stack.append(selected_cav_processed['single_past_lidar'])
                # mask = np.zeros(self.params['postprocess']['max_num'])
                # mask[:single_object_stack.shape[0]] = 1
                # single_mask_stack.append(mask)
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: N, k
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k frames to cav past 0 frame trans matrix, list of len=N, pastk_2_past0_tr_mats[i]: ndarray(k, 4, 4)
            pastk_2_past0_tr_mats.append(selected_cav_processed['pastk_2_past0_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        
            if self.is_generate_gt_flow:
                # for flow
                flow_gt.append(selected_cav_processed['flow_gt'])
                warp_mask.append(selected_cav_processed['warp_mask'])
        
        pastk_2_past0_tr_mats = np.stack(pastk_2_past0_tr_mats, axis=0) # N, k, 4, 4

        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) TODO: 这里面没有搞懂为什么不用 past_k_tr_mats

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack, 
             'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats})
            #  'times': self.times})

        if self.is_generate_gt_flow:
            flow_gt = np.vstack(flow_gt) # (N, H, W, 2)
            processed_data_dict['ego'].update({'flow_gt': flow_gt})
            warp_mask = np.vstack(warp_mask) # (N, C, H, W)
            processed_data_dict['ego'].update({'warp_mask': warp_mask})
        
        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            tmp = past_k_time_diffs_stack[self.k:].reshape(-1, self.k) # (N, k)
            
            # tmp_past_k_time_diffs = np.concatenate((tmp[:, :1] , (tmp[:, 1:] - tmp[:, :-1])), axis=1) # (N, k) # TODO: 不同的方法这个计算方式不一样
            tmp_past_k_time_diffs = tmp[:, :1] # (N, 1) irregular setting 下，只用最近的一个时间间隔
            
            avg_time_diff = sum(tmp_past_k_time_diffs.reshape(-1)) / tmp_past_k_time_diffs.reshape(-1).shape[0]
            processed_data_dict['ego'].update({'avg_time_delay':\
                avg_time_diff})

            tmp = past_k_sample_interval_stack[self.k:].reshape(-1, self.k) # (N, k)
            # avg_sample_interval = sum(tmp.reshape(-1)) / len(tmp.reshape(-1)) # TODO: 不同的方法这个计算方式不一样
            avg_sample_interval = sum(tmp[:, :1].reshape(-1)) / len(tmp[:, :1].reshape(-1)) # (N, 1) irregular setting 下，只用最近的一个时间间隔
            processed_data_dict['ego'].update({'avg_sample_interval':\
                avg_sample_interval})

            tmp_var = np.var(tmp, axis=1) # (N, )
            avg_var = np.mean(tmp_var)
            processed_data_dict['ego'].update({'avg_var':\
                avg_var})
            
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        if self.viz_bbx_flag:
            for id, cav in enumerate(cav_id_list):
                processed_data_dict[id] = {}
                processed_data_dict[id].update({
                    'single_lidar': single_lidar_stack[id],
                    'single_object_bbx_center': single_object_stack[id],
                    # 'single_object_bbx_mask': single_mask_stack[i],
                    'single_object_ids': single_object_id_stack[id],
                    'single_past_lidar': single_past_lidar_stack[id]
                })

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                "single_object_bbx_center": single_object_bbx_center,       # 用于viz single view
                "single_object_ids": single_object_ids,                # 用于viz single view
                'flow_gt':              # single view flow
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'pastk_2_past0_tr_mats':    # list of past k frames' transformation matrix to past 0 frame
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        if self.viz_bbx_flag:
            selected_cav_processed.update({'single_lidar': lidar_np[:, :3]})
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past_k to past_0 tansfomation matrix
        pastk_2_past0_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)
            # past_k trans past_0 matrix
            pastk_2_past0 = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], selected_cav_base['past_k'][0]['params']['lidar_pose'])
            pastk_2_past0_tr_mats.append(pastk_2_past0)
            
            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            if self.viz_bbx_flag:
                selected_cav_processed.update({'single_past_lidar': lidar_np[:, :3]})
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            # # 5. single view label
            # # past_i label at past_i single view
            # # opencood/data_utils/post_processor/base_postprocessor.py
            # object_bbx_center, object_bbx_mask, object_ids = \
            #     self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # # generate the anchor boxes
            # # opencood/data_utils/post_processor/voxel_postprocessor.py
            # anchor_box = self.anchor_box
            # single_view_label_dict = self.post_processor.generate_label(
            #         gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            #     )
            # past_k_label_dicts.append(single_view_label_dict)
        
        

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)
        pastk_2_past0_tr_mats = np.stack(pastk_2_past0_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        single_object_bbx_center, single_object_bbx_mask, single_object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=single_object_bbx_center, anchors=anchor_box, mask=single_object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update({
            "single_label_dict": label_dict,
            "single_object_bbx_center": single_object_bbx_center[single_object_bbx_mask == 1],
            "single_object_ids": single_object_ids,
            "curr_feature": curr_feature,
            'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
            'object_ids': object_ids,
            'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
            'past_k_tr_mats': past_k_tr_mats,
            'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats,
            'past_k_poses': past_k_poses,
            'past_k_features': past_k_features,
            'past_k_time_diffs': past_k_time_diffs,
            'past_k_sample_interval': past_k_sample_interval,
        #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
        #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
        #  'past_k_label_dicts': past_k_label_dicts,
            'if_no_point': if_no_point
            })

        if self.is_generate_gt_flow:
            # generate flow, from past_0 and curr
            prev_object_id_stack = {}
            prev_object_stack = {}
            for t_i in range(2):
                split_part = selected_cav_base['past_k'][0] if t_i == 0 else selected_cav_base['curr'] # TODO: 这里面的 prev 和 curr 可能反了
                object_bbx_center, object_bbx_mask, object_ids = \
                    self.generate_object_center([split_part], selected_cav_base['past_k'][0]['params']['lidar_pose'])
                prev_object_id_stack[t_i] = object_ids
                prev_object_stack[t_i] = object_bbx_center
            
            for t_i in range(2):
                unique_object_ids = list(set(prev_object_id_stack[t_i]))
                unique_indices = \
                    [prev_object_id_stack[t_i].index(x) for x in unique_object_ids]
                prev_object_stack[t_i] = np.vstack(prev_object_stack[t_i])
                prev_object_stack[t_i] = prev_object_stack[t_i][unique_indices]
                prev_object_id_stack[t_i] = unique_object_ids
            
            # TODO: generate_flow_map: yhu, generate_flow_map_szwei: szwei
            flow_map, warp_mask = generate_flow_map_szwei(prev_object_stack,
                                            prev_object_id_stack,
                                            self.params['preprocess']['cav_lidar_range'],
                                            self.params['preprocess']['args']['voxel_size'],
                                            past_k=1)

            selected_cav_processed.update({'flow_gt': flow_map})
            selected_cav_processed.update({'warp_mask': warp_mask})

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def dist_time(ts1, ts2, i = -1):
        """caculate the time interval between two timestamps

        Args:
            ts1 (string): time stamp at some time
            ts2 (string): current time stamp
            i (int, optional): past frame id, for debug use. Defaults to -1.
        
        Returns:
            time_diff (float): time interval (ts1 - ts2)
        """
        if not i==-1:
            return -i
        else:
            return (float(ts1) - float(ts2))

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'flow_gt': [N, C, H, W]
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        past_k_avg_time_var = []
        pastk_2_past0_tr_mats = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []
        
        if self.is_generate_gt_flow:
            # flow gt
            flow_gt_list = []
            warp_mask_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            try:
                past_k_avg_time_var.append(ego_dict['avg_var'])
            except KeyError:
                past_k_avg_time_var.append(-1.0)
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            pastk_2_past0_tr_mats.append(ego_dict['pastk_2_past0_tr_mats'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            if self.is_generate_gt_flow:
                flow_gt_list.append(ego_dict['flow_gt'])
                warp_mask_list.append(ego_dict['warp_mask'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        past_k_avg_time_var = np.array(past_k_avg_time_var)
        avg_time_var = float(sum(past_k_avg_time_var) / len(past_k_avg_time_var))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        pastk_2_past0_tr_mats = torch.from_numpy(np.vstack(pastk_2_past0_tr_mats))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        if self.is_generate_gt_flow:
            # for flow
            flow_gt = torch.from_numpy(np.vstack(flow_gt_list))
            label_torch_dict.update({'flow_gt': flow_gt})
            warp_mask = torch.from_numpy(np.vstack(warp_mask_list))
            label_torch_dict.update({'warp_mask': warp_mask})

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'pastk_2_past0_tr_mats': pastk_2_past0_tr_mats,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay,
                                   'avg_time_var': avg_time_var})
                                #    'times': time_consume})
        output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(self.anchor_box))})
        
        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})

        if self.viz_bbx_flag:
            single_lidar = []
            single_object_bbx_center = []
            single_object_mask = []
            single_object_ids = []
            single_past_lidar = []
            for i in range(len(batch[0].keys())-1):
                single_lidar.append(torch.from_numpy(batch[0][i]['single_lidar']))
                single_past_lidar.append(torch.from_numpy(batch[0][i]['single_past_lidar']))
                single_object_bbx_center.append(torch.from_numpy(batch[0][i]['single_object_bbx_center']))
                single_object_ids.append(batch[0][i]['single_object_ids'])
            output_dict['ego'].update({
                'single_lidar_list': single_lidar,
                'single_object_bbx_center': single_object_bbx_center,
                'single_object_ids': single_object_ids,
                'single_past_lidar_list': single_past_lidar
            })
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

    def generate_pred_bbx_frames(self, m_single, trans_mat_pastk_2_past0, past_time_diff, anchor_box):
        '''
        m_single : {
            psm_single, rm_single, (dm_single)
        }
        box_result : {
            'past_k_time_diff' : 
            [0] : {
                pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. (n, 8, 3)
                pred_box_center_tensor : (n, 7)
                scores: (n, )
            },
            ...
            [k-1] : { ... }
        }
        '''
        box_results = self.post_processor.single_post_process(m_single, trans_mat_pastk_2_past0, past_time_diff, anchor_box, self.k, self.num_roi_thres)
        return box_results

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_v2xset_irregular_flow_new.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/1/27
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
为了给v2vnet 等baseline方法训练 用的数据集
有一些baseline方法需要使用多个车之间的pairwise_t_matrix 在多帧的数据集里面忽略了这点 该类将多车之间的坐标系变换补充好
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# global, for debug use
illegal_path_list = set()

class IntermediateFusionDatasetV2XSETSingleframe(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    deep features to ego.
    """
    def __init__(self, params, visualize, train=True):

        self.times = []

        self.params = params
        self.visualize = visualize
        self.train = train

        self.pre_processor = None
        self.post_processor = None
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in params:          # number of time delay
            self.tau = params['time_delay'] 
        else:
            self.tau = 0

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 0

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 1

        # 控制 past-0 处是否有扰动
        self.is_no_shift = False
        if 'is_no_shift' in params and params['is_no_shift']:
            self.is_no_shift = True

        # 控制每个场景内不同车的采样间隔是否一致
        self.is_same_sample_interval = False
        if 'is_same_sample_interval' in params and params['is_same_sample_interval']:
            self.is_same_sample_interval = True

        # 控制是否采用完全 regular 的数据 （整数timestamp）
        self.is_ab_regular = False
        if 'is_ab_regular' in params and params['is_ab_regular']:
            self.is_ab_regular = True
        
        self.sample_interval_exp = int(self.binomial_n * self.binomial_p)

        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']
        
        print("Dataset dir:", root_dir)

        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        # first load all paths of different scenarios
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        '''
        scenario_database Structure: 
        {
            scenario_id : {
                cav_1 : {
                    'ego' : true / false , 
                    timestamp1 : {
                        yaml: path,
                        lidar: path, 
                        cameras: list of path
                    },
                    ...
                },
                ...
            }
        }
        '''
        self.scenario_database = OrderedDict()
        self.len_record = []

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            self.scenario_database.update({i: OrderedDict()})

            # # copy timestamps npy file
            # timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
            # time_annotations = np.load(timestamps_file)

            # at least 1 cav should show up
            cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
            assert len(cav_list) > 0

            # roadside unit data's id is always negative, so here we want to
            # make sure they will be in the end of the list as they shouldn't
            # be ego vehicle.
            if int(cav_list[0]) < 0:
                cav_list = cav_list[1:] + [cav_list[0]]

            # use the frame number as key, the full path as the values
            yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
            if len(yaml_files)==0:
                yaml_files = sorted([x 
                            for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                            x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

            regular_timestamps = self.extract_timestamps(yaml_files)

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):
                if j > self.max_cav - 1:
                    print('too many cavs')
                    break
                self.scenario_database[i][cav_id] = OrderedDict()

                cav_path = os.path.join(scenario_folder, cav_id)

                timestamps = regular_timestamps

                for timestamp in timestamps:
                    timestamp = "%06d" % float(timestamp)
                    self.scenario_database[i][cav_id][timestamp] = \
                        OrderedDict()

                    yaml_file = os.path.join(cav_path,
                                             timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                              timestamp + '.pcd')
                    # camera_files = self.load_camera_files(cav_path, timestamp)

                    self.scenario_database[i][cav_id][timestamp]['yaml'] = \
                        yaml_file
                    self.scenario_database[i][cav_id][timestamp]['lidar'] = \
                        lidar_file
                    # self.scenario_database[i][cav_id][timestamp]['camera0'] = \
                        # camera_files
                
                self.scenario_database[i][cav_id]['regular'] = self.scenario_database[i][cav_id]

                # Assume all cavs will have the same timestamps length. Thus
                # we only need to calculate for the first vehicle in the
                # scene.
                if j == 0:  # ego 
                    # we regard the agent with the minimum id as the ego
                    self.scenario_database[i][cav_id]['ego'] = True
                    # num_ego_timestamps = len(timestamps) - (self.tau + self.k - 1)		# 从第 tau+k 个往后, store 0 时刻的 time stamp
                    num_ego_timestamps = len(timestamps) - self.binomial_n * self.k * 3 # TODO:
                    if not self.len_record:
                        self.len_record.append(num_ego_timestamps)
                    else:
                        prev_last = self.len_record[-1]
                        self.len_record.append(prev_last + num_ego_timestamps)
                else:
                    self.scenario_database[i][cav_id]['ego'] = False
                    

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead. 
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        self.anchor_box = self.post_processor.generate_anchor_box()

        print("=== OPV2V-Irregular Multi-sweep dataset with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def extract_timestamps(self, yaml_files):
        """
        Given the list of the yaml files, extract the mocked timestamps.

        Parameters
        ----------
        yaml_files : list
            The full path of all yaml files of ego vehicle

        Returns
        -------
        timestamps : list
            The list containing timestamps only.
        """
        timestamps = []

        for file in yaml_files:
            res = file.split('/')[-1]
            if res.endswith('.yaml'):
                timestamp = res.replace('.yaml', '')
            elif res.endswith('.json'):
                timestamp = res.replace('.json', '')
            else:
                print("Woops! There is no processing method for file {}".format(res))
                sys.exit(1)
            timestamps.append(timestamp)

        return timestamps

    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		                # (k) totally
                        [0]:{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)
                        ...,		
                        [k-1] : {} (id-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {		                # (k) totally
                    'ego': false, 
                    curr : 	{
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {
                        [0] : {
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 	
                        [k-1]:{}  (id-\tau-(k-1))
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                }, 
                ...
            }
        """
        # we loop the accumulated length list to get the scenario index
        scenario_index = 0
        for i, ele in enumerate(self.len_record):
            if idx < ele:
                scenario_index = i
                break
        scenario_database = self.scenario_database[scenario_index]
        
        # 生成冻结分布函数
        bernoulliDist = stats.bernoulli(self.binomial_p) 

        data = OrderedDict()
        # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
        curr_timestamp_idx = idx if scenario_index == 0 else \
                        idx - self.len_record[scenario_index - 1]
        curr_timestamp_idx = curr_timestamp_idx + self.binomial_n * self.k * 3 # TODO:
        
        # load files for all CAVs
        for cav_id, cav_content in scenario_database.items():
            '''
            cav_content 
            {
                timestamp_1 : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                ...
                timestamp_n : {
                    yaml: path,
                    lidar: path, 
                    cameras: list of path
                },
                'regular' : {
                    timestamp_1 : {},
                    ...
                    timestamp_n : {}
                },
                'ego' : true / false , 
            },
            '''
            data[cav_id] = OrderedDict()
            
            # 1. set 'ego' flag
            data[cav_id]['ego'] = cav_content['ego']
            
            # 2. current frame, for co-perception lable use
            data[cav_id]['curr'] = {}

            timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
            
            # 2.1 load curr params
            # json is faster than yaml
            json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

            # TODO: debug use, uncomment: to use regular version yaml GT
            # time_new = str( '%06d' % int(float(timestamp_key)))
            # json_file = json_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # json_file = json_file.replace(timestamp_key, time_new)

            # json_file = cav_content['regular'][timestamp_key]['yaml'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2/dataset_irregular_v2")
            # data[cav_id]['curr']['params'] = \
            #                 load_yaml(json_file)

            # print(json_file)
            # store the scene name, which is convinient for debug
            scene_name = json_file.split('/')[-3]
            if os.path.exists(json_file):
                with open(json_file, "r") as f:
                    data[cav_id]['curr']['params'] = json.load(f)
            else:
                data[cav_id]['curr']['params'] = \
                            load_yaml(cav_content['regular'][timestamp_key]['yaml'])
            # 没有 lidar pose
            if not ('lidar_pose' in data[cav_id]['curr']['params']):
                tmp_ego_pose = np.array(data[cav_id]['curr']['params']['true_ego_pos'])
                tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                data[cav_id]['curr']['params']['lidar_pose'] = list(tmp_ego_pose)

            # 2.2 load curr lidar file
            # npy is faster than pcd
            npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("pcd", "npy")
            
            # TODO: debug use, uncomment: to use regular version lidar input
            # time_new = str( '%06d' % int(float(timestamp_key))) 
            # npy_file = npy_file.replace('OPV2V_irregular_npy','OPV2V_w_npy')
            # npy_file = npy_file.replace(timestamp_key, time_new)

            # npy_file = cav_content['regular'][timestamp_key]['lidar'].replace("OPV2V_irregular_npy", "OPV2V_Irregular_V2")
            # data[cav_id]['curr']['lidar_np'] = \
            #     pcd_utils.pcd_to_np(npy_file)
            # print(npy_file)

            if os.path.exists(npy_file): 
                data[cav_id]['curr']['lidar_np'] = np.load(npy_file)
            else:
                data[cav_id]['curr']['lidar_np'] = \
                        pcd_utils.pcd_to_np(cav_content['regular'][timestamp_key]['lidar'])

            # 2.3 store curr timestamp and time_diff
            data[cav_id]['curr']['timestamp'] = timestamp_key
            data[cav_id]['curr']['time_diff'] = 0.0
            data[cav_id]['curr']['sample_interval'] = 0

            # 3. past frames, for model input
            data[cav_id]['past_k'] = OrderedDict()
            latest_sample_stamp_idx = curr_timestamp_idx
            # past k frames, pose | lidar | label(for single view confidence map generator use)
            for i in range(self.k):
                # sample_interval
                if data[cav_id]['ego']:             # ego sample_interval = E(B(n, p))
                    if i == 0: # ego-past-0 与 ego-curr 是一样的
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    sample_interval = self.sample_interval_exp
                    if sample_interval == 0:
                        sample_interval = 1
                else:                               # non-ego sample_interval ~ B(n, p)
                    if self.sample_interval_exp==0 \
                        and self.is_no_shift \
                            and i == 0:
                        data[cav_id]['past_k'][i] = data[cav_id]['curr']
                        continue
                    if self.is_same_sample_interval:
                        sample_interval = self.sample_interval_exp
                    else:
                        # B(n, p)
                        trails = bernoulliDist.rvs(self.binomial_n)
                        sample_interval = sum(trails)
                    if sample_interval==0:
                        if i==0: # 检查past 0 的实际时间是否在curr 的后面
                            tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                            if self.dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                                sample_interval = 1
                        if i>0: # 过去的几帧不要重复
                            sample_interval = 1                

                # check the timestamp index
                data[cav_id]['past_k'][i] = {}
                latest_sample_stamp_idx -= sample_interval
                timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                # load the corresponding data into the dictionary
                # load param file: json is faster than yaml
                json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")

                if os.path.exists(json_file):
                    with open(json_file, "r") as f:
                        data[cav_id]['past_k'][i]['params'] = json.load(f)
                else:
                    data[cav_id]['past_k'][i]['params'] = \
                        load_yaml(cav_content[timestamp_key]['yaml'])
                # 没有 lidar pose
                if not ('lidar_pose' in data[cav_id]['past_k'][i]['params']):
                    tmp_ego_pose = np.array(data[cav_id]['past_k'][i]['params']['true_ego_pos'])
                    tmp_ego_pose += np.array([-0.5, 0, 1.9, 0, 0, 0])
                    data[cav_id]['past_k'][i]['params']['lidar_pose'] = list(tmp_ego_pose)

                # load lidar file: npy is faster than pcd
                npy_file = cav_content[timestamp_key]['lidar'].replace("pcd", "npy")
                if os.path.exists(npy_file):
                    data[cav_id]['past_k'][i]['lidar_np'] = np.load(npy_file)
                else:
                    data[cav_id]['past_k'][i]['lidar_np'] = \
                            pcd_utils.pcd_to_np(cav_content[timestamp_key]['lidar'])

                data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
                data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
                data[cav_id]['past_k'][i]['time_diff'] = \
                    self.dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])
            
            data[cav_id]['debug'] = {}
            data[cav_id]['debug']['scene'] = scene_name
            data[cav_id]['debug']['cav_id'] = cav_id

        return data

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        processed_data_dict : dict consist of all processed info, whose structure is:
        {
            'single_object_dict_stack': single_label_dict_stack,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, np.array of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
            'avg_sample_interval': float,
            'avg_time_delay': float

        }
        '''
        # TODO: debug use
        global illegal_path_list

        base_data_dict = self.retrieve_base_data(idx)
        ''' base_data_dict structure:
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            },
            cav_id_2 : { ... }
        }
        '''

        processed_data_dict = OrderedDict()
        processed_data_dict['ego'] = {}

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        too_far = []
        curr_lidar_pose_list = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = math.sqrt( \
                (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            curr_lidar_pose_list.append(selected_cav_base['curr']['params']['lidar_pose']) # 6dof pose
            cav_id_list.append(cav_id)  
        for cav_id in too_far: # filter those out of communicate range
            base_data_dict.pop(cav_id)
        
        single_label_dict_stack = []
        object_stack = []
        object_id_stack = []
        curr_pose_stack = []
        curr_feature_stack = []
        past_k_pose_stack = []
        past_k_features_stack = [] 
        past_k_tr_mats = []
        past_k_label_dicts_stack = []
        past_k_sample_interval_stack = []
        past_k_time_diffs_stack = []
        # avg_past_k_time_diffs = 0.0
        # avg_past_k_sample_interval = 0.0
        illegal_cav = []
        if self.visualize:
            projected_lidar_stack = []
        
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            ''' selected_cav_base:
            {
                'ego' : true,
                curr : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }
            '''
            selected_cav_processed = self.get_item_single_car(
                selected_cav_base,
                ego_lidar_pose,
                idx
            )
            ''' selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # curr lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换, 用于单体监督            cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature                 cav view + curr feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                # 'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                # 'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue

            
            if self.visualize:
                projected_lidar_stack.append(selected_cav_processed['projected_lidar'])
            
            # single view feature
            curr_feature_stack.append(selected_cav_processed['curr_feature'])
            # single view label
            single_label_dict_stack.append(selected_cav_processed['single_label_dict'])

            # curr ego view label
            object_stack.append(selected_cav_processed['object_bbx_center'])
            object_id_stack += selected_cav_processed['object_ids']

            # current pose: N, 6
            curr_pose_stack.append(selected_cav_processed['curr_pose']) 
            # features: N, k, 
            past_k_features_stack.append(selected_cav_processed['past_k_features'])
            # poses: N, k, 6
            past_k_pose_stack.append(selected_cav_processed['past_k_poses'])
            # time differences: N, k
            past_k_time_diffs_stack += selected_cav_processed['past_k_time_diffs']
            # sample intervals: N, k
            past_k_sample_interval_stack += selected_cav_processed['past_k_sample_interval']
            # past k frames to ego pose trans matrix, list of len=N, past_k_tr_mats[i]: ndarray(k, 4, 4)
            past_k_tr_mats.append(selected_cav_processed['past_k_tr_mats'])
            # past k label dict: N, k, object_num, 7
            # past_k_label_dicts_stack.append(selected_cav_processed['past_k_label_dicts'])
        # {pos: array[num_cav, k, 100, 252, 2], neg: array[num_cav, k, 100, 252, 2], target: array[num_cav, k, 100, 252, 2]}
        # past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts_stack)
        # self.times.append(time.time())

        # filter those cav who has no points left
        # then we can calculate get_pairwise_transformation
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)
            cav_id_list.remove(cav_id)

        merged_curr_feature_dict = self.merge_features_to_dict(curr_feature_stack)  # current 在各自view 下 feature
        
        single_label_dict = self.post_processor.collate_batch(single_label_dict_stack) # current 在各自view 下 label

        past_k_time_diffs_stack = np.array(past_k_time_diffs_stack)
        past_k_sample_interval_stack = np.array(past_k_sample_interval_stack)
        
        pairwise_t_matrix = \
            self.get_pairwise_transformation(base_data_dict, self.max_cav)
            # self.get_past_k_pairwise_transformation2ego(base_data_dict, 
            # ego_lidar_pose, self.max_cav) # np.tile(np.eye(4), (max_cav, self.k, 1, 1)) (L, k, 4, 4) TODO: 这里面没有搞懂为什么不用 past_k_tr_mats

        curr_lidar_poses = np.array(curr_pose_stack).reshape(-1, 6)  # (N_cav, 6)
        past_k_lidar_poses = np.array(past_k_pose_stack).reshape(-1, self.k, 6)  # (N, k, 6)

        # exclude all repetitive objects    
        unique_indices = \
            [object_id_stack.index(x) for x in set(object_id_stack)]
        try:
            object_stack = np.vstack(object_stack)
        except ValueError:
            # print("!!! vstack ValueError !!!")
            return None
        object_stack = object_stack[unique_indices]

        # make sure bounding boxes across all frames have the same number
        object_bbx_center = \
            np.zeros((self.params['postprocess']['max_num'], 7))
        mask = np.zeros(self.params['postprocess']['max_num'])
        object_bbx_center[:object_stack.shape[0], :] = object_stack
        mask[:object_stack.shape[0]] = 1

        # self.times.append(time.time())

        # merge preprocessed features from different cavs into the same dict
        cav_num = len(cav_id_list)
        # past_k_features_stack: list, len is num_cav. [i] is list, len is k. [cav_id][time_id] is Orderdict, {'voxel_features': array, ...}
        merged_feature_dict = self.merge_past_k_features_to_dict(past_k_features_stack)

        # generate the anchor boxes
        anchor_box = self.anchor_box

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=mask)

        # self.times.append(time.time())

        # self.times = (np.array(self.times[1:]) - np.array(self.times[:-1]))

        # self.times = np.hstack((self.times, time4data))

        processed_data_dict['ego'].update(
            {'single_object_dict_stack': single_label_dict,
             'curr_processed_lidar': merged_curr_feature_dict,
             'object_bbx_center': object_bbx_center,
             'object_bbx_mask': mask,
             'object_ids': [object_id_stack[i] for i in unique_indices],
             'anchor_box': anchor_box,
             'processed_lidar': merged_feature_dict,
             'label_dict': label_dict,
             'cav_num': cav_num,
             'pairwise_t_matrix': pairwise_t_matrix,
             'curr_lidar_poses': curr_lidar_poses,
             'past_k_lidar_poses': past_k_lidar_poses,
             'past_k_time_diffs': past_k_time_diffs_stack,
             'past_k_sample_interval': past_k_sample_interval_stack})
            #  'times': self.times})

        processed_data_dict['ego'].update({'sample_idx': idx,
                                            'cav_id_list': cav_id_list})
        try:
            processed_data_dict['ego'].update({'avg_sample_interval':\
                sum(past_k_sample_interval_stack[self.k:]) / len(past_k_sample_interval_stack[self.k:])})
            processed_data_dict['ego'].update({'avg_time_delay':\
                sum(past_k_time_diffs_stack[self.k:]) / len(past_k_time_diffs_stack[self.k:])})
        except ZeroDivisionError:
            # print("!!! ZeroDivisionError !!!")
            return None

        if self.visualize:
            processed_data_dict['ego'].update({'origin_lidar':
                np.vstack(projected_lidar_stack)})

        return processed_data_dict


    def get_item_single_car(self, selected_cav_base, ego_pose, idx):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true,
                'curr' : {
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                'past_k' : {		           # (k) totally
                    [0]:{
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},
                    ...,		
                    [k-1] : {}
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }
            }

        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.

        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                'projected_lidar':      # lidar in ego space, 用于viz
                'single_label_dict':	# single view label. 没有经过坐标变换,                      cav view + curr 的label
                'curr_feature':         # current feature, lidar预处理得到的feature
                'object_bbx_center':	# ego view label. np.ndarray. Shape is (max_num, 7).    ego view + curr 的label
                'object_ids':			# ego view label index. list. length is (max_num, 7).   ego view + curr 的label
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		    # list of past k frames' poses
                'past_k_features': 		    # list of past k frames' lidar
                'past_k_time_diffs': 	    # list of past k frames' time diff with current frame
                'past_k_tr_mats': 		    # list of past k frames' transformation matrix to current ego coordinate
                'past_k_sample_interval':   # list of past k frames' sample interval with later frame
                'avg_past_k_time_diffs':    # avg_past_k_time_diffs,
                'avg_past_k_sample_interval': # avg_past_k_sample_interval,
                'if_no_point':              # bool, 用于判断是否合法
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        if self.visualize:
            # trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['curr']['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            projected_lidar = \
                box_utils.project_points_by_matrix_torch(lidar_np[:, :3], transformation_matrix)
            selected_cav_processed.update({'projected_lidar': projected_lidar})

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # past k transfomation matrix
        past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample intervals
        past_k_sample_interval = []

        # for debug use
        # avg_past_k_time_diff = 0
        # avg_past_k_sample_interval = 0
        
        # past k label 
        # past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # 1. trans matrix
            transformation_matrix = \
                x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True

            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

            ################################################################
            # sizhewei
            # for past k frames' single view label
            ################################################################
            '''
            # 5. single view label
            # past_i label at past_i single view
            # opencood/data_utils/post_processor/base_postprocessor.py
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([selected_cav_base['past_k'][i]], selected_cav_base['past_k'][i]['params']['lidar_pose'])  
            # generate the anchor boxes
            # opencood/data_utils/post_processor/voxel_postprocessor.py
            anchor_box = self.anchor_box
            single_view_label_dict = self.post_processor.generate_label(
                    gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
                )
            past_k_label_dicts.append(single_view_label_dict)
            '''

        '''
        # past k merge
        past_k_label_dicts = self.post_processor.merge_label_to_dict(past_k_label_dicts)
        '''

        past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # avg_past_k_time_diffs = float(sum(past_k_time_diffs) / len(past_k_time_diffs))
        # avg_past_k_sample_interval = float(sum(past_k_sample_interval) / len(past_k_sample_interval))

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        
        # curr label at ego view
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
             'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
            #  'avg_past_k_time_diffs': avg_past_k_time_diffs,
            #  'avg_past_k_sample_interval': avg_past_k_sample_interval,
            #  'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        return selected_cav_processed

    @staticmethod
    def return_timestamp_key_async(cav_content, timestamp_index):
        """
        Given the timestamp index, return the correct timestamp key, e.g.
        2 --> '000078'.

        Parameters
        ----------
        scenario_database : OrderedDict
            The dictionary contains all contents in the current scenario.

        timestamp_index : int
            The index for timestamp.

        Returns
        -------
        timestamp_key : str
            The timestamp key saved in the cav dictionary.
        """
        # retrieve the correct index
        timestamp_key = list(cav_content.items())[timestamp_index][0]

        return timestamp_key

    @staticmethod
    def dist_time(ts1, ts2, i = -1):
        """caculate the time interval between two timestamps

        Args:
            ts1 (string): time stamp at some time
            ts2 (string): current time stamp
            i (int, optional): past frame id, for debug use. Defaults to -1.
        
        Returns:
            time_diff (float): time interval (ts1 - ts2)
        """
        if not i==-1:
            return -i
        else:
            return (float(ts1) - float(ts2))

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_diff = []
        past_k_sample_interval = []
        past_k_avg_time_delay = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        # for debug use:
        sum_time_diff = 0.0
        sum_sample_interval = 0.0
        # time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_diff.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_interval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            past_k_avg_time_delay.append(ego_dict['avg_time_delay']) # ego_dict['avg_sample_interval'] is float
            # avg_time_delay += ego_dict['avg_time_delay']
            # avg_sample_interval += ego_dict['avg_sample_interval']
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            # time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_diff = np.hstack(past_k_time_diff)
        past_k_time_diff= torch.from_numpy(past_k_time_diff)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_interval = np.hstack(past_k_sample_interval)
        past_k_sample_interval = torch.from_numpy(past_k_sample_interval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = float(sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval))

        past_k_avg_time_delay = np.array(past_k_avg_time_delay)
        avg_time_delay = float(sum(past_k_avg_time_delay) / len(past_k_avg_time_delay))

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for debug use: 
        # time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                   'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_diff,
                                   'past_k_sample_interval': past_k_sample_interval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'avg_time_delay': avg_time_delay})
                                #    'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        output_dict = self.collate_batch_train(batch)
        if output_dict is None:
            return None

        # check if anchor box in the batch
        if batch[0]['ego']['anchor_box'] is not None:
            output_dict['ego'].update({'anchor_box':
                torch.from_numpy(np.array(
                    batch[0]['ego'][
                        'anchor_box']))})

        # save the transformation matrix (4, 4) to ego vehicle
        # transformation is only used in post process (no use.)
        # we all predict boxes in ego coord.
        transformation_matrix_torch = \
            torch.from_numpy(np.identity(4)).float()
        transformation_matrix_clean_torch = \
            torch.from_numpy(np.identity(4)).float()

        output_dict['ego'].update({'transformation_matrix':
                                       transformation_matrix_torch,
                                    'transformation_matrix_clean':
                                       transformation_matrix_clean_torch,})

        output_dict['ego'].update({
            "sample_idx": batch[0]['ego']['sample_idx'],
            "cav_id_list": batch[0]['ego']['cav_id_list']
        })

        # output_dict['ego'].update({'veh_frame_id': batch[0]['ego']['veh_frame_id']})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                # lidar_pose = cav_content['curr']['params']['lidar_pose']
                lidar_pose = cav_content['past_k'][0]['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

# if __name__ == '__main__':   

opencood/data_utils/datasets/intermediate_fusion_dataset_v2xset_singleframe.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/4/15
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
two stage flow update framework

"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
from tqdm import tqdm
from tqdm.contrib import tenumerate
from tqdm.auto import trange
# import opencood.data_utils.post_processor as post_processor
# import opencood.utils.pcd_utils as pcd_utils
# from opencood.utils.keypoint_utils import bev_sample, get_keypoints
# from opencood.data_utils.datasets import basedataset
# from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
# from opencood.utils.pcd_utils import \
#     mask_points_by_range, mask_ego_points, shuffle_points, \
#     downsample_lidar_minimum
# from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
# from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
# from opencood.utils import box_utils
# from opencood.utils.flow_utils import generate_flow_map, generate_flow_map_szwei

from opencood.utils.box_utils import boxes_to_corners2d # for debug use
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# global, for debug use
illegal_path_list = set()


def extract_timestamps(yaml_files):
    """
    Given the list of the yaml files, extract the mocked timestamps.

    Parameters
    ----------
    yaml_files : list
        The full path of all yaml files of ego vehicle

    Returns
    -------
    timestamps : list
        The list containing timestamps only.
    """
    timestamps = []

    for file in yaml_files:
        res = file.split('/')[-1]
        if res.endswith('.yaml'):
            timestamp = res.replace('.yaml', '')
        elif res.endswith('.json'):
            timestamp = res.replace('.json', '')
        else:
            print("Woops! There is no processing method for file {}".format(res))
            sys.exit(1)
        timestamps.append(timestamp)

    return timestamps


def statistics_irv2v(split_name="validate"):
    # total box number
    box_total = 0
    # box number distribution
    box_num_dict = {}
    # speed distribution
    speed_dict = {}

    # 处理的数据集
    data_seperate_name = split_name
    # 文件路径
    root_dir = "/remote-home/share/OPV2V_irregular_npy/" + data_seperate_name

    # first load all paths of different scenarios
    scenario_folders = sorted([os.path.join(root_dir, x)
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])
    scenario_folders_name = sorted([x
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])

    print("### Processing Started! ###")

    # loop over all scenarios
    for (i, scenario_folder) in tenumerate(scenario_folders):
        start_time = time.time()
        # copy timestamps npy file
        timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
        time_annotations = np.load(timestamps_file)

        # save_scene_folder = os.path.join(save_dir, scenario_folders_name[i])
        # if not os.path.exists(save_scene_folder):
        #     os.makedirs(save_scene_folder)
        # save_timestamps_file = os.path.join(save_scene_folder, 'timestamps.npy')
        # if not os.path.exists(save_timestamps_file):
        #     shutil.copyfile(timestamps_file, save_timestamps_file)

        # iterate all cav in this scenario
        cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))], key=lambda y:int(y))
        assert len(cav_list) > 0

        # use the frame number as key, the full path as the values, store all json or yaml files in this scenario
        yaml_files = sorted([x
                    for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                    x.endswith(".json")], 
                    key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
        if len(yaml_files)==0:
            yaml_files = sorted([x 
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith('.yaml')], key=lambda y:float((y.split('/')[-1]).split('.yaml')[0]))

        start_timestamp = int(float(extract_timestamps(yaml_files)[0]))
        while(1):
            time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
            time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
            if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
                start_timestamp += 1
            else:
                break

        end_timestamp = int(float(extract_timestamps(yaml_files)[-1]))
        if start_timestamp%2 == 0:
            # even
            end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
        else:
            end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
        num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)
        regular_timestamps = [start_timestamp+2*i for i in range(num_timestamps)]

        for (j, cav_id) in tenumerate(cav_list):
            cav_path = os.path.join(scenario_folder, cav_id)
            if j==0: # ego
                timestamps = regular_timestamps
            else: 
                timestamps = list(time_annotations[j-1, :])

            for timestamp in timestamps:
                timestamp = "%.3f" % float(timestamp)
                
                # self.scenario_database[i][cav_id][timestamp] = OrderedDict()

                json_file = os.path.join(cav_path,
                                            timestamp + '.json')
                json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")
                
                with open(json_file, "r") as f:
                    curr_label = json.load(f)
            
                scene_box_num = len(curr_label['vehicles'].keys())
                # total box
                box_total += scene_box_num
                # box number distribution
                if scene_box_num in box_num_dict.keys():
                    box_num_dict[scene_box_num] += 1
                else:
                    box_num_dict[scene_box_num] = 1
                # speed distribution
                for car_idx, car_content in curr_label['vehicles'].items():
                    if int(car_content['speed']) in speed_dict.keys():
                        speed_dict[int(car_content['speed'])] += 1
                    else:
                        speed_dict[int(car_content['speed'])] = 1
    
    print(f"box_total = {box_total}")
    save_path = '/root/percp/OpenCOOD/opencood/'
    torch.save(box_num_dict, f"{save_path}box_num_dict_{split_name}.pt")
    torch.save(speed_dict, f"{save_path}speed_dict_{split_name}.pt")

def speed_hist():
    # import numpy as np
    import matplotlib.pyplot as plt
    plt.style.use('ggplot')
    # 假设你有一个名为 dict 的字典
    split_name_list = ['train', 'validate', 'test']
    dicts = []
    for split_name in split_name_list:
        speed_dict_path = f"/root/percp/OpenCOOD/opencood/speed_dict_{split_name}.pt"
        unit_dict = torch.load(speed_dict_path, map_location='cpu')
        dicts.append(unit_dict)

    result_dict = {}
    for unit_dict in dicts:
        for key, value in unit_dict.items():
            if key in result_dict.keys():
                result_dict[key] += value
            else:
                result_dict[key] = value

    # 按照键的大小进行排序，生成排序后的键列表
    sorted_keys = sorted(result_dict.keys())

    # 根据排序后的键列表，获取对应的值列表
    sorted_values = [result_dict[key] for key in sorted_keys]

    # 绘制直方图
    plt.bar(sorted_keys[1:], sorted_values[1:])

    avg_speed = np.dot(np.array(sorted_keys),np.array(sorted_values)) / np.sum(np.array(sorted_values[1:]))
    print(f'avg speed is {avg_speed}')
    print(f'max speed is {sorted_keys[-1]}')
    print(f'speed=0 has cav: {sorted_values[0]}')
    print(f'all cav num is {np.sum(np.array(sorted_values))}')

    # labels = []
    # for i, unit_dict in enumerate(dicts):
    #     # 按照键的大小进行排序，生成排序后的键列表
    #     sorted_keys = sorted(unit_dict.keys())

    #     # 根据排序后的键列表，获取对应的值列表
    #     sorted_values = [unit_dict[key] for key in sorted_keys]

    #     # 绘制直方图
    #     plt.bar(sorted_keys[1:], sorted_values[1:], label=split_name_list[i])

    #     avg_speed = np.dot(np.array(sorted_keys),np.array(sorted_values)) / np.sum(np.array(sorted_values[1:]))
    #     print(f'{split_name} avg speed is {avg_speed}')

    #     labels.append(split_name_list[i])

    # 设置图表标题和坐标轴标签
    # plt.title('speed distribution histogram.')

    fontsize = 20
    # 设置图表标题和坐标轴标签
    # plt.title('num of box distribution histogram.')
    plt.xticks(fontsize=16); plt.yticks(fontsize=16)
    plt.xlabel('Speed of the vehicle (km/h)', fontsize=fontsize, color='black')
    plt.ylabel('Number of vehicles', fontsize=fontsize, color='black')
    plt.tight_layout()

    # plt.legend(labels)

    # 展示图表
    plt.savefig(f'/root/percp/OpenCOOD/opencood/distribution-speed.png', dpi=600)

def box_hist():
    # import numpy as np
    import matplotlib.pyplot as plt
    plt.style.use('ggplot')
    # 假设你有一个名为 dict 的字典
    split_name_list = ['train', 'validate', 'test'] # ['test'] #
    dicts = []
    for split_name in split_name_list:
        speed_dict_path = f"/root/percp/OpenCOOD/opencood/box_num_dict_{split_name}.pt"
        unit_dict = torch.load(speed_dict_path, map_location='cpu')
        dicts.append(unit_dict)

    result_dict = {}
    for unit_dict in dicts:
        for key, value in unit_dict.items():
            if key in result_dict.keys():
                result_dict[key] += value
            else:
                result_dict[key] = value

    # 按照键的大小进行排序，生成排序后的键列表
    sorted_keys = sorted(result_dict.keys())

    # 根据排序后的键列表，获取对应的值列表
    sorted_values = [result_dict[key] for key in sorted_keys]

    # 绘制直方图
    plt.bar(sorted_keys, sorted_values)

    avg_speed = np.dot(np.array(sorted_keys),np.array(sorted_values)) / np.sum(np.array(sorted_values))
    print(f'avg num of box is {avg_speed}')
    print(f'max num is {sorted_keys[-1]}')
    print(f'total cav is : {np.sum(np.dot(np.array(sorted_keys), np.array(sorted_values)))}')

    # labels = []
    # for i, unit_dict in enumerate(dicts):
    #     # 按照键的大小进行排序，生成排序后的键列表
    #     sorted_keys = sorted(unit_dict.keys())

    #     # 根据排序后的键列表，获取对应的值列表
    #     sorted_values = [unit_dict[key] for key in sorted_keys]

    #     # 绘制直方图
    #     plt.bar(sorted_keys[1:], sorted_values[1:], label=split_name_list[i])

    #     avg_speed = np.dot(np.array(sorted_keys),np.array(sorted_values)) / np.sum(np.array(sorted_values[1:]))
    #     print(f'{split_name} avg speed is {avg_speed}')

    #     labels.append(split_name_list[i])

    fontsize = 20
    # 设置图表标题和坐标轴标签
    # plt.title('num of box distribution histogram.')
    plt.xticks(fontsize=16); plt.yticks(fontsize=16)
    plt.xlabel('Number of vehicles in the scene', fontsize=fontsize, color='black')
    plt.ylabel('Number of scenes', fontsize=fontsize, color='black')
    plt.tight_layout()
    # plt.legend(labels)

    # 展示图表
    plt.savefig(f'/root/percp/OpenCOOD/opencood/distribution-box.png', dpi=600)


if __name__ == '__main__':
    # statistics_irv2v()
    # speed_hist()
    box_hist()
    
opencood/data_utils/datasets/irv2v_statistics.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for late fusion
"""
import random
import math
from collections import OrderedDict

import numpy as np
import torch
from torch.utils.data import DataLoader

import opencood.data_utils.datasets
from opencood.data_utils.post_processor import build_postprocessor
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils import box_utils
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.transformation_utils import x1_to_x2
from opencood.utils.common_utils import read_json
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np
from opencood.utils.pose_utils import add_noise_data_dict

class LateFusionDataset(basedataset.BaseDataset):
    """
    This class is for intermediate fusion where each vehicle transmit the
    detection outputs to ego.
    """
    def __init__(self, params, visualize, train=True):
        super(LateFusionDataset, self).__init__(params, visualize, train)
        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = build_postprocessor(params['postprocess'], train)

        if "box_align" in params.keys():
            self.box_align = True
            self.stage1_result_path = params['box_align']['train_result'] if train else params['box_align']['val_result']
            self.stage1_result = read_json(self.stage1_result_path)
            self.box_align_args = params['box_align']['args']
        
        else:
            self.box_align = False

    def __getitem__(self, idx):
        base_data_dict = self.retrieve_base_data(idx)
        if self.train:
            reformat_data_dict = self.get_item_train(base_data_dict)
        else:
            reformat_data_dict = self.get_item_test(base_data_dict, idx)

        return reformat_data_dict

    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.
        The object_bbx_center is in ego coordinate.

        Notice: it is a wrap of postprocessor

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """
        return self.post_processor.generate_object_center(cav_contents,
                                                        reference_lidar_pose)

    def get_item_single_car(self, selected_cav_base):
        """
        Process a single CAV's information for the train/test pipeline.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
        """
        selected_cav_processed = {}

        # filter lidar
        lidar_np = selected_cav_base['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_points_by_range(lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
        # remove points that hit ego vehicle
        lidar_np = mask_ego_points(lidar_np)

        # generate the bounding box(n, 7) under the cav's space
        object_bbx_center, object_bbx_mask, object_ids = self.generate_object_center([selected_cav_base],
                                                    selected_cav_base[
                                                           'params'][
                                                           'lidar_pose_clean'])

        # data augmentation
        lidar_np, object_bbx_center, object_bbx_mask = \
            self.augment(lidar_np, object_bbx_center, object_bbx_mask)

        if self.visualize:
            selected_cav_processed.update({'origin_lidar': lidar_np})

        # pre-process the lidar to voxel/bev/downsampled lidar
        lidar_dict = self.pre_processor.preprocess(lidar_np)
        selected_cav_processed.update({'processed_lidar': lidar_dict})

        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()
        selected_cav_processed.update({'anchor_box': anchor_box})

        selected_cav_processed.update({'object_bbx_center': object_bbx_center,
                                       'object_bbx_mask': object_bbx_mask,
                                       'object_ids': object_ids})

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=object_bbx_mask)
        selected_cav_processed.update({'label_dict': label_dict})

        return selected_cav_processed

    def get_item_train(self, base_data_dict):
        processed_data_dict = OrderedDict()
        base_data_dict = add_noise_data_dict(base_data_dict, self.params['noise_setting'])
        # during training, we return a random cav's data
        # only one vehicle is in processed_data_dict
        if not self.visualize:
            selected_cav_id, selected_cav_base = \
                random.choice(list(base_data_dict.items()))
        else:
            selected_cav_id, selected_cav_base = \
                list(base_data_dict.items())[0]

        selected_cav_processed = self.get_item_single_car(selected_cav_base)
        processed_data_dict.update({'ego': selected_cav_processed})

        return processed_data_dict

    def get_item_test(self, base_data_dict, idx):
        """
            processed_data_dict.keys() = ['ego', "650", "659", ...]
        """
        base_data_dict = add_noise_data_dict(base_data_dict,self.params['noise_setting'])

        processed_data_dict = OrderedDict()
        ego_id = -1
        ego_lidar_pose = []
        cav_id_list = []
        lidar_pose_list = []

        # first find the ego vehicle's lidar pose
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['params']['lidar_pose']
                ego_lidar_pose_clean = cav_content['params']['lidar_pose_clean']
                break

        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            distance = \
                math.sqrt((selected_cav_base['params']['lidar_pose'][0] -
                           ego_lidar_pose[0]) ** 2 + (
                                  selected_cav_base['params'][
                                      'lidar_pose'][1] - ego_lidar_pose[
                                      1]) ** 2)
            if distance > self.params['comm_range']:
                continue
            cav_id_list.append(cav_id)
            lidar_pose_list.append(selected_cav_base['params']['lidar_pose'])

        '''
        ########## Added by Yifan Lu 2022.8.14 ##############
        # box align to correct pose.
        if self.box_align and str(idx) in self.stage1_result.keys():
            stage1_content = self.stage1_result[str(idx)]
            if stage1_content is not None:
                cav_id_list_stage1 = stage1_content['cav_id_list']
                
                pred_corners_list = stage1_content['pred_corner3d_np_list']
                pred_corners_list = [np.array(corners, dtype=np.float64) for corners in pred_corners_list]
                uncertainty_list = stage1_content['uncertainty_np_list']
                uncertainty_list = [np.array(uncertainty, dtype=np.float64) for uncertainty in uncertainty_list]
                stage1_lidar_pose_list = [base_data_dict[cav_id]['params']['lidar_pose'] for cav_id in cav_id_list_stage1]
                stage1_lidar_pose = np.array(stage1_lidar_pose_list)


                refined_pose = box_alignment_relative_sample_np(pred_corners_list,
                                                                stage1_lidar_pose, 
                                                                uncertainty_list=uncertainty_list, 
                                                                **self.box_align_args)
                stage1_lidar_pose[:,[0,1,4]] = refined_pose
                stage1_lidar_pose_refined_list = stage1_lidar_pose.tolist() # updated lidar_pose_list
                for cav_id, lidar_pose_refined in zip(cav_id_list_stage1, stage1_lidar_pose_refined_list):
                    if cav_id not in cav_id_list:
                        continue
                    idx_in_list = cav_id_list.index(cav_id)
                    lidar_pose_list[idx_in_list] = lidar_pose_refined
                    base_data_dict[cav_id]['params']['lidar_pose'] = lidar_pose_refined
            '''
            
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]
            # find the transformation matrix from current cav to ego.
            cav_lidar_pose = selected_cav_base['params']['lidar_pose']
            transformation_matrix = x1_to_x2(cav_lidar_pose, ego_lidar_pose)
            cav_lidar_pose_clean = selected_cav_base['params']['lidar_pose_clean']
            transformation_matrix_clean = x1_to_x2(cav_lidar_pose_clean, ego_lidar_pose_clean)

            selected_cav_processed = \
                self.get_item_single_car(selected_cav_base)
            selected_cav_processed.update({'transformation_matrix': transformation_matrix,
                                           'transformation_matrix_clean': transformation_matrix_clean})
            update_cav = "ego" if cav_id == ego_id else cav_id
            processed_data_dict.update({update_cav: selected_cav_processed})

        return processed_data_dict

    def collate_batch_test(self, batch):
        """
        Customized collate function for pytorch dataloader during testing
        for late fusion dataset.

        Parameters
        ----------
        batch : dict

        Returns
        -------
        batch : dict
            Reformatted batch.
        """
        # currently, we only support batch size of 1 during testing
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        batch = batch[0]

        output_dict = {}

        # for late fusion, we also need to stack the lidar for better
        # visualization
        if self.visualize:
            projected_lidar_list = []
            origin_lidar = []

        for cav_id, cav_content in batch.items():
            output_dict.update({cav_id: {}})
            # shape: (1, max_num, 7)
            object_bbx_center = \
                torch.from_numpy(np.array([cav_content['object_bbx_center']]))
            object_bbx_mask = \
                torch.from_numpy(np.array([cav_content['object_bbx_mask']]))
            object_ids = cav_content['object_ids']

            # the anchor box is the same for all bounding boxes usually, thus
            # we don't need the batch dimension.
            if cav_content['anchor_box'] is not None:
                output_dict[cav_id].update({'anchor_box':
                    torch.from_numpy(np.array(
                        cav_content[
                            'anchor_box']))})
            if self.visualize:
                transformation_matrix = cav_content['transformation_matrix']
                origin_lidar = [cav_content['origin_lidar']]

                if (self.params['only_vis_ego'] is False) or (cav_id=='ego'):
                    # print(cav_id)
                    import copy
                    projected_lidar = copy.deepcopy(cav_content['origin_lidar'])
                    projected_lidar[:, :3] = \
                        box_utils.project_points_by_matrix_torch(
                            projected_lidar[:, :3],
                            transformation_matrix)
                    projected_lidar_list.append(projected_lidar)

            # processed lidar dictionary
            processed_lidar_torch_dict = \
                self.pre_processor.collate_batch(
                    [cav_content['processed_lidar']])
            # label dictionary
            label_torch_dict = \
                self.post_processor.collate_batch([cav_content['label_dict']])

            # save the transformation matrix (4, 4) to ego vehicle
            transformation_matrix_torch = \
                torch.from_numpy(
                    np.array(cav_content['transformation_matrix'])).float()
            
            # late fusion training, no noise
            transformation_matrix_clean_torch = transformation_matrix_torch

            output_dict[cav_id].update({'object_bbx_center': object_bbx_center,
                                        'object_bbx_mask': object_bbx_mask,
                                        'processed_lidar': processed_lidar_torch_dict,
                                        'label_dict': label_torch_dict,
                                        'object_ids': object_ids,
                                        'transformation_matrix': transformation_matrix_torch,
                                        'transformation_matrix_clean': transformation_matrix_clean_torch})

            if self.visualize:
                origin_lidar = \
                    np.array(
                        downsample_lidar_minimum(pcd_np_list=origin_lidar))
                origin_lidar = torch.from_numpy(origin_lidar)
                output_dict[cav_id].update({'origin_lidar': origin_lidar})

        if self.visualize:
            projected_lidar_stack = [torch.from_numpy(
                np.vstack(projected_lidar_list))]
            output_dict['ego'].update({'origin_lidar': projected_lidar_stack})
            # output_dict['ego'].update({'projected_lidar_list': projected_lidar_list})

        return output_dict

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor


    def post_process_no_fusion(self, data_dict, output_dict_ego, return_uncertainty=False):
        data_dict_ego = OrderedDict()
        data_dict_ego['ego'] = data_dict['ego']
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        if return_uncertainty:
            pred_box_tensor, pred_score, uncertainty = \
                self.post_processor.post_process(data_dict_ego, output_dict_ego, return_uncertainty=True)
            return pred_box_tensor, pred_score, gt_box_tensor, uncertainty
        else:
            pred_box_tensor, pred_score = \
                self.post_processor.post_process(data_dict_ego, output_dict_ego)
            return pred_box_tensor, pred_score, gt_box_tensor

        

opencood/data_utils/datasets/late_fusion_dataset.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Qunahao Li

"""
Dataset class for DAIR-V2X dataset late fusion
"""
import random
import math
from collections import OrderedDict
import os
import os.path as osp
import opencood.data_utils.post_processor as post_processor
import numpy as np
import torch
from torch.utils.data import DataLoader
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
import opencood.data_utils.datasets
from opencood.data_utils.post_processor import build_postprocessor
from opencood.data_utils.datasets import late_fusion_dataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils import box_utils
from opencood.utils.transformation_utils import x1_to_x2
import json
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.transformation_utils import tfm_to_pose
from opencood.utils.transformation_utils import veh_side_rot_and_trans_to_trasnformation_matrix
from opencood.utils.transformation_utils import inf_side_rot_and_trans_to_trasnformation_matrix

def load_json(path):
    with open(path, mode="r") as f:
        data = json.load(f)
    return data

def build_idx_to_info(data):
    idx2info = {}
    for elem in data:
        if elem["pointcloud_path"] == "":
            continue
        idx = elem["pointcloud_path"].split("/")[-1].replace(".pcd", "")
        idx2info[idx] = elem
    return idx2info

def build_idx_to_co_info(data):
    idx2info = {}
    for elem in data:
        if elem["vehicle_pointcloud_path"] == "":
            continue
        idx = elem["vehicle_pointcloud_path"].split("/")[-1].replace(".pcd", "")
        idx2info[idx] = elem
    return idx2info

def build_inf_fid_to_veh_fid(data):
    inf_fid2veh_fid = {}
    for elem in data:
        veh_fid = elem["vehicle_pointcloud_path"].split("/")[-1].rstrip('.pcd')
        inf_fid = elem["infrastructure_pointcloud_path"].split("/")[-1].rstrip('.pcd')
        inf_fid2veh_fid[inf_fid] = veh_fid
    return inf_fid2veh_fid

def id_to_str(id, digits=6):
    result = ""
    for i in range(digits):
        result = str(id % 10) + result
        id //= 10
    return result

class LateFusionDatasetDAIR(late_fusion_dataset.LateFusionDataset):
    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        self.max_cav = 2

        if 'num_sweep_frames' in params:    # number of frames we use in LSTM
            self.k = params['num_sweep_frames']
        else:
            self.k = 1

        if 'binomial_n' in params:
            self.binomial_n = params['binomial_n']
        else:
            self.binomial_n = 10

        if 'binomial_p' in params:
            self.binomial_p = params['binomial_p']
        else:
            self.binomial_p = 0

        # 控制是否需要生成GT flow
        self.is_generate_gt_flow = False
        if 'is_generate_gt_flow' in params and params['is_generate_gt_flow']:
            self.is_generate_gt_flow = True

        self.viz_bbx_flag = False

        # if project first, cav's lidar will first be projected to
        # the ego's coordinate frame. otherwise, the feature will be
        # projected instead.
        assert 'proj_first' in params['fusion']['args']
        if params['fusion']['args']['proj_first']:
            self.proj_first = True
        else:
            self.proj_first = False


        assert 'clip_pc' in params['fusion']['args']
        if params['fusion']['args']['clip_pc']:
            self.clip_pc = True
        else:
            self.clip_pc = False
        
        if 'select_kp' in params:
            self.select_keypoint = params['select_kp']
        else:
            self.select_keypoint = None

        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)

        if self.train:
            split_dir = params['root_dir']
        else:
            split_dir = params['validate_dir']

        self.root_dir = params['data_dir']

        self.inf_idx2info = build_idx_to_info(
            load_json(osp.join(self.root_dir, "infrastructure-side/data_info.json"))
        )
        self.co_idx2info = build_idx_to_co_info(
            load_json(osp.join(self.root_dir, "cooperative/data_info.json"))
        )
        self.co_data = self.co_idx2info

        self.inf_fid2veh_fid = build_inf_fid_to_veh_fid(load_json(osp.join(self.root_dir, "cooperative/data_info.json"))
        )

        self.split_info = load_json(split_dir)
        self.data = []
        for veh_idx in self.split_info:
            if self.is_valid_id(veh_idx):
                self.data.append(veh_idx)


    def __len__(self):
        return len(self.data)
        # return len(self.split_info)

    def is_valid_id(self, veh_frame_id):
        """
        Written by sizhewei @ 2022/10/05
        Given veh_frame_id, determine whether there is a corresponding inf_frame that meets the k delay requirement.

        Parameters
        ----------
        veh_frame_id : 05d
            Vehicle frame id

        Returns
        -------
        bool valud
            True means there is a corresponding road-side frame.
        """
        # print('veh_frame_id: ',veh_frame_id,'\n')
        frame_info = {}
        
        frame_info = self.co_idx2info[veh_frame_id]
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")
        cur_inf_info = self.inf_idx2info[inf_frame_id]
        if (int(inf_frame_id) - self.binomial_n*self.k < int(cur_inf_info["batch_start_id"])):
            return False
        for i in range(self.binomial_n * self.k):
            delay_id = id_to_str(int(inf_frame_id) - i) 
            if delay_id not in self.inf_fid2veh_fid.keys():
                return False

        return True
    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.
        NOTICE!
        It is different from Intermediate Fusion and Early Fusion
        Parameters
        ----------
        idx : int
            Index given by dataloader.
        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
        """
        # veh_frame_id = self.split_info[idx]
        veh_frame_id = self.data[idx]
        frame_info = self.co_data[veh_frame_id]
        system_error_offset = frame_info["system_error_offset"]
        data = OrderedDict()

        data[0] = OrderedDict()
        data[0]['ego'] = True
        data[1] = OrderedDict()
        data[1]['ego'] = False
                
        data[0]['params'] = OrderedDict()
        data[0]['params']['vehicles'] = load_json(os.path.join(self.root_dir, 'vehicle-side/label/lidar/',veh_frame_id + '.json'))
        lidar_to_novatel_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/lidar_to_novatel/'+str(veh_frame_id)+'.json'))
        novatel_to_world_json_file = load_json(os.path.join(self.root_dir,'vehicle-side/calib/novatel_to_world/'+str(veh_frame_id)+'.json'))

        transformation_matrix = veh_side_rot_and_trans_to_trasnformation_matrix(lidar_to_novatel_json_file,novatel_to_world_json_file)

        data[0]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix)

        data[0]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["vehicle_pointcloud_path"]))

        data[1]['params'] = OrderedDict()
        inf_frame_id = frame_info['infrastructure_image_path'].split("/")[-1].replace(".jpg", "")

        data[1]['params']['vehicles'] = load_json(os.path.join(self.root_dir, 'infrastructure-side/label/virtuallidar/',inf_frame_id + '.json'))
        virtuallidar_to_world_json_file = load_json(os.path.join(self.root_dir,'infrastructure-side/calib/virtuallidar_to_world/'+str(inf_frame_id)+'.json'))
        transformation_matrix1 = inf_side_rot_and_trans_to_trasnformation_matrix(virtuallidar_to_world_json_file,system_error_offset)
        data[1]['params']['lidar_pose'] = tfm_to_pose(transformation_matrix1)

        data[1]['lidar_np'], _ = pcd_utils.read_pcd(os.path.join(self.root_dir,frame_info["infrastructure_pointcloud_path"]))
        return data

    ### rewrite generate_object_center ###
    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_dairv2x_single(cav_contents)

    ### rewrite post_process ###
    """
    We have to rewrite post_process for LateFusionDatasetDAIR
    because the object id can not used for identifying the same object
    
    here we will to use the IoU to determine it.
    """
    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx_by_iou(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    ### rewrite post_process_no_fusion ###
    """
    We have to rewrite post_process for LateFusionDatasetDAIR
    because the object id can not used for identifying the same object
    
    here we will to use the IoU to determine it.
    """
    def post_process_no_fusion(self, data_dict, output_dict_ego, return_uncertainty=False):
        data_dict_ego = OrderedDict()
        data_dict_ego['ego'] = data_dict['ego']
        gt_box_tensor = self.post_processor.generate_gt_bbx_by_iou(data_dict)

        if return_uncertainty:
            pred_box_tensor, pred_score, uncertainty = \
                self.post_processor.post_process(data_dict_ego, output_dict_ego, return_uncertainty=True)
            return pred_box_tensor, pred_score, gt_box_tensor, uncertainty
        else:
            pred_box_tensor, pred_score = \
                self.post_processor.post_process(data_dict_ego, output_dict_ego)
            return pred_box_tensor, pred_score, gt_box_tensor
opencood/data_utils/datasets/late_fusion_dataset_dair.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/1/27
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for Late fusion with past k frames on irregular OPV2V
TODO: train 部分因为暂时没有用到 可能存在bug 使用前需要检查
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset, intermediate_fusion_dataset_opv2v_irregular
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# from opencood.data_utils.datasets import build_dataset 

class LateFusionDatasetIrregular(intermediate_fusion_dataset_opv2v_irregular.IntermediateFusionDatasetIrregular):
    """
    This class is for late fusion where each vehicle transmit the
    detection outputs to ego.
    """
    def __init__(self, params, visualize, train=True):
        super(LateFusionDatasetIrregular, self).__init__(params, visualize, train)

        print("=== OPV2V-Irregular Multi-sweep dataset for late fusion with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        {
            # if train 
            # if test

        }
        '''
        # # debug use TODO:
        # if idx < 170 or idx > 270:
        #     return None
        base_data_dict = self.retrieve_base_data(idx)
        
        if self.train:
            reformat_data_dict = self.get_item_train(base_data_dict)
        else:
            reformat_data_dict = self.get_item_test(base_data_dict)

        return reformat_data_dict

    def get_item_train(self, base_data_dict):
        processed_data_dict = OrderedDict()
        # base_data_dict = add_noise_data_dict(base_data_dict, self.params['noise_setting'])
        # during training, we return a random cav's data
        # only one vehicle is in processed_data_dict
        if not self.visualize:
            selected_cav_id, selected_cav_base = \
                random.choice(list(base_data_dict.items()))
        else:
            selected_cav_id, selected_cav_base = \
                list(base_data_dict.items())[0]

        selected_cav_processed = self.get_item_single_car(selected_cav_base)
        processed_data_dict.update({'ego': selected_cav_processed})

        return processed_data_dict

    def get_item_test(self, base_data_dict):
        ''' 
        Fetch useful info from base_data_dict, filter out too far cav.
        Return a dict match the point_pillar model.
        
        Params:
        ------
        base_data_dict : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {		(id)			#      |       | label
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		# (k) totally
                        [0]:{		(id)			# pose | lidar | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)			# pose | lidar | label
                        ...,						# pose | lidar | label
                        [k-1] : {} (id-(k-1))		# pose | lidar | label
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {
                    'ego': false, 
                    curr : 	{		(id)			#      |       | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {		# (k) totally
                        [0] : {		(id - \tau)		# pose | lidar | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 						# pose | lidar | label
                        [k-1]:{}  (id-\tau-(k-1))	# pose | lidar | label
                    },
                }, 
                ...
            }

        Returns:
        ------ 
        {
            'ego' : {
                'transformation_matrix_curr':
                'transformation_matrix_past': 
                'if_no_point' : True ,
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego')
                },
                'past_k' : {
                    'origin_lidar' : 
                    'processed_lidar' : 
                    'anchor_box' : 
                    'object_bbx_center' : 
                    'object_bbx_mask' : 
                    'object_ids' : 
                    'label_dict' : 
                },
                'curr' : { ... }
            },
            cav_id: { ... }     
        }
        '''
        processed_data_dict = OrderedDict()

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        # scene_name = base_data_dict[ego_id]['debug']['scene_name']
        
        too_far = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = \
                math.sqrt( \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            cav_id_list.append(cav_id)  
        # filter those out of communicate range
        for cav_id in too_far:
            base_data_dict.pop(cav_id)

        if self.visualize:
            projected_lidar_stack = []

        illegal_cav = []
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]

            selected_cav_processed = self.get_item_single_car_test(
                selected_cav_base
            )
            '''
            selected_cav_processed : {
                'if_no_point' : False / True ,
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego')
                },
                'past_k' : {
                    'origin_lidar' : 
                    'processed_lidar' : 
                    'anchor_box' : 
                    'object_bbx_center' : 
                    'object_bbx_mask' : 
                    'object_ids' : 
                    'label_dict' : 
                },
                'curr' : {}
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                # illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                # illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue
            
            cav_lidar_pose_past = selected_cav_base['past_k'][0]['params']['lidar_pose']
            transformation_matrix_past = x1_to_x2(cav_lidar_pose_past, ego_lidar_pose)
            selected_cav_processed.update({'transformation_matrix_past': transformation_matrix_past})

            cav_lidar_pose_curr = selected_cav_base['curr']['params']['lidar_pose']
            transformation_matrix_curr = x1_to_x2(cav_lidar_pose_curr, ego_lidar_pose)
            selected_cav_processed.update({'transformation_matrix_curr': transformation_matrix_curr})
            
            update_cav = "ego" if cav_id == ego_id else cav_id
            processed_data_dict.update({update_cav: selected_cav_processed})
        
        return processed_data_dict

    def get_item_single_car_test(self, selected_cav_base):
        """
        Process a single CAV's information for the train/test pipeline.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information.
            Structure : {
                'ego' : true,
                curr : {		(id)			#      |       | label
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		# (k) totally
                    [0]:{		(id)			# pose | lidar | label
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)			# pose | lidar | label
                    ...,						# pose | lidar | label
                    [k-1] : {} (id-(k-1))		# pose | lidar | label
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }    
            }

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            Structure : {
                'if_no_point' : False / True ,
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego'),
                    'timestamp' : string ('past_k'-0-timestamp)
                },
                'past_k' : {
                    'origin_lidar' : 
                    'processed_lidar' : 
                    'anchor_box' : 
                    'object_bbx_center' : 
                    'object_bbx_mask' : 
                    'object_ids' : 
                    'label_dict' : 
                },
                'curr' : {
                    'origin_lidar' : 
                    'processed_lidar' : 
                    'anchor_box' : 
                    'object_bbx_center' : 
                    'object_bbx_mask' : 
                    'object_ids' : 
                    'label_dict' : 
                }
            }
        """
        selected_cav_processed = {}

        if_no_point = False
        # for past: lidar | feature | pose | label
        for part in ['past_k', 'curr']:
            processed_part = {}
            if part=='past_k':
                processing_base = selected_cav_base['past_k'][0]
            else:
                processing_base = selected_cav_base['curr']

            # filter lidar
            lidar_np = processing_base['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_points_by_range(lidar_np,
                                            self.params['preprocess'][
                                                'cav_lidar_range'])
            # remove points that hit ego vehicle
            lidar_np = mask_ego_points(lidar_np)
            
            # tag illegal situation
            if lidar_np.shape[0] == 0: # 没有点留下
                selected_cav_processed.update({'if_no_point': True})
                return selected_cav_processed

            # generate the bounding box(n, 7) under the cav's space
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([processing_base], processing_base['params']['lidar_pose'])  

            # data augmentation
            lidar_np, object_bbx_center, object_bbx_mask = \
                self.augment(lidar_np, object_bbx_center, object_bbx_mask) # TODO: check

            if self.visualize:
                processed_part.update({'origin_lidar': lidar_np})

            # pre-process the lidar to voxel/bev/downsampled lidar
            lidar_dict = self.pre_processor.preprocess(lidar_np)
            processed_part.update({'processed_lidar': lidar_dict})

            # generate the anchor boxes
            anchor_box = self.post_processor.generate_anchor_box()
            processed_part.update({'anchor_box': anchor_box})

            processed_part.update({'object_bbx_center': object_bbx_center,
                                        'object_bbx_mask': object_bbx_mask,
                                        'object_ids': object_ids})

            # generate targets label
            label_dict = \
                self.post_processor.generate_label(
                    gt_box_center=object_bbx_center,
                    anchors=anchor_box,
                    mask=object_bbx_mask)
            
            processed_part.update({'label_dict': label_dict})

            selected_cav_processed.update({part : processed_part})
 
        debug_part = {}
        # print(selected_cav_base['debug'])
        debug_part.update({'time_diff': selected_cav_base['past_k'][0]['time_diff'], 
                            'sample_interval': selected_cav_base['past_k'][0]['sample_interval'],
                            'scene_name': selected_cav_base['debug']['scene'],
                            'cav_id': selected_cav_base['debug']['cav_id'],
                            'timestamp': selected_cav_base['past_k'][0]['timestamp']})
        selected_cav_processed.update({'if_no_point': if_no_point,
                                        'debug': debug_part})

        return selected_cav_processed

    def get_item_single_car(self, selected_cav_base, idx=-1):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true / false,
                curr : {		(id)							#      |       | label
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		# (k) totally
                    [0]:{		(id) / (id-\tau-0)				# pose | lidar | label
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    [1] : {},	(id-1) / (id-\tau-1))			# pose | lidar | label
                    ...,										# pose | lidar | label
                    [k-1] : {} 	(id-(k-1)) / (id-\tau-(k-1))	# pose | lidar | label
                }	
            }, 
        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.
        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                single_label_dict:		# single view label
                curr_feature:           # curr_feature,
                object_bbx_center:		# ego view label. np.ndarray. Shape is (max_num, 7)
                object_ids:				# ego view label index. list. length is (max_num, 7)
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		# list of past k frames' poses
                'past_k_features': 		# list of past k frames' lidar
                'past_k_time_diffs': 	# list of past k frames' time diff with current frame
                # 'past_k_tr_mats': 		# list of past k frames' transformation matrix to current ego coordinate
                'past_k_label_dicts'    # [], [TBD] list of past k frames' label dict 
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # # past k transfomation matrix
        # past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample invervals
        past_k_sample_interval = []
        
        # past k label 
        past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # # 1. trans matrix
            # transformation_matrix = \
            #     x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            # past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True
            
            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

        # past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        selected_cav_processed.update({})
        
        # # curr label at ego view
        # object_bbx_center, object_bbx_mask, object_ids = \
        #     self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
            #  'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
             'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        return selected_cav_processed

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_interval = []
        past_k_sample_inverval = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        # for debug use:
        time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_interval.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_inverval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            # pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_interval = np.hstack(past_k_time_interval)
        past_k_time_interval = torch.from_numpy(past_k_time_interval)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_inverval = np.hstack(past_k_sample_inverval)
        past_k_sample_inverval = torch.from_numpy(past_k_sample_inverval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval)

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        # pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        # label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for debug use: 
        time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                #    'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_interval,
                                   'past_k_sample_inverval': past_k_sample_inverval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        """
        Parameters:
        -----------
        batch: list, len = batch_size
        [0] : {
            'ego' : {
                'transformation_matrix_curr':               cav['curr'] to ego['curr']
                'transformation_matrix_past':               cav['past'][0] to ego['curr']
                'if_no_point' : True ,
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego')
                },
                'past_k' : {
                    'origin_lidar' : 
                    'processed_lidar' : 
                    'anchor_box' : 
                    'object_bbx_center' : 
                    'object_bbx_mask' : 
                    'object_ids' : 
                    'label_dict' : 
                },
                'curr' : { ... }
            },
            cav_id: { ... }        
        } 

        Returns:
        ------
        Structure: {
            'ego' / cav_id : {
                'anchor_box' : ,
                'object_bbx_center':                    curr,
                'object_bbx_mask':                      curr,
                'processed_lidar':                      past_k 0 ,
                'label_dict':                           curr,
                'object_ids':                           curr,
                'transformation_matrix':                cav-past to ego-curr,
                'transformation_matrix_clean':          cav-curr to ego-curr,
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego')
                },
                'origin_lidar' :                        cav-curr in ego view
            },
            cav_id : { ... }
        }
        """
        if batch[0] is None:
            return None
        # currently, we only support batch size of 1 during testing
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        batch = batch[0]

        output_dict = {}

        # for late fusion, we also need to stack the lidar for better
        # visualization
        if self.visualize:
            projected_lidar_list = []
            origin_lidar = []

        for cav_id, cav_content in batch.items():
            output_dict.update({cav_id: {}})
            # shape: (1, max_num, 7)
            object_bbx_center = \
                torch.from_numpy(np.array([cav_content['curr']['object_bbx_center']]))
            object_bbx_mask = \
                torch.from_numpy(np.array([cav_content['curr']['object_bbx_mask']]))
            object_ids = cav_content['curr']['object_ids']

            # the anchor box is the same for all bounding boxes usually, thus
            # we don't need the batch dimension.
            if cav_content['past_k']['anchor_box'] is not None:
                output_dict[cav_id].update({'anchor_box':
                    torch.from_numpy(np.array(
                        cav_content['past_k']['anchor_box']))})
            if self.visualize:
                transformation_matrix = cav_content['transformation_matrix_curr']
                origin_lidar = [cav_content['curr']['origin_lidar']] # TODO: check

                if (self.params['only_vis_ego'] is False) or (cav_id=='ego'):
                    # print(cav_id)
                    import copy
                    projected_lidar = copy.deepcopy(cav_content['curr']['origin_lidar']) # TODO: check
                    projected_lidar[:, :3] = \
                        box_utils.project_points_by_matrix_torch(
                            projected_lidar[:, :3],
                            transformation_matrix)
                    projected_lidar_list.append(projected_lidar)

            # processed lidar dictionary
            processed_lidar_torch_dict = \
                self.pre_processor.collate_batch(
                    [cav_content['past_k']['processed_lidar']])
            # label dictionary
            label_torch_dict = \
                self.post_processor.collate_batch([cav_content['curr']['label_dict']]) # TODO: check

            # save the transformation matrix (4, 4) to ego vehicle
            transformation_matrix_torch = \
                torch.from_numpy(
                    np.array(cav_content['transformation_matrix_past'])).float()
            
            # late fusion training, no noise
            transformation_matrix_clean_torch = \
                torch.from_numpy(
                    np.array(cav_content['transformation_matrix_curr'])).float()

            output_dict[cav_id].update({'object_bbx_center': object_bbx_center,
                                        'object_bbx_mask': object_bbx_mask,
                                        'processed_lidar': processed_lidar_torch_dict,
                                        'label_dict': label_torch_dict,
                                        'object_ids': object_ids,
                                        'transformation_matrix': transformation_matrix_torch,
                                        'transformation_matrix_clean': transformation_matrix_clean_torch})

            if self.visualize:
                origin_lidar = \
                    np.array(
                        downsample_lidar_minimum(pcd_np_list=origin_lidar))
                origin_lidar = torch.from_numpy(origin_lidar)
                output_dict[cav_id].update({'origin_lidar': origin_lidar})

            output_dict[cav_id].update({'debug' : cav_content['debug']})

        if self.visualize:
            projected_lidar_stack = [torch.from_numpy(
                np.vstack(projected_lidar_list))]
            output_dict['ego'].update({'origin_lidar': projected_lidar_stack})
            # output_dict['ego'].update({'projected_lidar_list': projected_lidar_list})

        return output_dict

    def post_process_no_fusion(self, data_dict, output_dict_ego, return_uncertainty=False):
        data_dict_ego = OrderedDict()
        data_dict_ego['ego'] = data_dict['ego']
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        if return_uncertainty:
            pred_box_tensor, pred_score, uncertainty = \
                self.post_processor.post_process(data_dict_ego, output_dict_ego, return_uncertainty=True)
            return pred_box_tensor, pred_score, gt_box_tensor, uncertainty
        else:
            pred_box_tensor, pred_score = \
                self.post_processor.post_process(data_dict_ego, output_dict_ego)
            return pred_box_tensor, pred_score, gt_box_tensor
            
    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

if __name__ == '__main__':   
    
    def train_parser():
        parser = argparse.ArgumentParser(description="synthetic data generation")
        parser.add_argument("--hypes_yaml", "-y", type=str, required=True,
                            help='data generation yaml file needed ')
        parser.add_argument('--model_dir', default='',
                            help='Continued training path')
        parser.add_argument('--fusion_method', '-f', default="intermediate",
                            help='passed to inference.')
        opt = parser.parse_args()
        return opt
    
    opt = train_parser()
    hypes = yaml_utils.load_yaml(opt.hypes_yaml, opt)

    print('### Dataset Building ... ###')
    opencood_train_dataset = build_dataset(hypes, visualize=False, train=True)
opencood/data_utils/datasets/late_fusion_dataset_opv2v_irregular.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/1/27
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for Late fusion with past k frames on irregular OPV2V
TODO: train 部分因为暂时没有用到 可能存在bug 使用前需要检查
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset, intermediate_fusion_dataset_opv2v_irregular
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# from opencood.data_utils.datasets import build_dataset 

class LateFusionDatasetIrregular4Flow(intermediate_fusion_dataset_opv2v_irregular.IntermediateFusionDatasetIrregular):
    """
    This class is for late fusion where each vehicle transmit the
    detection outputs to ego.
    """
    def __init__(self, params, visualize, train=True):
        super(LateFusionDatasetIrregular4Flow, self).__init__(params, visualize, train)

        print("=== Late fusion with non-ego cavs' past {} frames collected initialized! Expectation of sample interval is {}. ### {} ###  samples totally! ===".format(self.k, self.binomial_n * self.binomial_p, self.len_record[-1]))

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        {
            # if train 
            # if test

        }
        '''
        # # debug use TODO:
        # if idx < 170 or idx > 270:
        #     return None
        base_data_dict = self.retrieve_base_data(idx)
        
        if self.train:
            reformat_data_dict = self.get_item_train(base_data_dict)
        else:
            reformat_data_dict = self.get_item_test(base_data_dict)

        return reformat_data_dict

    def get_item_train(self, base_data_dict):
        processed_data_dict = OrderedDict()
        # base_data_dict = add_noise_data_dict(base_data_dict, self.params['noise_setting'])
        # during training, we return a random cav's data
        # only one vehicle is in processed_data_dict
        if not self.visualize:
            selected_cav_id, selected_cav_base = \
                random.choice(list(base_data_dict.items()))
        else:
            selected_cav_id, selected_cav_base = \
                list(base_data_dict.items())[0]

        selected_cav_processed = self.get_item_single_car(selected_cav_base)
        processed_data_dict.update({'ego': selected_cav_processed})

        return processed_data_dict

    def get_item_test(self, base_data_dict):
        ''' 
        Fetch useful info from base_data_dict, filter out too far cav.
        Return a dict match the point_pillar model.
        
        Params:
        ------
        base_data_dict : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {		(id)			#      |       | label
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		# (k) totally
                        [0]:{		(id)			# pose | lidar | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        },
                        [1] : {},	(id-1)			# pose | lidar | label
                        ...,						# pose | lidar | label
                        [k-1] : {} (id-(k-1))		# pose | lidar | label
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {
                    'ego': false, 
                    curr : 	{		(id)			#      |       | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {		# (k) totally
                        [0] : {		(id - \tau)		# pose | lidar | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                            'sample_interval': int
                        }			
                        ..., 						# pose | lidar | label
                        [k-1]:{}  (id-\tau-(k-1))	# pose | lidar | label
                    },
                }, 
                ...
            }

        Returns:
        ------ 
        {
            'ego' : {
                'if_no_point' : bool,
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego'),
                    'timestamp' : string ('past_k'-0-timestamp)
                },
                'origin_lidar' : (if visualize) original lidar
                'curr_feature' : current processed feature at cav current view
                'object_bbx_center' : 
                'object_ids' :
                'label_dict' : 
                'trans_mat_cavpast0_2_egocurr' 
                'past_k_processed_lidar' : list, len is k
                'past_k_time_diff' :  list, len is k
                'past_k_trans_mat_cavpast_2_cavpast0' : np.array, [k, 4, 4]
                'transformation_matrix_curr': cavcurr to egocurr, [4, 4]
                'past0_pose' 
            },
            cav_id: { ... } 
        }
        '''
        processed_data_dict = OrderedDict()

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        # scene_name = base_data_dict[ego_id]['debug']['scene_name']
        
        too_far = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = \
                math.sqrt( \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            cav_id_list.append(cav_id)  
        # filter those out of communicate range
        for cav_id in too_far:
            base_data_dict.pop(cav_id)

        if self.visualize:
            projected_lidar_stack = []

        illegal_cav = []
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]

            selected_cav_processed = self.get_item_single_car_test(
                selected_cav_base,
                ego_lidar_pose
            )
            '''
            selected_cav_processed : {
                'if_no_point' : False / True ,
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego')
                },
                'trans_mat_cavpast0_2_egocurr' : ,
                'past_k' : {
                    [0] : {
                        'origin_lidar' : 
                        'processed_lidar' : 
                        'anchor_box' : 
                        'object_bbx_center' : 
                        'object_bbx_mask' : 
                        'object_ids' : 
                        'label_dict' : 
                        'trans_mat_cavpast_2_cavpast0' : 
                    }, 
                    [1] : { ... },
                    ...,
                    [k-1] : { ... }
                }, 
                'curr' : {}
            }
            selected_cav_processed: dict{
                'if_no_point' : bool,
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego'),
                    'timestamp' : string ('past_k'-0-timestamp)
                },
                'origin_lidar' : (if visualize) original lidar
                'curr_feature' : current processed feature at cav current view
                'object_bbx_center' : 
                'object_ids' :
                'anchor_box' :   
                'label_dict' : 
                'trans_mat_cavpast0_2_egocurr' 
                'past_k_processed_lidar' : list, len is k
                'past_k_time_diff' :  list, len is k
                'past_k_trans_mat_cavpast_2_cavpast0' : np.array, k, 4, 4 
                'past0_pose'
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                # illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                # illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue
            
            # cav_lidar_pose_past = selected_cav_base['past_k'][0]['params']['lidar_pose']
            # transformation_matrix_past = x1_to_x2(cav_lidar_pose_past, ego_lidar_pose)
            # selected_cav_processed.update({'transformation_matrix_past': transformation_matrix_past})

            cav_lidar_pose_curr = selected_cav_base['curr']['params']['lidar_pose']
            transformation_matrix_curr = x1_to_x2(cav_lidar_pose_curr, ego_lidar_pose)
            selected_cav_processed.update({'transformation_matrix_curr': transformation_matrix_curr})
            
            update_cav = "ego" if cav_id == ego_id else cav_id
            processed_data_dict.update({update_cav: selected_cav_processed})
        
        # filter out cav with no point:
        for cav_id in illegal_cav:
            base_data_dict.pop(cav_id)

        return processed_data_dict

    def get_item_single_car_test(self, selected_cav_base, ego_lidar_pose):
        """
        Process a single CAV's information for the train/test pipeline.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information.
            Structure : {
                'ego' : true,
                curr : {		(id)			#      |       | label
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		# (k) totally
                    [0]:{		(id)			# pose | lidar | label
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)			# pose | lidar | label
                    ...,						# pose | lidar | label
                    [k-1] : {} (id-(k-1))		# pose | lidar | label
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }    
            }

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            Structure : {
                'if_no_point' : False / True ,
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego'),
                    'timestamp' : string ('past_k'-0-timestamp)
                },
                'trans_mat_cavpast0_2_egocurr' : 
                'past_k' : {
                    [0] : {
                        'origin_lidar' : 
                        'processed_lidar' : 
                        'anchor_box' : 
                        'object_bbx_center' : 
                        'object_bbx_mask' : 
                        'object_ids' : 
                        'label_dict' : 
                        'time_diff' : 
                        'trans_mat_cavpast_2_cavpast0' : 
                    }
                    [1] : { ... },
                    ...
                    [k-1] : { ... }
                },
                'curr' : {
                    'origin_lidar' : 
                    'processed_lidar' : 
                    'anchor_box' : 
                    'object_bbx_center' : 
                    'object_bbx_mask' : 
                    'object_ids' : 
                    'label_dict' : 
                }
            }
        selected_cav_processed: dict{
            'if_no_point' : bool,
            'debug' : {
                'scene_name' : string,
                'cav_id' : string,
                'time_diff': float (0.0 if 'ego'),
                'sample_interval': int (0 if 'ego'),
                'timestamp' : string ('past_k'-0-timestamp)
            },
            'origin_lidar' : (if visualize) original lidar
            'curr_feature' : current processed feature at cav current view
            'object_bbx_center' : 
            'object_bbx_mask' :
            'object_ids' :
            'label_dict' : 
            'anchor_box' : 
            'trans_mat_cavpast0_2_egocurr'                  cav-past0 to ego-curr
            'trans_mat_cavpast0_2_cavcurr'                  cav-past0 to cav-curr
            'past_k_processed_lidar' : list, len is k
            'past_k_time_diff' :  list, len is k
            'past_k_trans_mat_cavpast_2_cavpast0' : np.array, k, 4, 4 .  cav-past to cav-past0
            'past0_pose'
        }
        """
        selected_cav_processed = {}
        ego_lidar_pose = ego_lidar_pose

        if_no_point = False
        
        ###################### for current frame ######################
        processed_part = {}
        processing_base = selected_cav_base['curr']

        # filter lidar
        lidar_np = processing_base['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_points_by_range(lidar_np,
                                        self.params['preprocess'][
                                            'cav_lidar_range'])
        # remove points that hit ego vehicle
        lidar_np = mask_ego_points(lidar_np)
        
        # tag illegal situation
        if lidar_np.shape[0] == 0: # 没有点留下
            selected_cav_processed.update({'if_no_point': True})
            return selected_cav_processed

        # generate the bounding box(n, 7) under the cav's space
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([processing_base], processing_base['params']['lidar_pose'])  

        # data augmentation
        lidar_np, object_bbx_center, object_bbx_mask = \
            self.augment(lidar_np, object_bbx_center, object_bbx_mask) # TODO: check

        if self.visualize:
            selected_cav_processed.update({'origin_lidar': lidar_np})

        # pre-process the lidar to voxel/bev/downsampled lidar
        lidar_dict = self.pre_processor.preprocess(lidar_np)
        # processed_part.update({'processed_lidar': lidar_dict})
        selected_cav_processed.update({'curr_feature': lidar_dict})

        # generate the anchor boxes
        anchor_box = self.post_processor.generate_anchor_box()
        processed_part.update({'anchor_box': anchor_box})

        selected_cav_processed.update({'object_bbx_center': object_bbx_center,
            'object_bbx_mask': object_bbx_mask,
            'object_ids': object_ids, 
            'anchor_box': anchor_box})

        # generate targets label
        label_dict = \
            self.post_processor.generate_label(
                gt_box_center=object_bbx_center,
                anchors=anchor_box,
                mask=object_bbx_mask)
        
        selected_cav_processed.update({'label_dict': label_dict})

        # selected_cav_processed.update({'curr' : processed_part})

        # trans mat : past to past 0 / past 0 to curr
        trans_mat_cavpast0_2_egocurr = x1_to_x2(selected_cav_base['past_k'][0]['params']['lidar_pose'], ego_lidar_pose)
        selected_cav_processed.update({'trans_mat_cavpast0_2_egocurr': trans_mat_cavpast0_2_egocurr})

        trans_mat_cavpast0_2_cavcurr = x1_to_x2(selected_cav_base['past_k'][0]['params']['lidar_pose'], \
             selected_cav_base['curr']['params']['lidar_pose'])
        selected_cav_processed.update({'trans_mat_cavpast0_2_cavcurr': trans_mat_cavpast0_2_cavcurr})

        ################### for past k frames ########################
        """
        'processed_lidar' : 
        'time_diff' : 
        'trans_mat_cavpast_2_cavpast0' :  
        """
        # if self.visualize:
        #     past_k_origin_lidar = []
        past_k_processed_lidar = []
        # past_k_anchor_box = []
        # past_k_object_bbx_center = []
        # past_k_object_bbx_mask = [] 
        # past_k_object_ids = []
        # past_k_label_dict = []
        past_k_time_diff = [] 
        past_k_trans_mat_cavpast_2_cavpast0 = [] 

        # selected_cav_processed['past_k'] = OrderedDict()
        for i in range(self.k):
            # processed_part = {}
            processing_base = selected_cav_base['past_k'][i]
            
            # trans mat : past to past 0 / past 0 to curr
            trans_mat_cavpast_2_cavpast0 = x1_to_x2(processing_base['params']['lidar_pose'], selected_cav_base['past_k'][0]['params']['lidar_pose'])
            # processed_part.update({'trans_mat_cavpast_2_cavpast0', trans_mat_cavpast_2_cavpast0})
            past_k_trans_mat_cavpast_2_cavpast0.append(trans_mat_cavpast_2_cavpast0)

            # filter lidar
            lidar_np = processing_base['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_points_by_range(lidar_np,
                                            self.params['preprocess'][
                                                'cav_lidar_range'])
            # remove points that hit ego vehicle
            lidar_np = mask_ego_points(lidar_np)
            
            # tag illegal situation
            if lidar_np.shape[0] == 0: # 没有点留下
                selected_cav_processed.update({'if_no_point': True})
                return selected_cav_processed

            # generate the bounding box(n, 7) under the cav's space
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([processing_base], processing_base['params']['lidar_pose'])  

            # data augmentation
            lidar_np, object_bbx_center, object_bbx_mask = \
                self.augment(lidar_np, object_bbx_center, object_bbx_mask) # TODO: check

            # if self.visualize:
            #     # processed_part.update({'origin_lidar': lidar_np})
            #     past_k_origin_lidar.append(lidar_np)

            # pre-process the lidar to voxel/bev/downsampled lidar
            lidar_dict = self.pre_processor.preprocess(lidar_np)
            # processed_part.update({'processed_lidar': lidar_dict})
            past_k_processed_lidar.append(lidar_dict)

            # generate the anchor boxes
            # anchor_box = self.post_processor.generate_anchor_box()
            # processed_part.update({'anchor_box': anchor_box})
            # past_k_anchor_box.append(anchor_box)

            # processed_part.update({'object_bbx_center': object_bbx_center,
            #                             'object_bbx_mask': object_bbx_mask,
            #                             'object_ids': object_ids})
            
            # processed_part.update({'time_diff': selected_cav_base['past_k'][i]['time_diff']})
            past_k_time_diff.append(selected_cav_base['past_k'][i]['time_diff'])

            # generate targets label
            # label_dict = \
            #     self.post_processor.generate_label(
            #         gt_box_center=object_bbx_center,
            #         anchors=anchor_box,
            #         mask=object_bbx_mask)
            
            # processed_part.update({'label_dict': label_dict})

            # selected_cav_processed['past_k'][i] = processed_part

        # past_k_time_diff = np.array(past_k_time_diff)  # len = k
        merged_past_k_processed_lidar = self.merge_cav_past_k_features_to_dict(past_k_processed_lidar)
        past_k_trans_mat_cavpast_2_cavpast0 = np.stack(past_k_trans_mat_cavpast_2_cavpast0, axis=0)  # k, 4, 4

        selected_cav_processed.update({
            'past_k_processed_lidar': merged_past_k_processed_lidar, 
            'past_k_time_diff': past_k_time_diff,
            'past_k_trans_mat_cavpast_2_cavpast0': past_k_trans_mat_cavpast_2_cavpast0})

        selected_cav_processed.update({
            'past0_pose' : selected_cav_base['past_k'][0]['params']['lidar_pose']
        })
        
        ############################################################################
        debug_part = {}
        # print(selected_cav_base['debug'])
        debug_part.update({'time_diff': selected_cav_base['past_k'][0]['time_diff'], 
                            'sample_interval': selected_cav_base['past_k'][0]['sample_interval'],
                            'scene_name': selected_cav_base['debug']['scene'],
                            'cav_id': selected_cav_base['debug']['cav_id'],
                            'timestamp': selected_cav_base['past_k'][0]['timestamp']})
        selected_cav_processed.update({'if_no_point': if_no_point,
                                        'debug': debug_part})

        return selected_cav_processed
    '''
    def get_item_single_car(self, selected_cav_base, idx=-1):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true / false,
                curr : {		(id)							#      |       | label
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		# (k) totally
                    [0]:{		(id) / (id-\tau-0)				# pose | lidar | label
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    [1] : {},	(id-1) / (id-\tau-1))			# pose | lidar | label
                    ...,										# pose | lidar | label
                    [k-1] : {} 	(id-(k-1)) / (id-\tau-(k-1))	# pose | lidar | label
                }	
            }, 
        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.
        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                single_label_dict:		# single view label
                curr_feature:           # curr_feature,
                object_bbx_center:		# ego view label. np.ndarray. Shape is (max_num, 7)
                object_ids:				# ego view label index. list. length is (max_num, 7)
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		# list of past k frames' poses
                'past_k_features': 		# list of past k frames' lidar
                'past_k_time_diffs': 	# list of past k frames' time diff with current frame
                # 'past_k_tr_mats': 		# list of past k frames' transformation matrix to current ego coordinate
                'past_k_label_dicts'    # [], [TBD] list of past k frames' label dict 
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # # past k transfomation matrix
        # past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample invervals
        past_k_sample_interval = []
        
        # past k label 
        past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # # 1. trans matrix
            # transformation_matrix = \
            #     x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            # past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True
            
            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

        # past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        selected_cav_processed.update({})
        
        # # curr label at ego view
        # object_bbx_center, object_bbx_mask, object_ids = \
        #     self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
            #  'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
             'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        return selected_cav_processed
    '''
    def merge_cav_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for time_id in range(self.k):
            for feature_name, feature in processed_feature_list[time_id].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_interval = []
        past_k_sample_inverval = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        # for debug use:
        time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_interval.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_inverval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            # pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_interval = np.hstack(past_k_time_interval)
        past_k_time_interval = torch.from_numpy(past_k_time_interval)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_inverval = np.hstack(past_k_sample_inverval)
        past_k_sample_inverval = torch.from_numpy(past_k_sample_inverval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval)

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        # pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        # label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for debug use: 
        time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                #    'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_interval,
                                   'past_k_sample_inverval': past_k_sample_inverval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        """
        Parameters:
        -----------
        batch: list, len = batch_size
        [0] : {
            'ego' / cav_id : {
                'if_no_point' : bool,
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego'),
                    'timestamp' : string ('past_k'-0-timestamp)
                },
                'origin_lidar' : (if visualize) original lidar
                'curr_feature' : current processed feature at cav current view
                'object_bbx_center' : 
                'object_ids' :
                'label_dict' : 
                'trans_mat_cavpast0_2_egocurr':  4, 4
                'trans_mat_cavpast0_2_cavcurr':  4, 4
                'past_k_processed_lidar' : list, len is k
                'past_k_time_diff' :  list, len is k
                'past_k_trans_mat_cavpast_2_cavpast0' : np.array, [k, 4, 4]
                'past0_pose'
                'transformation_matrix_curr': cavcurr to egocurr, [4, 4]
            },     
        } 

        Returns:
        ------
        Structure: {
            'ego' / cav_id : {
                'anchor_box' : ,
                'object_bbx_center':                    curr,
                'object_bbx_mask':                      curr,
                'object_ids':                           curr,
                'label_dict':                           curr,
                'processed_lidar':                      past_k 0 ,
                # 'transformation_matrix':                cav-past to ego-curr,
                'transformation_matrix_clean':          cav-curr to ego-curr, for viz
                'trans_mat_past_2_past0_torch' :    cav-past to cav-past0
                'trans_mat_cavpast0_2_cavcurr_torch':   cav-past0 to cav-curr
                'trans_mat_past0_2_curr_torch' :    cav-past0 to ego-curr
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego')
                },
                'origin_lidar' :                        cav-curr in ego view
            },
            cav_id : { ... }
        }
        """
        if batch[0] is None:
            return None
        # currently, we only support batch size of 1 during testing
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        batch = batch[0]

        output_dict = {}

        # for late fusion, we also need to stack the lidar for better
        # visualization
        if self.visualize:
            projected_lidar_list = []
            origin_lidar = []

        for cav_id, cav_content in batch.items():
            output_dict.update({cav_id: {}})
            # shape: (1, max_num, 7)
            object_bbx_center = \
                torch.from_numpy(np.array([cav_content['object_bbx_center']]))
            object_bbx_mask = \
                torch.from_numpy(np.array([cav_content['object_bbx_mask']]))
            object_ids = cav_content['object_ids']

            output_dict[cav_id].update({'past_k_time_diff': 
                torch.from_numpy(np.array(
                    cav_content['past_k_time_diff']))})
            
            # 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:
                output_dict[cav_id].update({'anchor_box':
                    torch.from_numpy(np.array(
                        cav_content['anchor_box']))})
            if self.visualize:
                transformation_matrix = cav_content['transformation_matrix_curr']
                origin_lidar = [cav_content['origin_lidar']] # TODO: check

                if (self.params['only_vis_ego'] is False) or (cav_id=='ego'):
                    # print(cav_id)
                    import copy
                    projected_lidar = copy.deepcopy(cav_content['origin_lidar']) # TODO: check
                    projected_lidar[:, :3] = \
                        box_utils.project_points_by_matrix_torch(
                            projected_lidar[:, :3],
                            transformation_matrix)
                    projected_lidar_list.append(projected_lidar)

            # label dictionary
            label_torch_dict = \
                self.post_processor.collate_batch([cav_content['label_dict']]) # TODO: check
            
            # processed lidar dictionary
            processed_lidar_torch_dict = \
                self.pre_processor.collate_batch(
                    cav_content['past_k_processed_lidar'])
            

            # save the transformation matrix (4, 4) to ego vehicle
            # transformation_matrix_torch = \
            #     torch.from_numpy(
            #         np.array(cav_content['transformation_matrix_past'])).float()
            
            trans_mat_past_2_past0_torch = \
                torch.from_numpy(
                    np.array(cav_content['past_k_trans_mat_cavpast_2_cavpast0'])
                ).float()

            trans_mat_cavpast0_2_cavcurr_torch = \
                torch.from_numpy(
                    np.array(cav_content['trans_mat_cavpast0_2_cavcurr'])
                ).float()

            trans_mat_cavpast0_2_eogcurr_torch = \
                torch.from_numpy(
                    np.array(cav_content['trans_mat_cavpast0_2_egocurr'])
                ).float()
            
            # late fusion training, no noise
            transformation_matrix_clean_torch = \
                torch.from_numpy(
                    np.array(cav_content['transformation_matrix_curr'])).float()

            output_dict[cav_id].update({'object_bbx_center': object_bbx_center,
                                        'object_bbx_mask': object_bbx_mask,
                                        'object_ids': object_ids,
                                        'label_dict': label_torch_dict,
                                        'processed_lidar': processed_lidar_torch_dict,
                                        # 'transformation_matrix': transformation_matrix_torch,
                                        'transformation_matrix_clean': transformation_matrix_clean_torch,
                                        'trans_mat_past_2_past0_torch':trans_mat_past_2_past0_torch,
                                        'trans_mat_cavpast0_2_cavcurr_torch': trans_mat_cavpast0_2_cavcurr_torch,
                                        'trans_mat_past0_2_curr_torch':trans_mat_cavpast0_2_eogcurr_torch})

            if self.visualize:
                origin_lidar = \
                    np.array(
                        downsample_lidar_minimum(pcd_np_list=origin_lidar))
                origin_lidar = torch.from_numpy(origin_lidar)
                output_dict[cav_id].update({'origin_lidar': origin_lidar})

            output_dict[cav_id].update({'debug' : cav_content['debug']})

        pairwise_t_matrix = torch.from_numpy(self.get_pairwise_transformation_w_time_diff(batch, self.max_cav))
        output_dict['ego'].update({'pairwise_t_matrix': pairwise_t_matrix})
        
        if self.visualize:
            projected_lidar_stack = [torch.from_numpy(
                np.vstack(projected_lidar_list))]
            output_dict['ego'].update({'origin_lidar': projected_lidar_stack})
            # output_dict['ego'].update({'projected_lidar_list': projected_lidar_list})

        return output_dict

    def post_process_no_fusion(self, data_dict, output_dict_ego, return_uncertainty=False):
        data_dict_ego = OrderedDict()
        data_dict_ego['ego'] = data_dict['ego']
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        if return_uncertainty:
            pred_box_tensor, pred_score, uncertainty = \
                self.post_processor.post_process(data_dict_ego, output_dict_ego, return_uncertainty=True)
            return pred_box_tensor, pred_score, gt_box_tensor, uncertainty
        else:
            pred_box_tensor, pred_score = \
                self.post_processor.post_process(data_dict_ego, output_dict_ego)
            return pred_box_tensor, pred_score, gt_box_tensor
            
    def generate_pred_bbx_frames(self, data_dict, output_dict):
        box_results = self.post_processor.flow_post_process(data_dict, output_dict, self.k)
        return box_results

    def generate_gt_bbx_cav_curr(self, data_dict):
        return self.post_processor.generate_gt_bbx_cav_curr(data_dict)

    def post_process_updated(self, data_dict, output_dict):
        return self.post_processor.post_process_fuse_updated_frame(data_dict, output_dict)

    def post_process_for_intermediate(self, data_dict, output_dict):
        pred_box_tensor, pred_score = \
            self.post_processor.post_process_for_intermediate(data_dict, output_dict)

        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor
    
    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)

        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation_w_time_diff(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents
        Note, this function caculate the trans_mat between cav-past-0 and ego-curr

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['past0_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

if __name__ == '__main__':   
    
    def train_parser():
        parser = argparse.ArgumentParser(description="synthetic data generation")
        parser.add_argument("--hypes_yaml", "-y", type=str, required=True,
                            help='data generation yaml file needed ')
        parser.add_argument('--model_dir', default='',
                            help='Continued training path')
        parser.add_argument('--fusion_method', '-f', default="intermediate",
                            help='passed to inference.')
        opt = parser.parse_args()
        return opt
    
    opt = train_parser()
    hypes = yaml_utils.load_yaml(opt.hypes_yaml, opt)

    print('### Dataset Building ... ###')
    opencood_train_dataset = build_dataset(hypes, visualize=False, train=True)
opencood/data_utils/datasets/late_fusion_dataset_opv2v_irregular_4flow.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: sizhewei @ 2023/1/27
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Dataset class for Late fusion with past k frames on regular OPV2V
TODO: train 部分因为暂时没有用到 可能存在bug 使用前需要检查
"""

from collections import OrderedDict
import os
import numpy as np
import torch
import math
import random
import copy
import sys
import time
import json
from scipy import stats
import opencood.data_utils.post_processor as post_processor
import opencood.utils.pcd_utils as pcd_utils
from opencood.utils.keypoint_utils import bev_sample, get_keypoints
from opencood.data_utils.datasets import basedataset, intermediate_fusion_dataset_opv2v_multisweep
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.utils.pose_utils import add_noise_data_dict, remove_z_axis
from opencood.utils.common_utils import read_json
from opencood.utils import box_utils
# from opencood.models.sub_modules.box_align_v2 import box_alignment_relative_sample_np

# from opencood.data_utils.datasets import build_dataset 

class LateFusionDatasetRegular(intermediate_fusion_dataset_opv2v_multisweep.IntermediateFusionDatasetMultisweep):
    """
    This class is for late fusion where each vehicle transmit the
    detection outputs to ego.
    """
    def __init__(self, params, visualize, train=True):
        super(LateFusionDatasetRegular, self).__init__(params, visualize, train)

        print("=== OPV2V-Regular Multi-sweep dataset for late fusion with non-ego cavs' past {} frames collected initialized! Time delay is {}. ### {} ###  samples totally! ===".format(self.k, self.tau, self.len_record[-1]))

    def __getitem__(self, idx):
        '''
        Returns:
        ------ 
        {
            # if train 
            # if test

        }
        '''
        # # debug use TODO:
        # if idx < 170 or idx > 270:
        #     return None
        base_data_dict = self.retrieve_base_data(idx)
        
        if self.train:
            reformat_data_dict = self.get_item_train(base_data_dict)
        else:
            reformat_data_dict = self.get_item_test(base_data_dict)

        return reformat_data_dict

    def get_item_train(self, base_data_dict):
        processed_data_dict = OrderedDict()
        # base_data_dict = add_noise_data_dict(base_data_dict, self.params['noise_setting'])
        # during training, we return a random cav's data
        # only one vehicle is in processed_data_dict
        if not self.visualize:
            selected_cav_id, selected_cav_base = \
                random.choice(list(base_data_dict.items()))
        else:
            selected_cav_id, selected_cav_base = \
                list(base_data_dict.items())[0]

        selected_cav_processed = self.get_item_single_car(selected_cav_base)
        processed_data_dict.update({'ego': selected_cav_processed})

        return processed_data_dict

    def get_item_test(self, base_data_dict):
        ''' 
        Fetch useful info from base_data_dict, filter out too far cav.
        Return a dict match the point_pillar model.
        
        Params:
        ------
        base_data_dict : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            Structure: 
            {
                cav_id_1 : {
                    'ego' : true,
                    curr : {		(id)			#      |       | label
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    past_k : {		# (k) totally
                        [0]:{		(id)			# pose | lidar | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float
                        },
                        [1] : {},	(id-1)			# pose | lidar | label
                        ...,						# pose | lidar | label
                        [k-1] : {} (id-(k-1))		# pose | lidar | label
                    },
                    'debug' : {                     # debug use
                        scene_name : string         
                        cav_id : string
                    }
                    
                }, 
                cav_id_2 : {
                    'ego': false, 
                    curr : 	{		(id)			#      |       | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string
                    },
                    past_k: {		# (k) totally
                        [0] : {		(id - \tau)		# pose | lidar | label
                            'params': (yaml),
                            'lidar_np': (numpy),
                            'timestamp': string,
                            'time_diff': float,
                        }			
                        ..., 						# pose | lidar | label
                        [k-1]:{}  (id-\tau-(k-1))	# pose | lidar | label
                    },
                }, 
                ...
            }

        Returns:
        ------ 
        {
            'ego' : {
                'transformation_matrix_curr':
                'transformation_matrix_past': 
                'if_no_point' : True ,
                'debug' : {
                    'scene_name' : string,
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego'),
                    'sample_interval': int (0 if 'ego')
                },
                'past_k' : {
                    'origin_lidar' : 
                    'processed_lidar' : 
                    'anchor_box' : 
                    'object_bbx_center' : 
                    'object_bbx_mask' : 
                    'object_ids' : 
                    'label_dict' : 
                },
                'curr' : { ... }
            },
            cav_id: { ... }     
        }
        '''
        processed_data_dict = OrderedDict()

        # first find the ego vehicle's lidar pose
        ego_id = -1
        ego_lidar_pose = []
        for cav_id, cav_content in base_data_dict.items():
            if cav_content['ego']:
                ego_id = cav_id
                ego_lidar_pose = cav_content['curr']['params']['lidar_pose']
                break	
        assert cav_id == list(base_data_dict.keys())[
            0], "The first element in the OrderedDict must be ego"
        assert ego_id != -1
        assert len(ego_lidar_pose) > 0

        # scene_name = base_data_dict[ego_id]['debug']['scene_name']
        
        too_far = []
        cav_id_list = []

        # loop over all CAVs to process information
        for cav_id, selected_cav_base in base_data_dict.items():
            # check if the cav is within the communication range with ego
            # for non-ego cav, we use the latest frame's pose
            distance = \
                math.sqrt( \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][0] - ego_lidar_pose[0]) ** 2 + \
                    (selected_cav_base['past_k'][0]['params']['lidar_pose'][1] - ego_lidar_pose[1]) ** 2)
            # if distance is too far, we will just skip this agent
            if distance > self.params['comm_range']:
                too_far.append(cav_id)
                continue
            cav_id_list.append(cav_id)  
        # filter those out of communicate range
        for cav_id in too_far:
            base_data_dict.pop(cav_id)

        if self.visualize:
            projected_lidar_stack = []

        illegal_cav = []
        for cav_id in cav_id_list:
            selected_cav_base = base_data_dict[cav_id]

            selected_cav_processed = self.get_item_single_car_test(
                selected_cav_base
            )
            '''
            selected_cav_processed : {
                'if_no_point' : False / True ,
                'debug' : {
                    'time_diff': float (0.0 if 'ego'),
                    'timestamp': string
                },
                'past_k' : {
                    'origin_lidar' : 
                    'processed_lidar' : 
                    'anchor_box' : 
                    'object_bbx_center' : 
                    'object_bbx_mask' : 
                    'object_ids' : 
                    'label_dict' : 
                },
                'curr' : { ... }
            }
            '''

            if selected_cav_processed['if_no_point']: # 把点的数量不合法的车排除
                illegal_cav.append(cav_id)
                # # 把出现不合法sample的 场景、车辆、时刻 记录下来:
                # illegal_path = os.path.join(base_data_dict[cav_id]['debug']['scene'], cav_id, base_data_dict[cav_id]['past_k'][0]['timestamp']+'.npy')
                # illegal_path_list.add(illegal_path)
                # print(illegal_path)
                continue
            
            cav_lidar_pose_past = selected_cav_base['past_k'][0]['params']['lidar_pose']
            transformation_matrix_past = x1_to_x2(cav_lidar_pose_past, ego_lidar_pose)
            selected_cav_processed.update({'transformation_matrix_past': transformation_matrix_past})

            cav_lidar_pose_curr = selected_cav_base['curr']['params']['lidar_pose']
            transformation_matrix_curr = x1_to_x2(cav_lidar_pose_curr, ego_lidar_pose)
            selected_cav_processed.update({'transformation_matrix_curr': transformation_matrix_curr})
            
            update_cav = "ego" if cav_id == ego_id else cav_id
            processed_data_dict.update({update_cav: selected_cav_processed})
        
        return processed_data_dict

    def get_item_single_car_test(self, selected_cav_base):
        """
        Process a single CAV's information for the train/test pipeline.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information.
            Structure : {
                'ego' : true / false , 
                curr : {		(id)			#      |       | label
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		# (k) totally
                    [0]:{		(id)			# pose | lidar | label
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string,
                        'time_diff': float,
                    },
                    [1] : {},	(id-1)			# pose | lidar | label
                    ...,						# pose | lidar | label
                    [k-1] : {} (id-(k-1))		# pose | lidar | label
                },
                'debug' : {                     # debug use
                    scene_name : string         
                    cav_id : string
                }    
            }

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            Structure : {
                'if_no_point' : False / True ,
                'debug' : {
                    'time_diff': float (0.0 if 'ego'),
                    'timestamp' : string ('past_k'-0-timestamp)
                },
                'past_k' : {
                    'origin_lidar' : 
                    'processed_lidar' : 
                    'anchor_box' : 
                    'object_bbx_center' : 
                    'object_bbx_mask' : 
                    'object_ids' : 
                    'label_dict' : 
                },
                'curr' : {
                    'origin_lidar' : 
                    'processed_lidar' : 
                    'anchor_box' : 
                    'object_bbx_center' : 
                    'object_bbx_mask' : 
                    'object_ids' : 
                    'label_dict' : 
                }
            }
        """
        selected_cav_processed = {}

        if_no_point = False
        # for past: lidar | feature | pose | label
        for part in ['past_k', 'curr']:
            processed_part = {}
            if part=='past_k':
                processing_base = selected_cav_base['past_k'][0] # 只用距离最近的过去帧
            else:
                processing_base = selected_cav_base['curr']

            # filter lidar
            lidar_np = processing_base['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_points_by_range(lidar_np,
                                            self.params['preprocess'][
                                                'cav_lidar_range'])
            # remove points that hit ego vehicle
            lidar_np = mask_ego_points(lidar_np)
            
            # tag illegal situation
            if lidar_np.shape[0] == 0: # 没有点留下
                selected_cav_processed.update({'if_no_point': True})
                return selected_cav_processed

            # generate the bounding box(n, 7) under the cav's space
            object_bbx_center, object_bbx_mask, object_ids = \
                self.generate_object_center([processing_base], processing_base['params']['lidar_pose'])  

            # data augmentation
            lidar_np, object_bbx_center, object_bbx_mask = \
                self.augment(lidar_np, object_bbx_center, object_bbx_mask) # TODO: check

            if self.visualize:
                processed_part.update({'origin_lidar': lidar_np})

            # pre-process the lidar to voxel/bev/downsampled lidar
            lidar_dict = self.pre_processor.preprocess(lidar_np)
            processed_part.update({'processed_lidar': lidar_dict})

            # generate the anchor boxes
            anchor_box = self.post_processor.generate_anchor_box()
            processed_part.update({'anchor_box': anchor_box})

            processed_part.update({'object_bbx_center': object_bbx_center,
                                        'object_bbx_mask': object_bbx_mask,
                                        'object_ids': object_ids})

            # generate targets label
            label_dict = \
                self.post_processor.generate_label(
                    gt_box_center=object_bbx_center,
                    anchors=anchor_box,
                    mask=object_bbx_mask)
            
            processed_part.update({'label_dict': label_dict})

            selected_cav_processed.update({part : processed_part})
 
        debug_part = {}
        # print(selected_cav_base['debug'])
        debug_part.update({'time_diff': selected_cav_base['past_k'][0]['time_diff'], 
                            'timestamp': selected_cav_base['past_k'][0]['timestamp']})
        selected_cav_processed.update({'if_no_point': if_no_point,
                                        'debug': debug_part})

        return selected_cav_processed

    def get_item_single_car(self, selected_cav_base, idx=-1):
        """
        Project the lidar and bbx to ego space first, and then do clipping.

        Parameters
        ----------
        selected_cav_base : dict
            The dictionary contains a single CAV's raw information, 
            structure: {
                'ego' : true / false,
                curr : {		(id)							#      |       | label
                    'params': (yaml),
                    'lidar_np': (numpy),
                    'timestamp': string
                },
                past_k : {		# (k) totally
                    [0]:{		(id) / (id-\tau-0)				# pose | lidar | label
                        'params': (yaml),
                        'lidar_np': (numpy),
                        'timestamp': string
                    },
                    [1] : {},	(id-1) / (id-\tau-1))			# pose | lidar | label
                    ...,										# pose | lidar | label
                    [k-1] : {} 	(id-(k-1)) / (id-\tau-(k-1))	# pose | lidar | label
                }	
            }, 
        ego_pose : list, length 6
            The ego vehicle lidar pose under world coordinate.
        idx: int,
            debug use.

        Returns
        -------
        selected_cav_processed : dict
            The dictionary contains the cav's processed information.
            {
                single_label_dict:		# single view label
                curr_feature:           # curr_feature,
                object_bbx_center:		# ego view label. np.ndarray. Shape is (max_num, 7)
                object_ids:				# ego view label index. list. length is (max_num, 7)
                'curr_pose':			# current pose, list, len = 6
                'past_k_poses': 		# list of past k frames' poses
                'past_k_features': 		# list of past k frames' lidar
                'past_k_time_diffs': 	# list of past k frames' time diff with current frame
                # 'past_k_tr_mats': 		# list of past k frames' transformation matrix to current ego coordinate
                'past_k_label_dicts'    # [], [TBD] list of past k frames' label dict 
            }
        """
        selected_cav_processed = {}

        # curr lidar feature
        lidar_np = selected_cav_base['curr']['lidar_np']
        lidar_np = shuffle_points(lidar_np)
        lidar_np = mask_ego_points(lidar_np) # remove points that hit itself

        lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
        curr_feature = self.pre_processor.preprocess(lidar_np)
        
        # # past k transfomation matrix
        # past_k_tr_mats = []
        # past k lidars
        past_k_features = []
        # past k poses
        past_k_poses = []
        # past k timestamps
        past_k_time_diffs = []
        # past k sample invervals
        past_k_sample_interval = []
        
        # past k label 
        past_k_label_dicts = [] # todo 这个部分可以删掉

        # 判断点的数量是否合法
        if_no_point = False

        # past k frames [trans matrix], [lidar feature], [pose], [time interval]
        for i in range(self.k):
            # # 1. trans matrix
            # transformation_matrix = \
            #     x1_to_x2(selected_cav_base['past_k'][i]['params']['lidar_pose'], ego_pose) # T_ego_cav, np.ndarray
            # past_k_tr_mats.append(transformation_matrix)

            # 2. lidar feature
            lidar_np = selected_cav_base['past_k'][i]['lidar_np']
            lidar_np = shuffle_points(lidar_np)
            lidar_np = mask_ego_points(lidar_np) # remove points that hit itself
            lidar_np = mask_points_by_range(lidar_np, self.params['preprocess']['cav_lidar_range'])
            processed_features = self.pre_processor.preprocess(lidar_np)
            past_k_features.append(processed_features)

            if lidar_np.shape[0] == 0: # 没有点留下
                if_no_point = True
            
            # 3. pose
            past_k_poses.append(selected_cav_base['past_k'][i]['params']['lidar_pose'])

            # 4. time interval and sample interval
            past_k_time_diffs.append(selected_cav_base['past_k'][i]['time_diff'])
            past_k_sample_interval.append(selected_cav_base['past_k'][i]['sample_interval'])

        # past_k_tr_mats = np.stack(past_k_tr_mats, axis=0) # (k, 4, 4)

        # curr label at single view
        # opencood/data_utils/post_processor/base_postprocessor.py
        object_bbx_center, object_bbx_mask, object_ids = \
            self.generate_object_center([selected_cav_base['curr']], selected_cav_base['curr']['params']['lidar_pose'])  
        # generate the anchor boxes
        # opencood/data_utils/post_processor/voxel_postprocessor.py
        anchor_box = self.anchor_box
        label_dict = self.post_processor.generate_label(
                gt_box_center=object_bbx_center, anchors=anchor_box, mask=object_bbx_mask
            )
        selected_cav_processed.update({})
        
        # # curr label at ego view
        # object_bbx_center, object_bbx_mask, object_ids = \
        #     self.generate_object_center([selected_cav_base['curr']], ego_pose)
            
        selected_cav_processed.update(
            {"single_label_dict": label_dict,
             "curr_feature": curr_feature,
             'object_bbx_center': object_bbx_center[object_bbx_mask == 1],
             'object_ids': object_ids,
             'curr_pose': selected_cav_base['curr']['params']['lidar_pose'],
            #  'past_k_tr_mats': past_k_tr_mats,
             'past_k_poses': past_k_poses,
             'past_k_features': past_k_features,
             'past_k_time_diffs': past_k_time_diffs,
             'past_k_sample_interval': past_k_sample_interval,
             'past_k_label_dicts': past_k_label_dicts,
             'if_no_point': if_no_point
             })

        return selected_cav_processed

    def merge_past_k_features_to_dict(self, processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for cav_id in range(len(processed_feature_list)):
            for time_id in range(self.k):
                for feature_name, feature in processed_feature_list[cav_id][time_id].items():
                    if feature_name not in merged_feature_dict:
                        merged_feature_dict[feature_name] = []
                    if isinstance(feature, list):
                        merged_feature_dict[feature_name] += feature
                    else:
                        merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    @staticmethod
    def merge_features_to_dict(processed_feature_list):
        """
        Merge the preprocessed features from different cavs to the same
        dictionary.

        Parameters
        ----------
        processed_feature_list : list
            A list of dictionary containing all processed features from
            different cavs.

        Returns
        -------
        merged_feature_dict: dict
            key: feature names, value: list of features.
        """

        merged_feature_dict = OrderedDict()

        for i in range(len(processed_feature_list)):
            for feature_name, feature in processed_feature_list[i].items():
                if feature_name not in merged_feature_dict:
                    merged_feature_dict[feature_name] = []
                if isinstance(feature, list):
                    merged_feature_dict[feature_name] += feature
                else:
                    merged_feature_dict[feature_name].append(feature) # merged_feature_dict['coords'] = [f1,f2,f3,f4]
        return merged_feature_dict

    def collate_batch_train(self, batch):
        '''
        Parameters:
        ----------
        batch[i]['ego'] structure:
        {   
            'single_object_dict_stack': single_label_dict_stack,
            'curr_processed_lidar': merged_curr_feature_dict,
            'object_bbx_center': object_bbx_center,
            'object_bbx_mask': mask,
            'object_ids': [object_id_stack[i] for i in unique_indices],
            'anchor_box': anchor_box,
            'processed_lidar': merged_feature_dict,
            'label_dict': label_dict,
            'cav_num': cav_num,
            'pairwise_t_matrix': pairwise_t_matrix,
            'curr_lidar_poses': curr_lidar_poses,
            'past_k_lidar_poses': past_k_lidar_poses,
            'sample_idx': idx,
            'cav_id_list': cav_id_list,
            'past_k_time_diffs': past_k_time_diffs_stack, 
                list of len(\sum_i^num_cav k_i), k_i represents the num of past frames of cav_i
        }
        '''
        for i in range(len(batch)):
            if batch[i] is None:
                return None
        # Intermediate fusion is different the other two
        output_dict = {'ego': {}}

        single_object_label = []

        pos_equal_one_single = []
        neg_equal_one_single = []
        targets_single = []

        object_bbx_center = []
        object_bbx_mask = []
        object_ids = []
        curr_processed_lidar_list = []
        processed_lidar_list = []
        # used to record different scenario
        record_len = []
        label_dict_list = []
        curr_lidar_pose_list = []
        past_k_lidar_pose_list = []
        past_k_label_list = []
        # store the time interval of each feature map
        past_k_time_interval = []
        past_k_sample_inverval = []
        past_k_avg_sample_interval = []
        # pairwise transformation matrix
        pairwise_t_matrix_list = []

        # for debug use:
        time_consume = np.zeros_like(batch[0]['ego']['times'])

        if self.visualize:
            origin_lidar = []
        
        for i in range(len(batch)):
            ego_dict = batch[i]['ego']
            single_object_label.append(ego_dict['single_object_dict_stack'])

            pos_equal_one_single.append(ego_dict['single_object_dict_stack']['pos_equal_one'])
            neg_equal_one_single.append(ego_dict['single_object_dict_stack']['neg_equal_one'])
            targets_single.append(ego_dict['single_object_dict_stack']['targets'])

            curr_processed_lidar_list.append(ego_dict['curr_processed_lidar'])
            object_bbx_center.append(ego_dict['object_bbx_center'])
            object_bbx_mask.append(ego_dict['object_bbx_mask'])
            object_ids.append(ego_dict['object_ids'])
            curr_lidar_pose_list.append(ego_dict['curr_lidar_poses']) # ego_dict['curr_lidar_pose'] is np.ndarray [N,6]
            past_k_lidar_pose_list.append(ego_dict['past_k_lidar_poses']) # ego_dict['past_k_lidar_pose'] is np.ndarray [N,k,6]
            past_k_time_interval.append(ego_dict['past_k_time_diffs']) # ego_dict['past_k_time_diffs'] is np.array(), len=nxk
            past_k_sample_inverval.append(ego_dict['past_k_sample_interval']) # ego_dict['past_k_sample_interval'] is np.array(), len=nxk
            past_k_avg_sample_interval.append(ego_dict['avg_sample_interval']) # ego_dict['avg_sample_interval'] is float
            processed_lidar_list.append(ego_dict['processed_lidar']) # different cav_num, ego_dict['processed_lidar'] is list.
            record_len.append(ego_dict['cav_num'])
            label_dict_list.append(ego_dict['label_dict'])
            # pairwise_t_matrix_list.append(ego_dict['pairwise_t_matrix'])
            # past_k_label_list.append(ego_dict['past_k_label_dicts'])

            time_consume += ego_dict['times']
            if self.visualize:
                origin_lidar.append(ego_dict['origin_lidar'])
        
        
        # single_object_label = self.post_processor.collate_batch(single_object_label)
        single_object_label = { "pos_equal_one": torch.cat(pos_equal_one_single, dim=0),
                                "neg_equal_one": torch.cat(neg_equal_one_single, dim=0),
                                 "targets": torch.cat(targets_single, dim=0)}
                                 
        # collate past k single view label from different batch [B, cav_num, k, 100, 252, 2]...
        past_k_single_label_torch_dict = self.post_processor.collate_batch(past_k_label_list)
        
        # collate past k time interval from different batch, (B, )
        past_k_time_interval = np.hstack(past_k_time_interval)
        past_k_time_interval = torch.from_numpy(past_k_time_interval)

        # collate past k sample interval from different batch, (B, )
        past_k_sample_inverval = np.hstack(past_k_sample_inverval)
        past_k_sample_inverval = torch.from_numpy(past_k_sample_inverval)

        past_k_avg_sample_interval = np.array(past_k_avg_sample_interval)
        avg_sample_interval = sum(past_k_avg_sample_interval) / len(past_k_avg_sample_interval)

        # convert to numpy, (B, max_num, 7)
        object_bbx_center = torch.from_numpy(np.array(object_bbx_center))
        # （B, max_num)
        object_bbx_mask = torch.from_numpy(np.array(object_bbx_mask))
        
        curr_merged_feature_dict = self.merge_features_to_dict(curr_processed_lidar_list)
        curr_processed_lidar_torch_dict = \
            self.pre_processor.collate_batch(curr_merged_feature_dict)
        # processed_lidar_list: list, len is 6. [batch_i] is OrderedDict, 3 keys: {'voxel_features': , ...}
        # example: {'voxel_features':[np.array([1,2,3]]),
        # np.array([3,5,6]), ...]}
        merged_feature_dict = self.merge_features_to_dict(processed_lidar_list)
        # [sum(record_len), C, H, W]
        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))
        # [[N1, 6], [N2, 6]...] -> [[N1+N2+...], 6]
        curr_lidar_pose = torch.from_numpy(np.concatenate(curr_lidar_pose_list, axis=0))
        # [[N1, k, 6], [N2, k, 6]...] -> [(N1+N2+...), k, 6]
        past_k_lidar_pose = torch.from_numpy(np.concatenate(past_k_lidar_pose_list, axis=0))
        label_torch_dict = \
            self.post_processor.collate_batch(label_dict_list)

        # (B, max_cav, k, 4, 4)
        # pairwise_t_matrix = torch.from_numpy(np.array(pairwise_t_matrix_list))

        # add pairwise_t_matrix to label dict
        # label_torch_dict['pairwise_t_matrix'] = pairwise_t_matrix
        label_torch_dict['record_len'] = record_len

        # for debug use: 
        time_consume = torch.from_numpy(time_consume)

        # object id is only used during inference, where batch size is 1.
        # so here we only get the first element.
        output_dict['ego'].update({'single_object_label': single_object_label,
                                   'curr_processed_lidar': curr_processed_lidar_torch_dict,
                                   '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,
                                   'single_past_dict': past_k_single_label_torch_dict,
                                   'object_ids': object_ids[0],
                                #    'pairwise_t_matrix': pairwise_t_matrix,
                                   'curr_lidar_pose': curr_lidar_pose,
                                   'past_lidar_pose': past_k_lidar_pose,
                                   'past_k_time_interval': past_k_time_interval,
                                   'past_k_sample_inverval': past_k_sample_inverval,
                                   'avg_sample_interval': avg_sample_interval,
                                   'times': time_consume})

        if self.visualize:
            origin_lidar = \
                np.array(downsample_lidar_minimum(pcd_np_list=origin_lidar))
            origin_lidar = torch.from_numpy(origin_lidar)
            output_dict['ego'].update({'origin_lidar': origin_lidar})
            
        # if self.params['preprocess']['core_method'] == 'SpVoxelPreprocessor' and \
        #     (output_dict['ego']['processed_lidar']['voxel_coords'][:, 0].max().int().item() + 1) != record_len.sum().int().item():
        #     return None

        return output_dict

    def collate_batch_test(self, batch):
        """
        Parameters:
        -----------
        batch: list, len = batch_size
        [0] : {
            'ego' : {
                'transformation_matrix_curr':               cav['curr'] to ego['curr']
                'transformation_matrix_past':               cav['past'][0] to ego['curr']
                'if_no_point' : True ,
                'debug' : {
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego')
                },
                'past_k' : {
                    'origin_lidar' : 
                    'processed_lidar' : 
                    'anchor_box' : 
                    'object_bbx_center' : 
                    'object_bbx_mask' : 
                    'object_ids' : 
                    'label_dict' : 
                },
                'curr' : { ... }
            },
            cav_id: { ... }        
        } 

        Returns:
        ------
        Structure: {
            'ego' / cav_id : {
                'anchor_box' : ,
                'object_bbx_center':                    curr,
                'object_bbx_mask':                      curr,
                'processed_lidar':                      past_k 0 ,
                'label_dict':                           curr,
                'object_ids':                           curr,
                'transformation_matrix':                cav-past to ego-curr,
                'transformation_matrix_clean':          cav-curr to ego-curr,
                'debug' : {
                    'cav_id' : string,
                    'time_diff': float (0.0 if 'ego')
                },
                'origin_lidar' :                        cav-curr in ego view
            },
            cav_id : { ... }
        }
        """
        if batch[0] is None:
            return None
        # currently, we only support batch size of 1 during testing
        assert len(batch) <= 1, "Batch size 1 is required during testing!"
        batch = batch[0]

        output_dict = {}

        # for late fusion, we also need to stack the lidar for better
        # visualization
        if self.visualize:
            projected_lidar_list = []
            origin_lidar = []

        for cav_id, cav_content in batch.items():
            output_dict.update({cav_id: {}})
            # shape: (1, max_num, 7)
            object_bbx_center = \
                torch.from_numpy(np.array([cav_content['curr']['object_bbx_center']]))
            object_bbx_mask = \
                torch.from_numpy(np.array([cav_content['curr']['object_bbx_mask']]))
            object_ids = cav_content['curr']['object_ids']

            # the anchor box is the same for all bounding boxes usually, thus
            # we don't need the batch dimension.
            if cav_content['past_k']['anchor_box'] is not None:
                output_dict[cav_id].update({'anchor_box':
                    torch.from_numpy(np.array(
                        cav_content['past_k']['anchor_box']))})
            if self.visualize:
                transformation_matrix = cav_content['transformation_matrix_curr']
                origin_lidar = [cav_content['curr']['origin_lidar']] # TODO: check

                if (self.params['only_vis_ego'] is False) or (cav_id=='ego'):
                    # print(cav_id)
                    import copy
                    projected_lidar = copy.deepcopy(cav_content['curr']['origin_lidar']) # TODO: check
                    projected_lidar[:, :3] = \
                        box_utils.project_points_by_matrix_torch(
                            projected_lidar[:, :3],
                            transformation_matrix)
                    projected_lidar_list.append(projected_lidar)

            # processed lidar dictionary
            processed_lidar_torch_dict = \
                self.pre_processor.collate_batch(
                    [cav_content['past_k']['processed_lidar']])
            # label dictionary
            label_torch_dict = \
                self.post_processor.collate_batch([cav_content['curr']['label_dict']]) # TODO: check

            # save the transformation matrix (4, 4) to ego vehicle
            transformation_matrix_torch = \
                torch.from_numpy(
                    np.array(cav_content['transformation_matrix_past'])).float()
            
            # late fusion training, no noise
            transformation_matrix_clean_torch = \
                torch.from_numpy(
                    np.array(cav_content['transformation_matrix_curr'])).float()

            output_dict[cav_id].update({'object_bbx_center': object_bbx_center,
                                        'object_bbx_mask': object_bbx_mask,
                                        'processed_lidar': processed_lidar_torch_dict,
                                        'label_dict': label_torch_dict,
                                        'object_ids': object_ids,
                                        'transformation_matrix': transformation_matrix_torch,
                                        'transformation_matrix_clean': transformation_matrix_clean_torch})

            if self.visualize:
                origin_lidar = \
                    np.array(
                        downsample_lidar_minimum(pcd_np_list=origin_lidar))
                origin_lidar = torch.from_numpy(origin_lidar)
                output_dict[cav_id].update({'origin_lidar': origin_lidar})

            output_dict[cav_id].update({'debug' : cav_content['debug']})

        if self.visualize:
            projected_lidar_stack = [torch.from_numpy(
                np.vstack(projected_lidar_list))]
            output_dict['ego'].update({'origin_lidar': projected_lidar_stack})
            # output_dict['ego'].update({'projected_lidar_list': projected_lidar_list})

        return output_dict

    def post_process_no_fusion(self, data_dict, output_dict_ego, return_uncertainty=False):
        data_dict_ego = OrderedDict()
        data_dict_ego['ego'] = data_dict['ego']
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        if return_uncertainty:
            pred_box_tensor, pred_score, uncertainty = \
                self.post_processor.post_process(data_dict_ego, output_dict_ego, return_uncertainty=True)
            return pred_box_tensor, pred_score, gt_box_tensor, uncertainty
        else:
            pred_box_tensor, pred_score = \
                self.post_processor.post_process(data_dict_ego, output_dict_ego)
            return pred_box_tensor, pred_score, gt_box_tensor
            
    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box_tensor : torch.Tensor
            The tensor of prediction bounding box after NMS.
        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)
        gt_box_tensor = self.post_processor.generate_gt_bbx(data_dict)

        return pred_box_tensor, pred_score, gt_box_tensor

    def get_pairwise_transformation(self, base_data_dict, max_cav):
        """
        Get pair-wise transformation matrix accross different agents.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The pairwise transformation matrix across each cav.
            shape: (L, L, 4, 4), L is the max cav number in a scene
            pairwise_t_matrix[i, j] is Tji, i_to_j
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                lidar_pose = cav_content['curr']['params']['lidar_pose']
                t_list.append(x_to_world(lidar_pose))  # Twx

            for i in range(len(t_list)):
                for j in range(len(t_list)):
                    # identity matrix to self
                    if i != j:
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix
    
    def get_past_k_pairwise_transformation2ego(self, base_data_dict, ego_pose, max_cav):
        """
        Get transformation matrixes accross different agents to curr ego at all past timestamps.

        Parameters
        ----------
        base_data_dict : dict
            Key : cav id, item: transformation matrix to ego, lidar points.
        
        ego_pose : list
            ego pose

        max_cav : int
            The maximum number of cav, default 5

        Return
        ------
        pairwise_t_matrix : np.array
            The transformation matrix each cav to curr ego at past k frames.
            shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
            pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
        """
        pairwise_t_matrix = np.tile(np.eye(4), (max_cav, self.k, 1, 1)) # (L, k, 4, 4)

        if self.proj_first:
            # if lidar projected to ego first, then the pairwise matrix
            # becomes identity
            # no need to warp again in fusion time.

            # pairwise_t_matrix[:, :] = np.identity(4)
            return pairwise_t_matrix
        else:
            t_list = []

            # save all transformation matrix in a list in order first.
            for cav_id, cav_content in base_data_dict.items():
                past_k_poses = []
                for time_id in range(self.k):
                    past_k_poses.append(x_to_world(cav_content['past_k'][time_id]['params']['lidar_pose']))
                t_list.append(past_k_poses) # Twx
            
            ego_pose = x_to_world(ego_pose)
            for i in range(len(t_list)): # different cav
                if i!=0 :
                    for j in range(len(t_list[i])): # different time
                        # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                        # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                        t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                        pairwise_t_matrix[i, j] = t_matrix

        return pairwise_t_matrix

if __name__ == '__main__':   
    
    def train_parser():
        parser = argparse.ArgumentParser(description="synthetic data generation")
        parser.add_argument("--hypes_yaml", "-y", type=str, required=True,
                            help='data generation yaml file needed ')
        parser.add_argument('--model_dir', default='',
                            help='Continued training path')
        parser.add_argument('--fusion_method', '-f', default="intermediate",
                            help='passed to inference.')
        opt = parser.parse_args()
        return opt
    
    opt = train_parser()
    hypes = yaml_utils.load_yaml(opt.hypes_yaml, opt)

    print('### Dataset Building ... ###')
    opencood_train_dataset = build_dataset(hypes, visualize=False, train=True)
opencood/data_utils/datasets/late_fusion_dataset_opv2v_regular.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu

"""
Dataset class for V2X-Sim2.0 dataset late fusion
"""
import random
import math
from collections import OrderedDict

import numpy as np
import torch
from torch.utils.data import DataLoader
import pickle

from opencood.data_utils.augmentor.data_augmentor import DataAugmentor
import opencood.data_utils.post_processor as post_processor
from opencood.utils import box_utils
from opencood.data_utils.datasets import basedataset
from opencood.data_utils.datasets import late_fusion_dataset
from opencood.data_utils.pre_processor import build_preprocessor
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import \
    mask_points_by_range, mask_ego_points, shuffle_points, \
    downsample_lidar_minimum
from opencood.utils.pose_utils import add_noise_data_dict
from opencood.utils.transformation_utils import x1_to_x2, x_to_world, tfm_to_pose
from opencood.utils.common_utils import read_json
from icecream import ic


class LateFusionDatasetV2X(late_fusion_dataset.LateFusionDataset):
    """
        IntermediateFusionDatasetV2X is children of IntermediateFusionDataset
        It works almost the same as IntermediateFusionDataset,
        But it 
        rewrite __init__, __len__, retrieve_base_data , which is from BaseDataset
        rewrite generate_object_center, which is from IntermediateFusionDataset
    """

    ### rewrite __init__ ###
    def __init__(self, params, visualize, train=True):
        self.params = params
        self.visualize = visualize
        self.train = train
        self.data_augmentor = DataAugmentor(params['data_augment'],
                                            train)
        
        # same as IntermediateFusionDataset
        self.proj_first = True
        if 'proj_first' in params['fusion']['args'] and \
            not params['fusion']['args']['proj_first']:
            self.proj_first = False

        self.pre_processor = build_preprocessor(params['preprocess'],
                                                train)
        self.post_processor = post_processor.build_postprocessor(
            params['postprocess'],
            train)
        
        if "box_align" in params.keys():
            self.box_align = True
            self.stage1_result_path = params['box_align']['train_result'] if train else params['box_align']['val_result']
            self.stage1_result = read_json(self.stage1_result_path)
            self.box_align_args = params['box_align']['args']
        
        else:
            self.box_align = False

        if self.train:
            root_dir = params['root_dir']
        else:
            root_dir = params['validate_dir']

        print("Dataset dir:", root_dir)

        # v2x: read the pickle file
        with open(root_dir, 'rb') as f:
            dataset_infos = pickle.load(f)  # dataset_infos is a list 

        self.agent_start = eval(min([i[-1] for i in dataset_infos[0].keys() if i.startswith("lidar_pose")]))


        if 'train_params' not in params or\
                'max_cav' not in params['train_params']:
            self.max_cav = 5
        else:
            self.max_cav = params['train_params']['max_cav']

        self.keyframe_database = OrderedDict()
        self.len_record = len(dataset_infos)

        # loop over all keyframe.
        # data_info is one sample.
        np.random.seed(303)
        for (i, data_info) in enumerate(dataset_infos):
            self.keyframe_database.update({i: OrderedDict()})

            # at least 1 cav should show up
            cav_num = data_info['agent_num']
            assert cav_num > 0

            select_ego = 1

            # in one keyframe, loop all agent
            for cav_id in range(self.agent_start, self.agent_start + cav_num):

                self.keyframe_database[i][cav_id] = OrderedDict()
                self.keyframe_database[i][cav_id]['lidar'] = data_info[f'lidar_path_{cav_id}']  # maybe add camera in the future
                self.keyframe_database[i][cav_id]['params'] = OrderedDict()
                self.keyframe_database[i][cav_id]['params']['lidar_pose'] = tfm_to_pose(data_info[f"lidar_pose_{cav_id}"]) # tfm in data_info, turn to [x,y,z,roll,yaw,pitch]
                self.keyframe_database[i][cav_id]['params']['vehicles'] = data_info[f'labels_{cav_id}']['gt_boxes_global']
                self.keyframe_database[i][cav_id]['params']['object_ids'] = data_info[f'labels_{cav_id}']['gt_object_ids'].tolist()

                # randomly select ego is better
                if cav_id == select_ego:
                    self.keyframe_database[i][cav_id]['ego'] = True
                    
                else:
                    self.keyframe_database[i][cav_id]['ego'] = False

    ### rewrite __len__ ###
    def __len__(self):
        return self.len_record

    ### rewrite retrieve_base_data ###
    def retrieve_base_data(self, idx):
        """
        Given the index, return the corresponding data.

        Parameters
        ----------
        idx : int
            Index given by dataloader.

        Returns
        -------
        data : dict
            The dictionary contains loaded yaml params and lidar data for
            each cav.
            lidar_np: (N, 4)
        """
        # we loop the accumulated length list to see get the scenario index
        keyframe = self.keyframe_database[idx]

        data = OrderedDict()
        ego_idx = -1
        # load files for all CAVs
        for cav_id, cav_content in keyframe.items():
            data[cav_id] = OrderedDict()
            data[cav_id]['ego'] = cav_content['ego']
            if cav_content['ego']:
                ego_idx = cav_id
            data[cav_id]['params'] = cav_content['params'] # lidar_pose, vehicles(gt_boxes), object_id(token)

            # load the corresponding data into the dictionary
            nbr_dims = 4 # x,y,z,intensity
            scan = np.fromfile(cav_content['lidar'], dtype='float32')
            points = scan.reshape((-1, 5))[:, :nbr_dims] 
            data[cav_id]['lidar_np'] = points

        # in fact we want to put ego in the first place
        # we will exchange two agent.
        if ego_idx != 0:
            if self.agent_start == 0:  # v2x-sim1.0
                data[ego_idx], data[0] = data[0], data[ego_idx]
            elif self.agent_start == 1:  # v2x-sim2.0
                data[0] = data[ego_idx]
                data.move_to_end(0, last=False)
                data.pop(ego_idx)

        return data

    ### rewrite generate_object_center ###
    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Notice: it is a wrap of postprocessor function

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        return self.post_processor.generate_object_center_v2x(cav_contents,
                                                        reference_lidar_pose)

    def get_item_train(self, base_data_dict):
        """
            We only choose ego for v2x-sim 2.0.

            Because agent have been shuffled
        """
        base_data_dict = add_noise_data_dict(base_data_dict, self.params['noise_setting'])
        processed_data_dict = OrderedDict()
        selected_cav_base = base_data_dict[0] # 0 is ego
        selected_cav_processed = self.get_item_single_car(selected_cav_base)
        processed_data_dict.update({'ego': selected_cav_processed})

        return processed_data_dict


opencood/data_utils/datasets/late_fusion_dataset_v2x.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

from opencood.data_utils.post_processor.voxel_postprocessor import VoxelPostprocessor
from opencood.data_utils.post_processor.bev_postprocessor import BevPostprocessor
from opencood.data_utils.post_processor.ciassd_postprocessor import CiassdPostprocessor
from opencood.data_utils.post_processor.fpvrcnn_postprocessor import FpvrcnnPostprocessor
from opencood.data_utils.post_processor.simple_voxel_postprocessor import SimpleVoxelPostprocessor
from opencood.data_utils.post_processor.uncertainty_voxel_postprocessor import UncertaintyVoxelPostprocessor
from opencood.data_utils.post_processor.voxel_postprocessor_flow import VoxelPostprocessorFlow

__all__ = {
    'VoxelPostprocessor': VoxelPostprocessor,
    'BevPostprocessor': BevPostprocessor,
    'CiassdPostprocessor': CiassdPostprocessor,
    'FpvrcnnPostprocessor': FpvrcnnPostprocessor,
    'SimpleVoxelPostprocessor' : SimpleVoxelPostprocessor,
    'UncertaintyVoxelPostprocessor': UncertaintyVoxelPostprocessor,
    'VoxelPostprocessorFlow': VoxelPostprocessorFlow
}


def build_postprocessor(anchor_cfg, train):
    process_method_name = anchor_cfg['core_method']
    assert process_method_name in ['VoxelPostprocessor', 'BevPostprocessor', 'CiassdPostprocessor', 'FpvrcnnPostprocessor', 'SimpleVoxelPostprocessor', 'UncertaintyVoxelPostprocessor', 'VoxelPostprocessorFlow']
    anchor_generator = __all__[process_method_name](
        anchor_params=anchor_cfg,
        train=train
    )

    return anchor_generator
opencood/data_utils/post_processor/__init__.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Template for AnchorGenerator
"""

import numpy as np
import torch

from opencood.utils import box_utils
from opencood.utils import common_utils


class BasePostprocessor(object):
    """
    Template for Anchor generator.

    Parameters
    ----------
    anchor_params : dict
        The dictionary containing all anchor-related parameters.
    train : bool
        Indicate train or test mode.

    Attributes
    ----------
    bbx_dict : dictionary
        Contain all objects information across the cav, key: id, value: bbx
        coordinates (1, 7)
    """

    def __init__(self, anchor_params, train=True):
        self.params = anchor_params
        self.bbx_dict = {}
        self.train = train

    def generate_anchor_box(self):
        # needs to be overloaded
        return None

    def generate_label(self, *argv):
        return None

    def generate_gt_bbx_cav_curr(self, cav_content):
        """
        The base postprocessor will generate 3d groundtruth bounding box.

        For early and intermediate fusion,
            data_dict only contains ego.

        For late fusion,
            data_dcit contains all cavs, so we need transformation matrix.
            To generate gt boxes, transformation_matrix should be clean

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model for one cav.

        Returns
        -------
        gt_box3d_tensor : torch.Tensor
            The groundtruth bounding box tensor, shape (N, 8, 3).
        """
        gt_box3d_list = []
        # used to avoid repetitive bounding box
        object_id_list = []

        # for cav_id, cav_content in data_dict.items():
        # used to project gt bounding box to ego space
        # object_bbx_center is clean.
        # transformation_matrix = cav_content['transformation_matrix_clean']

        object_bbx_center = cav_content['object_bbx_center']
        object_bbx_mask = cav_content['object_bbx_mask']
        object_ids = cav_content['object_ids']
        object_bbx_center = object_bbx_center[object_bbx_mask == 1]

        # convert center to corner
        object_bbx_corner = \
            box_utils.boxes_to_corners_3d(object_bbx_center,
                                            self.params['order'])
        # projected_object_bbx_corner = \
        #     box_utils.project_box3d(object_bbx_corner.float(),
        #                             transformation_matrix)

        projected_object_bbx_corner = object_bbx_corner
        gt_box3d_list.append(projected_object_bbx_corner)
        # append the corresponding ids
        object_id_list += object_ids

        # gt bbx 3d
        gt_box3d_list = torch.vstack(gt_box3d_list)
        # some of the bbx may be repetitive, use the id list to filter
        gt_box3d_selected_indices = \
            [object_id_list.index(x) for x in set(object_id_list)]
        gt_box3d_tensor = gt_box3d_list[gt_box3d_selected_indices]

        # filter the gt_box to make sure all bbx are in the range
        mask = \
            box_utils.get_mask_for_boxes_within_range_torch(gt_box3d_tensor, self.params['gt_range'])
        gt_box3d_tensor = gt_box3d_tensor[mask, :, :]

        return gt_box3d_tensor
    
    def generate_single_bbx(self, cav_content):
        gt_box3d_list = []
        # used to avoid repetitive bounding box
        object_id_list = []

        object_bbx_center = cav_content['object_bbx_center']
        object_ids = cav_content['object_ids']

        # convert center to corner
        object_bbx_corner = \
            box_utils.boxes_to_corners_3d(object_bbx_center,
                                            self.params['order'])

        gt_box3d_list.append(projected_object_bbx_corner)
        # append the corresponding ids
        object_id_list += object_ids

        # gt bbx 3d
        gt_box3d_list = torch.vstack(gt_box3d_list)
        # some of the bbx may be repetitive, use the id list to filter
        gt_box3d_selected_indices = \
            [object_id_list.index(x) for x in set(object_id_list)]
        gt_box3d_tensor = gt_box3d_list[gt_box3d_selected_indices]

        # filter the gt_box to make sure all bbx are in the range
        mask = \
            box_utils.get_mask_for_boxes_within_range_torch(gt_box3d_tensor, self.params['gt_range'])
        gt_box3d_tensor = gt_box3d_tensor[mask, :, :]
        return 0

    def generate_gt_bbx(self, data_dict):
        """
        The base postprocessor will generate 3d groundtruth bounding box.

        For early and intermediate fusion,
            data_dict only contains ego.

        For late fusion,
            data_dcit contains all cavs, so we need transformation matrix.
            To generate gt boxes, transformation_matrix should be clean

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        Returns
        -------
        gt_box3d_tensor : torch.Tensor
            The groundtruth bounding box tensor, shape (N, 8, 3).
        """
        gt_box3d_list = []
        # used to avoid repetitive bounding box
        object_id_list = []

        for cav_id, cav_content in data_dict.items():
            # used to project gt bounding box to ego space
            # object_bbx_center is clean.
            transformation_matrix = cav_content['transformation_matrix_clean']

            object_bbx_center = cav_content['object_bbx_center']
            object_bbx_mask = cav_content['object_bbx_mask']
            object_ids = cav_content['object_ids']
            object_bbx_center = object_bbx_center[object_bbx_mask == 1]

            # convert center to corner
            object_bbx_corner = \
                box_utils.boxes_to_corners_3d(object_bbx_center,
                                              self.params['order'])
            projected_object_bbx_corner = \
                box_utils.project_box3d(object_bbx_corner.float(),
                                        transformation_matrix)
            gt_box3d_list.append(projected_object_bbx_corner)
            # append the corresponding ids
            object_id_list += object_ids

        # gt bbx 3d
        gt_box3d_list = torch.vstack(gt_box3d_list)
        # some of the bbx may be repetitive, use the id list to filter
        gt_box3d_selected_indices = \
            [object_id_list.index(x) for x in set(object_id_list)]
        gt_box3d_tensor = gt_box3d_list[gt_box3d_selected_indices]

        # filter the gt_box to make sure all bbx are in the range
        mask = \
            box_utils.get_mask_for_boxes_within_range_torch(gt_box3d_tensor, self.params['gt_range'])
        gt_box3d_tensor = gt_box3d_tensor[mask, :, :]

        return gt_box3d_tensor

    def generate_gt_bbx_by_iou(self, data_dict):
        """
        This function is only used by LateFusionDatasetDAIR

        LateFusionDatasetDAIR's label are from veh-side and inf-side
        and do not have unique object id.

        So we will filter the same object by IoU

        The base postprocessor will generate 3d groundtruth bounding box.

        For early and intermediate fusion,
            data_dict only contains ego.

        For late fusion,
            data_dcit contains all cavs, so we need transformation matrix.
            To generate gt boxes, transformation_matrix should be clean

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        Returns
        -------
        gt_box3d_tensor : torch.Tensor
            The groundtruth bounding box tensor, shape (N, 8, 3).
        """
        gt_box3d_list = []

        for cav_id, cav_content in data_dict.items():
            # used to project gt bounding box to ego space
            # object_bbx_center is clean.
            transformation_matrix = cav_content['transformation_matrix_clean']

            object_bbx_center = cav_content['object_bbx_center']
            object_bbx_mask = cav_content['object_bbx_mask']
            object_ids = cav_content['object_ids']
            object_bbx_center = object_bbx_center[object_bbx_mask == 1]

            # convert center to corner
            object_bbx_corner = \
                box_utils.boxes_to_corners_3d(object_bbx_center,
                                              self.params['order'])
            projected_object_bbx_corner = \
                box_utils.project_box3d(object_bbx_corner.float(),
                                        transformation_matrix)
            gt_box3d_list.append(projected_object_bbx_corner)

        # if only ego agent
        if len(data_dict) == 1:
            gt_box3d_tensor = torch.vstack(gt_box3d_list)
        # both veh-side and inf-side label
        else:
            veh_corners_np = gt_box3d_list[0].cpu().numpy()
            inf_corners_np = gt_box3d_list[1].cpu().numpy()
            inf_polygon_list = list(common_utils.convert_format(inf_corners_np))
            veh_polygon_list = list(common_utils.convert_format(veh_corners_np))
            iou_thresh = 0.05 


            gt_from_inf = []
            for i in range(len(inf_polygon_list)):
                inf_polygon = inf_polygon_list[i]
                ious = common_utils.compute_iou(inf_polygon, veh_polygon_list)
                if (ious > iou_thresh).any():
                    continue
                gt_from_inf.append(inf_corners_np[i])
            
            if len(gt_from_inf):
                gt_from_inf = np.stack(gt_from_inf)
                gt_box3d = np.vstack([veh_corners_np, gt_from_inf])
            else:
                gt_box3d = veh_corners_np

            gt_box3d_tensor = torch.from_numpy(gt_box3d).to(device=gt_box3d_list[0].device)

        # filter the gt_box to make sure all bbx are in the range
        mask = \
            box_utils.get_mask_for_boxes_within_range_torch(gt_box3d_tensor, self.params['gt_range'])
        gt_box3d_tensor = gt_box3d_tensor[mask, :, :]

        return gt_box3d_tensor

    def generate_object_center(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            in fact it is used in get_item_single_car, so the list length is 1

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """
        from opencood.data_utils.datasets import GT_RANGE_OPV2V

        tmp_object_dict = {}
        for cav_content in cav_contents:
            tmp_object_dict.update(cav_content['params']['vehicles'])

        output_dict = {}
        filter_range = self.params['anchor_args']['cav_lidar_range'] # if self.train else GT_RANGE_OPV2V

        box_utils.project_world_objects(tmp_object_dict,
                                        output_dict,
                                        reference_lidar_pose,
                                        filter_range,
                                        self.params['order'])

        object_np = np.zeros((self.params['max_num'], 7))
        mask = np.zeros(self.params['max_num'])
        object_ids = []

        for i, (object_id, object_bbx) in enumerate(output_dict.items()):
            object_np[i] = object_bbx[0, :]
            mask[i] = 1
            object_ids.append(object_id)

        return object_np, mask, object_ids


    def generate_object_center_v2x(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.
            In fact, only the ego vehile needs to generate object center

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """
        # from opencood.data_utils.datasets import GT_RANGE

        assert len(cav_contents) == 1
        
        """
        In old version, we only let ego agent return gt box.
        Other agent return empty.

        But it's not suitable for late fusion.
        Also, we should filter out boxes that don't have any lidar point hits.

        Thankfully, 'lidar_np' is in cav_contents[0].keys()
        """


        gt_boxes = cav_contents[0]['params']['vehicles'] # notice [N,10], 10 includes [x,y,z,dx,dy,dz,w,a,b,c]
        object_ids = cav_contents[0]['params']['object_ids']
        lidar_np = cav_contents[0]['lidar_np']
        
        tmp_object_dict = {"gt_boxes": gt_boxes, "object_ids":object_ids}

        output_dict = {}
        filter_range = self.params['anchor_args']['cav_lidar_range'] # v2x we don't use GT_RANGE.

        box_utils.project_world_objects_v2x(tmp_object_dict,
                                        output_dict,
                                        reference_lidar_pose,
                                        filter_range,
                                        self.params['order'],
                                        lidar_np=lidar_np)

        object_np = np.zeros((self.params['max_num'], 7))
        mask = np.zeros(self.params['max_num'])
        object_ids = []


        for i, (object_id, object_bbx) in enumerate(output_dict.items()):
            object_np[i] = object_bbx[0, :]
            mask[i] = 1
            object_ids.append(object_id)

        return object_np, mask, object_ids

    def generate_object_center_dairv2x(self,
                               cav_contents,
                               reference_lidar_pose):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        # tmp_object_dict = {}
        tmp_object_list = []
        cav_content = cav_contents[0]
        tmp_object_list = cav_content['params']['vehicles']

        output_dict = {}
        filter_range = self.params['anchor_args']['cav_lidar_range']


        box_utils.project_world_objects_dairv2x(tmp_object_list,
                                        output_dict,
                                        reference_lidar_pose,
                                        filter_range,
                                        self.params['order'])

        object_np = np.zeros((self.params['max_num'], 7))
        mask = np.zeros(self.params['max_num'])
        object_ids = []

        for i, (object_id, object_bbx) in enumerate(output_dict.items()):
            object_np[i] = object_bbx[0, :]
            mask[i] = 1
            object_ids.append(object_id)

        return object_np, mask, object_ids


    def generate_object_center_dairv2x_single(self,
                               cav_contents, suffix=""):
        """
        Retrieve all objects in a format of (n, 7), where 7 represents
        x, y, z, l, w, h, yaw or x, y, z, h, w, l, yaw.

        Parameters
        ----------
        cav_contents : list
            List of dictionary, save all cavs' information.

        reference_lidar_pose : list
            The final target lidar pose with length 6.

        Returns
        -------
        object_np : np.ndarray
            Shape is (max_num, 7).
        mask : np.ndarray
            Shape is (max_num,).
        object_ids : list
            Length is number of bbx in current sample.
        """

        # tmp_object_dict = {}
        tmp_object_list = []
        cav_content = cav_contents[0]
        tmp_object_list = cav_content['params'][f'vehicles{suffix}'] 

        output_dict = {}
        filter_range = self.params['anchor_args']['cav_lidar_range']


        box_utils.load_single_objects_dairv2x(tmp_object_list,
                                        output_dict,
                                        filter_range,
                                        self.params['order'])

        object_np = np.zeros((self.params['max_num'], 7))
        mask = np.zeros(self.params['max_num'])
        object_ids = []

        for i, (object_id, object_bbx) in enumerate(output_dict.items()):
            object_np[i] = object_bbx[0, :]
            mask[i] = 1
            object_ids.append(object_id)

        return object_np, mask, object_ids

opencood/data_utils/post_processor/base_postprocessor.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang haxiang@g.ucla.edu
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Anchor-free 2d Generator
"""

import numpy as np
import torch
import torch.nn.functional as F

from opencood.utils.transformation_utils import dist_to_continuous
from opencood.data_utils.post_processor.base_postprocessor \
    import BasePostprocessor
from opencood.utils import box_utils
from opencood.visualization import vis_utils


class BevPostprocessor(BasePostprocessor):
    def __init__(self, anchor_params, train):
        super(BevPostprocessor, self).__init__(anchor_params, train)
        # self.geometry_param = anchor_params["geometry"]
        self.geometry_param = anchor_params["geometry_param"]

        # TODO
        # Hard coded for now. Need to calculate for our own training dataset
        self.target_mean = np.array([0.008, 0.001, 0.202, 0.2, 0.43, 1.368])
        self.target_std_dev = np.array([0.866, 0.5, 0.954, 0.668, 0.09, 0.111])

    def generate_anchor_box(self):
        return None

    def generate_label(self, **kwargs):
        """
        Generate targets for training.

        Parameters
        ----------
        kwargs : list
            gt_box_center:(max_num, 7)

        Returns
        -------
        label_dict : dict
            Dictionary that contains all target related info.
        """
        assert self.params['order'] == 'lwh', \
            'Currently BEV only support lwh bbx order.'
        # (max_num, 7)
        gt_box_center = kwargs['gt_box_center']

        # (max_num)
        masks = kwargs['mask']

        # (n, 7)
        gt_box_center_valid = gt_box_center[masks == 1]
        # (n, 4, 3)
        bev_corners = box_utils.boxes_to_corners2d(gt_box_center_valid,
                                                   self.params['order'])

        n = gt_box_center_valid.shape[0]
        # (n, 4, 2)
        bev_corners = bev_corners[:, :, :2]
        yaw = gt_box_center_valid[:, -1]
        x, y = gt_box_center_valid[:, 0], gt_box_center_valid[:, 1]
        dx, dy = gt_box_center_valid[:, 3], gt_box_center_valid[:, 4]
        # (n, 6)
        reg_targets = np.column_stack([np.cos(yaw), np.sin(yaw), x, y, dx, dy])

        # target label map including classification and regression targets
        label_map = np.zeros(self.geometry_param["label_shape"])
        self.update_label_map(label_map, bev_corners, reg_targets)
        label_map = self.normalize_targets(label_map)
        label_dict = {
            # (7, label_shape[0], label_shape[1])
            "label_map": np.transpose(label_map, (2, 0, 1)).astype(np.float32),
            "bev_corners": bev_corners
        }
        return label_dict

    def update_label_map(self, label_map, bev_corners, reg_targets):
        """
        Update label_map based on bbx and regression targets.

        Parameters
        ----------
        label_map : numpy.array
            Targets array for classification and regression tasks with
            the shape of label_shape. (H, W, 7). 

        bev_corners : numpy.array
            The bbx corners in lidar frame with shape (n, 4, 2)

        reg_targets : numpy.array
            Array containing the regression targets information. It need to be
            further processed.

        """
        res = self.geometry_param["res"]
        downsample_rate = self.geometry_param["downsample_rate"]

        bev_origin = np.array([self.geometry_param["L1"],
                               self.geometry_param["W1"]]).reshape(1, -1)

        # discretized bbx corner representations -- (n, 4, 2)
        # bev_corners is real coordinate
        # bev_corners_dist is pixel coordinate
        bev_corners_dist = (bev_corners - bev_origin) / res / downsample_rate
        # generate the coordinates of m
        x = np.arange(self.geometry_param["label_shape"][0]) # H (x in lidar coordinate)
        y = np.arange(self.geometry_param["label_shape"][1]) # W (y in lidar coordinate)
        xx, yy = np.meshgrid(x, y)

        # (label_shape[0]*label_shape[1], 2)
        points = np.concatenate([xx.reshape(-1, 1), yy.reshape(-1, 1)],
                                axis=-1) # pixel 
        bev_origin_dist = bev_origin / res / downsample_rate

        # loop over each bbx, find the points within the bbx.
        for i in range(bev_corners.shape[0]):
            reg_target = reg_targets[i, :]

            # find discredited points in bbx
            points_in_box = \
                box_utils.get_points_in_rotated_box(points,
                                                    bev_corners_dist[i, ...])
            # convert points to continuous space
            points_continuous = dist_to_continuous(points_in_box,
                                                   bev_origin_dist,
                                                   res,
                                                   downsample_rate)
            actual_reg_target = np.repeat(reg_target.reshape(1, -1),
                                          points_continuous.shape[0],
                                          axis=0)
            # build learning targets
            actual_reg_target[:, 2:4] = \
                actual_reg_target[:, 2:4] - points_continuous
            actual_reg_target[:, 4:] = np.log(actual_reg_target[:, 4:])

            # update label map
            label_map[points_in_box[:, 0], points_in_box[:, 1], 0] = 1.0
            label_map[points_in_box[:, 0], points_in_box[:, 1], 1:] = \
                actual_reg_target

    def normalize_targets(self, label_map):
        """
        Normalize label_map

        Parameters
        ----------
        label_map : numpy.array
            Targets array for classification and regression tasks with the
            shape of label_shape.

        Returns
        -------
        label_map: numpy.array
            Nromalized label_map.

        """
        label_map[..., 1:] = \
            (label_map[..., 1:] - self.target_mean) / self.target_std_dev
        return label_map

    def denormalize_reg_map(self, reg_map):
        """
        Denormalize the regression map

        Parameters
        ----------
        reg_map : np.ndarray / torch.Tensor
            Regression output mapwith the shape of (label_shape[0],
            label_shape[1], 6).

        Returns
        -------
        reg_map : np.ndarray / torch.Tensor
            Denormalized regression map.

        """
        if isinstance(reg_map, np.ndarray):
            target_mean = self.target_mean
            target_std_dev = self.target_std_dev

        else:
            target_mean = \
                torch.from_numpy(self.target_mean).to(reg_map.device)
            target_std_dev = \
                torch.from_numpy(self.target_std_dev).to(reg_map.device)
        reg_map = reg_map * target_std_dev + target_mean
        return reg_map

    @staticmethod
    def collate_batch(label_batch_list):
        """
        Customized collate function for target label generation.

        Parameters
        ----------
        label_batch_list : list
            The list of dictionary  that contains all labels for several
            frames.

        Returns
        -------
        processed_batch : dict
            Reformatted labels in torch tensor.
        """
        label_map_list = [x["label_map"][np.newaxis, ...] for x in
                          label_batch_list]
        processed_batch = {
            # (batch_size, 7, label_shape[0], label_shape[1])
            "label_map": torch.from_numpy(np.concatenate(label_map_list,
                                                         axis=0)),
            "bev_corners": [torch.from_numpy(x["bev_corners"]) for x in
                            label_batch_list]
        }
        return processed_batch

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D bounding box.
        Step1: convert each cav's output to bounding box format
        Step2: project the bounding boxes to ego space.
        Step:3 NMS

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box2d_tensor : torch.Tensor
            The prediction bounding box tensor after NMS.

        gt_box2d_tensor : torch.Tensor
            The groundtruth bounding box tensor.
        """

        # the final bounding box list
        pred_box2d_list = []
        pred_score_list = []

        for cav_id, cav_content in data_dict.items():
            assert cav_id in 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)
                # 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)
                box3d = F.pad(corners2d, (0, 1))
                # (number of high confidence bbx, 4, 2)
                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)

        if len(pred_box2d_list):
            pred_box2ds = torch.cat(pred_box2d_list, dim=0)
            pred_scores = torch.cat(pred_score_list, dim=0)
        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, self.params['gt_range'])
        pred_box2ds = pred_box2ds[mask, :, :]
        pred_scores = pred_scores[mask]
        assert pred_scores.shape[0] == pred_box2ds.shape[0]
        return pred_box2ds, pred_scores

    def reg_map_to_bbx_corners(self, reg_map, mask):
        """
        Construct bbx from the regression output of the model.

        Parameters
        ----------
        reg_map : torch.Tensor
            Regression output of neural networks.

        mask : torch.Tensor
            Masks used to filter bbx.

        Returns
        -------
        corners : torch.Tensor
            Bbx output with shape (N, 4, 2).

        """

        assert len(reg_map.shape) == 3, \
            "only support shape of label_shape i.e. (*, *, 6)"
        device = reg_map.device

        cos_t, sin_t, x, y, log_dx, log_dy = \
            [tt.squeeze(-1) for tt in torch.chunk(reg_map, 6, dim=-1)]
        yaw = torch.atan2(sin_t, cos_t)
        dx, dy = log_dx.exp(), log_dy.exp()

        grid_size = self.geometry_param["res"] * \
                    self.geometry_param["downsample_rate"]
        grid_x = torch.arange(self.geometry_param["L1"],
                              self.geometry_param["L2"],
                              grid_size, dtype=torch.float32, device=device)
        grid_y = torch.arange(self.geometry_param["W1"],
                              self.geometry_param["W2"],
                              grid_size,
                              dtype=torch.float32,
                              device=device)

        xx, yy = torch.meshgrid([grid_x, grid_y])
        center_x = xx + x
        center_y = yy + y

        bbx2d = torch.stack([center_x, center_y, dx, dy, yaw], dim=-1)
        bbx2d = bbx2d[mask, :]
        corners = box_utils.boxes2d_to_corners2d(bbx2d)

        return corners

    def post_process_debug(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D bounding box for debug purpose.
        Step1: convert each cav's output to bounding box format
        Step2: project the bounding boxes to ego space.
        Step:3 NMS

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box2d_tensor : torch.Tensor
            The prediction bounding box tensor after NMS.
        gt_box2d_tensor : torch.Tensor
            The groundtruth bounding box tensor.
        """
        # the final bounding box list
        pred_box2d_list = []
        pred_score_list = []

        # the transformation matrix to ego space
        transformation_matrix = data_dict['transformation_matrix']

        # classification probability -- (label_shape[0], label_shape[1])
        prob = output_dict['cls'].squeeze(0).squeeze(0)
        prob = torch.sigmoid(prob)

        # regression map -- (label_shape[0], label_shape[1], 6)
        reg_map = output_dict['reg'].squeeze(0).permute(1, 2, 0)
        reg_map = self.denormalize_reg_map(reg_map)

        threshold = 0.5
        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)
            # 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)
            box3d = F.pad(corners2d, (0, 1))

            # (number of high confidence bbx, 4, 2)
            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)

        pred_box2ds = torch.cat(pred_box2d_list, dim=0)
        pred_scores = torch.cat(pred_score_list, dim=0)

        keep_index = box_utils.nms_rotated(pred_box2ds,
                                           pred_scores,
                                           self.params['nms_thresh'])
        pred_box2ds = pred_box2ds[keep_index]

        # filter out the prediction out of the range.
        mask = box_utils.get_mask_for_boxes_within_range_torch(pred_box2ds, self.params['gt_range'])
        pred_box2ds = pred_box2ds[mask, :, :]
        return pred_box2ds

    @staticmethod
    def visualize(pred_box_tensor, gt_tensor, pcd, show_vis, save_path,
                  dataset=None):
        """
        Visualize the BEV 2D prediction, ground truth with point cloud together.

        Parameters
        ----------
        pred_box_tensor : torch.Tensor
            (N, 8, 3) prediction.

        gt_tensor : torch.Tensor
            (N, 8, 3) groundtruth bbx

        pcd : torch.Tensor
            PointCloud, (N, 4).

        show_vis : bool
            Whether to show visualization.

        save_path : str
            Save the visualization results to given path.

        dataset : BaseDataset
            opencood dataset object.
        """
        assert dataset is not None, "dataset argument can't be None"
        vis_utils.visualize_single_sample_output_bev(pred_box_tensor,
                                                     gt_tensor,
                                                     pcd,
                                                     dataset,
                                                     show_vis,
                                                     save_path)

opencood/data_utils/post_processor/bev_postprocessor.py
MediaBrain-SJTU/CoBEVFlow
"""
3D Anchor Generator for Voxel
"""
import math
import sys

import numpy as np
import torch
import torch.nn.functional as F

from opencood.data_utils.post_processor.voxel_postprocessor \
    import VoxelPostprocessor
from opencood.utils import box_utils


class CiassdPostprocessor(VoxelPostprocessor):
    def __init__(self, anchor_params, train):
        super(CiassdPostprocessor, self).__init__(anchor_params, train)
        self.train = train
        self.anchor_num = self.params['anchor_args']['num']

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.
        Step1: convert each cav's output to bounding box format
        Step2: project the bounding boxes to ego space.
        Step:3 NMS

        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box3d_tensor : torch.Tensor
            The prediction bounding box tensor after NMS.
        gt_box3d_tensor : torch.Tensor
            The groundtruth bounding box tensor.
        """
        # the final bounding box list
        global batch_num_box_count
        pred_box3d_original_list = []
        pred_box3d_list = []
        pred_box2d_list = []

        for cav_id, cav_content in data_dict.items():
            assert cav_id in output_dict
            # the transformation matrix to ego space
            if 'transformation_matrix' in cav_content:
                transformation_matrix = cav_content['transformation_matrix']
            else:
                transformation_matrix = torch.from_numpy(np.identity(4)).float().\
                    to(cav_content['anchor_box'].device)

            # (H, W, anchor_num, 7)
            anchor_box = cav_content['anchor_box']

            # prediction result
            preds_dict = output_dict[cav_id]['preds_dict_stage1']

            # preds
            prob = preds_dict['cls_preds']
            prob = torch.sigmoid(prob.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)
            dir = preds_dict['dir_cls_preds'].permute(0, 2, 3, 1).contiguous().reshape(1, -1, 2) # [N, H*W*2, 2]

            # 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']) # [N, H, W, 2]
            batch_num_box_count = [int(m.sum()) for m in mask]
            mask = mask.view(1, -1) # [1,N*H*W*2]
            mask_reg = mask.unsqueeze(2).repeat(1, 1, 7)

            # during validation/testing, the batch size should be 1
            if not self.train:
                assert batch_box3d.shape[0] == 1

            boxes3d = torch.masked_select(batch_box3d.view(-1, 7), mask_reg[0]).view(-1, 7)
            scores = torch.masked_select(prob.view(-1), mask[0])

            dir_labels = torch.max(dir, dim=-1)[1]  # indices. shape [N, H*W*2].  value 0 or 1
            dir_labels = dir_labels[mask] # sum(mask==1)
            # top_labels = torch.zeros([scores.shape[0]], dtype=torch.long).cuda()
            if scores.shape[0] != 0:
                iou = (iou + 1) * 0.5
                scores = scores * torch.pow(iou.masked_select(mask), 4)
                # correct_direction
                top_labels = (boxes3d[..., -1] > 0) ^ (dir_labels.byte() == 1)
                boxes3d[..., -1] += torch.where(top_labels, torch.tensor(np.pi).type_as(boxes3d),
                                                  torch.tensor(0.0).type_as(boxes3d))
                pred_box3d_original_list.append(boxes3d.detach())

            # 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'])
                # (N, 8, 3)
                projected_boxes3d = box_utils.project_box3d(boxes3d_corner, transformation_matrix)
                # convert 3d bbx to 2d, (N,4)
                projected_boxes2d = box_utils.corner_to_standup_box_torch(projected_boxes3d)
                # (N, 5)
                boxes2d_score = torch.cat((projected_boxes2d, scores.unsqueeze(1)), dim=1)

                pred_box2d_list.append(boxes2d_score)
                pred_box3d_list.append(projected_boxes3d)

        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)
        pred_box3d_original = torch.vstack(pred_box3d_original_list)

        if not self.train:
            # 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
        else:
            cur_idx = 0
            batch_pred_boxes3d = []
            batch_scores = []
            for n in batch_num_box_count:
                cur_boxes = pred_box3d_tensor[cur_idx:cur_idx+n]
                cur_scores = scores[cur_idx:cur_idx+n]
                # nms
                keep_index = box_utils.nms_rotated(cur_boxes,
                                                   cur_scores,
                                                   self.params['nms_thresh']
                                                   )
                cur_boxes = pred_box3d_original[cur_idx:cur_idx+n] # [:, [0, 1, 2, 5, 4, 3, 6]] # hwl -> lwh
                batch_pred_boxes3d.append(cur_boxes[keep_index])
                batch_scores.append(cur_scores[keep_index])
                cur_idx += n

            return batch_pred_boxes3d, batch_scores

opencood/data_utils/post_processor/ciassd_postprocessor.py
MediaBrain-SJTU/CoBEVFlow
"""
3D Anchor Generator for Voxel
"""
import numpy as np
import torch

from opencood.data_utils.post_processor.voxel_postprocessor \
    import VoxelPostprocessor
from opencood.utils import box_utils
from opencood.utils import common_utils


class FpvrcnnPostprocessor(VoxelPostprocessor):
    def __init__(self, anchor_params, train):
        super(FpvrcnnPostprocessor, self).__init__(anchor_params, train)

    def post_process(self, data_dict, output_dict, stage1=False):
        if stage1:
            return self.post_process_stage1(data_dict, output_dict)
        else:
            return self.post_process_stage2(data_dict)

    def post_process_stage1(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.
        No NMS


        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box3d_tensor : torch.Tensor
            The prediction bounding box tensor after NMS.
        gt_box3d_tensor : torch.Tensor
            The groundtruth bounding box tensor.
        """
        # the final bounding box list
        pred_box3d_original_list = []
        pred_box3d_list = []
        pred_box2d_list = []

        for cav_id, cav_content in data_dict.items():
            assert cav_id in output_dict
            # the transformation matrix to ego space
            if 'transformation_matrix' in cav_content:
                transformation_matrix = cav_content['transformation_matrix']
            else:
                transformation_matrix = torch.from_numpy(np.identity(4)).float().\
                    to(cav_content['anchor_box'].device)


            # (H, W, anchor_num, 7)
            anchor_box = cav_content['anchor_box']

            # prediction result
            preds_dict = output_dict[cav_id]['preds_dict_stage1']

            # preds
            prob = preds_dict['cls_preds']
            prob = torch.sigmoid(prob.permute(0, 2, 3, 1).contiguous())
            reg = preds_dict['box_preds']  # .permute(0, 2, 3, 1).contiguous()
            iou = torch.sigmoid(preds_dict['iou_preds'].permute(0, 2, 3, 1).contiguous()).reshape(1, -1)
            dir = preds_dict['dir_cls_preds'].permute(0, 2, 3, 1).contiguous().reshape(1, -1, 2)

            # convert regression map back to bounding box
            # (N, W*L*anchor_num, 7)
            # reg: (N, anchor_num*7, H, W), permute inside delta_to_boxes3d

            batch_box3d = self.delta_to_boxes3d(reg, anchor_box) # hwl
            mask = torch.gt(prob, self.params['target_args']['score_threshold'])
            batch_num_box_count = [int(m.sum()) for m in mask]
            mask = mask.view(1, -1)
            mask_reg = mask.unsqueeze(2).repeat(1, 1, 7)

            # during validation/testing, the batch size should be 1
            if not self.train:
                assert batch_box3d.shape[0] == 1

            boxes3d = torch.masked_select(batch_box3d.view(-1, 7), mask_reg[0]).view(-1, 7) # hwl. right
            scores = torch.masked_select(prob.view(-1), mask[0])


            dir_labels = torch.max(dir, dim=-1)[1]
            dir_labels = dir_labels[mask]
            # top_labels = torch.zeros([scores.shape[0]], dtype=torch.long).cuda()
            # process per agent
            if scores.shape[0] != 0:
                iou = torch.clamp(iou, min=0.0, max=1.0)
                iou = (iou + 1) * 0.5
                scores = scores * torch.pow(iou.masked_select(mask), 4)
                # correct_direction
                top_labels = (boxes3d[..., -1] > 0) ^ (dir_labels.byte() == 1)
                boxes3d[..., -1] += torch.where(top_labels, torch.tensor(np.pi).type_as(boxes3d),
                                                  torch.tensor(0.0).type_as(boxes3d))

                # filter invalid boxes
                keep_idx = torch.logical_and((boxes3d[:, 3:6] > 1).all(dim=1), (boxes3d[:, 3:6] < 10).all(dim=1))
                idx_start = 0
                count = []
                for i, n in enumerate(batch_num_box_count):
                    count.append(int(keep_idx[idx_start:idx_start+n].sum()))
                batch_num_box_count = count
                boxes3d = boxes3d[keep_idx] # hwl
                scores = scores[keep_idx]


                # if the number of boxes is too huge, this would consume a lot of memory in the second stage
                # therefore, randomly select some boxes if the box number is too big at the beginning of the training
                if len(boxes3d) > 300:
                    keep_idx = torch.multinomial(scores, 300)
                    idx_start = 0
                    count = []
                    for i, n in enumerate(batch_num_box_count):
                        count.append(int(torch.logical_and(keep_idx>=idx_start, keep_idx<idx_start + n).sum()))
                    batch_num_box_count = count
                    boxes3d = boxes3d[keep_idx] 
                    scores = scores[keep_idx]

                pred_box3d_original_list.append(boxes3d.detach()) # hwl
            
            # 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'])
                # (N, 8, 3)
                projected_boxes3d = box_utils.project_box3d(boxes3d_corner, transformation_matrix)
                # convert 3d bbx to 2d, (N,4)
                projected_boxes2d = box_utils.corner_to_standup_box_torch(projected_boxes3d)
                # (N, 5)
                boxes2d_score = torch.cat((projected_boxes2d, scores.unsqueeze(1)), dim=1)

                pred_box2d_list.append(boxes2d_score)
                pred_box3d_list.append(projected_boxes3d)


        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)

        pred_box3d_original = torch.vstack(pred_box3d_original_list)

        cur_idx = 0
        batch_pred_boxes3d = []
        batch_scores = []
        for n in batch_num_box_count:
            cur_boxes = pred_box3d_tensor[cur_idx:cur_idx+n]
            cur_scores = scores[cur_idx:cur_idx+n]
            # nms
            keep_index = box_utils.nms_rotated(cur_boxes,
                                               cur_scores,
                                               self.params['nms_thresh']
                                               )
            cur_boxes = pred_box3d_original[cur_idx:cur_idx+n][:, [0, 1, 2, 5, 4, 3, 6]] 
            batch_pred_boxes3d.append(cur_boxes[keep_index])
            batch_scores.append(cur_scores[keep_index])
            cur_idx += n
            # print("stage1 cur_boxes:",cur_boxes) # hwl

        return batch_pred_boxes3d, batch_scores

    def post_process_stage2(self, data_dict):
        from opencood.pcdet_utils.iou3d_nms.iou3d_nms_utils import nms_gpu
        if 'fpvrcnn_out' not in data_dict['ego'].keys():
            return None, None
        output_dict = data_dict['ego']['fpvrcnn_out']
        label_dict = data_dict['ego']['rcnn_label_dict']
        rcnn_cls = output_dict['rcnn_cls'].sigmoid().view(-1)
        rcnn_iou = output_dict['rcnn_iou'].view(-1)
        rcnn_score = rcnn_cls * rcnn_iou**4
        rcnn_reg = output_dict['rcnn_reg'].view(-1, 7)
        rois_anchor = label_dict['rois_anchor'] # rois_anchor is hwl order
        rois = label_dict['rois']
        roi_center = rois[:, 0:3]
        roi_ry = rois[:, 6] % (2 * np.pi)
        boxes_local = box_utils.box_decode(rcnn_reg, rois_anchor)

        # print("label_dict['gt_of_rois']:",label_dict['gt_of_rois'])      # hwl order
        # print("rois_anchor:", label_dict['rois_anchor'])                 # hwl order

        # boxes_local = rcnn_reg + rois_anchor
        detections = common_utils.rotate_points_along_z(
            points=boxes_local.view(-1, 1, boxes_local.shape[-1]), angle=roi_ry.view(-1)
        ).view(-1, boxes_local.shape[-1])
        detections[:, :3] = detections[:, :3] + roi_center
        detections[:, 6] = detections[:, 6] + roi_ry
        mask = rcnn_score >= 0.1
        detections = detections[mask]
        scores = rcnn_score[mask]

        # gt_boxes = label_dict['gt_of_rois_src'][mask]
        mask = nms_gpu(detections, scores, thresh=0.01)[0]
        boxes3d = detections[mask][:, [0, 1, 2, 5, 4, 3, 6]] # lwh -> hwl

        projected_boxes3d = None
        if len(boxes3d) != 0:
            # (N, 8, 3)
            boxes3d_corner = \
                box_utils.boxes_to_corners_3d(boxes3d,
                                              order=self.params['order'])
            # (N, 8, 3)
            projected_boxes3d = \
                box_utils.project_box3d(boxes3d_corner,
                                        data_dict['ego']['transformation_matrix'])

        ## Added by Yifan Lu, filter box outside of GT range
        if projected_boxes3d is None:
            return None, None
        scores = scores[mask]
        cav_range = self.params['gt_range']
        mask = box_utils.get_mask_for_boxes_within_range_torch(projected_boxes3d, cav_range)
        projected_boxes3d = projected_boxes3d[mask]


        return projected_boxes3d, scores[mask]

opencood/data_utils/post_processor/fpvrcnn_postprocessor.py
MediaBrain-SJTU/CoBEVFlow
# Author: Yifan Lu

"""
3D Anchor Generator for Voxel
"""
import numpy as np
import torch
import torch.nn.functional as F
import math

from opencood.data_utils.post_processor.voxel_postprocessor \
    import VoxelPostprocessor
from opencood.utils import box_utils
from opencood.utils import common_utils


class SimpleVoxelPostprocessor():
    def __init__(self, params, train=True):
        self.params = params
        self.anchor_num = params['anchor_args']['num']

    def generate_anchor_box(self):
        W = self.params['anchor_args']['W']
        H = self.params['anchor_args']['H']

        l = self.params['anchor_args']['l']
        w = self.params['anchor_args']['w']
        h = self.params['anchor_args']['h']
        r = self.params['anchor_args']['r']

        assert self.anchor_num == len(r)
        r = [math.radians(ele) for ele in r]

        vh = self.params['anchor_args']['vh'] # voxel_size
        vw = self.params['anchor_args']['vw']

        xrange = [self.params['anchor_args']['cav_lidar_range'][0],
                  self.params['anchor_args']['cav_lidar_range'][3]]
        yrange = [self.params['anchor_args']['cav_lidar_range'][1],
                  self.params['anchor_args']['cav_lidar_range'][4]]

        if 'feature_stride' in self.params['anchor_args']:
            feature_stride = self.params['anchor_args']['feature_stride']
        else:
            feature_stride = 2

        x = np.linspace(xrange[0] + vw, xrange[1] - vw, W // feature_stride)
        y = np.linspace(yrange[0] + vh, yrange[1] - vh, H // feature_stride)

        cx, cy = np.meshgrid(x, y)
        cx = np.tile(cx[..., np.newaxis], self.anchor_num) # center
        cy = np.tile(cy[..., np.newaxis], self.anchor_num)
        cz = np.ones_like(cx) * -1.0

        w = np.ones_like(cx) * w
        l = np.ones_like(cx) * l
        h = np.ones_like(cx) * h

        r_ = np.ones_like(cx)
        for i in range(self.anchor_num):
            r_[..., i] = r[i]

        if self.params['order'] == 'hwl': # pointpillar
            anchors = np.stack([cx, cy, cz, h, w, l, r_], axis=-1) # (50, 176, 2, 7)

        elif self.params['order'] == 'lhw':
            anchors = np.stack([cx, cy, cz, l, h, w, r_], axis=-1)
        else:
            sys.exit('Unknown bbx order.')

        return anchors

    def post_process_stage1(self, stage1_output_dict, anchor_box):
        """
        Process the outputs of the model to 2D/3D bounding box.
        Step1: convert each cav's output to bounding box format
        Step2: (project the bounding boxes to ego space). 
            For stage1, we do not project to ego. Latter we project to world coord.
        Step:3 NMS

        Parameters
        ---------
            stage1_output_dict :
                psm: torch.Size([12, 2, 100, 352])
                rm: torch.Size([12, 14, 100, 352])
            anchor_box: torch.Size()

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box3d_tensor : torch.Tensor
            The prediction bounding box tensor after NMS.
        gt_box3d_tensor : torch.Tensor
            The groundtruth bounding box tensor.
        """
        # the final bounding box list
        psm = stage1_output_dict['psm']
        rm = stage1_output_dict['rm']
        
        prob = psm  # [N, anchor_num , H, W]
        prob = F.sigmoid(prob.permute(0, 2, 3, 1).contiguous())  # [N, H, W, anchor_num]

        # regression map
        reg = rm  # [N, anchor_num * 7, H, W]

        # convert regression map back to bounding box
        batch_box3d = self.delta_to_boxes3d(reg, anchor_box)
        mask = torch.gt(prob, self.params['target_args']['score_threshold'])
        batch_num_box_count = [int(m.sum()) for m in mask]
        mask = mask.view(1, -1)
        mask_reg = mask.unsqueeze(2).repeat(1, 1, 7)


        boxes3d = torch.masked_select(batch_box3d.view(-1, 7), mask_reg[0]).view(-1, 7) 
        scores = torch.masked_select(prob.view(-1), mask[0])

        # convert output to bounding box
        if len(boxes3d) != 0:
            # save origianl format box. [N, 7]
            pred_box3d_original = boxes3d.detach()
            # (N, 8, 3)
            boxes3d_corner = box_utils.boxes_to_corners_3d(boxes3d, order=self.params['order'])
            # (N, 8, 3)
            pred_corners_tensor = boxes3d_corner  # box_utils.project_box3d(boxes3d_corner, transformation_matrix)
            # convert 3d bbx to 2d, (N,4)
            projected_boxes2d = box_utils.corner_to_standup_box_torch(pred_corners_tensor)
            # (N, 5)
            pred_box2d_score_tensor = torch.cat((projected_boxes2d, scores.unsqueeze(1)), dim=1)
            scores = pred_box2d_score_tensor[:, -1]

        else:
             return None, None, None

        # divide boxes to each cav

        cur_idx = 0
        batch_pred_corners3d = [] # [[N1, 8, 3], [N2, 8, 3], ...]
        batch_pred_boxes3d = [] # [[N1, 7], [N2, 7], ...]
        batch_scores = []
        for n in batch_num_box_count:
            cur_corners = pred_corners_tensor[cur_idx: cur_idx+n]
            cur_boxes = pred_box3d_original[cur_idx: cur_idx+n]
            cur_scores = scores[cur_idx:cur_idx+n]
            # nms
            keep_index = box_utils.nms_rotated(cur_corners,
                                               cur_scores,
                                               self.params['nms_thresh']
                                               )
            batch_pred_corners3d.append(cur_corners[keep_index])
            batch_pred_boxes3d.append(cur_boxes[keep_index])
            batch_scores.append(cur_scores[keep_index])
            cur_idx += n

        return batch_pred_corners3d, batch_pred_boxes3d, batch_scores



    @staticmethod
    def delta_to_boxes3d(deltas, anchors):
        """
        Convert the output delta to 3d bbx.

        Parameters
        ----------
        deltas : torch.Tensor
            (N, W, L, 14)?? should be (N, 14, H, W)
        anchors : torch.Tensor
            (W, L, 2, 7) -> xyzhwlr

        Returns
        -------
        box3d : torch.Tensor
            (N, W*L*2, 7)
        """
        # batch size
        N = deltas.shape[0]
        deltas = deltas.permute(0, 2, 3, 1).contiguous().view(N, -1, 7)
        boxes3d = torch.zeros_like(deltas)

        if deltas.is_cuda:
            anchors = anchors.cuda()
            boxes3d = boxes3d.cuda()

        # (W*L*2, 7)
        anchors_reshaped = anchors.view(-1, 7).float()
        # the diagonal of the anchor 2d box, (W*L*2)
        anchors_d = torch.sqrt(
            anchors_reshaped[:, 4] ** 2 + anchors_reshaped[:, 5] ** 2)
        anchors_d = anchors_d.repeat(N, 2, 1).transpose(1, 2)
        anchors_reshaped = anchors_reshaped.repeat(N, 1, 1)

        # Inv-normalize to get xyz
        boxes3d[..., [0, 1]] = torch.mul(deltas[..., [0, 1]], anchors_d) + \
                               anchors_reshaped[..., [0, 1]]
        boxes3d[..., [2]] = torch.mul(deltas[..., [2]],
                                      anchors_reshaped[..., [3]]) + \
                            anchors_reshaped[..., [2]]
        # hwl
        boxes3d[..., [3, 4, 5]] = torch.exp(
            deltas[..., [3, 4, 5]]) * anchors_reshaped[..., [3, 4, 5]]
        # yaw angle
        boxes3d[..., 6] = deltas[..., 6] + anchors_reshaped[..., 6]

        return boxes3d
opencood/data_utils/post_processor/simple_voxel_postprocessor.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu


"""
3D Anchor Generator for Voxel
"""
import math
import sys

import numpy as np
import torch
from torch.nn.functional import sigmoid
import torch.nn.functional as F

from opencood.data_utils.post_processor.base_postprocessor \
    import BasePostprocessor
from opencood.utils import box_utils
from opencood.utils.box_overlaps import bbox_overlaps
from opencood.visualization import vis_utils
from opencood.utils.common_utils import limit_period


class UncertaintyVoxelPostprocessor(BasePostprocessor):
    def __init__(self, anchor_params, train):
        super(UncertaintyVoxelPostprocessor, self).__init__(anchor_params, train)
        self.anchor_num = self.params['anchor_args']['num']

    def generate_anchor_box(self):
        W = self.params['anchor_args']['W']
        H = self.params['anchor_args']['H']

        l = self.params['anchor_args']['l']
        w = self.params['anchor_args']['w']
        h = self.params['anchor_args']['h']
        r = self.params['anchor_args']['r']

        assert self.anchor_num == len(r)
        r = [math.radians(ele) for ele in r]

        vh = self.params['anchor_args']['vh'] # voxel_size
        vw = self.params['anchor_args']['vw']

        xrange = [self.params['anchor_args']['cav_lidar_range'][0],
                  self.params['anchor_args']['cav_lidar_range'][3]]
        yrange = [self.params['anchor_args']['cav_lidar_range'][1],
                  self.params['anchor_args']['cav_lidar_range'][4]]

        if 'feature_stride' in self.params['anchor_args']:
            feature_stride = self.params['anchor_args']['feature_stride']
        else:
            feature_stride = 2

        x = np.linspace(xrange[0] + vw, xrange[1] - vw, W // feature_stride)
        y = np.linspace(yrange[0] + vh, yrange[1] - vh, H // feature_stride)

        cx, cy = np.meshgrid(x, y)
        cx = np.tile(cx[..., np.newaxis], self.anchor_num) # center
        cy = np.tile(cy[..., np.newaxis], self.anchor_num)
        cz = np.ones_like(cx) * -1.0

        w = np.ones_like(cx) * w
        l = np.ones_like(cx) * l
        h = np.ones_like(cx) * h

        r_ = np.ones_like(cx)
        for i in range(self.anchor_num):
            r_[..., i] = r[i]

        if self.params['order'] == 'hwl': # pointpillar
            anchors = np.stack([cx, cy, cz, h, w, l, r_], axis=-1) # (50, 176, 2, 7)

        elif self.params['order'] == 'lhw':
            anchors = np.stack([cx, cy, cz, l, h, w, r_], axis=-1)
        else:
            sys.exit('Unknown bbx order.')

        return anchors

    def generate_label(self, **kwargs):
        """
        Generate targets for training.

        Parameters
        ----------
        argv : list
            gt_box_center:(max_num, 7), anchor:(H, W, anchor_num, 7)

        Returns
        -------
        label_dict : dict
            Dictionary that contains all target related info.
        """
        assert self.params['order'] == 'hwl', 'Currently Voxel only support' \
                                              'hwl bbx order.'
        # (max_num, 7)
        gt_box_center = kwargs['gt_box_center']
        # (H, W, anchor_num, 7)
        anchors = kwargs['anchors']
        # (max_num)
        masks = kwargs['mask']

        # (H, W)
        feature_map_shape = anchors.shape[:2]

        # (H*W*anchor_num, 7)
        anchors = anchors.reshape(-1, 7)
        # normalization factor, (H * W * anchor_num)
        anchors_d = np.sqrt(anchors[:, 4] ** 2 + anchors[:, 5] ** 2)

        # (H, W, 2)
        pos_equal_one = np.zeros((*feature_map_shape, self.anchor_num))
        neg_equal_one = np.zeros((*feature_map_shape, self.anchor_num))
        # (H, W, self.anchor_num * 7)
        targets = np.zeros((*feature_map_shape, self.anchor_num * 7))

        # (n, 7)
        gt_box_center_valid = gt_box_center[masks == 1]
        # (n, 8, 3)
        gt_box_corner_valid = \
            box_utils.boxes_to_corners_3d(gt_box_center_valid,
                                          self.params['order'])
        # (H*W*anchor_num, 8, 3)
        anchors_corner = \
            box_utils.boxes_to_corners_3d(anchors,
                                          order=self.params['order'])
        # (H*W*anchor_num, 4)
        anchors_standup_2d = \
            box_utils.corner2d_to_standup_box(anchors_corner)
        # (n, 4)
        gt_standup_2d = \
            box_utils.corner2d_to_standup_box(gt_box_corner_valid)

        # (H*W*anchor_n)
        iou = bbox_overlaps(
            np.ascontiguousarray(anchors_standup_2d).astype(np.float32),
            np.ascontiguousarray(gt_standup_2d).astype(np.float32),
        )

        # the anchor boxes has the largest iou across
        # shape: (n)
        id_highest = np.argmax(iou.T, axis=1)
        # [0, 1, 2, ..., n-1]
        id_highest_gt = np.arange(iou.T.shape[0])
        # make sure all highest iou is larger than 0
        mask = iou.T[id_highest_gt, id_highest] > 0
        id_highest, id_highest_gt = id_highest[mask], id_highest_gt[mask]


        # find anchors iou > params['pos_iou']
        id_pos, id_pos_gt = \
            np.where(iou >
                     self.params['target_args']['pos_threshold'])
        #  find anchors iou  params['neg_iou']
        id_neg = np.where(np.sum(iou <
                                 self.params['target_args']['neg_threshold'],
                                 axis=1) == iou.shape[1])[0]
        id_pos = np.concatenate([id_pos, id_highest])
        id_pos_gt = np.concatenate([id_pos_gt, id_highest_gt])
        id_pos, index = np.unique(id_pos, return_index=True)
        id_pos_gt = id_pos_gt[index]
        id_neg.sort()

        # cal the target and set the equal one
        index_x, index_y, index_z = np.unravel_index(
            id_pos, (*feature_map_shape, self.anchor_num))
        pos_equal_one[index_x, index_y, index_z] = 1

        # calculate the targets
        targets[index_x, index_y, np.array(index_z) * 7] = \
            (gt_box_center[id_pos_gt, 0] - anchors[id_pos, 0]) / anchors_d[
                id_pos]
        targets[index_x, index_y, np.array(index_z) * 7 + 1] = \
            (gt_box_center[id_pos_gt, 1] - anchors[id_pos, 1]) / anchors_d[
                id_pos]
        targets[index_x, index_y, np.array(index_z) * 7 + 2] = \
            (gt_box_center[id_pos_gt, 2] - anchors[id_pos, 2]) / anchors[
                id_pos, 3]
        targets[index_x, index_y, np.array(index_z) * 7 + 3] = np.log(
            gt_box_center[id_pos_gt, 3] / anchors[id_pos, 3])
        targets[index_x, index_y, np.array(index_z) * 7 + 4] = np.log(
            gt_box_center[id_pos_gt, 4] / anchors[id_pos, 4])
        targets[index_x, index_y, np.array(index_z) * 7 + 5] = np.log(
            gt_box_center[id_pos_gt, 5] / anchors[id_pos, 5])
        targets[index_x, index_y, np.array(index_z) * 7 + 6] = (
                gt_box_center[id_pos_gt, 6] - anchors[id_pos, 6])


        index_x, index_y, index_z = np.unravel_index(
            id_neg, (*feature_map_shape, self.anchor_num))
        neg_equal_one[index_x, index_y, index_z] = 1

        # to avoid a box be pos/neg in the same time
        index_x, index_y, index_z = np.unravel_index(
            id_highest, (*feature_map_shape, self.anchor_num))
        neg_equal_one[index_x, index_y, index_z] = 0


        label_dict = {'pos_equal_one': pos_equal_one,
                      'neg_equal_one': neg_equal_one,
                      'targets': targets}

        return label_dict

    @staticmethod
    def collate_batch(label_batch_list):
        """
        Customized collate function for target label generation.

        Parameters
        ----------
        label_batch_list : list
            The list of dictionary  that contains all labels for several
            frames.

        Returns
        -------
        target_batch : dict
            Reformatted labels in torch tensor.
        """
        pos_equal_one = []
        neg_equal_one = []
        targets = []

        for i in range(len(label_batch_list)):
            pos_equal_one.append(label_batch_list[i]['pos_equal_one'])
            neg_equal_one.append(label_batch_list[i]['neg_equal_one'])
            targets.append(label_batch_list[i]['targets'])

        pos_equal_one = \
            torch.from_numpy(np.array(pos_equal_one))
        neg_equal_one = \
            torch.from_numpy(np.array(neg_equal_one))
        targets = \
            torch.from_numpy(np.array(targets))

        return {'targets': targets,
                'pos_equal_one': pos_equal_one,
                'neg_equal_one': neg_equal_one}
    
    def post_process_stage1(self, stage1_output_dict, anchor_box):
        """
        Process the outputs of the model to 2D/3D bounding box.
        Step1: convert each cav's output to bounding box format
        Step2: (project the bounding boxes to ego space). 
            For stage1, we do not project to ego. Latter we project to world coord.
        Step:3 NMS

        Parameters
        ---------
            stage1_output_dict:
                psm: torch.Size([12, 2, 100, 352])
                rm: torch.Size([12, 14, 100, 352])
                sm: torch.Size([12, 4, 100, 352])
            anchor_box: torch.Size()

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box3d_tensor : torch.Tensor
            The prediction bounding box tensor after NMS.
        gt_box3d_tensor : torch.Tensor
            The groundtruth bounding box tensor.
        """
        psm = stage1_output_dict['psm']
        rm = stage1_output_dict['rm']
        sm = stage1_output_dict['sm']

        dm = None
        if 'dm' in stage1_output_dict.keys():
            dm = stage1_output_dict['dm']

        # the final bounding box list
        uncertainty_dim = sm.shape[1] // psm.shape[1]

        prob = psm  # [N, anchor_num , H, W]
        prob = F.sigmoid(prob.permute(0, 2, 3, 1).contiguous())  # [N, H, W, anchor_num]

        # regression map
        reg = rm  # [N, anchor_num * 7, H, W]

        # sm
        sm = sm.permute(0,2,3,1).contiguous() #[N, H, W, anchor_num * 2]


        # convert regression map back to bounding box
        batch_box3d = self.delta_to_boxes3d(reg, anchor_box)  # (N, W*L*2, 7)
        mask = torch.gt(prob, self.params['target_args']['score_threshold'])
        batch_num_box_count = [int(m.sum()) for m in mask]
        mask = mask.view(1, -1)
        mask_reg = mask.unsqueeze(2).repeat(1, 1, 7)
        mask_sm = mask.unsqueeze(2).repeat(1, 1, uncertainty_dim)
        


        boxes3d = torch.masked_select(batch_box3d.view(-1, 7), mask_reg[0]).view(-1, 7) 
        uncertainty = torch.masked_select(sm.view(-1,uncertainty_dim), mask_sm[0]).view(-1,uncertainty_dim) # [N*H*W*#anchor_num, 2] -> [num_select, 2]
        scores = torch.masked_select(prob.view(-1), mask[0])
        if dm is not None and len(boxes3d) != 0:
            dir_offset = self.params['dir_args']['dir_offset']
            num_bins = self.params['dir_args']['num_bins']

            dir_cls_preds = dm.permute(0, 2, 3, 1).contiguous().reshape(1, -1, num_bins) # [1, N*H*W*2, 2]
            dir_cls_preds = dir_cls_preds[mask]
            # if rot_gt > 0, then the label is 1, then the regression target is [0, 1]
            dir_labels = torch.max(dir_cls_preds, dim=-1)[1]  # indices. shape [1, N*H*W*2].  value 0 or 1. If value is 1, then rot_gt > 0
            
            period = (2 * np.pi / num_bins) # pi
            dir_rot = limit_period(
                boxes3d[..., 6] - dir_offset, 0, period
            ) # 限制在0到pi之间
            boxes3d[..., 6] = dir_rot + dir_offset + period * dir_labels.to(boxes3d.dtype) # 转化0.25pi到2.5pi
            boxes3d[..., 6] = limit_period(boxes3d[..., 6], 0.5, 2 * np.pi) # limit to [-pi, pi]

        # convert output to bounding box
        if len(boxes3d) != 0:
            # save origianl format box. [N, 7]
            pred_box3d_original = boxes3d.detach()
            # (N, 8, 3)
            boxes3d_corner = box_utils.boxes_to_corners_3d(boxes3d, order=self.params['order'])
            # (N, 8, 3)
            pred_corners_tensor = boxes3d_corner  # box_utils.project_box3d(boxes3d_corner, transformation_matrix)
            # convert 3d bbx to 2d, (N,4)
            projected_boxes2d = box_utils.corner_to_standup_box_torch(pred_corners_tensor)
            # (N, 5)
            pred_box2d_score_tensor = torch.cat((projected_boxes2d, scores.unsqueeze(1)), dim=1)
            scores = pred_box2d_score_tensor[:, -1]

        else:
             return None, None, None

        # divide boxes to each cav

        cur_idx = 0
        batch_pred_corners3d = [] # [[N1, 8, 3], [N2, 8, 3], ...]
        batch_pred_boxes3d = [] # [[N1, 7], [N2, 7], ...]
        batch_uncertainty = [] # [[N1, 2], [N2, 2], ...]
        batch_scores = []
        for n in batch_num_box_count:
            cur_corners = pred_corners_tensor[cur_idx: cur_idx+n]
            cur_boxes = pred_box3d_original[cur_idx: cur_idx+n]
            cur_scores = scores[cur_idx:cur_idx+n]
            cur_uncertainty = uncertainty[cur_idx: cur_idx+n]
            # nms
            keep_index = box_utils.nms_rotated(cur_corners,
                                               cur_scores,
                                               self.params['nms_thresh']
                                               )
            batch_pred_corners3d.append(cur_corners[keep_index])
            batch_pred_boxes3d.append(cur_boxes[keep_index])
            batch_scores.append(cur_scores[keep_index])
            batch_uncertainty.append(cur_uncertainty[keep_index])
            cur_idx += n

        return batch_pred_corners3d, batch_pred_boxes3d, batch_uncertainty


    def post_process(self, data_dict, output_dict, return_uncertainty=False):
        """
        Process the outputs of the model to 2D/3D bounding box.
        Step1: convert each cav's output to bounding box format
        Step2: project the bounding boxes to ego space.
        Step:3 NMS

        For early and intermediate fusion,
            data_dict only contains ego.

        For late fusion,
            data_dcit contains all cavs, so we need transformation matrix.


        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box3d_tensor : torch.Tensor
            The prediction bounding box tensor after NMS.
        gt_box3d_tensor : torch.Tensor
            The groundtruth bounding box tensor.
        """
        # the final bounding box list
        pred_box3d_list = []
        pred_box2d_list = []
        uncertainty_list = []
        for cav_id, cav_content in data_dict.items():
            if cav_id not in output_dict:
                continue
            # the transformation matrix to ego space
            transformation_matrix = cav_content['transformation_matrix'] # no clean

            # (H, W, anchor_num, 7)
            anchor_box = cav_content['anchor_box']

            # classification probability
            uncertainty_dim = output_dict[cav_id]['sm'].shape[1] // output_dict[cav_id]['psm'].shape[1]
            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']

            # uncertainty map
            sm = output_dict[cav_id]['sm'].permute(0, 2, 3, 1).contiguous()
            sm = sm.view(sm.shape[0], -1, uncertainty_dim) # [N, H*W*#anchor_num, 2]

            # convert regression map back to bounding box
            batch_box3d = self.delta_to_boxes3d(reg, anchor_box) # (N, H*W*#anchor_num, 7)
            mask = \
                torch.gt(prob, self.params['target_args']['score_threshold'])
            mask = mask.view(1, -1)
            mask_reg = mask.unsqueeze(2).repeat(1, 1, 7)
            mask_sm = mask.unsqueeze(2).repeat(1, 1, uncertainty_dim)
            
            # 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])
            uncertainty = torch.masked_select(sm[0], mask_sm[0]).view(-1, uncertainty_dim)


            # adding dir classifier
            if 'dm' in output_dict[cav_id].keys() and len(boxes3d) != 0:
                dir_offset = self.params['dir_args']['dir_offset']
                num_bins = self.params['dir_args']['num_bins']


                dm  = output_dict[cav_id]['dm'] # [N, H, W, 4]
                dir_cls_preds = dm.permute(0, 2, 3, 1).contiguous().reshape(1, -1, num_bins) # [1, N*H*W*2, 2]
                dir_cls_preds = dir_cls_preds[mask]
                # if rot_gt > 0, then the label is 1, then the regression target is [0, 1]
                dir_labels = torch.max(dir_cls_preds, dim=-1)[1]  # indices. shape [1, N*H*W*2].  value 0 or 1. If value is 1, then rot_gt > 0
                
                period = (2 * np.pi / num_bins) # pi
                dir_rot = limit_period(
                    boxes3d[..., 6] - dir_offset, 0, period
                ) # 限制在0到pi之间
                boxes3d[..., 6] = dir_rot + dir_offset + period * dir_labels.to(dir_cls_preds.dtype) # 转化0.25pi到2.5pi
                boxes3d[..., 6] = limit_period(boxes3d[..., 6], 0.5, 2 * np.pi) # limit to [-pi, pi]


            # 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'])
                # (N, 8, 3)
                projected_boxes3d = \
                    box_utils.project_box3d(boxes3d_corner,
                                            transformation_matrix)
                # convert 3d bbx to 2d, (N,4)
                projected_boxes2d = \
                    box_utils.corner_to_standup_box_torch(projected_boxes3d)
                # (N, 5)
                boxes2d_score = \
                    torch.cat((projected_boxes2d, scores.unsqueeze(1)), dim=1)

                pred_box2d_list.append(boxes2d_score)
                pred_box3d_list.append(projected_boxes3d)
                uncertainty_list.append(uncertainty)


        if len(pred_box2d_list) ==0 or len(pred_box3d_list) == 0:
            if return_uncertainty:
                return None, None, None
            return None, None
        # shape: (N, 5)
        pred_box2d_list = torch.vstack(pred_box2d_list)
        uncertainty_list = torch.vstack(uncertainty_list)
        uncertainty = uncertainty_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]
        uncertainty = uncertainty[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]
        uncertainty = uncertainty[keep_index]

        # filter out the prediction out of the range.
        mask = \
            box_utils.get_mask_for_boxes_within_range_torch(pred_box3d_tensor, self.params['gt_range'])
        pred_box3d_tensor = pred_box3d_tensor[mask, :, :]
        scores = scores[mask]
        uncertainty = uncertainty[mask]

        assert scores.shape[0] == pred_box3d_tensor.shape[0]
        
        if return_uncertainty:
            return pred_box3d_tensor, scores, uncertainty
        return pred_box3d_tensor, scores

    @staticmethod
    def delta_to_boxes3d(deltas, anchors):
        """
        Convert the output delta to 3d bbx.

        Parameters
        ----------
        deltas : torch.Tensor
            (N, W, L, 14)?? should be (N, 14, H, W)
        anchors : torch.Tensor
            (W, L, 2, 7) -> xyzhwlr

        Returns
        -------
        box3d : torch.Tensor
            (N, W*L*2, 7)
        """
        # batch size
        N = deltas.shape[0]
        deltas = deltas.permute(0, 2, 3, 1).contiguous().view(N, -1, 7)
        boxes3d = torch.zeros_like(deltas)

        if deltas.is_cuda:
            anchors = anchors.cuda()
            boxes3d = boxes3d.cuda()

        # (W*L*2, 7)
        anchors_reshaped = anchors.view(-1, 7).float()
        # the diagonal of the anchor 2d box, (W*L*2)
        anchors_d = torch.sqrt(
            anchors_reshaped[:, 4] ** 2 + anchors_reshaped[:, 5] ** 2)
        anchors_d = anchors_d.repeat(N, 2, 1).transpose(1, 2)
        anchors_reshaped = anchors_reshaped.repeat(N, 1, 1)

        # Inv-normalize to get xyz
        boxes3d[..., [0, 1]] = torch.mul(deltas[..., [0, 1]], anchors_d) + \
                               anchors_reshaped[..., [0, 1]]
        boxes3d[..., [2]] = torch.mul(deltas[..., [2]],
                                      anchors_reshaped[..., [3]]) + \
                            anchors_reshaped[..., [2]]
        # hwl
        boxes3d[..., [3, 4, 5]] = torch.exp(
            deltas[..., [3, 4, 5]]) * anchors_reshaped[..., [3, 4, 5]]
        # yaw angle
        boxes3d[..., 6] = deltas[..., 6] + anchors_reshaped[..., 6]

        return boxes3d

    @staticmethod
    def visualize(pred_box_tensor, gt_tensor, pcd, show_vis, save_path, dataset=None):
        """
        Visualize the prediction, ground truth with point cloud together.

        Parameters
        ----------
        pred_box_tensor : torch.Tensor
            (N, 8, 3) prediction.

        gt_tensor : torch.Tensor
            (N, 8, 3) groundtruth bbx

        pcd : torch.Tensor
            PointCloud, (N, 4).

        show_vis : bool
            Whether to show visualization.

        save_path : str
            Save the visualization results to given path.

        dataset : BaseDataset
            opencood dataset object.

        """
        vis_utils.visualize_single_sample_output_gt(pred_box_tensor,
                                                    gt_tensor,
                                                    pcd,
                                                    show_vis,
                                                    save_path)

opencood/data_utils/post_processor/uncertainty_voxel_postprocessor.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, OpenPCDet
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
3D Anchor Generator for Voxel
"""
import math
import sys

import numpy as np
import torch
from torch.nn.functional import sigmoid
import torch.nn.functional as F

from opencood.data_utils.post_processor.base_postprocessor \
    import BasePostprocessor
from opencood.utils import box_utils
from opencood.utils.box_overlaps import bbox_overlaps
from opencood.visualization import vis_utils
from opencood.utils.common_utils import limit_period

from collections import OrderedDict


class VoxelPostprocessor(BasePostprocessor):
    def __init__(self, anchor_params, train):
        super(VoxelPostprocessor, self).__init__(anchor_params, train)
        self.anchor_num = self.params['anchor_args']['num']

    def generate_anchor_box(self):
        W = self.params['anchor_args']['W']
        H = self.params['anchor_args']['H']

        l = self.params['anchor_args']['l']
        w = self.params['anchor_args']['w']
        h = self.params['anchor_args']['h']
        r = self.params['anchor_args']['r']

        assert self.anchor_num == len(r)
        r = [math.radians(ele) for ele in r]

        vh = self.params['anchor_args']['vh'] # voxel_size
        vw = self.params['anchor_args']['vw']

        xrange = [self.params['anchor_args']['cav_lidar_range'][0],
                  self.params['anchor_args']['cav_lidar_range'][3]]
        yrange = [self.params['anchor_args']['cav_lidar_range'][1],
                  self.params['anchor_args']['cav_lidar_range'][4]]

        if 'feature_stride' in self.params['anchor_args']:
            feature_stride = self.params['anchor_args']['feature_stride']
        else:
            feature_stride = 2


        x = np.linspace(xrange[0] + vw, xrange[1] - vw, W // feature_stride)
        y = np.linspace(yrange[0] + vh, yrange[1] - vh, H // feature_stride)


        cx, cy = np.meshgrid(x, y)
        cx = np.tile(cx[..., np.newaxis], self.anchor_num) # center
        cy = np.tile(cy[..., np.newaxis], self.anchor_num)
        cz = np.ones_like(cx) * -1.0

        w = np.ones_like(cx) * w
        l = np.ones_like(cx) * l
        h = np.ones_like(cx) * h

        r_ = np.ones_like(cx)
        for i in range(self.anchor_num):
            r_[..., i] = r[i]

        if self.params['order'] == 'hwl': # pointpillar
            anchors = np.stack([cx, cy, cz, h, w, l, r_], axis=-1) # (50, 176, 2, 7)

        elif self.params['order'] == 'lhw':
            anchors = np.stack([cx, cy, cz, l, h, w, r_], axis=-1)
        else:
            sys.exit('Unknown bbx order.')

        return anchors

    def generate_label(self, **kwargs):
        """
        Generate targets for training.

        Parameters
        ----------
        argv : list
            gt_box_center:(max_num, 7), anchor:(H, W, anchor_num, 7)

        Returns
        -------
        label_dict : dict
            Dictionary that contains all target related info.
        """
        assert self.params['order'] == 'hwl', 'Currently Voxel only support' \
                                              'hwl bbx order.'
        # (max_num, 7)
        gt_box_center = kwargs['gt_box_center']
        # (H, W, anchor_num, 7)
        anchors = kwargs['anchors']
        # (max_num)
        masks = kwargs['mask']

        # (H, W)
        feature_map_shape = anchors.shape[:2]

        # (H*W*anchor_num, 7)
        anchors = anchors.reshape(-1, 7)
        # normalization factor, (H * W * anchor_num)
        anchors_d = np.sqrt(anchors[:, 4] ** 2 + anchors[:, 5] ** 2)

        # (H, W, 2)
        pos_equal_one = np.zeros((*feature_map_shape, self.anchor_num))
        neg_equal_one = np.zeros((*feature_map_shape, self.anchor_num))
        # (H, W, self.anchor_num * 7)
        targets = np.zeros((*feature_map_shape, self.anchor_num * 7))

        # (n, 7)
        gt_box_center_valid = gt_box_center[masks == 1]
        # (n, 8, 3)
        gt_box_corner_valid = \
            box_utils.boxes_to_corners_3d(gt_box_center_valid,
                                          self.params['order'])
        # (H*W*anchor_num, 8, 3)
        anchors_corner = \
            box_utils.boxes_to_corners_3d(anchors,
                                          order=self.params['order'])
        # (H*W*anchor_num, 4)
        anchors_standup_2d = \
            box_utils.corner2d_to_standup_box(anchors_corner)
        # (n, 4)
        gt_standup_2d = \
            box_utils.corner2d_to_standup_box(gt_box_corner_valid)

        # (H*W*anchor_n)
        iou = bbox_overlaps(
            np.ascontiguousarray(anchors_standup_2d).astype(np.float32),
            np.ascontiguousarray(gt_standup_2d).astype(np.float32),
        )

        # the anchor boxes has the largest iou across
        # shape: (n)
        id_highest = np.argmax(iou.T, axis=1)
        # [0, 1, 2, ..., n-1]
        id_highest_gt = np.arange(iou.T.shape[0])
        # make sure all highest iou is larger than 0
        mask = iou.T[id_highest_gt, id_highest] > 0
        id_highest, id_highest_gt = id_highest[mask], id_highest_gt[mask]


        # find anchors iou > params['pos_iou']
        id_pos, id_pos_gt = \
            np.where(iou >
                     self.params['target_args']['pos_threshold'])
        #  find anchors iou  params['neg_iou']
        id_neg = np.where(np.sum(iou <
                                 self.params['target_args']['neg_threshold'],
                                 axis=1) == iou.shape[1])[0]
        id_pos = np.concatenate([id_pos, id_highest])
        id_pos_gt = np.concatenate([id_pos_gt, id_highest_gt])
        id_pos, index = np.unique(id_pos, return_index=True)
        id_pos_gt = id_pos_gt[index]
        id_neg.sort()

        # cal the target and set the equal one
        index_x, index_y, index_z = np.unravel_index(
            id_pos, (*feature_map_shape, self.anchor_num))
        pos_equal_one[index_x, index_y, index_z] = 1

        # calculate the targets
        targets[index_x, index_y, np.array(index_z) * 7] = \
            (gt_box_center[id_pos_gt, 0] - anchors[id_pos, 0]) / anchors_d[
                id_pos]
        targets[index_x, index_y, np.array(index_z) * 7 + 1] = \
            (gt_box_center[id_pos_gt, 1] - anchors[id_pos, 1]) / anchors_d[
                id_pos]
        targets[index_x, index_y, np.array(index_z) * 7 + 2] = \
            (gt_box_center[id_pos_gt, 2] - anchors[id_pos, 2]) / anchors[
                id_pos, 3]
        targets[index_x, index_y, np.array(index_z) * 7 + 3] = np.log(
            gt_box_center[id_pos_gt, 3] / anchors[id_pos, 3])
        targets[index_x, index_y, np.array(index_z) * 7 + 4] = np.log(
            gt_box_center[id_pos_gt, 4] / anchors[id_pos, 4])
        targets[index_x, index_y, np.array(index_z) * 7 + 5] = np.log(
            gt_box_center[id_pos_gt, 5] / anchors[id_pos, 5])
        targets[index_x, index_y, np.array(index_z) * 7 + 6] = (
                gt_box_center[id_pos_gt, 6] - anchors[id_pos, 6])

        index_x, index_y, index_z = np.unravel_index(
            id_neg, (*feature_map_shape, self.anchor_num))
        neg_equal_one[index_x, index_y, index_z] = 1

        # to avoid a box be pos/neg in the same time
        index_x, index_y, index_z = np.unravel_index(
            id_highest, (*feature_map_shape, self.anchor_num))
        neg_equal_one[index_x, index_y, index_z] = 0


        label_dict = {'pos_equal_one': pos_equal_one,
                      'neg_equal_one': neg_equal_one,
                      'targets': targets}

        return label_dict
    
    @staticmethod
    def merge_label_to_dict(label_k_list):
        """
        Customized collate function for target label generation.

        Parameters
        ----------
        label_batch_list : list
            The list of dictionary  that contains all labels for several
            frames.

        Returns
        -------
        target_batch : dict
            Reformatted labels in torch tensor.
        """
        pos_equal_one = []
        neg_equal_one = []
        targets = []

        for i in range(len(label_k_list)):
            pos_equal_one.append(label_k_list[i]['pos_equal_one'])
            neg_equal_one.append(label_k_list[i]['neg_equal_one'])
            targets.append(label_k_list[i]['targets'])

        return {'targets': targets,
                'pos_equal_one': pos_equal_one,
                'neg_equal_one': neg_equal_one}

    @staticmethod
    def collate_batch(label_batch_list):
        """
        Customized collate function for target label generation.

        Parameters
        ----------
        label_batch_list : list
            The list of dictionary  that contains all labels for several
            frames.

        Returns
        -------
        target_batch : dict
            Reformatted labels in torch tensor.
        """
        pos_equal_one = []
        neg_equal_one = []
        targets = []

        for i in range(len(label_batch_list)):
            pos_equal_one.append(label_batch_list[i]['pos_equal_one'])
            neg_equal_one.append(label_batch_list[i]['neg_equal_one'])
            targets.append(label_batch_list[i]['targets'])

        pos_equal_one = \
            torch.from_numpy(np.array(pos_equal_one))
        neg_equal_one = \
            torch.from_numpy(np.array(neg_equal_one))
        targets = \
            torch.from_numpy(np.array(targets))

        return {'targets': targets,
                'pos_equal_one': pos_equal_one,
                'neg_equal_one': neg_equal_one}

    """
    Added by Sizhewei @ 2023-04-15
    Generate box results on each cav's past_0 time
    """
    def bandwidth_filter(self, input, num_box):
        if input['scores'].shape[0] <= num_box:
            return input
        
        topk_idx = torch.argsort(input['scores'], descending=True)[:num_box]
        output = {}
        output['scores'] = input['scores'][topk_idx]
        output['pred_box_3dcorner_tensor']  = input['pred_box_3dcorner_tensor'][topk_idx]
        output['pred_box_center_tensor'] = input['pred_box_center_tensor'][topk_idx]

        return output

    def single_post_process(self, m_single, trans_mat_pastk_2_past0, past_time_diff, anchor_box, num_sweeps=2, num_roi_thres=-1):
        """
        Process the outputs of the model to 2D/3D bounding box.
        Step1: convert each cav's output to bounding box format
        Step2: project the bounding boxes to ego space.
        Step:3 NMS

        For early and intermediate fusion,
            data_dict only contains ego.

        For late fusion,
            data_dcit contains all cavs, so we need transformation matrix.


        Parameters
        ----------
        data_dict : dict
            
        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        {
            'past_k_time_diff' : (k, )
            [0] : {
                pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. (n, 8, 3)
                pred_box_center_tensor : (n, 7)
                scores: (n, )
            },
            ...
            [k-1] : { ... }
        }
        """
        self.k = num_sweeps
        self.num_roi_thres = num_roi_thres
        
        psm_single = m_single['psm_single']
        rm_single = m_single['rm_single']

        use_dir_flag = False
        if 'dm_single' in m_single.keys():
            use_dir_flag = True
            dm_single = m_single['dm_single']

        box_results = OrderedDict()
        # for cav_id, cav_content in data_dict.items():

        box_results.update({
            'past_k_time_diff': past_time_diff
        })
        # the transformation matrix to ego space
        # transformation_matrix = cav_content['transformation_matrix'] # no clean
        transformation_matrix = trans_mat_pastk_2_past0.to(torch.float32)

        # classification probability
        prob = psm_single  # k,2,100,352
        prob = torch.sigmoid(prob.permute(0, 2, 3, 1))
        prob = prob.reshape(self.k, -1) # (k, HxWx2)

        # regression map
        reg = rm_single # (k,14,100,352)

        # convert regression map back to bounding box
        batch_box3d = self.delta_to_boxes3d(reg, anchor_box) # (k, HxWx2, 7)
        mask = \
            torch.gt(prob, self.params['target_args']['score_threshold']) 
        mask = mask.view(self.k, -1) # (k, HxWx2)
        mask_reg = mask.unsqueeze(2).repeat(1, 1, 7) # (k, HxWx2, 7)

        # during validation/testing, the batch size should be 1
        # assert batch_box3d.shape[0] == 1

        if use_dir_flag:
            dir_offset = self.params['dir_args']['dir_offset']
            num_bins = self.params['dir_args']['num_bins']
            dm = dm_single # [N, H, W, 4]
        
        for i in range(self.k):  # 遍历所有的时间戳
            box_results[i] = OrderedDict()
            unit_trans_mat = transformation_matrix[i]
            boxes3d = torch.masked_select(batch_box3d[i],
                                        mask_reg[i]).view(-1, 7)
            scores = torch.masked_select(prob[i], mask[i])

            ########### adding dir classifier
            if use_dir_flag and len(boxes3d)!=0:
                dir_cls_preds = dm[i:i+1].permute(0, 2, 3, 1).contiguous().reshape(1, -1, num_bins) # [1, N*H*W*2, 2]
                dir_cls_preds = dir_cls_preds[mask[i:i+1]]
                # if rot_gt > 0, then the label is 1, then the regression target is [0, 1]
                dir_labels = torch.max(dir_cls_preds, dim=-1)[1]  # indices. shape [1, N*H*W*2].  value 0 or 1. If value is 1, then rot_gt > 0p
                
                period = (2 * np.pi / num_bins) # pi
                dir_rot = limit_period(
                    boxes3d[..., 6] - dir_offset, 0, period
                ) # 限制在0到pi之间
                boxes3d[..., 6] = dir_rot + dir_offset + period * dir_labels.to(dir_cls_preds.dtype) # 转化0.25pi到2.5pi
                boxes3d[..., 6] = limit_period(boxes3d[..., 6], 0.5, 2 * np.pi) # limit to [-pi, pi]
            ###########################################

            # 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'])
                # STEP 2
                # (N, 8, 3)
                projected_boxes3d = \
                    box_utils.project_box3d(boxes3d_corner,
                                            unit_trans_mat)

                # backup_projected_boxes3d = projected_boxes3d.clone()
                # backup_scores = scores.clone()

                # convert 3d bbx to 2d, (N,4)
                projected_boxes2d = \
                    box_utils.corner_to_standup_box_torch(projected_boxes3d)
                # (N, 5)
                boxes2d_score = \
                    torch.cat((projected_boxes2d, scores.unsqueeze(1)), dim=1)

                # pred_box2d_list.append(boxes2d_score)
                # pred_box3d_list.append(projected_boxes3d)

                # scores = boxes2d_score[:, -1]

                # remove large bbx
                keep_index_1 = box_utils.remove_large_pred_bbx(projected_boxes3d)
                keep_index_2 = box_utils.remove_bbx_abnormal_z(projected_boxes3d)
                keep_index = torch.logical_and(keep_index_1, keep_index_2)

                projected_boxes3d = projected_boxes3d[keep_index]
                scores = scores[keep_index]

                # STEP3
                # nms
                keep_index = box_utils.nms_rotated(projected_boxes3d,
                                                scores,
                                                self.params['nms_thresh']
                                                )
                pred_box3d_tensor = projected_boxes3d[keep_index]
                # select cooresponding score
                scores = scores[keep_index]

                # filter out the prediction out of the range.
                range_mask = \
                    box_utils.get_mask_for_boxes_within_range_torch(pred_box3d_tensor, self.params['gt_range'])
                pred_box_3dcorner_tensor = pred_box3d_tensor[range_mask, :, :]
                scores = scores[range_mask]
                pred_box_center_tensor = box_utils.corner_to_center_torch(pred_box_3dcorner_tensor, self.params['order'])

                assert scores.shape[0] == pred_box_3dcorner_tensor.shape[0]

                # if scores.shape[0] == 0: # 没有bbx剩下，那么保存之前剩的框
                #     pred_box_3dcorner_tensor = backup_projected_boxes3d
                #     scores = backup_scores
                #     pred_box_center_tensor = box_utils.corner_to_center_torch(pred_box_3dcorner_tensor, self.params['order'])

                try:
                    self.num_roi_thres = self.num_roi_thres
                except KeyError:
                    print('Note! There is no num_roi_thres in the config file.')
                    self.num_roi_thres = -1
                if self.num_roi_thres != -1:
                    original_results = {
                        'pred_box_3dcorner_tensor': pred_box_3dcorner_tensor, 
                        'pred_box_center_tensor': pred_box_center_tensor,
                        'scores': scores
                    }
                    sorted_box_results = self.bandwidth_filter(original_results, self.num_roi_thres)
                    box_results[i].update({
                        'pred_box_3dcorner_tensor': sorted_box_results['pred_box_3dcorner_tensor'], 
                        'pred_box_center_tensor': sorted_box_results['pred_box_center_tensor'],
                        'scores': sorted_box_results['scores']
                    })
                else: 
                    box_results[i].update({
                        'pred_box_3dcorner_tensor': pred_box_3dcorner_tensor, 
                        'pred_box_center_tensor': pred_box_center_tensor,
                        'scores': scores
                    })

            else:
                box_results[i].update({
                    'pred_box_3dcorner_tensor': torch.Tensor(0, 8, 3).to(scores.device), 
                    'pred_box_center_tensor': torch.Tensor(0, 7).to(scores.device),
                    'scores': torch.Tensor(0).to(scores.device)
                })

        return box_results

    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.
        Step1: convert each cav's output to bounding box format
        Step2: project the bounding boxes to ego space.
        Step:3 NMS

        For early and intermediate fusion,
            data_dict only contains ego.

        For late fusion,
            data_dcit contains all cavs, so we need transformation matrix.


        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box3d_tensor : torch.Tensor
            The prediction bounding box tensor after NMS.
        gt_box3d_tensor : torch.Tensor
            The groundtruth bounding box tensor.
        """
        # the final bounding box list
        pred_box3d_list = []
        pred_box2d_list = []
        for cav_id, cav_content in data_dict.items():
            assert cav_id in output_dict
            # the transformation matrix to ego space
            transformation_matrix = cav_content['transformation_matrix'] # no clean

            # (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
            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])

            # adding dir classifier
            if 'dm' in output_dict[cav_id].keys() and len(boxes3d) !=0:
                dir_offset = self.params['dir_args']['dir_offset']
                num_bins = self.params['dir_args']['num_bins']
                dm  = output_dict[cav_id]['dm'] # [N, H, W, 4]
                dir_cls_preds = dm.permute(0, 2, 3, 1).contiguous().reshape(1, -1, num_bins) # [1, N*H*W*2, 2]
                dir_cls_preds = dir_cls_preds[mask]
                # if rot_gt > 0, then the label is 1, then the regression target is [0, 1]
                dir_labels = torch.max(dir_cls_preds, dim=-1)[1]  # indices. shape [1, N*H*W*2].  value 0 or 1. If value is 1, then rot_gt > 0
                
                period = (2 * np.pi / num_bins) # pi
                dir_rot = limit_period(
                    boxes3d[..., 6] - dir_offset, 0, period
                ) # 限制在0到pi之间
                boxes3d[..., 6] = dir_rot + dir_offset + period * dir_labels.to(dir_cls_preds.dtype) # 转化0.25pi到2.5pi
                boxes3d[..., 6] = limit_period(boxes3d[..., 6], 0.5, 2 * np.pi) # limit to [-pi, pi]

            # 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'])
                
                # STEP 2
                # (N, 8, 3)
                projected_boxes3d = \
                    box_utils.project_box3d(boxes3d_corner,
                                            transformation_matrix)
                # convert 3d bbx to 2d, (N,4)
                projected_boxes2d = \
                    box_utils.corner_to_standup_box_torch(projected_boxes3d)
                # (N, 5)
                boxes2d_score = \
                    torch.cat((projected_boxes2d, scores.unsqueeze(1)), dim=1)

                pred_box2d_list.append(boxes2d_score)
                pred_box3d_list.append(projected_boxes3d)

        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]

        # STEP3
        # 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, self.params['gt_range'])
        pred_box3d_tensor = pred_box3d_tensor[mask, :, :]
        scores = scores[mask]

        assert scores.shape[0] == pred_box3d_tensor.shape[0]

        return pred_box3d_tensor, scores

    @staticmethod
    def delta_to_boxes3d(deltas, anchors):
        """
        Convert the output delta to 3d bbx.

        Parameters
        ----------
        deltas : torch.Tensor
            (N, W, L, 14)?? should be (N, 14, H, W)
        anchors : torch.Tensor
            (W, L, 2, 7) -> xyzhwlr

        Returns
        -------
        box3d : torch.Tensor
            (N, W*L*2, 7)
        """
        # batch size
        N = deltas.shape[0]
        deltas = deltas.permute(0, 2, 3, 1).contiguous().view(N, -1, 7)
        boxes3d = torch.zeros_like(deltas)

        if deltas.is_cuda:
            anchors = anchors.cuda()
            boxes3d = boxes3d.cuda()

        # (W*L*2, 7)
        anchors_reshaped = anchors.view(-1, 7).float()
        # the diagonal of the anchor 2d box, (W*L*2)
        anchors_d = torch.sqrt(
            anchors_reshaped[:, 4] ** 2 + anchors_reshaped[:, 5] ** 2)
        anchors_d = anchors_d.repeat(N, 2, 1).transpose(1, 2)
        anchors_reshaped = anchors_reshaped.repeat(N, 1, 1)

        # Inv-normalize to get xyz
        boxes3d[..., [0, 1]] = torch.mul(deltas[..., [0, 1]], anchors_d) + \
                               anchors_reshaped[..., [0, 1]]
        boxes3d[..., [2]] = torch.mul(deltas[..., [2]],
                                      anchors_reshaped[..., [3]]) + \
                            anchors_reshaped[..., [2]]
        # hwl
        boxes3d[..., [3, 4, 5]] = torch.exp(
            deltas[..., [3, 4, 5]]) * anchors_reshaped[..., [3, 4, 5]]
        # yaw angle
        boxes3d[..., 6] = deltas[..., 6] + anchors_reshaped[..., 6]

        return boxes3d

    @staticmethod
    def visualize(pred_box_tensor, gt_tensor, pcd, show_vis, save_path, dataset=None):
        """
        Visualize the prediction, ground truth with point cloud together.

        Parameters
        ----------
        pred_box_tensor : torch.Tensor
            (N, 8, 3) prediction.

        gt_tensor : torch.Tensor
            (N, 8, 3) groundtruth bbx

        pcd : torch.Tensor
            PointCloud, (N, 4).

        show_vis : bool
            Whether to show visualization.

        save_path : str
            Save the visualization results to given path.

        dataset : BaseDataset
            opencood dataset object.

        """
        vis_utils.visualize_single_sample_output_gt(pred_box_tensor,
                                                    gt_tensor,
                                                    pcd,
                                                    show_vis,
                                                    save_path)

opencood/data_utils/post_processor/voxel_postprocessor.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, OpenPCDet
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
3D Anchor Generator for Voxel
"""
import math
import sys

import numpy as np
import torch
from torch.nn.functional import sigmoid
import torch.nn.functional as F

from opencood.data_utils.post_processor.base_postprocessor \
    import BasePostprocessor
from opencood.utils import box_utils
from opencood.utils.box_overlaps import bbox_overlaps
from opencood.visualization import vis_utils
from opencood.utils.common_utils import limit_period

from collections import OrderedDict

class VoxelPostprocessorFlow(BasePostprocessor):
    def __init__(self, anchor_params, train):
        super(VoxelPostprocessorFlow, self).__init__(anchor_params, train)
        self.anchor_num = self.params['anchor_args']['num']

    def generate_anchor_box(self):
        W = self.params['anchor_args']['W']
        H = self.params['anchor_args']['H']

        l = self.params['anchor_args']['l']
        w = self.params['anchor_args']['w']
        h = self.params['anchor_args']['h']
        r = self.params['anchor_args']['r']

        assert self.anchor_num == len(r)
        r = [math.radians(ele) for ele in r]

        vh = self.params['anchor_args']['vh'] # voxel_size
        vw = self.params['anchor_args']['vw']

        xrange = [self.params['anchor_args']['cav_lidar_range'][0],
                  self.params['anchor_args']['cav_lidar_range'][3]]
        yrange = [self.params['anchor_args']['cav_lidar_range'][1],
                  self.params['anchor_args']['cav_lidar_range'][4]]

        if 'feature_stride' in self.params['anchor_args']:
            feature_stride = self.params['anchor_args']['feature_stride']
        else:
            feature_stride = 2


        x = np.linspace(xrange[0] + vw, xrange[1] - vw, W // feature_stride)
        y = np.linspace(yrange[0] + vh, yrange[1] - vh, H // feature_stride)


        cx, cy = np.meshgrid(x, y)
        cx = np.tile(cx[..., np.newaxis], self.anchor_num) # center
        cy = np.tile(cy[..., np.newaxis], self.anchor_num)
        cz = np.ones_like(cx) * -1.0

        w = np.ones_like(cx) * w
        l = np.ones_like(cx) * l
        h = np.ones_like(cx) * h

        r_ = np.ones_like(cx)
        for i in range(self.anchor_num):
            r_[..., i] = r[i]

        if self.params['order'] == 'hwl': # pointpillar
            anchors = np.stack([cx, cy, cz, h, w, l, r_], axis=-1) # (50, 176, 2, 7)

        elif self.params['order'] == 'lhw':
            anchors = np.stack([cx, cy, cz, l, h, w, r_], axis=-1)
        else:
            sys.exit('Unknown bbx order.')

        return anchors

    def generate_label(self, **kwargs):
        """
        Generate targets for training.

        Parameters
        ----------
        argv : list
            gt_box_center:(max_num, 7), anchor:(H, W, anchor_num, 7)

        Returns
        -------
        label_dict : dict
            Dictionary that contains all target related info.
        """
        assert self.params['order'] == 'hwl', 'Currently Voxel only support' \
                                              'hwl bbx order.'
        # (max_num, 7)
        gt_box_center = kwargs['gt_box_center']
        # (H, W, anchor_num, 7)
        anchors = kwargs['anchors']
        # (max_num)
        masks = kwargs['mask']

        # (H, W)
        feature_map_shape = anchors.shape[:2]

        # (H*W*anchor_num, 7)
        anchors = anchors.reshape(-1, 7)
        # normalization factor, (H * W * anchor_num)
        anchors_d = np.sqrt(anchors[:, 4] ** 2 + anchors[:, 5] ** 2)

        # (H, W, 2)
        pos_equal_one = np.zeros((*feature_map_shape, self.anchor_num))
        neg_equal_one = np.zeros((*feature_map_shape, self.anchor_num))
        # (H, W, self.anchor_num * 7)
        targets = np.zeros((*feature_map_shape, self.anchor_num * 7))

        # (n, 7)
        gt_box_center_valid = gt_box_center[masks == 1]
        # (n, 8, 3)
        gt_box_corner_valid = \
            box_utils.boxes_to_corners_3d(gt_box_center_valid,
                                          self.params['order'])
        # (H*W*anchor_num, 8, 3)
        anchors_corner = \
            box_utils.boxes_to_corners_3d(anchors,
                                          order=self.params['order'])
        # (H*W*anchor_num, 4)
        anchors_standup_2d = \
            box_utils.corner2d_to_standup_box(anchors_corner)
        # (n, 4)
        gt_standup_2d = \
            box_utils.corner2d_to_standup_box(gt_box_corner_valid)

        # (H*W*anchor_n)
        iou = bbox_overlaps(
            np.ascontiguousarray(anchors_standup_2d).astype(np.float32),
            np.ascontiguousarray(gt_standup_2d).astype(np.float32),
        )

        # the anchor boxes has the largest iou across
        # shape: (n)
        id_highest = np.argmax(iou.T, axis=1)
        # [0, 1, 2, ..., n-1]
        id_highest_gt = np.arange(iou.T.shape[0])
        # make sure all highest iou is larger than 0
        mask = iou.T[id_highest_gt, id_highest] > 0
        id_highest, id_highest_gt = id_highest[mask], id_highest_gt[mask]


        # find anchors iou > params['pos_iou']
        id_pos, id_pos_gt = \
            np.where(iou >
                     self.params['target_args']['pos_threshold'])
        #  find anchors iou  params['neg_iou']
        id_neg = np.where(np.sum(iou <
                                 self.params['target_args']['neg_threshold'],
                                 axis=1) == iou.shape[1])[0]
        id_pos = np.concatenate([id_pos, id_highest])
        id_pos_gt = np.concatenate([id_pos_gt, id_highest_gt])
        id_pos, index = np.unique(id_pos, return_index=True)
        id_pos_gt = id_pos_gt[index]
        id_neg.sort()

        # cal the target and set the equal one
        index_x, index_y, index_z = np.unravel_index(
            id_pos, (*feature_map_shape, self.anchor_num))
        pos_equal_one[index_x, index_y, index_z] = 1

        # calculate the targets
        targets[index_x, index_y, np.array(index_z) * 7] = \
            (gt_box_center[id_pos_gt, 0] - anchors[id_pos, 0]) / anchors_d[
                id_pos]
        targets[index_x, index_y, np.array(index_z) * 7 + 1] = \
            (gt_box_center[id_pos_gt, 1] - anchors[id_pos, 1]) / anchors_d[
                id_pos]
        targets[index_x, index_y, np.array(index_z) * 7 + 2] = \
            (gt_box_center[id_pos_gt, 2] - anchors[id_pos, 2]) / anchors[
                id_pos, 3]
        targets[index_x, index_y, np.array(index_z) * 7 + 3] = np.log(
            gt_box_center[id_pos_gt, 3] / anchors[id_pos, 3])
        targets[index_x, index_y, np.array(index_z) * 7 + 4] = np.log(
            gt_box_center[id_pos_gt, 4] / anchors[id_pos, 4])
        targets[index_x, index_y, np.array(index_z) * 7 + 5] = np.log(
            gt_box_center[id_pos_gt, 5] / anchors[id_pos, 5])
        targets[index_x, index_y, np.array(index_z) * 7 + 6] = (
                gt_box_center[id_pos_gt, 6] - anchors[id_pos, 6])

        index_x, index_y, index_z = np.unravel_index(
            id_neg, (*feature_map_shape, self.anchor_num))
        neg_equal_one[index_x, index_y, index_z] = 1

        # to avoid a box be pos/neg in the same time
        index_x, index_y, index_z = np.unravel_index(
            id_highest, (*feature_map_shape, self.anchor_num))
        neg_equal_one[index_x, index_y, index_z] = 0


        label_dict = {'pos_equal_one': pos_equal_one,
                      'neg_equal_one': neg_equal_one,
                      'targets': targets}

        return label_dict
    
    @staticmethod
    def merge_label_to_dict(label_k_list):
        """
        Customized collate function for target label generation.

        Parameters
        ----------
        label_batch_list : list
            The list of dictionary  that contains all labels for several
            frames.

        Returns
        -------
        target_batch : dict
            Reformatted labels in torch tensor.
        """
        pos_equal_one = []
        neg_equal_one = []
        targets = []

        for i in range(len(label_k_list)):
            pos_equal_one.append(label_k_list[i]['pos_equal_one'])
            neg_equal_one.append(label_k_list[i]['neg_equal_one'])
            targets.append(label_k_list[i]['targets'])

        return {'targets': targets,
                'pos_equal_one': pos_equal_one,
                'neg_equal_one': neg_equal_one}

    @staticmethod
    def collate_batch(label_batch_list):
        """
        Customized collate function for target label generation.

        Parameters
        ----------
        label_batch_list : list
            The list of dictionary  that contains all labels for several
            frames.

        Returns
        -------
        target_batch : dict
            Reformatted labels in torch tensor.
        """
        pos_equal_one = []
        neg_equal_one = []
        targets = []

        for i in range(len(label_batch_list)):
            pos_equal_one.append(label_batch_list[i]['pos_equal_one'])
            neg_equal_one.append(label_batch_list[i]['neg_equal_one'])
            targets.append(label_batch_list[i]['targets'])

        pos_equal_one = \
            torch.from_numpy(np.array(pos_equal_one))
        neg_equal_one = \
            torch.from_numpy(np.array(neg_equal_one))
        targets = \
            torch.from_numpy(np.array(targets))

        return {'targets': targets,
                'pos_equal_one': pos_equal_one,
                'neg_equal_one': neg_equal_one}

    def flow_post_process(self, data_dict, output_dict, num_sweeps=2):
        """
        Process the outputs of the model to 2D/3D bounding box.
        Step1: convert each cav's output to bounding box format
        Step2: project the bounding boxes to ego space.
        Step:3 NMS

        For early and intermediate fusion,
            data_dict only contains ego.

        For late fusion,
            data_dcit contains all cavs, so we need transformation matrix.


        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        box_result : dict for each cav at each time
            {
                'ego' / cav_id : {
                    'past_k_time_diff' : 
                    ['comp'] {
                         pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                         pred_box_center_tensor : n, 7
                         scores: (n, )
                    }
                }
            }
        """
        self.k = num_sweeps

        box_results = OrderedDict()
        for cav_id, cav_content in data_dict.items():
            assert cav_id in output_dict
            box_results[cav_id] = OrderedDict()
            box_results[cav_id].update({
                'past_k_time_diff' : data_dict[cav_id]['past_k_time_diff']
            })
            # the transformation matrix to ego space
            # transformation_matrix = cav_content['transformation_matrix'] # no clean
            transformation_matrix = cav_content['trans_mat_past_2_past0_torch']

            # (H, W, anchor_num, 7)
            anchor_box = cav_content['anchor_box']

            # classification probability
            prob = output_dict[cav_id]['psm']  # k,2,100,352
            prob = torch.sigmoid(prob.permute(0, 2, 3, 1))
            prob = prob.reshape(self.k, -1) # (k, HxWx2)

            # regression map
            reg = output_dict[cav_id]['rm'] # (k,14,100,352)

            # convert regression map back to bounding box
            batch_box3d = self.delta_to_boxes3d(reg, anchor_box) # (k, HxWx2, 7)
            mask = \
                torch.gt(prob, self.params['target_args']['score_threshold']) 
            mask = mask.view(self.k, -1) # (k, HxWx2)
            mask_reg = mask.unsqueeze(2).repeat(1, 1, 7) # (k, HxWx2, 7)

            # during validation/testing, the batch size should be 1
            # assert batch_box3d.shape[0] == 1
            
            for i in range(self.k):  # 遍历所有的时间戳
                box_results[cav_id][i] = OrderedDict()
                unit_trans_mat = transformation_matrix[i]
                boxes3d = torch.masked_select(batch_box3d[i],
                                            mask_reg[i]).view(-1, 7)
                scores = torch.masked_select(prob[i], mask[i])

                # 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'])
                    # STEP 2
                    # (N, 8, 3)
                    projected_boxes3d = \
                        box_utils.project_box3d(boxes3d_corner,
                                                unit_trans_mat)

                    # backup_projected_boxes3d = projected_boxes3d.clone()
                    # backup_scores = scores.clone()

                    # convert 3d bbx to 2d, (N,4)
                    projected_boxes2d = \
                        box_utils.corner_to_standup_box_torch(projected_boxes3d)
                    # (N, 5)
                    boxes2d_score = \
                        torch.cat((projected_boxes2d, scores.unsqueeze(1)), dim=1)

                    # pred_box2d_list.append(boxes2d_score)
                    # pred_box3d_list.append(projected_boxes3d)

                    # scores = boxes2d_score[:, -1]

                    # remove large bbx
                    keep_index_1 = box_utils.remove_large_pred_bbx(projected_boxes3d)
                    keep_index_2 = box_utils.remove_bbx_abnormal_z(projected_boxes3d)
                    keep_index = torch.logical_and(keep_index_1, keep_index_2)

                    projected_boxes3d = projected_boxes3d[keep_index]
                    scores = scores[keep_index]

                    # STEP3
                    # nms
                    keep_index = box_utils.nms_rotated(projected_boxes3d,
                                                    scores,
                                                    self.params['nms_thresh']
                                                    )
                    pred_box3d_tensor = projected_boxes3d[keep_index]
                    # select cooresponding score
                    scores = scores[keep_index]

                    # filter out the prediction out of the range.
                    range_mask = \
                        box_utils.get_mask_for_boxes_within_range_torch(pred_box3d_tensor, self.params['gt_range'])
                    pred_box_3dcorner_tensor = pred_box3d_tensor[range_mask, :, :]
                    scores = scores[range_mask]
                    pred_box_center_tensor = box_utils.corner_to_center_torch(pred_box_3dcorner_tensor, self.params['order'])

                    assert scores.shape[0] == pred_box_3dcorner_tensor.shape[0]

                    # if scores.shape[0] == 0: # 没有bbx剩下，那么保存之前剩的框
                    #     pred_box_3dcorner_tensor = backup_projected_boxes3d
                    #     scores = backup_scores
                    #     pred_box_center_tensor = box_utils.corner_to_center_torch(pred_box_3dcorner_tensor, self.params['order'])

                    box_results[cav_id][i].update({
                        'pred_box_3dcorner_tensor': pred_box_3dcorner_tensor, 
                        'pred_box_center_tensor': pred_box_center_tensor,
                        'scores': scores
                    })

                else:
                    box_results[cav_id][i].update({
                        'pred_box_3dcorner_tensor': torch.Tensor(0, 8, 3).to(scores.device), 
                        'pred_box_center_tensor': torch.Tensor(0, 7).to(scores.device),
                        'scores': torch.Tensor(0).to(scores.device)
                    })

        return box_results

    def post_process_fuse_updated_frame(self, data_dict, output_dict):
        """
        Process the flow-updated box from cavs to ego-cav view.
        Step2: project the bounding boxes to ego space.
        Step3: NMS

        For early and intermediate fusion,
            data_dict only contains ego.

        For late fusion,
            data_dcit contains all cavs, so we need transformation matrix.


        Parameters
        ----------
        data_dict : 
            {

            }

        output_dict : dict for each cav at each time
            {
                'ego' / cav_id : {
                    'past_k_time_diff' : 
                    ['comp'] {
                         pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                         pred_box_center_tensor : n, 7
                         scores: (n, )
                    }
                }
            }

        Returns
        -------
        pred_box3d_tensor : bbx fused from different cavs at ego-curr view
            torch.tensor(N, 8, 3)
        
        scores : torch.tensor(N, )
        
        """
        # the final bounding box list
        pred_box3d_list = []
        delay_box3d_list = []
        scores_list = []

        for cav_id, cav_content in output_dict.items():
            assert cav_id in data_dict
            transformation_matrix = data_dict[cav_id]['trans_mat_past0_2_curr_torch']  # 4, 4

            box_3d = cav_content['comp']['pred_box_3dcorner_tensor']
            scores = cav_content['comp']['scores']
            box_3d_delay = cav_content[0]['pred_box_3dcorner_tensor']
            if box_3d.shape[0] != 0:
                projected_boxes3d = box_utils.project_box3d(box_3d, transformation_matrix)
                pred_box3d_list.append(projected_boxes3d)
                scores_list += scores
            if box_3d_delay.shape[0] != 0:
                projected_boxes3d_delay = box_utils.project_box3d(box_3d_delay, transformation_matrix)
                delay_box3d_list.append(projected_boxes3d_delay)

        pred_box3d_tensor = torch.vstack(pred_box3d_list)
        delay_box3d_tensor = torch.vstack(delay_box3d_list)
        scores = torch.hstack(scores_list)
        assert pred_box3d_tensor.shape[0] == scores.shape[0]
        
        # STEP3
        # nms
        keep_index = box_utils.nms_rotated(pred_box3d_tensor,
                                           scores,
                                           self.params['nms_thresh']
                                           )

        pred_box3d_tensor = pred_box3d_tensor[keep_index]
        delay_box3d_tensor = delay_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, self.params['gt_range'])
        pred_box3d_tensor = pred_box3d_tensor[mask, :, :]
        scores = scores[mask]

        delay_box3d_tensor = delay_box3d_tensor[mask, :, :]

        assert scores.shape[0] == pred_box3d_tensor.shape[0]

        
        return pred_box3d_tensor, scores, delay_box3d_tensor


        for cav_id, cav_content in data_dict.items():
            assert cav_id in output_dict
            # the transformation matrix to ego space
            transformation_matrix = cav_content['transformation_matrix'] # no clean

            # (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
            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])

            # adding dir classifier
            if 'dm' in output_dict[cav_id].keys() and len(boxes3d) !=0:
                dir_offset = self.params['dir_args']['dir_offset']
                num_bins = self.params['dir_args']['num_bins']


                dm  = output_dict[cav_id]['dm'] # [N, H, W, 4]
                dir_cls_preds = dm.permute(0, 2, 3, 1).contiguous().reshape(1, -1, num_bins) # [1, N*H*W*2, 2]
                dir_cls_preds = dir_cls_preds[mask]
                # if rot_gt > 0, then the label is 1, then the regression target is [0, 1]
                dir_labels = torch.max(dir_cls_preds, dim=-1)[1]  # indices. shape [1, N*H*W*2].  value 0 or 1. If value is 1, then rot_gt > 0
                
                period = (2 * np.pi / num_bins) # pi
                dir_rot = limit_period(
                    boxes3d[..., 6] - dir_offset, 0, period
                ) # 限制在0到pi之间
                boxes3d[..., 6] = dir_rot + dir_offset + period * dir_labels.to(dir_cls_preds.dtype) # 转化0.25pi到2.5pi
                boxes3d[..., 6] = limit_period(boxes3d[..., 6], 0.5, 2 * np.pi) # limit to [-pi, pi]

            # 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'])
                
                # STEP 2
                # (N, 8, 3)
                projected_boxes3d = \
                    box_utils.project_box3d(boxes3d_corner,
                                            transformation_matrix)
                # convert 3d bbx to 2d, (N,4)
                projected_boxes2d = \
                    box_utils.corner_to_standup_box_torch(projected_boxes3d)
                # (N, 5)
                boxes2d_score = \
                    torch.cat((projected_boxes2d, scores.unsqueeze(1)), dim=1)

                pred_box2d_list.append(boxes2d_score)
                pred_box3d_list.append(projected_boxes3d)

        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]

        # STEP3
        # 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, self.params['gt_range'])
        pred_box3d_tensor = pred_box3d_tensor[mask, :, :]
        scores = scores[mask]

        assert scores.shape[0] == pred_box3d_tensor.shape[0]

        return pred_box3d_tensor, scores

    def post_process_for_intermediate(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.
        Step1: convert each cav's output to bounding box format
        Step2: project the bounding boxes to ego space.
        Step3: NMS

        For early and intermediate fusion,
            data_dict only contains ego.

        For late fusion,
            data_dcit contains all cavs, so we need transformation matrix.


        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box3d_tensor : torch.Tensor
            The prediction bounding box tensor after NMS.
        gt_box3d_tensor : torch.Tensor
            The groundtruth bounding box tensor.
        """
        # the final bounding box list
        # pred_box3d_list = []
        # pred_box2d_list = []
        # for cav_id, cav_content in data_dict.items():
        #     assert cav_id in output_dict
        #     # the transformation matrix to ego space
        #     transformation_matrix = cav_content['transformation_matrix'] # no clean

        # (H, W, anchor_num, 7)
        anchor_box = data_dict['ego']['anchor_box']

        # classification probability
        prob = output_dict['psm']
        prob = torch.sigmoid(prob.permute(0, 2, 3, 1))
        prob = prob.reshape(1, -1)

        # regression map
        reg = output_dict['rm']

        # convert regression map back to bounding box
        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'])

            # predicted 3d bbx
            pred_box3d_tensor = boxes3d_corner
            # 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]

            # STEP3
            # 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, self.params['gt_range'])
            pred_box3d_tensor = pred_box3d_tensor[mask, :, :]
            scores = scores[mask]

            assert scores.shape[0] == pred_box3d_tensor.shape[0]
            return pred_box3d_tensor, scores

        else:
            pred_box3d_tensor = torch.Tensor(0, 8, 3).to(boxes3d.device)
            scores = torch.Tensor(0).to(boxes3d.device)
            return pred_box3d_tensor, scores
    
    def post_process(self, data_dict, output_dict):
        """
        Process the outputs of the model to 2D/3D bounding box.
        Step1: convert each cav's output to bounding box format
        Step2: project the bounding boxes to ego space.
        Step3: NMS

        For early and intermediate fusion,
            data_dict only contains ego.

        For late fusion,
            data_dcit contains all cavs, so we need transformation matrix.


        Parameters
        ----------
        data_dict : dict
            The dictionary containing the origin input data of model.

        output_dict :dict
            The dictionary containing the output of the model.

        Returns
        -------
        pred_box3d_tensor : torch.Tensor
            The prediction bounding box tensor after NMS.
        gt_box3d_tensor : torch.Tensor
            The groundtruth bounding box tensor.
        """
        # the final bounding box list
        pred_box3d_list = []
        pred_box2d_list = []
        for cav_id, cav_content in data_dict.items():
            assert cav_id in output_dict
            # the transformation matrix to ego space
            transformation_matrix = cav_content['transformation_matrix'] # no clean

            # (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
            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])

            # adding dir classifier
            if 'dm' in output_dict[cav_id].keys() and len(boxes3d) !=0:
                dir_offset = self.params['dir_args']['dir_offset']
                num_bins = self.params['dir_args']['num_bins']


                dm  = output_dict[cav_id]['dm'] # [N, H, W, 4]
                dir_cls_preds = dm.permute(0, 2, 3, 1).contiguous().reshape(1, -1, num_bins) # [1, N*H*W*2, 2]
                dir_cls_preds = dir_cls_preds[mask]
                # if rot_gt > 0, then the label is 1, then the regression target is [0, 1]
                dir_labels = torch.max(dir_cls_preds, dim=-1)[1]  # indices. shape [1, N*H*W*2].  value 0 or 1. If value is 1, then rot_gt > 0
                
                period = (2 * np.pi / num_bins) # pi
                dir_rot = limit_period(
                    boxes3d[..., 6] - dir_offset, 0, period
                ) # 限制在0到pi之间
                boxes3d[..., 6] = dir_rot + dir_offset + period * dir_labels.to(dir_cls_preds.dtype) # 转化0.25pi到2.5pi
                boxes3d[..., 6] = limit_period(boxes3d[..., 6], 0.5, 2 * np.pi) # limit to [-pi, pi]

            # 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'])
                
                # STEP 2
                # (N, 8, 3)
                projected_boxes3d = \
                    box_utils.project_box3d(boxes3d_corner,
                                            transformation_matrix)
                # convert 3d bbx to 2d, (N,4)
                projected_boxes2d = \
                    box_utils.corner_to_standup_box_torch(projected_boxes3d)
                # (N, 5)
                boxes2d_score = \
                    torch.cat((projected_boxes2d, scores.unsqueeze(1)), dim=1)

                pred_box2d_list.append(boxes2d_score)
                pred_box3d_list.append(projected_boxes3d)

        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]

        # STEP3
        # 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, self.params['gt_range'])
        pred_box3d_tensor = pred_box3d_tensor[mask, :, :]
        scores = scores[mask]

        assert scores.shape[0] == pred_box3d_tensor.shape[0]

        return pred_box3d_tensor, scores
    
    @staticmethod
    def delta_to_boxes3d(deltas, anchors):
        """
        Convert the output delta to 3d bbx.

        Parameters
        ----------
        deltas : torch.Tensor
            (N, W, L, 14)?? should be (N, 14, H, W)
        anchors : torch.Tensor
            (W, L, 2, 7) -> xyzhwlr

        Returns
        -------
        box3d : torch.Tensor
            (N, W*L*2, 7)
        """
        # batch size
        N = deltas.shape[0]
        deltas = deltas.permute(0, 2, 3, 1).contiguous().view(N, -1, 7)
        boxes3d = torch.zeros_like(deltas)

        if deltas.is_cuda:
            anchors = anchors.cuda()
            boxes3d = boxes3d.cuda()

        # (W*L*2, 7)
        anchors_reshaped = anchors.view(-1, 7).float()
        # the diagonal of the anchor 2d box, (W*L*2)
        anchors_d = torch.sqrt(
            anchors_reshaped[:, 4] ** 2 + anchors_reshaped[:, 5] ** 2)
        anchors_d = anchors_d.repeat(N, 2, 1).transpose(1, 2)
        anchors_reshaped = anchors_reshaped.repeat(N, 1, 1)

        # Inv-normalize to get xyz
        boxes3d[..., [0, 1]] = torch.mul(deltas[..., [0, 1]], anchors_d) + \
                               anchors_reshaped[..., [0, 1]]
        boxes3d[..., [2]] = torch.mul(deltas[..., [2]],
                                      anchors_reshaped[..., [3]]) + \
                            anchors_reshaped[..., [2]]
        # hwl
        boxes3d[..., [3, 4, 5]] = torch.exp(
            deltas[..., [3, 4, 5]]) * anchors_reshaped[..., [3, 4, 5]]
        # yaw angle
        boxes3d[..., 6] = deltas[..., 6] + anchors_reshaped[..., 6]

        return boxes3d

    @staticmethod
    def visualize(pred_box_tensor, gt_tensor, pcd, show_vis, save_path, dataset=None):
        """
        Visualize the prediction, ground truth with point cloud together.

        Parameters
        ----------
        pred_box_tensor : torch.Tensor
            (N, 8, 3) prediction.

        gt_tensor : torch.Tensor
            (N, 8, 3) groundtruth bbx

        pcd : torch.Tensor
            PointCloud, (N, 4).

        show_vis : bool
            Whether to show visualization.

        save_path : str
            Save the visualization results to given path.

        dataset : BaseDataset
            opencood dataset object.

        """
        vis_utils.visualize_single_sample_output_gt(pred_box_tensor,
                                                    gt_tensor,
                                                    pcd,
                                                    show_vis,
                                                    save_path)

opencood/data_utils/post_processor/voxel_postprocessor_flow.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
Convert lidar to bev
"""

import numpy as np
import torch
from opencood.data_utils.pre_processor.base_preprocessor import \
    BasePreprocessor


class BevPreprocessor(BasePreprocessor):
    def __init__(self, preprocess_params, train):
        super(BevPreprocessor, self).__init__(preprocess_params, train)
        self.lidar_range = self.params['cav_lidar_range']
        self.geometry_param = preprocess_params["geometry_param"]

    def preprocess(self, pcd_raw):
        """
        Preprocess the lidar points to BEV representations.

        Parameters
        ----------
        pcd_raw : np.ndarray
            The raw lidar.

        Returns
        -------
        data_dict : the structured output dictionary.
        """
        bev = np.zeros(self.geometry_param['input_shape'], dtype=np.float32)
        intensity_map_count = np.zeros((bev.shape[0], bev.shape[1]),
                                       dtype=np.int)
        bev_origin = np.array(
            [self.geometry_param["L1"], self.geometry_param["W1"],
             self.geometry_param["H1"]]).reshape(1, -1)

        indices = ((pcd_raw[:, :3] - bev_origin) / self.geometry_param[
            "res"]).astype(int)

        # if any point hit this voxel, set the voxel to 1
        for i in range(indices.shape[0]):
            bev[indices[i, 0], indices[i, 1], indices[i, 2]] = 1
            bev[indices[i, 0], indices[i, 1], -1] += pcd_raw[i, 3] # intensity
            intensity_map_count[indices[i, 0], indices[i, 1]] += 1
        divide_mask = intensity_map_count != 0
        bev[divide_mask, -1] = np.divide(bev[divide_mask, -1],
                                         intensity_map_count[divide_mask])

        data_dict = {
            "bev_input": np.transpose(bev, (2, 0, 1)) # (C,H,W)
        }
        return data_dict

    @staticmethod
    def collate_batch_list(batch):
        """
        Customized pytorch data loader collate function.

        Parameters
        ----------
        batch : list
            List of dictionary. Each dictionary represent a single frame.

        Returns
        -------
        processed_batch : dict
            Updated lidar batch.
        """
        bev_input_list = [
            x["bev_input"][np.newaxis, ...] for x in batch
        ]
        processed_batch = {
            "bev_input": torch.from_numpy(
                np.concatenate(bev_input_list, axis=0))
        }
        return processed_batch

    @staticmethod
    def collate_batch_dict(batch):
        """
        Customized pytorch data loader collate function.

        Parameters
        ----------
        batch : dict
            Dict of list. Each element represents a CAV.

        Returns
        -------
        processed_batch : dict
            Updated lidar batch.
        """
        bev_input_list = [
            x[np.newaxis, ...] for x in batch["bev_input"]
        ]
        processed_batch = {
            "bev_input": torch.from_numpy(
                np.concatenate(bev_input_list, axis=0))
        }
        return processed_batch

    def collate_batch(self, batch):
        """
        Customized pytorch data loader collate function.

        Parameters
        ----------
        batch : list / dict
            Batched data.
        Returns
        -------
        processed_batch : dict
            Updated lidar batch.
        """
        if isinstance(batch, list):
            return self.collate_batch_list(batch)
        elif isinstance(batch, dict):
            return self.collate_batch_dict(batch)
        else:
            raise NotImplemented

opencood/data_utils/pre_processor/bev_preprocessor.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, OpenPCDet
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Transform points to voxels using sparse conv library
"""
import sys

import numpy as np
import torch
from icecream import ic

from opencood.data_utils.pre_processor.base_preprocessor import \
    BasePreprocessor


class SpVoxelPreprocessor(BasePreprocessor):
    def __init__(self, preprocess_params, train):
        super(SpVoxelPreprocessor, self).__init__(preprocess_params,
                                                  train)
        try:
            from spconv.utils import VoxelGeneratorV2 as VoxelGenerator
        except:
            from spconv.utils import VoxelGenerator   

        self.lidar_range = self.params['cav_lidar_range']
        self.voxel_size = self.params['args']['voxel_size']
        self.max_points_per_voxel = self.params['args']['max_points_per_voxel']

        if train:
            self.max_voxels = self.params['args']['max_voxel_train']
        else:
            self.max_voxels = self.params['args']['max_voxel_test']
        
        # whether there are more than on frames in the past
        if 'past_k' in preprocess_params and preprocess_params['past_k'] > 0:
            self.sweep = True
        else:
            self.sweep = False

        grid_size = (np.array(self.lidar_range[3:6]) -
                     np.array(self.lidar_range[0:3])) / np.array(self.voxel_size)
        self.grid_size = np.round(grid_size).astype(np.int64)

        # use sparse conv library to generate voxel
        self.voxel_generator = VoxelGenerator(
            voxel_size=self.voxel_size,
            point_cloud_range=self.lidar_range,
            max_num_points=self.max_points_per_voxel,
            max_voxels=self.max_voxels
        )

    def preprocess(self, pcd_np):
        data_dict = {}
        voxel_output = self.voxel_generator.generate(pcd_np)
        if isinstance(voxel_output, dict):
            voxels, coordinates, num_points = \
                voxel_output['voxels'], voxel_output['coordinates'], \
                voxel_output['num_points_per_voxel']
        else:
            voxels, coordinates, num_points = voxel_output

        data_dict['voxel_features'] = voxels
        data_dict['voxel_coords'] = coordinates
        data_dict['voxel_num_points'] = num_points


        return data_dict

    def collate_batch(self, batch):
        """
        Customized pytorch data loader collate function.

        Parameters
        ----------
        batch : list or dict
            List or dictionary.

        Returns
        -------
        processed_batch : dict
            Updated lidar batch.
        """
        if isinstance(batch, list):
            return self.collate_batch_list(batch)
        elif isinstance(batch, dict):
            return self.collate_batch_dict(batch)
        else:
            sys.exit('Batch has too be a list or a dictionarn')

    @staticmethod
    def collate_batch_list(batch):
        """
        Customized pytorch data loader collate function.

        Parameters
        ----------
        batch : list
            List of dictionary. Each dictionary represent a single frame.

        Returns
        -------
        processed_batch : dict
            Updated lidar batch.
        """
        voxel_features = []
        voxel_num_points = []
        voxel_coords = []

        for i in range(len(batch)):
            voxel_features.append(batch[i]['voxel_features'])
            voxel_num_points.append(batch[i]['voxel_num_points'])
            coords = batch[i]['voxel_coords']
            voxel_coords.append(
                np.pad(coords, ((0, 0), (1, 0)),
                       mode='constant', constant_values=i))

        voxel_num_points = torch.from_numpy(np.concatenate(voxel_num_points))
        voxel_features = torch.from_numpy(np.concatenate(voxel_features))
        voxel_coords = torch.from_numpy(np.concatenate(voxel_coords))

        return {'voxel_features': voxel_features,
                'voxel_coords': voxel_coords,
                'voxel_num_points': voxel_num_points}

    @staticmethod
    def collate_batch_dict(batch: dict):
        """
        Collate batch if the batch is a dictionary,
        eg: {'voxel_features': [feature1, feature2...., feature n]}

        Parameters
        ----------
        batch : dict

        Returns
        -------
        processed_batch : dict
            Updated lidar batch.
        """
        voxel_features = \
            torch.from_numpy(np.concatenate(batch['voxel_features']))
        voxel_num_points = \
            torch.from_numpy(np.concatenate(batch['voxel_num_points']))
        coords = batch['voxel_coords']
        voxel_coords = []

        for i in range(len(coords)):
            voxel_coords.append(
                np.pad(coords[i], ((0, 0), (1, 0)),
                       mode='constant', constant_values=i))
        voxel_coords = torch.from_numpy(np.concatenate(voxel_coords))

        return {'voxel_features': voxel_features,
                'voxel_coords': voxel_coords,
                'voxel_num_points': voxel_num_points}

opencood/data_utils/pre_processor/sp_voxel_preprocessor.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
Convert lidar to voxel. This class was manually designed, and we recommend
to use sp_voxel_preprocessor.
"""
import sys

import numpy as np
import torch

from opencood.data_utils.pre_processor.base_preprocessor import \
    BasePreprocessor


class VoxelPreprocessor(BasePreprocessor):
    def __init__(self, preprocess_params, train):
        super(VoxelPreprocessor, self).__init__(preprocess_params, train)
        # TODO: add intermediate lidar range later
        self.lidar_range = self.params['cav_lidar_range']

        self.vw = self.params['args']['vw']
        self.vh = self.params['args']['vh']
        self.vd = self.params['args']['vd']
        self.T = self.params['args']['T']

    def preprocess(self, pcd_np):
        """
        Preprocess the lidar points by  voxelization.

        Parameters
        ----------
        pcd_np : np.ndarray
            The raw lidar.

        Returns
        -------
        data_dict : the structured output dictionary.
        """
        data_dict = {}

        # calculate the voxel coordinates
        voxel_coords = ((pcd_np[:, :3] -
                         np.floor(np.array([self.lidar_range[0],
                                            self.lidar_range[1],
                                            self.lidar_range[2]])) / (
                             self.vw, self.vh, self.vd))).astype(np.int32)

        # convert to  (D, H, W) as the paper
        voxel_coords = voxel_coords[:, [2, 1, 0]]
        voxel_coords, inv_ind, voxel_counts = np.unique(voxel_coords, axis=0,
                                                        return_inverse=True,
                                                        return_counts=True)

        voxel_features = []

        for i in range(len(voxel_coords)):
            voxel = np.zeros((self.T, 7), dtype=np.float32)
            pts = pcd_np[inv_ind == i]
            if voxel_counts[i] > self.T:
                pts = pts[:self.T, :]
                voxel_counts[i] = self.T

            # augment the points
            voxel[:pts.shape[0], :] = np.concatenate((pts, pts[:, :3] -
                                                      np.mean(pts[:, :3], 0)),
                                                     axis=1)
            voxel_features.append(voxel)

        data_dict['voxel_features'] = np.array(voxel_features)
        data_dict['voxel_coords'] = voxel_coords

        return data_dict

    def collate_batch(self, batch):
        """
        Customized pytorch data loader collate function.

        Parameters
        ----------
        batch : list or dict
            List or dictionary.

        Returns
        -------
        processed_batch : dict
            Updated lidar batch.
        """

        if isinstance(batch, list):
            return self.collate_batch_list(batch)
        elif isinstance(batch, dict):
            return self.collate_batch_dict(batch)
        else:
            sys.exit('Batch has too be a list or a dictionary')

    @staticmethod
    def collate_batch_list(batch):
        """
        Customized pytorch data loader collate function.

        Parameters
        ----------
        batch : list
            List of dictionary. Each dictionary represent a single frame.

        Returns
        -------
        processed_batch : dict
            Updated lidar batch.
        """
        voxel_features = []
        voxel_coords = []

        for i in range(len(batch)):
            voxel_features.append(batch[i]['voxel_features'])
            coords = batch[i]['voxel_coords']
            voxel_coords.append(
                np.pad(coords, ((0, 0), (1, 0)),
                       mode='constant', constant_values=i))

        voxel_features = torch.from_numpy(np.concatenate(voxel_features))
        voxel_coords = torch.from_numpy(np.concatenate(voxel_coords))

        return {'voxel_features': voxel_features,
                'voxel_coords': voxel_coords}

    @staticmethod
    def collate_batch_dict(batch: dict):
        """
        Collate batch if the batch is a dictionary,
        eg: {'voxel_features': [feature1, feature2...., feature n]}

        Parameters
        ----------
        batch : dict

        Returns
        -------
        processed_batch : dict
            Updated lidar batch.
        """
        voxel_features = \
            torch.from_numpy(np.concatenate(batch['voxel_features']))
        coords = batch['voxel_coords']
        voxel_coords = []

        for i in range(len(coords)):
            voxel_coords.append(
                np.pad(coords[i], ((0, 0), (1, 0)),
                       mode='constant', constant_values=i))
        voxel_coords = torch.from_numpy(np.concatenate(voxel_coords))

        return {'voxel_features': voxel_features,
                'voxel_coords': voxel_coords}

opencood/data_utils/pre_processor/voxel_preprocessor.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import re
import yaml
import os
import math

import numpy as np


def load_yaml(file, opt=None, config_suffix=''):
    """
    Load yaml file and return a dictionary.

    Parameters
    ----------
    file : string
        yaml file path.

    opt : argparser
         Argparser.
    Returns
    -------
    param : dict
        A dictionary that contains defined parameters.
    """
    if opt and opt.model_dir:
        if config_suffix:
            file = os.path.join(opt.model_dir, f'config_{config_suffix}.yaml')
        else:
            file = os.path.join(opt.model_dir, 'config.yaml')

    stream = open(file, 'r')
    loader = yaml.Loader
    loader.add_implicit_resolver(
        u'tag:yaml.org,2002:float',
        re.compile(u'''^(?:
         [-+]?(?:[0-9][0-9_]*)\\.[0-9_]*(?:[eE][-+]?[0-9]+)?
        |[-+]?(?:[0-9][0-9_]*)(?:[eE][-+]?[0-9]+)
        |\\.[0-9_]+(?:[eE][-+][0-9]+)?
        |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\\.[0-9_]*
        |[-+]?\\.(?:inf|Inf|INF)
        |\\.(?:nan|NaN|NAN))$''', re.X),
        list(u'-+0123456789.'))
    param = yaml.load(stream, Loader=loader)
    if "yaml_parser" in param:
        param = eval(param["yaml_parser"])(param)

    return param


def load_voxel_params(param):
    """
    Based on the lidar range and resolution of voxel, calcuate the anchor box
    and target resolution.

    Parameters
    ----------
    param : dict
        Original loaded parameter dictionary.

    Returns
    -------
    param : dict
        Modified parameter dictionary with new attribute `anchor_args[W][H][L]`
    """
    anchor_args = param['postprocess']['anchor_args']
    cav_lidar_range = anchor_args['cav_lidar_range']
    voxel_size = param['preprocess']['args']['voxel_size']

    vw = voxel_size[0]
    vh = voxel_size[1]
    vd = voxel_size[2]

    anchor_args['vw'] = vw
    anchor_args['vh'] = vh
    anchor_args['vd'] = vd

    anchor_args['W'] = int((cav_lidar_range[3] - cav_lidar_range[0]) / vw)
    anchor_args['H'] = int((cav_lidar_range[4] - cav_lidar_range[1]) / vh)
    anchor_args['D'] = int((cav_lidar_range[5] - cav_lidar_range[2]) / vd)

    param['postprocess'].update({'anchor_args': anchor_args})

    # sometimes we just want to visualize the data without implementing model
    if 'model' in param:
        param['model']['args']['W'] = anchor_args['W']
        param['model']['args']['H'] = anchor_args['H']
        param['model']['args']['D'] = anchor_args['D']

    return param


def load_point_pillar_params(param):
    """
    Based on the lidar range and resolution of voxel, calcuate the anchor box
    and target resolution.

    Parameters
    ----------
    param : dict
        Original loaded parameter dictionary.

    Returns
    -------
    param : dict
        Modified parameter dictionary with new attribute.
    """
    cav_lidar_range = param['preprocess']['cav_lidar_range']
    voxel_size = param['preprocess']['args']['voxel_size']

    grid_size = (np.array(cav_lidar_range[3:6]) - np.array(
        cav_lidar_range[0:3])) / \
                np.array(voxel_size)
    grid_size = np.round(grid_size).astype(np.int64)
    param['model']['args']['point_pillar_scatter']['grid_size'] = grid_size

    anchor_args = param['postprocess']['anchor_args']

    vw = voxel_size[0]
    vh = voxel_size[1]
    vd = voxel_size[2]

    anchor_args['vw'] = vw
    anchor_args['vh'] = vh
    anchor_args['vd'] = vd

    anchor_args['W'] = math.ceil((cav_lidar_range[3] - cav_lidar_range[0]) / vw) # W is image width, but along with x axis in lidar coordinate
    anchor_args['H'] = math.ceil((cav_lidar_range[4] - cav_lidar_range[1]) / vh) # H is image height
    anchor_args['D'] = math.ceil((cav_lidar_range[5] - cav_lidar_range[2]) / vd)

    param['postprocess'].update({'anchor_args': anchor_args})

    return param


def load_second_params(param):
    """
    Based on the lidar range and resolution of voxel, calcuate the anchor box
    and target resolution.

    Parameters
    ----------
    param : dict
        Original loaded parameter dictionary.

    Returns
    -------
    param : dict
        Modified parameter dictionary with new attribute.
    """
    cav_lidar_range = param['preprocess']['cav_lidar_range']
    voxel_size = param['preprocess']['args']['voxel_size']

    grid_size = (np.array(cav_lidar_range[3:6]) - np.array(
        cav_lidar_range[0:3])) / \
                np.array(voxel_size)
    grid_size = np.round(grid_size).astype(np.int64)
    param['model']['args']['grid_size'] = grid_size

    anchor_args = param['postprocess']['anchor_args']

    vw = voxel_size[0]
    vh = voxel_size[1]
    vd = voxel_size[2]

    anchor_args['vw'] = vw
    anchor_args['vh'] = vh
    anchor_args['vd'] = vd

    anchor_args['W'] = int((cav_lidar_range[3] - cav_lidar_range[0]) / vw)
    anchor_args['H'] = int((cav_lidar_range[4] - cav_lidar_range[1]) / vh)
    anchor_args['D'] = int((cav_lidar_range[5] - cav_lidar_range[2]) / vd)

    param['postprocess'].update({'anchor_args': anchor_args})

    return param


def load_bev_params(param):
    """
    Load bev related geometry parameters s.t. boundary, resolutions, input
    shape, target shape etc.

    Parameters
    ----------
    param : dict
        Original loaded parameter dictionary.

    Returns
    -------
    param : dict
        Modified parameter dictionary with new attribute `geometry_param`.

    """
    res = param["preprocess"]["args"]["res"]
    L1, W1, H1, L2, W2, H2 = param["preprocess"]["cav_lidar_range"]
    downsample_rate = param["preprocess"]["args"]["downsample_rate"]

    def f(low, high, r):
        return int((high - low) / r)

    input_shape = (
        int((f(L1, L2, res))),
        int((f(W1, W2, res))),
        int((f(H1, H2, res)) + 1)
    )
    label_shape = (
        int(input_shape[0] / downsample_rate),
        int(input_shape[1] / downsample_rate),
        7
    )
    geometry_param = {
        'L1': L1,
        'L2': L2,
        'W1': W1,
        'W2': W2,
        'H1': H1,
        'H2': H2,
        "downsample_rate": downsample_rate,
        "input_shape": input_shape,
        "label_shape": label_shape,
        "res": res
    }
    param["preprocess"]["geometry_param"] = geometry_param
    param["postprocess"]["geometry_param"] = geometry_param
    param["model"]["args"]["geometry_param"] = geometry_param
    return param


def save_yaml(data, save_name):
    """
    Save the dictionary into a yaml file.

    Parameters
    ----------
    data : dict
        The dictionary contains all data.

    save_name : string
        Full path of the output yaml file.
    """

    with open(save_name, 'w') as outfile:
        yaml.dump(data, outfile, default_flow_style=False)



def load_point_pillar_params_stage1(param):
    """
    Based on the lidar range and resolution of voxel, calcuate the anchor box
    and target resolution.

    Parameters
    ----------
    param : dict
        Original loaded parameter dictionary.

    Returns
    -------
    param : dict
        Modified parameter dictionary with new attribute.
    """
    cav_lidar_range = param['preprocess']['cav_lidar_range']
    voxel_size = param['preprocess']['args']['voxel_size']

    grid_size = (np.array(cav_lidar_range[3:6]) - np.array(
        cav_lidar_range[0:3])) / \
                np.array(voxel_size)
    grid_size = np.round(grid_size).astype(np.int64)
    param['box_align_pre_calc']['stage1_model_config']['point_pillar_scatter']['grid_size'] = grid_size

    anchor_args = param['box_align_pre_calc']['stage1_postprocessor_config']['anchor_args']

    vw = voxel_size[0]
    vh = voxel_size[1]
    vd = voxel_size[2]

    anchor_args['vw'] = vw
    anchor_args['vh'] = vh
    anchor_args['vd'] = vd

    anchor_args['W'] = int((cav_lidar_range[3] - cav_lidar_range[0]) / vw) # W is image width, but along with x axis in lidar coordinate
    anchor_args['H'] = int((cav_lidar_range[4] - cav_lidar_range[1]) / vh) # H is image height
    anchor_args['D'] = int((cav_lidar_range[5] - cav_lidar_range[2]) / vd)

    param['box_align_pre_calc']['stage1_postprocessor_config'].update({'anchor_args': anchor_args})

    return param

opencood/hypes_yaml/yaml_utils.py
MediaBrain-SJTU/CoBEVFlow
import torch
import torch.nn as nn

from opencood.data_utils.post_processor.voxel_postprocessor import VoxelPostprocessor
from opencood.pcdet_utils.iou3d_nms.iou3d_nms_utils import aligned_boxes_iou3d_gpu


class CiassdLoss(nn.Module):
    def __init__(self, args):
        super(CiassdLoss, self).__init__()
        self.pos_cls_weight = args['pos_cls_weight']
        self.encode_rad_error_by_sin = args['encode_rad_error_by_sin']
        self.cls = args['cls']
        self.reg = args['reg']
        self.iou = args['iou']
        self.dir = args['dir']
        self.loss_dict = {}
        ##
        self.num_cls = 2
        self.box_codesize = 7

    def forward(self, output_dict, label_dict):
        """
        Parameters
        ----------
        output_dict : dict
        target_dict : dict
        """
        preds_dict = output_dict['preds_dict_stage1']
        if 'stage1' in label_dict.keys():
            target_dict = label_dict['stage1']
        else:
            target_dict = label_dict
        if 'record_len' in output_dict:
            batch_size = int(output_dict['record_len'].sum())
        else:
            batch_size = output_dict['batch_size']

        # ########
        # pred = torch.sigmoid(preds_dict['cls_preds'][0]).sum(dim=0).cpu().detach().numpy()
        # tagt_pos = target_dict['pos_equal_one'][0].sum(dim=-1).cpu().detach().numpy()
        # tagt_neg = target_dict['neg_equal_one'][0].sum(dim=-1).cpu().detach().numpy()
        # import matplotlib.pyplot as plt
        #
        # fig = plt.figure(figsize=(18, 8))
        # ax1 = fig.add_subplot(3, 1, 1)
        # ax1.imshow(pred, cmap='cool')
        # ax2 = fig.add_subplot(3, 1, 2)
        # ax2.imshow(tagt_pos, cmap='cool')
        # ax3 = fig.add_subplot(3, 1, 3)
        # ax3.imshow(tagt_neg, cmap='cool')
        # plt.show()
        # plt.close()
        # #########

        cls_labls = target_dict['pos_equal_one'].view(batch_size, -1,  self.num_cls - 1)
        positives = cls_labls > 0
        negatives = target_dict['neg_equal_one'].view(batch_size, -1,  self.num_cls - 1) > 0
        cared = torch.logical_or(positives, negatives)
        cls_labls = cls_labls * cared.type_as(cls_labls)
        # num_normalizer = cared.sum(1, keepdim=True)
        pos_normalizer = positives.sum(1, keepdim=True).float()

        # cls loss
        cls_preds = preds_dict["cls_preds"].permute(0, 2, 3, 1).contiguous() \
                    .view(batch_size, -1,  self.num_cls - 1)
        cls_weights = positives * self.pos_cls_weight + negatives * 1.0
        cls_weights /= torch.clamp(pos_normalizer, min=1.0)
        cls_loss = sigmoid_focal_loss(cls_preds, cls_labls, weights=cls_weights, **self.cls)
        cls_loss_reduced = cls_loss.sum() * self.cls['weight'] / batch_size

        # reg loss
        reg_weights = positives / torch.clamp(pos_normalizer, min=1.0)
        reg_preds = preds_dict['box_preds'].permute(0, 2, 3, 1).contiguous().view(batch_size, -1, self.box_codesize)
        reg_targets = target_dict['targets'].view(batch_size, -1, self.box_codesize)
        if self.encode_rad_error_by_sin:
            reg_preds, reg_targets = add_sin_difference(reg_preds, reg_targets)
        reg_loss = weighted_smooth_l1_loss(reg_preds, reg_targets, weights=reg_weights, sigma=self.reg['sigma'])
        reg_loss_reduced = reg_loss.sum() * self.reg['weight'] / batch_size

        from icecream import ic
        # dir loss
        dir_targets = get_direction_target(reg_targets, output_dict['anchor_box'])
        dir_logits = preds_dict["dir_cls_preds"].permute(0, 2, 3, 1).contiguous().view(batch_size, -1, 2)
        dir_weights = (cls_labls > 0).type_as(dir_logits).view(batch_size, -1)
        dir_weights /= torch.clamp(dir_weights.sum(-1, keepdim=True), min=1.0)  # [8, 70400], averaged in sample.
        dir_loss = softmax_cross_entropy_with_logits(dir_logits.view(-1, self.num_cls), dir_targets.view(-1, self.num_cls))
        dir_loss = dir_loss.view(dir_weights.shape) * dir_weights
        dir_loss_reduced = dir_loss.sum() * self.dir['weight'] / batch_size

        # iou loss
        iou_preds = preds_dict["iou_preds"].permute(0, 2, 3, 1).contiguous()
        pos_pred_mask = reg_weights.squeeze(dim=-1) > 0 # (4, 70400)
        iou_pos_preds = iou_preds.view(batch_size, -1)[pos_pred_mask]
        boxes3d_pred = VoxelPostprocessor.delta_to_boxes3d(preds_dict['box_preds'].permute(0, 2, 3, 1).contiguous().detach(),
                                                           output_dict['anchor_box'])[pos_pred_mask]
        boxes3d_tgt = VoxelPostprocessor.delta_to_boxes3d(target_dict['targets'],
                                                          output_dict['anchor_box'])[pos_pred_mask]
        iou_weights = reg_weights[pos_pred_mask].view(-1)
        iou_pos_targets = aligned_boxes_iou3d_gpu(boxes3d_pred.float()[:, [0, 1, 2, 5, 4, 3, 6]],
                                                  boxes3d_tgt.float()[:, [0, 1, 2, 5, 4, 3, 6]]).detach().squeeze()
        iou_pos_targets = 2 * iou_pos_targets.view(-1) - 1
        iou_loss = weighted_smooth_l1_loss(iou_pos_preds, iou_pos_targets, weights=iou_weights, sigma=self.iou['sigma'])

        iou_loss_reduced = iou_loss.sum() * self.iou['weight'] / batch_size

        loss = cls_loss_reduced + reg_loss_reduced + dir_loss_reduced + iou_loss_reduced

        self.loss_dict.update({
            'total_loss': loss,
            'cls_loss': cls_loss_reduced,
            'reg_loss': reg_loss_reduced,
            'dir_loss': dir_loss_reduced,
            'iou_loss': iou_loss_reduced
        })

        return loss

    def logging(self, epoch, batch_id, batch_len, writer = None):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        total_loss = self.loss_dict['total_loss']
        reg_loss = self.loss_dict['reg_loss']
        cls_loss = self.loss_dict['cls_loss']
        dir_loss = self.loss_dict['dir_loss']
        iou_loss = self.loss_dict['iou_loss']
        if (batch_id + 1) % 10 == 0:
            print("[epoch %d][%d/%d], || Loss: %.4f || Cls: %.4f"
                  " || Loc: %.4f || Dir: %.4f || Iou: %.4f" % (
                      epoch, batch_id + 1, batch_len,
                      total_loss.item(), cls_loss.item(), reg_loss.item(), dir_loss.item(), iou_loss.item()))
        if writer is not None:
            writer.add_scalar('Regression_loss', reg_loss.item(),
                            epoch*batch_len + batch_id)
            writer.add_scalar('Confidence_loss', cls_loss.item(),
                            epoch*batch_len + batch_id)
            writer.add_scalar('Direction_loss', dir_loss.item(),
                            epoch*batch_len + batch_id)
            writer.add_scalar('Iou_loss', iou_loss.item(),
                            epoch*batch_len + batch_id)


def add_sin_difference(boxes1, boxes2):
    rad_pred_encoding = torch.sin(boxes1[..., -1:]) * torch.cos(boxes2[..., -1:])   # ry -> sin(pred_ry)*cos(gt_ry)
    rad_gt_encoding = torch.cos(boxes1[..., -1:]) * torch.sin(boxes2[..., -1:])     # ry -> cos(pred_ry)*sin(gt_ry)
    res_boxes1 = torch.cat([boxes1[..., :-1], rad_pred_encoding], dim=-1)
    res_boxes2 = torch.cat([boxes2[..., :-1], rad_gt_encoding], dim=-1)
    return res_boxes1, res_boxes2


def get_direction_target(reg_targets, anchors, one_hot=True, dir_offset=0.0):
    """
    Generate targets for bounding box direction classification.

    Parameters
    ----------
    anchors: torch.Tensor
        shape as (H*W*2, 7) or (H, W, 2, 7)
    reg_targets: torch.Tensor
        shape as (B, H*W*2, 7)

    Returns
    -------
    dir_cls_targets : torch.Tensor
        [batch_size, w*h*num_anchor_per_pos, 2]
    """
    batch_size = reg_targets.shape[0]
    anchors = anchors.view(1,  -1, anchors.shape[-1]).repeat(batch_size, 1, 1)
    rot_gt = reg_targets[..., -1] + anchors[..., -1]  # [4, 70400]
    dir_cls_targets = ((rot_gt - dir_offset) > 0).long()  # [4, 70400]
    if one_hot:
        dir_cls_targets = one_hot_f(dir_cls_targets, 2, dtype=anchors.dtype)
    return dir_cls_targets


def one_hot_f(tensor, depth, dim=-1, on_value=1.0, dtype=torch.float32):
    tensor_onehot = torch.zeros(*list(tensor.shape), depth, dtype=dtype, device=tensor.device) # [4, 70400, 2]
    tensor_onehot.scatter_(dim, tensor.unsqueeze(dim).long(), on_value)                        # [4, 70400, 2]
    return tensor_onehot


def sigmoid_focal_loss(preds, targets, weights=None, **kwargs):
    assert 'gamma' in kwargs and 'alpha' in kwargs
    # sigmoid cross entropy with logits
    # more details: https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits
    per_entry_cross_ent = torch.clamp(preds, min=0) - preds * targets.type_as(preds)
    per_entry_cross_ent += torch.log1p(torch.exp(-torch.abs(preds)))
    # focal loss
    prediction_probabilities = torch.sigmoid(preds)
    p_t = (targets * prediction_probabilities) + ((1 - targets) * (1 - prediction_probabilities))
    modulating_factor = torch.pow(1.0 - p_t, kwargs['gamma'])
    alpha_weight_factor = targets * kwargs['alpha'] + (1 - targets) * (1 - kwargs['alpha'])

    loss = modulating_factor * alpha_weight_factor * per_entry_cross_ent
    if weights is not None:
        loss *= weights
    return loss


def softmax_cross_entropy_with_logits(logits, labels):
    param = list(range(len(logits.shape)))
    transpose_param = [0] + [param[-1]] + param[1:-1]
    logits = logits.permute(*transpose_param)
    loss_ftor = torch.nn.CrossEntropyLoss(reduction="none")
    loss = loss_ftor(logits, labels.max(dim=-1)[1])
    return loss


def weighted_smooth_l1_loss(preds, targets, sigma=3.0, weights=None):
    diff = preds - targets
    abs_diff = torch.abs(diff)
    abs_diff_lt_1 = torch.le(abs_diff, 1 / (sigma ** 2)).type_as(abs_diff)
    loss = abs_diff_lt_1 * 0.5 * torch.pow(abs_diff * sigma, 2) + \
               (abs_diff - 0.5 / (sigma ** 2)) * (1.0 - abs_diff_lt_1)
    if weights is not None:
        loss *= weights
    return loss
opencood/loss/ciassd_loss.py
MediaBrain-SJTU/CoBEVFlow
# sizhewei @ 2023/04/17
# for flowmap loss
#

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

class FlowMapLoss(nn.Module):
    def __init__(self, args):
        super(FlowMapLoss, self).__init__()
    def forward(self, output_dict, foo=None):
        self.loss = output_dict['flow_recon_loss']
        return self.loss
    def logging(self, epoch, batch_id, batch_len, writer = None):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        total_loss = self.loss
        print_msg = ("[epoch %d][%d/%d], || Loss: %.4f " 
        % (epoch, batch_id + 1, batch_len,total_loss.item()))
        print(print_msg)
        if writer is not None:
            writer.add_scalar('loss', total_loss.item(), epoch * batch_len + batch_id)
opencood/loss/flowmap_loss.py
MediaBrain-SJTU/CoBEVFlow
import torch
from torch import nn
import numpy as np
from opencood.loss.ciassd_loss import CiassdLoss, weighted_smooth_l1_loss


class FpvrcnnLoss(nn.Module):
    def __init__(self, args):
        super(FpvrcnnLoss, self).__init__()
        self.ciassd_loss = CiassdLoss(args['stage1'])
        self.cls = args['stage2']['cls']
        self.reg = args['stage2']['reg']
        self.iou = args['stage2']['iou']
        self.loss_dict = {}

    def forward(self, output_dict, label_dict):
        """
        Parameters
        ----------
        output_dict : dict
        target_dict : dict
        """
        ciassd_loss = self.ciassd_loss(output_dict, label_dict)

        # only update ciassd if no bbox is detected in the first stage
        if 'fpvrcnn_out' not in output_dict:
            self.loss_dict = {
                'loss': ciassd_loss,
            }
            return ciassd_loss

        # rcnn out
        rcnn_cls = output_dict['fpvrcnn_out']['rcnn_cls'].view(1, -1, 1)
        rcnn_iou = output_dict['fpvrcnn_out']['rcnn_iou'].view(1, -1, 1)
        rcnn_reg = output_dict['fpvrcnn_out']['rcnn_reg'].view(1, -1, 7)

        tgt_cls = output_dict['rcnn_label_dict']['cls_tgt'].view(1, -1, 1)
        tgt_iou = output_dict['rcnn_label_dict']['iou_tgt'].view(1, -1, 1)
        tgt_reg = output_dict['rcnn_label_dict']['reg_tgt'].view(1, -1, 7)

        pos_norm = tgt_cls.sum()
        # cls loss
        loss_cls = weighted_sigmoid_binary_cross_entropy(rcnn_cls, tgt_cls)

        # iou loss
        # TODO: also count the negative samples
        loss_iou = weighted_smooth_l1_loss(rcnn_iou, tgt_iou,
                                           weights=tgt_cls).mean()

        # regression loss
        # Target resampling : Generate a weights mask to force the regressor concentrate on low iou predictions
        # sample 50% with iou>0.7 and 50% < 0.7
        weights = torch.ones(tgt_iou.shape, device=tgt_iou.device)
        weights[tgt_cls == 0] = 0
        neg = torch.logical_and(tgt_iou < 0.7, tgt_cls != 0)
        pos = torch.logical_and(tgt_iou >= 0.7, tgt_cls != 0)
        num_neg = int(neg.sum(dim=1))
        num_pos = int(pos.sum(dim=1))
        num_pos_smps = max(num_neg, 2)
        pos_indices = torch.where(pos)[1]
        not_selsected = torch.randperm(num_pos)[:num_pos - num_pos_smps]
        # not_selsected_indices = pos_indices[not_selsected]
        weights[:, pos_indices[not_selsected]] = 0
        loss_reg = weighted_smooth_l1_loss(rcnn_reg, tgt_reg,
                                           weights=weights / max(weights.sum(),
                                                                 1)).sum()

        loss_cls_reduced = loss_cls * self.cls['weight']
        loss_iou_reduced = loss_iou * self.iou['weight']
        loss_reg_reduced = loss_reg * self.reg['weight']

        # if torch.isnan(loss_reg_reduced):
        #     print('debug')

        rcnn_loss = loss_cls_reduced + loss_iou_reduced + loss_reg_reduced
        loss = rcnn_loss + ciassd_loss

        self.loss_dict.update({
            'loss': loss,
            'rcnn_loss': rcnn_loss,
            'cls_loss': loss_cls_reduced,
            'iou_loss': loss_iou_reduced,
            'reg_loss': loss_reg_reduced,
        })

        return loss

    def logging(self, epoch, batch_id, batch_len, writer=None):
        """
        Print out  the loss function for current iteration.
        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        ciassd_loss_dict = self.ciassd_loss.loss_dict
        ciassd_total_loss = ciassd_loss_dict['total_loss']
        reg_loss = ciassd_loss_dict['reg_loss']
        cls_loss = ciassd_loss_dict['cls_loss']
        dir_loss = ciassd_loss_dict['dir_loss']
        iou_loss = ciassd_loss_dict['iou_loss']

        if (batch_id + 1) % 10 == 0:
            str_to_print = "[epoch %d][%d/%d], || Loss: %.4f || Ciassd: %.4f " \
                           "|| Cls1: %.4f || Loc1: %.4f || Dir1: %.4f || Iou1: %.4f" % (
                               epoch, batch_id + 1, batch_len, self.loss_dict['loss'],
                               ciassd_total_loss.item(), cls_loss.item(), reg_loss.item(),
                               dir_loss.item(), iou_loss.item(),
                               )
            if 'rcnn_loss' in self.loss_dict:
                str_to_print += " || Rcnn: %.4f || Cls2: %.4f || Loc2: %.4f || Iou2: %.4f" % (
                        self.loss_dict['rcnn_loss'],
                        self.loss_dict['cls_loss'].item(),
                        self.loss_dict['reg_loss'].item(),
                        self.loss_dict['iou_loss'].item(),
                    )
            print(str_to_print)
            
        if writer:
            writer.add_scalar('Ciassd_regression_loss', reg_loss.item(),
                            epoch * batch_len + batch_id)
            writer.add_scalar('Ciassd_Confidence_loss', cls_loss.item(),
                            epoch * batch_len + batch_id)
            writer.add_scalar('Ciassd_Direction_loss', dir_loss.item(),
                            epoch * batch_len + batch_id)
            writer.add_scalar('Ciassd_Iou_loss', iou_loss.item(),
                            epoch * batch_len + batch_id)
            writer.add_scalar('Ciassd_loss', ciassd_total_loss.item(),
                            epoch * batch_len + batch_id)
            if 'rcnn_loss' in self.loss_dict:
                writer.add_scalar('Rcnn_regression_loss',
                                self.loss_dict['reg_loss'].item(),
                                epoch * batch_len + batch_id)
                writer.add_scalar('Rcnn_Confidence_loss',
                                self.loss_dict['cls_loss'].item(),
                                epoch * batch_len + batch_id)
                writer.add_scalar('Rcnn_Iou_loss',
                                self.loss_dict['iou_loss'].item(),
                                epoch * batch_len + batch_id)
                writer.add_scalar('Rcnn_loss', self.loss_dict['rcnn_loss'].item(),
                                epoch * batch_len + batch_id)
                writer.add_scalar('Total_loss', self.loss_dict['loss'].item(),
                                epoch * batch_len + batch_id)


def weighted_sigmoid_binary_cross_entropy(preds, tgts, weights=None,
                                          class_indices=None):
    if weights is not None:
        weights = weights.unsqueeze(-1)
    if class_indices is not None:
        weights *= (
            indices_to_dense_vector(class_indices, preds.shape[2])
                .view(1, 1, -1)
                .type_as(preds)
        )
    per_entry_cross_ent = nn.functional.binary_cross_entropy_with_logits(preds,
                                                                         tgts,
                                                                         weights)
    return per_entry_cross_ent


def indices_to_dense_vector(
        indices, size, indices_value=1.0, default_value=0, dtype=np.float32
):
    """Creates dense vector with indices set to specific value and rest to zeros.
    This function exists because it is unclear if it is safe to use
        tf.sparse_to_dense(indices, [size], 1, validate_indices=False)
    with indices which are not ordered.
    This function accepts a dynamic size (e.g. tf.shape(tensor)[0])
    Args:
        indices: 1d Tensor with integer indices which are to be set to
            indices_values.
        size: scalar with size (integer) of output Tensor.
        indices_value: values of elements specified by indices in the output vector
        default_value: values of other elements in the output vector.
        dtype: data type.
    Returns:
        dense 1D Tensor of shape [size] with indices set to indices_values and the
            rest set to default_value.
    """
    dense = torch.zeros(size).fill_(default_value)
    dense[indices] = indices_value

    return dense
opencood/loss/fpvrcnn_loss.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


from functools import reduce

import torch
import torch.nn as nn
import torch.nn.functional as F


class PixorLoss(nn.Module):
    def __init__(self, args):
        super(PixorLoss, self).__init__()
        self.alpha = args["alpha"]
        self.beta = args["beta"]
        self.loss_dict = {}

    def forward(self, output_dict, target_dict):
        """
        Compute loss for pixor network
        Parameters
        ----------
        output_dict : dict
           The dictionary that contains the output.

        target_dict : dict
           The dictionary that contains the target.

        Returns
        -------
        total_loss : torch.Tensor
            Total loss.

        """
        targets = target_dict["label_map"]
        cls_preds, loc_preds = output_dict["cls"], output_dict["reg"]

        cls_targets, loc_targets = targets.split([1, 6], dim=1)
        pos_count = cls_targets.sum()
        neg_count = (cls_targets == 0).sum()
        w1, w2 = neg_count / (pos_count + neg_count), pos_count / (
                    pos_count + neg_count)
        weights = torch.ones_like(cls_preds.reshape(-1))
        weights[cls_targets.reshape(-1) == 1] = w1
        weights[cls_targets.reshape(-1) == 0] = w2
        # cls_targets = cls_targets.float()
        # cls_loss = F.binary_cross_entropy_with_logits(input=cls_preds.reshape(-1), target=cls_targets.reshape(-1), weight=weights,
        #                                               reduction='mean')
        cls_loss = F.binary_cross_entropy_with_logits(
            input=cls_preds, target=cls_targets,
            reduction='mean')
        pos_pixels = cls_targets.sum()

        loc_loss = F.smooth_l1_loss(cls_targets * loc_preds,
                                    cls_targets * loc_targets,
                                    reduction='sum')
        loc_loss = loc_loss / pos_pixels if pos_pixels > 0 else loc_loss

        total_loss = self.alpha * cls_loss + self.beta * loc_loss

        self.loss_dict.update({'total_loss': total_loss,
                               'reg_loss': loc_loss,
                               'cls_loss': cls_loss})

        return total_loss

    def logging(self, epoch, batch_id, batch_len, writer = None):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        total_loss = self.loss_dict['total_loss']
        reg_loss = self.loss_dict['reg_loss']
        cls_loss = self.loss_dict['cls_loss']

        print("[epoch %d][%d/%d], || Loss: %.4f || cls Loss: %.4f"
              " || reg Loss: %.4f" % (
                  epoch, batch_id + 1, batch_len,
                  total_loss.item(), cls_loss.item(), reg_loss.item()))
                  
        if not writer is None:
            writer.add_scalar('Regression_loss', reg_loss.item(),
                            epoch * batch_len + batch_id)
            writer.add_scalar('Confidence_loss', cls_loss.item(),
                            epoch * batch_len + batch_id)


def test():
    torch.manual_seed(0)
    loss = PixorLoss(None)
    pred = torch.sigmoid(torch.randn(1, 7, 2, 3))
    label = torch.zeros(1, 7, 2, 3)
    loss = loss(pred, label)
    print(loss)


if __name__ == "__main__":
    test()

opencood/loss/pixor_loss.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu
# Add direction classification loss
# The originally point_pillar_loss.py, can not determine if the box heading is opposite to the GT.


import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from opencood.utils.common_utils import limit_period


class WeightedSmoothL1Loss(nn.Module):
    """
    Code-wise Weighted Smooth L1 Loss modified based on fvcore.nn.smooth_l1_loss
    https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/smooth_l1_loss.py
                  | 0.5 * x ** 2 / beta   if abs(x) < beta
    smoothl1(x) = |
                  | abs(x) - 0.5 * beta   otherwise,
    where x = input - target.
    """
    def __init__(self, beta: float = 1.0 / 9.0, code_weights: list = None):
        """
        Args:
            beta: Scalar float.
                L1 to L2 change point.
                For beta values < 1e-5, L1 loss is computed.
            code_weights: (#codes) float list if not None.
                Code-wise weights.
        """
        super(WeightedSmoothL1Loss, self).__init__()
        self.beta = beta
        if code_weights is not None:
            self.code_weights = np.array(code_weights, dtype=np.float32)
            self.code_weights = torch.from_numpy(self.code_weights).cuda()

    @staticmethod
    def smooth_l1_loss(diff, beta):
        if beta < 1e-5:
            loss = torch.abs(diff)
        else:
            n = torch.abs(diff)
            loss = torch.where(n < beta, 0.5 * n ** 2 / beta, n - 0.5 * beta)

        return loss

    def forward(self, input: torch.Tensor,
                target: torch.Tensor, weights: torch.Tensor = None):
        """
        Args:
            input: (B, #anchors, #codes) float tensor.
                Ecoded predicted locations of objects.
            target: (B, #anchors, #codes) float tensor.
                Regression targets.
            weights: (B, #anchors) float tensor if not None.

            #anchors = H * W * anchor_num

        Returns:
            loss: (B, #anchors) float tensor.
                Weighted smooth l1 loss without reduction.
        """
        target = torch.where(torch.isnan(target), input, target)  # ignore nan targets

        diff = input - target
        loss = self.smooth_l1_loss(diff, self.beta)

        # anchor-wise weighting
        if weights is not None:
            assert weights.shape[0] == loss.shape[0] and weights.shape[1] == loss.shape[1]
            loss = loss * weights.unsqueeze(-1)

        return loss


class PointPillarDirLoss(nn.Module):
    def __init__(self, args):
        super(PointPillarDirLoss, self).__init__()
        self.reg_loss_func = WeightedSmoothL1Loss()
        self.alpha = 0.25
        self.gamma = 2.0

        self.cls_weight = args['cls_weight']
        self.reg_coe = args['reg']

        self.dir_weight = args['dir_args']['dir_weight']
        self.dir_offset = args['dir_args']['args']['dir_offset']
        self.num_bins = args['dir_args']['args']['num_bins']
        anchor_yaw = np.deg2rad(np.array(args['dir_args']['anchor_yaw']))  # for direction classification
        self.anchor_yaw_map = torch.from_numpy(anchor_yaw).view(1,-1,1)  # [1,2,1]
        self.anchor_num = self.anchor_yaw_map.shape[1]


        self.loss_dict = {}

    def forward(self, output_dict, target_dict):
        """
        Parameters
        ----------
        output_dict : dict
        target_dict : dict
        """
        rm = output_dict['rm']  # [B, 14, 50, 176]
        psm = output_dict['psm'] # [B, 2, 50, 176]
        dm = output_dict['dm']
        targets = target_dict['targets']

        cls_preds = psm.permute(0, 2, 3, 1).contiguous() # N, C, H, W -> N, H, W, C

        box_cls_labels = target_dict['pos_equal_one']  # [B, 50, 176, 2]
        box_cls_labels = box_cls_labels.view(psm.shape[0], -1).contiguous()  # [B, 50*176*2]

        positives = box_cls_labels > 0
        negatives = box_cls_labels == 0
        negative_cls_weights = negatives * 1.0
        cls_weights = (negative_cls_weights + 1.0 * positives).float()
        reg_weights = positives.float()

        pos_normalizer = positives.sum(1, keepdim=True).float()
        reg_weights /= torch.clamp(pos_normalizer, min=1.0) # [N, H*W*anchor_num]
        cls_weights /= torch.clamp(pos_normalizer, min=1.0)
        cls_targets = box_cls_labels
        cls_targets = cls_targets.unsqueeze(dim=-1)

        cls_targets = cls_targets.squeeze(dim=-1)
        one_hot_targets = torch.zeros(
            *list(cls_targets.shape), 2,
            dtype=cls_preds.dtype, device=cls_targets.device
        )
        one_hot_targets.scatter_(-1, cls_targets.unsqueeze(dim=-1).long(), 1.0)
        cls_preds = cls_preds.view(psm.shape[0], -1, 1)
        one_hot_targets = one_hot_targets[..., 1:]

        cls_loss_src = self.cls_loss_func(cls_preds,
                                          one_hot_targets,
                                          weights=cls_weights)  # [N, M]
        cls_loss = cls_loss_src.sum() / psm.shape[0]
        conf_loss = cls_loss * self.cls_weight

        # regression
        rm = rm.permute(0, 2, 3, 1).contiguous()
        rm = rm.view(rm.size(0), -1, 7)
        targets = targets.view(targets.size(0), -1, 7)
        box_preds_sin, reg_targets_sin = self.add_sin_difference(rm,
                                                                 targets)
        loc_loss_src =\
            self.reg_loss_func(box_preds_sin,
                               reg_targets_sin,
                               weights=reg_weights)
        reg_loss = loc_loss_src.sum() / rm.shape[0]
        reg_loss *= self.reg_coe

        ######## direction ##########
        dir_targets = self.get_direction_target(targets)
        N =  output_dict["dm"].shape[0]
        dir_logits = output_dict["dm"].permute(0, 2, 3, 1).contiguous().view(N, -1, 2) # [N, H*W*#anchor, 2]


        dir_loss = softmax_cross_entropy_with_logits(dir_logits.view(-1, self.anchor_num), dir_targets.view(-1, self.anchor_num)) 

        dir_loss = dir_loss.view(dir_logits.shape[:2]) * reg_weights # [N, H*W*anchor_num]

        dir_loss = dir_loss.sum() * self.dir_weight / N


        total_loss = reg_loss + conf_loss + dir_loss

        self.loss_dict.update({'total_loss': total_loss,
                               'reg_loss': reg_loss,
                               'conf_loss': conf_loss,
                               'dir_loss': dir_loss})

        return total_loss

    def cls_loss_func(self, input: torch.Tensor,
                      target: torch.Tensor,
                      weights: torch.Tensor):
        """
        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets
            weights: (B, #anchors) float tensor.
                Anchor-wise weights.

        Returns:
            weighted_loss: (B, #anchors, #classes) float tensor after weighting.
        """
        pred_sigmoid = torch.sigmoid(input)
        alpha_weight = target * self.alpha + (1 - target) * (1 - self.alpha)
        pt = target * (1.0 - pred_sigmoid) + (1.0 - target) * pred_sigmoid
        focal_weight = alpha_weight * torch.pow(pt, self.gamma)

        bce_loss = self.sigmoid_cross_entropy_with_logits(input, target)

        loss = focal_weight * bce_loss

        if weights.shape.__len__() == 2 or \
                (weights.shape.__len__() == 1 and target.shape.__len__() == 2):
            weights = weights.unsqueeze(-1)

        assert weights.shape.__len__() == loss.shape.__len__()

        return loss * weights

    @staticmethod
    def sigmoid_cross_entropy_with_logits(input: torch.Tensor, target: torch.Tensor):
        """ PyTorch Implementation for tf.nn.sigmoid_cross_entropy_with_logits:
            max(x, 0) - x * z + log(1 + exp(-abs(x))) in
            https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits

        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets

        Returns:
            loss: (B, #anchors, #classes) float tensor.
                Sigmoid cross entropy loss without reduction
        """
        loss = torch.clamp(input, min=0) - input * target + \
               torch.log1p(torch.exp(-torch.abs(input)))
        return loss

    @staticmethod
    def add_sin_difference(boxes1, boxes2, dim=6):
        assert dim != -1
        rad_pred_encoding = torch.sin(boxes1[..., dim:dim + 1]) * \
                            torch.cos(boxes2[..., dim:dim + 1])
        rad_tg_encoding = torch.cos(boxes1[..., dim:dim + 1]) * \
                          torch.sin(boxes2[..., dim:dim + 1])

        boxes1 = torch.cat([boxes1[..., :dim], rad_pred_encoding,
                            boxes1[..., dim + 1:]], dim=-1)
        boxes2 = torch.cat([boxes2[..., :dim], rad_tg_encoding,
                            boxes2[..., dim + 1:]], dim=-1)
        return boxes1, boxes2

    def get_direction_target(self, reg_targets):
        """
        Args:
            reg_targets:  [N, H * W * #anchor_num, 7]
                The last term is (theta_gt - theta_a)
        
        Returns:
            dir_targets:
                theta_gt: [N, H * W * #anchor_num, NUM_BIN] 
                NUM_BIN = 2
        """
        # (1, 2, 1)
        H_times_W_times_anchor_num = reg_targets.shape[1]
        anchor_map = self.anchor_yaw_map.repeat(1, H_times_W_times_anchor_num//self.anchor_num, 1).to(reg_targets.device) # [1, H * W * #anchor_num, 1]
        rot_gt = reg_targets[..., -1] + anchor_map[..., -1] # [N, H*W*anchornum]
        offset_rot = limit_period(rot_gt - self.dir_offset, 0, 2 * np.pi)
        dir_cls_targets = torch.floor(offset_rot / (2 * np.pi / self.num_bins)).long()  # [N, H*W*anchornum]
        dir_cls_targets = torch.clamp(dir_cls_targets, min=0, max=self.num_bins - 1)
        # one_hot:
        # if rot_gt > 0, then the label is 1, then the regression target is [0, 1]
        dir_cls_targets = one_hot_f(dir_cls_targets, self.num_bins)
        return dir_cls_targets



    def logging(self, epoch, batch_id, batch_len, writer = None):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        total_loss = self.loss_dict['total_loss']
        reg_loss = self.loss_dict['reg_loss']
        conf_loss = self.loss_dict['conf_loss']
        dir_loss = self.loss_dict['dir_loss']

        print("[epoch %d][%d/%d], || Loss: %.4f || Conf Loss: %.4f"
              " || Loc Loss: %.4f || Dir Loss: %.4f" % (
                  epoch, batch_id + 1, batch_len,
                  total_loss.item(), conf_loss.item(), reg_loss.item(), dir_loss.item()))

        if not writer is None:
            writer.add_scalar('Regression_loss', reg_loss.item(),
                            epoch*batch_len + batch_id)
            writer.add_scalar('Confidence_loss', conf_loss.item(),
                            epoch*batch_len + batch_id)
            writer.add_scalar('Dir_loss', dir_loss.item(),
                            epoch*batch_len + batch_id)

def one_hot_f(tensor, num_bins, dim=-1, on_value=1.0, dtype=torch.float32):
    tensor_onehot = torch.zeros(*list(tensor.shape), num_bins, dtype=dtype, device=tensor.device) 
    tensor_onehot.scatter_(dim, tensor.unsqueeze(dim).long(), on_value)                    
    return tensor_onehot

def softmax_cross_entropy_with_logits(logits, labels):
    param = list(range(len(logits.shape)))
    transpose_param = [0] + [param[-1]] + param[1:-1]
    logits = logits.permute(*transpose_param)
    loss_ftor = torch.nn.CrossEntropyLoss(reduction="none")
    loss = loss_ftor(logits, labels.max(dim=-1)[1])
    return loss

opencood/loss/point_pillar_dir_loss.py
MediaBrain-SJTU/CoBEVFlow
# Author: Yifan Lu, Yiming Li, Dekun Ma
# add KD loss

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class WeightedSmoothL1Loss(nn.Module):
    """
    Code-wise Weighted Smooth L1 Loss modified based on fvcore.nn.smooth_l1_loss
    https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/smooth_l1_loss.py
                  | 0.5 * x ** 2 / beta   if abs(x) < beta
    smoothl1(x) = |
                  | abs(x) - 0.5 * beta   otherwise,
    where x = input - target.
    """
    def __init__(self, beta: float = 1.0 / 9.0, code_weights: list = None):
        """
        Args:
            beta: Scalar float.
                L1 to L2 change point.
                For beta values < 1e-5, L1 loss is computed.
            code_weights: (#codes) float list if not None.
                Code-wise weights.
        """
        super(WeightedSmoothL1Loss, self).__init__()
        self.beta = beta
        if code_weights is not None:
            self.code_weights = np.array(code_weights, dtype=np.float32)
            self.code_weights = torch.from_numpy(self.code_weights).cuda()

    @staticmethod
    def smooth_l1_loss(diff, beta):
        if beta < 1e-5:
            loss = torch.abs(diff)
        else:
            n = torch.abs(diff)
            loss = torch.where(n < beta, 0.5 * n ** 2 / beta, n - 0.5 * beta)

        return loss

    def forward(self, input: torch.Tensor,
                target: torch.Tensor, weights: torch.Tensor = None):
        """
        Args:
            input: (B, #anchors, #codes) float tensor.
                Ecoded predicted locations of objects.
            target: (B, #anchors, #codes) float tensor.
                Regression targets.
            weights: (B, #anchors) float tensor if not None.

        Returns:
            loss: (B, #anchors) float tensor.
                Weighted smooth l1 loss without reduction.
        """
        target = torch.where(torch.isnan(target), input, target)  # ignore nan targets

        diff = input - target
        loss = self.smooth_l1_loss(diff, self.beta)

        # anchor-wise weighting
        if weights is not None:
            assert weights.shape[0] == loss.shape[0] and weights.shape[1] == loss.shape[1]
            loss = loss * weights.unsqueeze(-1)

        return loss


class PointPillarDiscoNetLoss(nn.Module):
    def __init__(self, args):
        super(PointPillarDiscoNetLoss, self).__init__()
        self.reg_loss_func = WeightedSmoothL1Loss()
        self.alpha = 0.25
        self.gamma = 2.0
        self.decoder_kd = args['decoder_kd'] # whether kd for decoded feature


        self.cls_weight = args['cls_weight']
        self.reg_coe = args['reg']
        self.kd_weight = args['kd_weight']
        self.loss_dict = {}

    def forward(self, output_dict, target_dict):
        """
        Parameters
        ----------
        output_dict : dict
        target_dict : dict
        """
        rm = output_dict['rm']  # [B, 14, 50, 176]
        psm = output_dict['psm'] # [B, 2, 50, 176]
        feature = output_dict['feature']

        targets = target_dict['targets']

        teacher_rm = output_dict['teacher_rm']
        teather_psm = output_dict['teacher_psm']
        teacher_feature = output_dict['teacher_feature']

        ### KL loss ###
        kl_loss_mean = nn.KLDivLoss(size_average=True, reduce=True)

        N, C, H, W = teacher_feature.shape
        teacher_feature = teacher_feature.permute(0,2,3,1).reshape(N*H*W, C)
        student_feature = feature.permute(0,2,3,1).reshape(N*H*W, C)
        kd_loss_feature = kl_loss_mean(
                F.log_softmax(student_feature, dim=1), F.softmax(teacher_feature, dim=1)
            )
        
        kd_loss = kd_loss_feature

        if self.decoder_kd:
            N, C, H, W = teacher_rm.shape
            teacher_rm = teacher_rm.permute(0,2,3,1).reshape(N*H*W, C)
            student_rm = rm.permuate(0,2,3,1).reshape(N*H*W, C)
            kd_loss_rm = kl_loss_mean(
                    F.log_softmax(student_rm, dim=1), F.softmax(teacher_rm, dim=1)
                )

            N, C, H, W = teacher_psm.shape
            teacher_psm = teather_psm.permute(0,2,3,1).reshape(N*H*W, C)
            student_psm = psm.permuate(0,2,3,1).reshape(N*H*W, C)
            kd_loss_psm = kl_loss_mean(
                    F.log_softmax(student_psm, dim=1), F.softmax(teacher_psm, dim=1)
                )

            kd_loss += kd_loss_rm + kd_loss_psm


        cls_preds = psm.permute(0, 2, 3, 1).contiguous() # N, C, H, W -> N, H, W, C

        box_cls_labels = target_dict['pos_equal_one']  # [B, 50, 176, 2]
        box_cls_labels = box_cls_labels.view(psm.shape[0], -1).contiguous()

        positives = box_cls_labels > 0
        negatives = box_cls_labels == 0
        negative_cls_weights = negatives * 1.0
        cls_weights = (negative_cls_weights + 1.0 * positives).float()
        reg_weights = positives.float()

        pos_normalizer = positives.sum(1, keepdim=True).float()
        reg_weights /= torch.clamp(pos_normalizer, min=1.0)
        cls_weights /= torch.clamp(pos_normalizer, min=1.0)
        cls_targets = box_cls_labels
        cls_targets = cls_targets.unsqueeze(dim=-1)

        cls_targets = cls_targets.squeeze(dim=-1)
        one_hot_targets = torch.zeros(
            *list(cls_targets.shape), 2,
            dtype=cls_preds.dtype, device=cls_targets.device
        )
        one_hot_targets.scatter_(-1, cls_targets.unsqueeze(dim=-1).long(), 1.0)
        cls_preds = cls_preds.view(psm.shape[0], -1, 1)
        one_hot_targets = one_hot_targets[..., 1:]

        cls_loss_src = self.cls_loss_func(cls_preds,
                                          one_hot_targets,
                                          weights=cls_weights)  # [N, M]
        cls_loss = cls_loss_src.sum() / psm.shape[0]
        conf_loss = cls_loss * self.cls_weight

        # regression
        rm = rm.permute(0, 2, 3, 1).contiguous()
        rm = rm.view(rm.size(0), -1, 7)
        targets = targets.view(targets.size(0), -1, 7)
        box_preds_sin, reg_targets_sin = self.add_sin_difference(rm,
                                                                 targets)
        loc_loss_src =\
            self.reg_loss_func(box_preds_sin,
                               reg_targets_sin,
                               weights=reg_weights)
        reg_loss = loc_loss_src.sum() / rm.shape[0]
        reg_loss *= self.reg_coe

        kd_loss =  self.kd_weight * kd_loss

        total_loss = reg_loss + conf_loss + kd_loss

        self.loss_dict.update({'total_loss': total_loss,
                               'reg_loss': reg_loss,
                               'conf_loss': conf_loss,
                               'kd_loss': kd_loss})

        return total_loss

    def cls_loss_func(self, input: torch.Tensor,
                      target: torch.Tensor,
                      weights: torch.Tensor):
        """
        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets
            weights: (B, #anchors) float tensor.
                Anchor-wise weights.

        Returns:
            weighted_loss: (B, #anchors, #classes) float tensor after weighting.
        """
        pred_sigmoid = torch.sigmoid(input)
        alpha_weight = target * self.alpha + (1 - target) * (1 - self.alpha)
        pt = target * (1.0 - pred_sigmoid) + (1.0 - target) * pred_sigmoid
        focal_weight = alpha_weight * torch.pow(pt, self.gamma)

        bce_loss = self.sigmoid_cross_entropy_with_logits(input, target)

        loss = focal_weight * bce_loss

        if weights.shape.__len__() == 2 or \
                (weights.shape.__len__() == 1 and target.shape.__len__() == 2):
            weights = weights.unsqueeze(-1)

        assert weights.shape.__len__() == loss.shape.__len__()

        return loss * weights

    @staticmethod
    def sigmoid_cross_entropy_with_logits(input: torch.Tensor, target: torch.Tensor):
        """ PyTorch Implementation for tf.nn.sigmoid_cross_entropy_with_logits:
            max(x, 0) - x * z + log(1 + exp(-abs(x))) in
            https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits

        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets

        Returns:
            loss: (B, #anchors, #classes) float tensor.
                Sigmoid cross entropy loss without reduction
        """
        loss = torch.clamp(input, min=0) - input * target + \
               torch.log1p(torch.exp(-torch.abs(input)))
        return loss

    @staticmethod
    def add_sin_difference(boxes1, boxes2, dim=6):
        assert dim != -1
        rad_pred_encoding = torch.sin(boxes1[..., dim:dim + 1]) * \
                            torch.cos(boxes2[..., dim:dim + 1])
        rad_tg_encoding = torch.cos(boxes1[..., dim:dim + 1]) * \
                          torch.sin(boxes2[..., dim:dim + 1])

        boxes1 = torch.cat([boxes1[..., :dim], rad_pred_encoding,
                            boxes1[..., dim + 1:]], dim=-1)
        boxes2 = torch.cat([boxes2[..., :dim], rad_tg_encoding,
                            boxes2[..., dim + 1:]], dim=-1)
        return boxes1, boxes2


    def logging(self, epoch, batch_id, batch_len, writer = None):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        total_loss = self.loss_dict['total_loss']
        reg_loss = self.loss_dict['reg_loss']
        conf_loss = self.loss_dict['conf_loss']
        kd_loss = self.loss_dict['kd_loss']

        print("[epoch %d][%d/%d], || Loss: %.4f || Conf Loss: %.4f"
              " || Loc Loss: %.4f || Kd Loss %.4f"  % (
                  epoch, batch_id + 1, batch_len,
                  total_loss.item(), conf_loss.item(), reg_loss.item(), kd_loss.item()))

        if not writer is None:
            writer.add_scalar('Regression_loss', reg_loss.item(),
                            epoch*batch_len + batch_id)
            writer.add_scalar('Confidence_loss', conf_loss.item(),
                            epoch*batch_len + batch_id)
            writer.add_scalar('Kd_loss', kd_loss.item(),
                            epoch*batch_len + batch_id)
opencood/loss/point_pillar_disconet_loss.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: OpenPCDet, Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class WeightedSmoothL1Loss(nn.Module):
    """
    Code-wise Weighted Smooth L1 Loss modified based on fvcore.nn.smooth_l1_loss
    https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/smooth_l1_loss.py
                  | 0.5 * x ** 2 / beta   if abs(x) < beta
    smoothl1(x) = |
                  | abs(x) - 0.5 * beta   otherwise,
    where x = input - target.
    """
    def __init__(self, beta: float = 1.0 / 9.0, code_weights: list = None):
        """
        Args:
            beta: Scalar float.
                L1 to L2 change point.
                For beta values < 1e-5, L1 loss is computed.
            code_weights: (#codes) float list if not None.
                Code-wise weights.
        """
        super(WeightedSmoothL1Loss, self).__init__()
        self.beta = beta
        if code_weights is not None:
            self.code_weights = np.array(code_weights, dtype=np.float32)
            self.code_weights = torch.from_numpy(self.code_weights).cuda()

    @staticmethod
    def smooth_l1_loss(diff, beta):
        if beta < 1e-5:
            loss = torch.abs(diff)
        else:
            n = torch.abs(diff)
            loss = torch.where(n < beta, 0.5 * n ** 2 / beta, n - 0.5 * beta)

        return loss

    def forward(self, input: torch.Tensor,
                target: torch.Tensor, weights: torch.Tensor = None):
        """
        Args:
            input: (B, #anchors, #codes) float tensor.
                Ecoded predicted locations of objects.
            target: (B, #anchors, #codes) float tensor.
                Regression targets.
            weights: (B, #anchors) float tensor if not None.

        Returns:
            loss: (B, #anchors) float tensor.
                Weighted smooth l1 loss without reduction.
        """
        target = torch.where(torch.isnan(target), input, target)  # ignore nan targets

        diff = input - target
        loss = self.smooth_l1_loss(diff, self.beta)

        # anchor-wise weighting
        if weights is not None:
            assert weights.shape[0] == loss.shape[0] and weights.shape[1] == loss.shape[1]
            loss = loss * weights.unsqueeze(-1)

        return loss


class PointPillarLoss(nn.Module):
    def __init__(self, args):
        super(PointPillarLoss, self).__init__()
        self.reg_loss_func = WeightedSmoothL1Loss()
        self.alpha = 0.25
        self.gamma = 2.0

        if 'dir_args' in args.keys():
            self.use_dir = True
            self.dir_weight = args['dir_args']['dir_weight']
            self.dir_offset = args['dir_args']['args']['dir_offset']
            self.num_bins = args['dir_args']['args']['num_bins']
            anchor_yaw = np.deg2rad(np.array(args['dir_args']['anchor_yaw']))  # for direction classification
            self.anchor_yaw_map = torch.from_numpy(anchor_yaw).view(1,-1,1)  # [1,2,1]
            self.anchor_num = self.anchor_yaw_map.shape[1]

        else:
            self.use_dir =False

        self.cls_weight = args['cls_weight']
        self.reg_coe = args['reg']
        self.loss_dict = {}

    def forward(self, output_dict, target_dict):
        """
        Parameters
        ----------
        output_dict : dict
        target_dict : dict
        """
        rm = output_dict['rm']  # [B, 14, 50, 176]
        psm = output_dict['psm'] # [B, 2, 50, 176]
        targets = target_dict['targets']

        cls_preds = psm.permute(0, 2, 3, 1).contiguous() # N, C, H, W -> N, H, W, C

        box_cls_labels = target_dict['pos_equal_one']  # [B, 50, 176, 2]
        box_cls_labels = box_cls_labels.view(psm.shape[0], -1).contiguous()

        positives = box_cls_labels > 0
        negatives = box_cls_labels == 0
        negative_cls_weights = negatives * 1.0
        cls_weights = (negative_cls_weights + 1.0 * positives).float()
        reg_weights = positives.float()

        pos_normalizer = positives.sum(1, keepdim=True).float()
        reg_weights /= torch.clamp(pos_normalizer, min=1.0)
        cls_weights /= torch.clamp(pos_normalizer, min=1.0)
        cls_targets = box_cls_labels
        cls_targets = cls_targets.unsqueeze(dim=-1)

        cls_targets = cls_targets.squeeze(dim=-1)
        one_hot_targets = torch.zeros(
            *list(cls_targets.shape), 2,
            dtype=cls_preds.dtype, device=cls_targets.device
        )
        one_hot_targets.scatter_(-1, cls_targets.unsqueeze(dim=-1).long(), 1.0)
        cls_preds = cls_preds.view(psm.shape[0], -1, 1)
        one_hot_targets = one_hot_targets[..., 1:]

        cls_loss_src = self.cls_loss_func(cls_preds,
                                          one_hot_targets,
                                          weights=cls_weights)  # [N, M]
        cls_loss = cls_loss_src.sum() / psm.shape[0]
        conf_loss = cls_loss * self.cls_weight

        # regression
        rm = rm.permute(0, 2, 3, 1).contiguous()
        rm = rm.view(rm.size(0), -1, 7)
        targets = targets.view(targets.size(0), -1, 7)
        box_preds_sin, reg_targets_sin = self.add_sin_difference(rm,
                                                                 targets)
        loc_loss_src =\
            self.reg_loss_func(box_preds_sin,
                               reg_targets_sin,
                               weights=reg_weights)
        reg_loss = loc_loss_src.sum() / rm.shape[0]
        reg_loss *= self.reg_coe

        ######## direction ##########
        if self.use_dir:
            dir_targets = self.get_direction_target(targets)
            N =  output_dict["dm"].shape[0]
            dir_logits = output_dict["dm"].permute(0, 2, 3, 1).contiguous().view(N, -1, 2) # [N, H*W*#anchor, 2]


            dir_loss = softmax_cross_entropy_with_logits(dir_logits.view(-1, self.anchor_num), dir_targets.view(-1, self.anchor_num)) 
            dir_loss = dir_loss.view(dir_logits.shape[:2]) * reg_weights # [N, H*W*anchor_num]
            dir_loss = dir_loss.sum() * self.dir_weight / N


        total_loss = reg_loss + conf_loss

        self.loss_dict.update({'total_loss': total_loss,
                               'reg_loss': reg_loss,
                               'conf_loss': conf_loss})

        if self.use_dir:
            total_loss += dir_loss
            self.loss_dict.update({'dir_loss': dir_loss})

        return total_loss

    def cls_loss_func(self, input: torch.Tensor,
                      target: torch.Tensor,
                      weights: torch.Tensor):
        """
        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets
            weights: (B, #anchors) float tensor.
                Anchor-wise weights.

        Returns:
            weighted_loss: (B, #anchors, #classes) float tensor after weighting.
        """
        pred_sigmoid = torch.sigmoid(input)
        alpha_weight = target * self.alpha + (1 - target) * (1 - self.alpha)
        pt = target * (1.0 - pred_sigmoid) + (1.0 - target) * pred_sigmoid
        focal_weight = alpha_weight * torch.pow(pt, self.gamma)

        bce_loss = self.sigmoid_cross_entropy_with_logits(input, target)

        loss = focal_weight * bce_loss

        if weights.shape.__len__() == 2 or \
                (weights.shape.__len__() == 1 and target.shape.__len__() == 2):
            weights = weights.unsqueeze(-1)

        assert weights.shape.__len__() == loss.shape.__len__()

        return loss * weights

    @staticmethod
    def sigmoid_cross_entropy_with_logits(input: torch.Tensor, target: torch.Tensor):
        """ PyTorch Implementation for tf.nn.sigmoid_cross_entropy_with_logits:
            max(x, 0) - x * z + log(1 + exp(-abs(x))) in
            https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits

        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets

        Returns:
            loss: (B, #anchors, #classes) float tensor.
                Sigmoid cross entropy loss without reduction
        """
        loss = torch.clamp(input, min=0) - input * target + \
               torch.log1p(torch.exp(-torch.abs(input)))
        return loss

    @staticmethod
    def add_sin_difference(boxes1, boxes2, dim=6):
        assert dim != -1
        rad_pred_encoding = torch.sin(boxes1[..., dim:dim + 1]) * \
                            torch.cos(boxes2[..., dim:dim + 1])
        rad_tg_encoding = torch.cos(boxes1[..., dim:dim + 1]) * \
                          torch.sin(boxes2[..., dim:dim + 1])

        boxes1 = torch.cat([boxes1[..., :dim], rad_pred_encoding,
                            boxes1[..., dim + 1:]], dim=-1)
        boxes2 = torch.cat([boxes2[..., :dim], rad_tg_encoding,
                            boxes2[..., dim + 1:]], dim=-1)
        return boxes1, boxes2


    def logging(self, epoch, batch_id, batch_len, writer = None):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        total_loss = self.loss_dict['total_loss']
        reg_loss = self.loss_dict['reg_loss']
        conf_loss = self.loss_dict['conf_loss']
        

        print_msg = ("[epoch %d][%d/%d], || Loss: %.4f || Conf Loss: %.4f"
                    " || Loc Loss: %.4f" % (
                        epoch, batch_id + 1, batch_len,
                        total_loss.item(), conf_loss.item(), reg_loss.item()))
        
        if self.use_dir:
            dir_loss = self.loss_dict['dir_loss']
            print_msg += " || Dir Loss: %.4f" % dir_loss.item()

        print(print_msg)

        if not writer is None:
            writer.add_scalar('Regression_loss', reg_loss.item(),
                            epoch*batch_len + batch_id)
            writer.add_scalar('Confidence_loss', conf_loss.item(),
                            epoch*batch_len + batch_id)
                            
            if self.use_dir:
                writer.add_scalar('dir_loss', dir_loss.item(),
                            epoch*batch_len + batch_id)

def softmax_cross_entropy_with_logits(logits, labels):
    param = list(range(len(logits.shape)))
    transpose_param = [0] + [param[-1]] + param[1:-1]
    logits = logits.permute(*transpose_param)
    loss_ftor = torch.nn.CrossEntropyLoss(reduction="none")
    loss = loss_ftor(logits, labels.max(dim=-1)[1])
    return loss

opencood/loss/point_pillar_loss.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu


import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class WeightedSmoothL1Loss(nn.Module):
    """
    Code-wise Weighted Smooth L1 Loss modified based on fvcore.nn.smooth_l1_loss
    https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/smooth_l1_loss.py
                  | 0.5 * x ** 2 / beta   if abs(x) < beta
    smoothl1(x) = |
                  | abs(x) - 0.5 * beta   otherwise,
    where x = input - target.
    """
    def __init__(self, beta: float = 1.0 / 9.0, code_weights: list = None):
        """
        Args:
            beta: Scalar float.
                L1 to L2 change point.
                For beta values < 1e-5, L1 loss is computed.
            code_weights: (#codes) float list if not None.
                Code-wise weights.
        """
        super(WeightedSmoothL1Loss, self).__init__()
        self.beta = beta
        if code_weights is not None:
            self.code_weights = np.array(code_weights, dtype=np.float32)
            self.code_weights = torch.from_numpy(self.code_weights).cuda()

    @staticmethod
    def smooth_l1_loss(diff, beta):
        if beta < 1e-5:
            loss = torch.abs(diff)
        else:
            n = torch.abs(diff)
            loss = torch.where(n < beta, 0.5 * n ** 2 / beta, n - 0.5 * beta)

        return loss

    def forward(self, input: torch.Tensor,
                target: torch.Tensor, weights: torch.Tensor = None):
        """
        Args:
            input: (B, #anchors, #codes) float tensor.
                Ecoded predicted locations of objects.
            target: (B, #anchors, #codes) float tensor.
                Regression targets.
            weights: (B, #anchors) float tensor if not None.

        Returns:
            loss: (B, #anchors) float tensor.
                Weighted smooth l1 loss without reduction.
        """
        target = torch.where(torch.isnan(target), input, target)  # ignore nan targets

        diff = input - target
        loss = self.smooth_l1_loss(diff, self.beta)

        # anchor-wise weighting
        if weights is not None:
            assert weights.shape[0] == loss.shape[0] and weights.shape[1] == loss.shape[1]
            loss = loss * weights.unsqueeze(-1)

        return loss


class PointPillarMashLoss(nn.Module):
    def __init__(self, args):
        super(PointPillarMashLoss, self).__init__()
        self.reg_loss_func = WeightedSmoothL1Loss()
        self.grid_loss_func = nn.CrossEntropyLoss()
        self.alpha = 0.25
        self.gamma = 2.0

        self.cls_weight = args['cls_weight']
        self.grid_weight = args['grid_weight']
        self.reg_coe = args['reg']
        self.H = args['H']
        self.W = args['W']
        self.downsample_rate = args['downsample_rate']
        self.discrete_ratio = args['voxel_size'][0]

        self.loss_dict = {}

    def forward(self, output_dict, target_dict):
        """
        Parameters
        ----------
        output_dict : dict
        target_dict : dict
        """
        rm = output_dict['rm']  # [B, 14, 50, 176]
        psm = output_dict['psm'] # [B, 2, 50, 176]
        targets = target_dict['targets']

        cls_preds = psm.permute(0, 2, 3, 1).contiguous() # N, C, H, W -> N, H, W, C

        box_cls_labels = target_dict['pos_equal_one']  # [B, 50, 176, 2]
        box_cls_labels = box_cls_labels.view(psm.shape[0], -1).contiguous()

        positives = box_cls_labels > 0
        negatives = box_cls_labels == 0
        negative_cls_weights = negatives * 1.0
        cls_weights = (negative_cls_weights + 1.0 * positives).float()
        reg_weights = positives.float()

        pos_normalizer = positives.sum(1, keepdim=True).float()
        reg_weights /= torch.clamp(pos_normalizer, min=1.0)
        cls_weights /= torch.clamp(pos_normalizer, min=1.0)
        cls_targets = box_cls_labels
        cls_targets = cls_targets.unsqueeze(dim=-1)

        cls_targets = cls_targets.squeeze(dim=-1)
        one_hot_targets = torch.zeros(
            *list(cls_targets.shape), 2,
            dtype=cls_preds.dtype, device=cls_targets.device
        )
        one_hot_targets.scatter_(-1, cls_targets.unsqueeze(dim=-1).long(), 1.0)
        cls_preds = cls_preds.view(psm.shape[0], -1, 1)
        one_hot_targets = one_hot_targets[..., 1:]

        cls_loss_src = self.cls_loss_func(cls_preds,
                                          one_hot_targets,
                                          weights=cls_weights)  # [N, M]
        cls_loss = cls_loss_src.sum() / psm.shape[0]
        conf_loss = cls_loss * self.cls_weight

        # regression
        rm = rm.permute(0, 2, 3, 1).contiguous()
        rm = rm.view(rm.size(0), -1, 7)
        targets = targets.view(targets.size(0), -1, 7)
        box_preds_sin, reg_targets_sin = self.add_sin_difference(rm,
                                                                 targets)
        loc_loss_src =\
            self.reg_loss_func(box_preds_sin,
                               reg_targets_sin,
                               weights=reg_weights)
        reg_loss = loc_loss_src.sum() / rm.shape[0]
        reg_loss *= self.reg_coe

        total_loss = reg_loss + conf_loss
        self.loss_dict.update({'total_loss': total_loss.item(),
                               'reg_loss': reg_loss.item(),
                               'conf_loss': conf_loss.item()})

        corr_vol = output_dict['corr_vol']
        # grid loss
        if corr_vol is not None:
            pairwise_t_matrix = target_dict['pairwise_t_matrix']
            record_len = target_dict['record_len']

            pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]
            pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * self.H / self.W
            pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * self.W / self.H
            pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * self.W) * 2
            pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * self.H) * 2
            
            B = len(record_len)
            t_matrix_list = []
            for b in range(B):
                if(record_len[b] == 1):
                    continue
                t_matrix_list.append(pairwise_t_matrix[b, 0, 1:record_len[b]]) # [N-1, 2, 3]

            t_matrix = torch.cat(t_matrix_list, 0) 
            N_ = t_matrix.shape[0]  # N_ is sum_i{Ni-1}
            
            grid_gt = F.affine_grid(t_matrix, (N_, 1, self.H, self.W)) # (N_, H, W, 2)
            X = grid_gt[...,0]  # (N_, H, W)
            Y = grid_gt[...,1]  # (N_, H, W)
            X_idx = ((X / 2 + 0.5) * self.W).to(torch.long)  # (N_, H, W)
            Y_idx = ((Y / 2 + 0.5) * self.H).to(torch.long)  # (N_, H, W)

            idx = (Y_idx * self.W + X_idx)  # (N_, H, W)

            # out of boundary
            mask = torch.gt(grid_gt, -1) * torch.lt(grid_gt, 1)  # (N_, H, W,2)
            mask = mask[...,0] * mask[...,1]  # (N_, H, W)
            mask = ~ mask 
            idx[mask] = self.H * self.W # No matching, set to empty, last dimension
            corr_vol_gt = idx

            grid_loss = self.grid_loss_func(corr_vol, corr_vol_gt)
            grid_loss *= self.grid_weight

            total_loss += grid_loss
            self.loss_dict.update({"total_loss": total_loss.item(),
                                    "grid_loss": grid_loss.item()})

        

        return total_loss

    def cls_loss_func(self, input: torch.Tensor,
                      target: torch.Tensor,
                      weights: torch.Tensor):
        """
        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets
            weights: (B, #anchors) float tensor.
                Anchor-wise weights.

        Returns:
            weighted_loss: (B, #anchors, #classes) float tensor after weighting.
        """
        pred_sigmoid = torch.sigmoid(input)
        alpha_weight = target * self.alpha + (1 - target) * (1 - self.alpha)
        pt = target * (1.0 - pred_sigmoid) + (1.0 - target) * pred_sigmoid
        focal_weight = alpha_weight * torch.pow(pt, self.gamma)

        bce_loss = self.sigmoid_cross_entropy_with_logits(input, target)

        loss = focal_weight * bce_loss

        if weights.shape.__len__() == 2 or \
                (weights.shape.__len__() == 1 and target.shape.__len__() == 2):
            weights = weights.unsqueeze(-1)

        assert weights.shape.__len__() == loss.shape.__len__()

        return loss * weights

    @staticmethod
    def sigmoid_cross_entropy_with_logits(input: torch.Tensor, target: torch.Tensor):
        """ PyTorch Implementation for tf.nn.sigmoid_cross_entropy_with_logits:
            max(x, 0) - x * z + log(1 + exp(-abs(x))) in
            https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits

        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets

        Returns:
            loss: (B, #anchors, #classes) float tensor.
                Sigmoid cross entropy loss without reduction
        """
        loss = torch.clamp(input, min=0) - input * target + \
               torch.log1p(torch.exp(-torch.abs(input)))
        return loss

    @staticmethod
    def add_sin_difference(boxes1, boxes2, dim=6):
        assert dim != -1
        rad_pred_encoding = torch.sin(boxes1[..., dim:dim + 1]) * \
                            torch.cos(boxes2[..., dim:dim + 1])
        rad_tg_encoding = torch.cos(boxes1[..., dim:dim + 1]) * \
                          torch.sin(boxes2[..., dim:dim + 1])

        boxes1 = torch.cat([boxes1[..., :dim], rad_pred_encoding,
                            boxes1[..., dim + 1:]], dim=-1)
        boxes2 = torch.cat([boxes2[..., :dim], rad_tg_encoding,
                            boxes2[..., dim + 1:]], dim=-1)
        return boxes1, boxes2


    def logging(self, epoch, batch_id, batch_len, writer = None):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        total_loss = self.loss_dict['total_loss']
        reg_loss = self.loss_dict['reg_loss']
        conf_loss = self.loss_dict['conf_loss']
        grid_loss = self.loss_dict['grid_loss']

        print("[epoch %d][%d/%d], || Loss: %.4f || Conf Loss: %.4f"
              " || Loc Loss: %.4f || Grid Loss: %.4f" % (
                  epoch, batch_id + 1, batch_len,
                  total_loss, conf_loss, reg_loss, grid_loss))

        if not writer is None:
            writer.add_scalar('Regression_loss', reg_loss,
                            epoch*batch_len + batch_id)
            writer.add_scalar('Confidence_loss', conf_loss,
                            epoch*batch_len + batch_id)
            writer.add_scalar('Grid_loss', grid_loss,
                            epoch*batch_len + batch_id)

opencood/loss/point_pillar_mash_loss.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: OpenPCDet, Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Point pillar loss for time compensation.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from opencood.utils.common_utils import limit_period

class WeightedSmoothL1Loss(nn.Module):
    """
    Code-wise Weighted Smooth L1 Loss modified based on fvcore.nn.smooth_l1_loss
    https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/smooth_l1_loss.py
                  | 0.5 * x ** 2 / beta   if abs(x) < beta
    smoothl1(x) = |
                  | abs(x) - 0.5 * beta   otherwise,
    where x = input - target.
    """
    def __init__(self, beta: float = 1.0 / 9.0, code_weights: list = None):
        """
        Args:
            beta: Scalar float.
                L1 to L2 change point.
                For beta values < 1e-5, L1 loss is computed.
            code_weights: (#codes) float list if not None.
                Code-wise weights.
        """
        super(WeightedSmoothL1Loss, self).__init__()
        self.beta = beta
        if code_weights is not None:
            self.code_weights = np.array(code_weights, dtype=np.float32)
            self.code_weights = torch.from_numpy(self.code_weights).cuda()

    @staticmethod
    def smooth_l1_loss(diff, beta):
        if beta < 1e-5:
            loss = torch.abs(diff)
        else:
            n = torch.abs(diff)
            loss = torch.where(n < beta, 0.5 * n ** 2 / beta, n - 0.5 * beta)

        return loss

    def forward(self, input: torch.Tensor,
                target: torch.Tensor, weights: torch.Tensor = None):
        """
        Args:
            input: (B, #anchors, #codes) float tensor.
                Ecoded predicted locations of objects.
            target: (B, #anchors, #codes) float tensor.
                Regression targets.
            weights: (B, #anchors) float tensor if not None.

        Returns:
            loss: (B, #anchors) float tensor.
                Weighted smooth l1 loss without reduction.
        """
        target = torch.where(torch.isnan(target), input, target)  # ignore nan targets

        diff = input - target
        loss = self.smooth_l1_loss(diff, self.beta)

        # anchor-wise weighting
        if weights is not None:
            assert weights.shape[0] == loss.shape[0] and weights.shape[1] == loss.shape[1]
            loss = loss * weights.unsqueeze(-1)

        return loss


class PointPillarTcLoss(nn.Module):
    def __init__(self, args):
        super(PointPillarTcLoss, self).__init__()
        self.loss_dict = {}

        self.backbone_fix = False
        self.use_dir = False
        if 'backbone_fix' in args.keys() and args['backbone_fix']:
            self.backbone_fix = True
            self.flow_loss = nn.SmoothL1Loss(reduction='none')
            self.state_loss = nn.BCEWithLogitsLoss()
        else:
            self.reg_loss_func = WeightedSmoothL1Loss()
            self.alpha = 0.25
            self.gamma = 2.0
            self.cls_weight = args['cls_weight']
            self.reg_coe = args['reg']

            if 'dir_args' in args.keys():
                self.use_dir = True
                self.dir_weight = args['dir_args']['dir_weight']
                self.dir_offset = args['dir_args']['args']['dir_offset']
                self.num_bins = args['dir_args']['args']['num_bins']
                anchor_yaw = np.deg2rad(np.array(args['dir_args']['anchor_yaw']))  # for direction classification
                self.anchor_yaw_map = torch.from_numpy(anchor_yaw).view(1,-1,1)  # [1,2,1]
                self.anchor_num = self.anchor_yaw_map.shape[1]
            else:
                self.use_dir =False

    def delete_ego_label(self, target_dict, record_len):
        '''
        将 target_dict 中每一项的ego部分去掉 

        Parameters: 
        -----------
        target_dict: dict
            {
                'pos_equal_one': # [\sum_{B} N* , H, W, 2]
                'neg_equal_one': # [\sum_{B} N* , H, W, 2]
                'targets':       # [\sum_{B} N* , H, W, 14]
            }
        record_len: list
            len = B, 存放每个batch有多少辆车

        Returns:
        --------
        updated_dict: dict
            {
                'pos_equal_one': # [\sum_{B} n* , H, W, 2]   n* = N* - 1
                'neg_equal_one': # [\sum_{B} n* , H, W, 2]
                'targets':       # [\sum_{B} n* , H, W, 14]
            }
        '''
        updated_dict = {}

        split_loc = torch.cumsum(record_len, dim=0)
        for key, values in target_dict.items():
            deleted_label = []
            batch_list = torch.tensor_split(values, split_loc[:-1].cpu())
            for unit_content in batch_list:
                deleted_label.append(unit_content[1:])
            updated_dict[key] = torch.cat(deleted_label, dim=0)
        return updated_dict


    def forward(self, output_dict, target_dict, mode=None, suffix=""):
        """
        Parameters
        ----------
        output_dict : dict
        target_dict : dict
        """
        if self.backbone_fix:
            total_loss = 0
            ###### flow loss ######
            if f'flow_preds{suffix}' in output_dict and output_dict[f'flow_preds{suffix}'].shape[0] != 0:
                h, w = output_dict[f'flow_preds{suffix}'].shape[-2], output_dict[f'flow_preds{suffix}'].shape[-1]
                output_dict[f'flow_preds{suffix}'] = output_dict[f'flow_preds{suffix}'].reshape(1, -1, h, w)
                target_dict[f'flow_gt{suffix}'] = target_dict[f'flow_gt{suffix}'].reshape(1, -1, h, w)
                flow_loss = self.flow_loss(output_dict[f'flow_preds{suffix}'],
                                        target_dict[f'flow_gt{suffix}'])
                valid_flow_mask = ((torch.abs(target_dict[f'flow_gt{suffix}'].reshape(-1, 2, h, w)).max(dim=1)[0] > 0.49) * 1.0).unsqueeze(1)
                flow_loss = (flow_loss * valid_flow_mask).sum() / (valid_flow_mask.sum() + 1e-6)
                # flow_loss *= self.flow_weight

                state_loss = self.state_loss(output_dict[f'state_preds{suffix}'],
                                            valid_flow_mask)

                total_loss = flow_loss + state_loss
                self.loss_dict.update({'total_loss': total_loss.item(),
                                    'flow_loss': flow_loss.item(),
                                    'state_loss': state_loss.item()})
        else:
            if mode=='curr':
                rm = output_dict['rm_curr']
                psm = output_dict['psm_curr']
            elif mode=='latency':
                rm = output_dict['rm_latency']
                psm = output_dict['psm_latency']
            elif mode == 'single':
                psm = output_dict['psm_nonego_single']              # [\sum_{B}n, 2, H, W]
                rm = output_dict['rm_nonego_single']                # [\sum_{B}n, 14, H, W]
                target_dict = self.delete_ego_label(target_dict, output_dict['record_len'])
            else:
                rm = output_dict['rm']  # [B, 14, 50, 176]
                psm = output_dict['psm'] # [B, 2, 50, 176]
                if self.use_dir:
                    dm = output_dict["dm"] # [B, 4, 50, 176]

            # rename variable 
            if f'psm{suffix}' in output_dict:
                psm = output_dict[f'psm{suffix}']
            if f'rm{suffix}' in output_dict:
                rm = output_dict[f'rm{suffix}']
            if self.use_dir:
                if f'dm{suffix}' in output_dict:
                    dm = output_dict[f'dm{suffix}']
            
            targets = target_dict['targets']                # B, H, W, 14
            box_cls_labels = target_dict['pos_equal_one']   # B, H, W, 2

            cls_preds = psm.permute(0, 2, 3, 1).contiguous() # N, C, H, W -> N, H, W, C
            box_cls_labels = box_cls_labels.view(psm.shape[0], -1).contiguous()     # B, HxWx2

            positives = box_cls_labels > 0
            negatives = box_cls_labels == 0
            negative_cls_weights = negatives * 1.0
            cls_weights = (negative_cls_weights + 1.0 * positives).float()
            reg_weights = positives.float()

            pos_normalizer = positives.sum(1, keepdim=True).float()
            reg_weights /= torch.clamp(pos_normalizer, min=1.0)
            cls_weights /= torch.clamp(pos_normalizer, min=1.0)
            cls_targets = box_cls_labels
            cls_targets = cls_targets.unsqueeze(dim=-1)

            cls_targets = cls_targets.squeeze(dim=-1)
            one_hot_targets = torch.zeros(
                *list(cls_targets.shape), 2,
                dtype=cls_preds.dtype, device=cls_targets.device
            )
            one_hot_targets.scatter_(-1, cls_targets.unsqueeze(dim=-1).long(), 1.0)
            cls_preds = cls_preds.view(psm.shape[0], -1, 1)
            one_hot_targets = one_hot_targets[..., 1:]

            cls_loss_src = self.cls_loss_func(cls_preds,
                                            one_hot_targets,
                                            weights=cls_weights)  # [N, M]
            cls_loss = cls_loss_src.sum() / psm.shape[0]
            conf_loss = cls_loss * self.cls_weight

            # regression
            rm = rm.permute(0, 2, 3, 1).contiguous()
            rm = rm.view(rm.size(0), -1, 7)
            targets = targets.view(targets.size(0), -1, 7)
            box_preds_sin, reg_targets_sin = self.add_sin_difference(rm,
                                                                    targets)
            loc_loss_src =\
                self.reg_loss_func(box_preds_sin,
                                reg_targets_sin,
                                weights=reg_weights)
            reg_loss = loc_loss_src.sum() / rm.shape[0]
            reg_loss *= self.reg_coe

            ######## direction ##########
            if self.use_dir:
                dir_targets = self.get_direction_target(targets)
                N =  dm.shape[0]
                dir_logits = dm.permute(0, 2, 3, 1).contiguous().view(N, -1, 2) # [N, H*W*#anchor, 2]

                dir_loss = softmax_cross_entropy_with_logits(dir_logits.view(-1, self.anchor_num), dir_targets.view(-1, self.anchor_num)) 
                dir_loss = dir_loss.view(dir_logits.shape[:2]) * reg_weights # [N, H*W*anchor_num]
                dir_loss = dir_loss.sum() * self.dir_weight / N

            total_loss = reg_loss + conf_loss
            
            self.loss_dict.update({'total_loss': total_loss,
                                'reg_loss': reg_loss,
                                'conf_loss': conf_loss})

            if self.use_dir:
                total_loss += dir_loss
                self.loss_dict.update({'dir_loss': dir_loss})

        return total_loss

    def cls_loss_func(self, input: torch.Tensor,
                      target: torch.Tensor,
                      weights: torch.Tensor):
        """
        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets
            weights: (B, #anchors) float tensor.
                Anchor-wise weights.

        Returns:
            weighted_loss: (B, #anchors, #classes) float tensor after weighting.
        """
        pred_sigmoid = torch.sigmoid(input)
        alpha_weight = target * self.alpha + (1 - target) * (1 - self.alpha)
        pt = target * (1.0 - pred_sigmoid) + (1.0 - target) * pred_sigmoid
        focal_weight = alpha_weight * torch.pow(pt, self.gamma)

        bce_loss = self.sigmoid_cross_entropy_with_logits(input, target)

        loss = focal_weight * bce_loss

        if weights.shape.__len__() == 2 or \
                (weights.shape.__len__() == 1 and target.shape.__len__() == 2):
            weights = weights.unsqueeze(-1)

        assert weights.shape.__len__() == loss.shape.__len__()

        return loss * weights

    @staticmethod
    def sigmoid_cross_entropy_with_logits(input: torch.Tensor, target: torch.Tensor):
        """ PyTorch Implementation for tf.nn.sigmoid_cross_entropy_with_logits:
            max(x, 0) - x * z + log(1 + exp(-abs(x))) in
            https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits

        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets

        Returns:
            loss: (B, #anchors, #classes) float tensor.
                Sigmoid cross entropy loss without reduction
        """
        loss = torch.clamp(input, min=0) - input * target + \
               torch.log1p(torch.exp(-torch.abs(input)))
        return loss

    @staticmethod
    def add_sin_difference(boxes1, boxes2, dim=6):
        assert dim != -1
        rad_pred_encoding = torch.sin(boxes1[..., dim:dim + 1]) * \
                            torch.cos(boxes2[..., dim:dim + 1])
        rad_tg_encoding = torch.cos(boxes1[..., dim:dim + 1]) * \
                          torch.sin(boxes2[..., dim:dim + 1])

        boxes1 = torch.cat([boxes1[..., :dim], rad_pred_encoding,
                            boxes1[..., dim + 1:]], dim=-1)
        boxes2 = torch.cat([boxes2[..., :dim], rad_tg_encoding,
                            boxes2[..., dim + 1:]], dim=-1)
        return boxes1, boxes2

    def get_direction_target(self, reg_targets):
        """
        Args:
            reg_targets:  [N, H * W * #anchor_num, 7]
                The last term is (theta_gt - theta_a)
        
        Returns:
            dir_targets:
                theta_gt: [N, H * W * #anchor_num, NUM_BIN] 
                NUM_BIN = 2
        """
        # (1, 2, 1)
        H_times_W_times_anchor_num = reg_targets.shape[1]
        anchor_map = self.anchor_yaw_map.repeat(1, H_times_W_times_anchor_num//self.anchor_num, 1).to(reg_targets.device) # [1, H * W * #anchor_num, 1]
        rot_gt = reg_targets[..., -1] + anchor_map[..., -1] # [N, H*W*anchornum]
        offset_rot = limit_period(rot_gt - self.dir_offset, 0, 2 * np.pi)
        dir_cls_targets = torch.floor(offset_rot / (2 * np.pi / self.num_bins)).long()  # [N, H*W*anchornum]
        dir_cls_targets = torch.clamp(dir_cls_targets, min=0, max=self.num_bins - 1)
        # one_hot:
        # if rot_gt > 0, then the label is 1, then the regression target is [0, 1]
        dir_cls_targets = one_hot_f(dir_cls_targets, self.num_bins)
        return dir_cls_targets



    def logging(self, epoch, batch_id, batch_len, writer = None, suffix=""):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        flow_loss = self.loss_dict.get('flow_loss', 0)
        state_loss = self.loss_dict.get('state_loss', 0)
        total_loss = self.loss_dict.get('total_loss', 0)
        reg_loss = self.loss_dict.get('reg_loss', 0)
        conf_loss = self.loss_dict.get('conf_loss', 0)
            
        print_msg = ("[epoch %d][%d/%d]%s, || Loss: %.4f || Conf Loss: %.4f"
                    " || Loc Loss: %.4f || Flow Loss: %.4f || State Loss: %.4f" % (
                        epoch, batch_id + 1, batch_len, suffix, 
                        total_loss, conf_loss, reg_loss, flow_loss, state_loss))
        
        if self.use_dir:
            dir_loss = self.loss_dict['dir_loss']
            print_msg += " || Dir Loss: %.4f" % dir_loss.item()

        print(print_msg)

        if not writer is None:
            writer.add_scalar('Regression_loss'+suffix, reg_loss,
                            epoch*batch_len + batch_id)
            writer.add_scalar('Confidence_loss'+suffix, conf_loss,
                            epoch*batch_len + batch_id)
            writer.add_scalar('Flow_recon_loss'+suffix, flow_loss,
                            epoch*batch_len + batch_id)
            writer.add_scalar('State_loss'+suffix, state_loss,
                            epoch*batch_len + batch_id)
                            
            if self.use_dir:
                writer.add_scalar('dir_loss'+suffix, dir_loss,
                            epoch*batch_len + batch_id)

def one_hot_f(tensor, num_bins, dim=-1, on_value=1.0, dtype=torch.float32):
    tensor_onehot = torch.zeros(*list(tensor.shape), num_bins, dtype=dtype, device=tensor.device) 
    tensor_onehot.scatter_(dim, tensor.unsqueeze(dim).long(), on_value)                    
    return tensor_onehot


def softmax_cross_entropy_with_logits(logits, labels):
    param = list(range(len(logits.shape)))
    transpose_param = [0] + [param[-1]] + param[1:-1]
    logits = logits.permute(*transpose_param)
    loss_ftor = torch.nn.CrossEntropyLoss(reduction="none")
    loss = loss_ftor(logits, labels.max(dim=-1)[1])
    return loss

opencood/loss/point_pillar_tc_loss.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: OpenPCDet, Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
Point pillar loss for time compensation.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


class WeightedSmoothL1Loss(nn.Module):
    """
    Code-wise Weighted Smooth L1 Loss modified based on fvcore.nn.smooth_l1_loss
    https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/smooth_l1_loss.py
                  | 0.5 * x ** 2 / beta   if abs(x) < beta
    smoothl1(x) = |
                  | abs(x) - 0.5 * beta   otherwise,
    where x = input - target.
    """
    def __init__(self, beta: float = 1.0 / 9.0, code_weights: list = None):
        """
        Args:
            beta: Scalar float.
                L1 to L2 change point.
                For beta values < 1e-5, L1 loss is computed.
            code_weights: (#codes) float list if not None.
                Code-wise weights.
        """
        super(WeightedSmoothL1Loss, self).__init__()
        self.beta = beta
        if code_weights is not None:
            self.code_weights = np.array(code_weights, dtype=np.float32)
            self.code_weights = torch.from_numpy(self.code_weights).cuda()

    @staticmethod
    def smooth_l1_loss(diff, beta):
        if beta < 1e-5:
            loss = torch.abs(diff)
        else:
            n = torch.abs(diff)
            loss = torch.where(n < beta, 0.5 * n ** 2 / beta, n - 0.5 * beta)

        return loss

    def forward(self, input: torch.Tensor,
                target: torch.Tensor, weights: torch.Tensor = None):
        """
        Args:
            input: (B, #anchors, #codes) float tensor.
                Ecoded predicted locations of objects.
            target: (B, #anchors, #codes) float tensor.
                Regression targets.
            weights: (B, #anchors) float tensor if not None.

        Returns:
            loss: (B, #anchors) float tensor.
                Weighted smooth l1 loss without reduction.
        """
        target = torch.where(torch.isnan(target), input, target)  # ignore nan targets

        diff = input - target
        loss = self.smooth_l1_loss(diff, self.beta)

        # anchor-wise weighting
        if weights is not None:
            assert weights.shape[0] == loss.shape[0] and weights.shape[1] == loss.shape[1]
            loss = loss * weights.unsqueeze(-1)

        return loss


class PointPillarTcLossSyncnet(nn.Module):
    def __init__(self, args):
        super(PointPillarTcLossSyncnet, self).__init__()
        self.loss_dict = {}

        self.backbone_fix = False
        self.use_dir = False
        if 'backbone_fix' in args.keys() and args['backbone_fix']:
            self.backbone_fix = True
            
        else:
            self.reg_loss_func = WeightedSmoothL1Loss()
            self.alpha = 0.25
            self.gamma = 2.0
            self.cls_weight = args['cls_weight']
            self.reg_coe = args['reg']

            if 'dir_args' in args.keys():
                self.use_dir = True
                self.dir_weight = args['dir_args']['dir_weight']
                self.dir_offset = args['dir_args']['args']['dir_offset']
                self.num_bins = args['dir_args']['args']['num_bins']
                anchor_yaw = np.deg2rad(np.array(args['dir_args']['anchor_yaw']))  # for direction classification
                self.anchor_yaw_map = torch.from_numpy(anchor_yaw).view(1,-1,1)  # [1,2,1]
                self.anchor_num = self.anchor_yaw_map.shape[1]
            else:
                self.use_dir =False

    def delete_ego_label(self, target_dict, record_len):
        '''
        将 target_dict 中每一项的ego部分去掉 

        Parameters: 
        -----------
        target_dict: dict
            {
                'pos_equal_one': # [\sum_{B} N* , H, W, 2]
                'neg_equal_one': # [\sum_{B} N* , H, W, 2]
                'targets':       # [\sum_{B} N* , H, W, 14]
            }
        record_len: list
            len = B, 存放每个batch有多少辆车

        Returns:
        --------
        updated_dict: dict
            {
                'pos_equal_one': # [\sum_{B} n* , H, W, 2]   n* = N* - 1
                'neg_equal_one': # [\sum_{B} n* , H, W, 2]
                'targets':       # [\sum_{B} n* , H, W, 14]
            }
        '''
        updated_dict = {}

        split_loc = torch.cumsum(record_len, dim=0)
        for key, values in target_dict.items():
            deleted_label = []
            batch_list = torch.tensor_split(values, split_loc[:-1].cpu())
            for unit_content in batch_list:
                deleted_label.append(unit_content[1:])
            updated_dict[key] = torch.cat(deleted_label, dim=0)
        return updated_dict


    def forward(self, output_dict, target_dict, mode=None, suffix=""):
        """
        Parameters
        ----------
        output_dict : dict
        target_dict : dict
        """
        if self.backbone_fix:
            total_loss = 0
        else:
            if mode=='curr':
                rm = output_dict['rm_curr']
                psm = output_dict['psm_curr']
            elif mode=='latency':
                rm = output_dict['rm_latency']
                psm = output_dict['psm_latency']
            elif mode == 'single':
                psm = output_dict['psm_nonego_single']              # [\sum_{B}n, 2, H, W]
                rm = output_dict['rm_nonego_single']                # [\sum_{B}n, 14, H, W]
                target_dict = self.delete_ego_label(target_dict, output_dict['record_len'])
            else:
                rm = output_dict['rm']  # [B, 14, 50, 176]
                psm = output_dict['psm'] # [B, 2, 50, 176]

            # rename variable 
            if f'psm{suffix}' in output_dict:
                psm = output_dict[f'psm{suffix}']
            if f'rm{suffix}' in output_dict:
                rm = output_dict[f'rm{suffix}']
            
            targets = target_dict['targets']                # B, H, W, 14
            box_cls_labels = target_dict['pos_equal_one']   # B, H, W, 2

            cls_preds = psm.permute(0, 2, 3, 1).contiguous() # N, C, H, W -> N, H, W, C
            box_cls_labels = box_cls_labels.view(psm.shape[0], -1).contiguous()     # B, HxWx2

            positives = box_cls_labels > 0
            negatives = box_cls_labels == 0
            negative_cls_weights = negatives * 1.0
            cls_weights = (negative_cls_weights + 1.0 * positives).float()
            reg_weights = positives.float()

            pos_normalizer = positives.sum(1, keepdim=True).float()
            reg_weights /= torch.clamp(pos_normalizer, min=1.0)
            cls_weights /= torch.clamp(pos_normalizer, min=1.0)
            cls_targets = box_cls_labels
            cls_targets = cls_targets.unsqueeze(dim=-1)

            cls_targets = cls_targets.squeeze(dim=-1)
            one_hot_targets = torch.zeros(
                *list(cls_targets.shape), 2,
                dtype=cls_preds.dtype, device=cls_targets.device
            )
            one_hot_targets.scatter_(-1, cls_targets.unsqueeze(dim=-1).long(), 1.0)
            cls_preds = cls_preds.view(psm.shape[0], -1, 1)
            one_hot_targets = one_hot_targets[..., 1:]

            cls_loss_src = self.cls_loss_func(cls_preds,
                                            one_hot_targets,
                                            weights=cls_weights)  # [N, M]
            cls_loss = cls_loss_src.sum() / psm.shape[0]
            conf_loss = cls_loss * self.cls_weight

            # regression
            rm = rm.permute(0, 2, 3, 1).contiguous()
            rm = rm.view(rm.size(0), -1, 7)
            targets = targets.view(targets.size(0), -1, 7)
            box_preds_sin, reg_targets_sin = self.add_sin_difference(rm,
                                                                    targets)
            loc_loss_src =\
                self.reg_loss_func(box_preds_sin,
                                reg_targets_sin,
                                weights=reg_weights)
            reg_loss = loc_loss_src.sum() / rm.shape[0]
            reg_loss *= self.reg_coe

            ######## direction ##########
            if self.use_dir:
                dir_targets = self.get_direction_target(targets)
                N =  output_dict["dm"].shape[0]
                dir_logits = output_dict["dm"].permute(0, 2, 3, 1).contiguous().view(N, -1, 2) # [N, H*W*#anchor, 2]

                dir_loss = softmax_cross_entropy_with_logits(dir_logits.view(-1, self.anchor_num), dir_targets.view(-1, self.anchor_num)) 
                dir_loss = dir_loss.view(dir_logits.shape[:2]) * reg_weights # [N, H*W*anchor_num]
                dir_loss = dir_loss.sum() * self.dir_weight / N

            total_loss = reg_loss + conf_loss
            
            self.loss_dict.update({'total_loss': total_loss,
                                'reg_loss': reg_loss,
                                'conf_loss': conf_loss})

            if self.use_dir:
                total_loss += dir_loss
                self.loss_dict.update({'dir_loss': dir_loss})

        ###### flow loss ######
        total_loss += output_dict['recon_loss']
        self.loss_dict.update({'total_loss': total_loss.item()})

        return total_loss

    def cls_loss_func(self, input: torch.Tensor,
                      target: torch.Tensor,
                      weights: torch.Tensor):
        """
        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets
            weights: (B, #anchors) float tensor.
                Anchor-wise weights.

        Returns:
            weighted_loss: (B, #anchors, #classes) float tensor after weighting.
        """
        pred_sigmoid = torch.sigmoid(input)
        alpha_weight = target * self.alpha + (1 - target) * (1 - self.alpha)
        pt = target * (1.0 - pred_sigmoid) + (1.0 - target) * pred_sigmoid
        focal_weight = alpha_weight * torch.pow(pt, self.gamma)

        bce_loss = self.sigmoid_cross_entropy_with_logits(input, target)

        loss = focal_weight * bce_loss

        if weights.shape.__len__() == 2 or \
                (weights.shape.__len__() == 1 and target.shape.__len__() == 2):
            weights = weights.unsqueeze(-1)

        assert weights.shape.__len__() == loss.shape.__len__()

        return loss * weights

    @staticmethod
    def sigmoid_cross_entropy_with_logits(input: torch.Tensor, target: torch.Tensor):
        """ PyTorch Implementation for tf.nn.sigmoid_cross_entropy_with_logits:
            max(x, 0) - x * z + log(1 + exp(-abs(x))) in
            https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits

        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets

        Returns:
            loss: (B, #anchors, #classes) float tensor.
                Sigmoid cross entropy loss without reduction
        """
        loss = torch.clamp(input, min=0) - input * target + \
               torch.log1p(torch.exp(-torch.abs(input)))
        return loss

    @staticmethod
    def add_sin_difference(boxes1, boxes2, dim=6):
        assert dim != -1
        rad_pred_encoding = torch.sin(boxes1[..., dim:dim + 1]) * \
                            torch.cos(boxes2[..., dim:dim + 1])
        rad_tg_encoding = torch.cos(boxes1[..., dim:dim + 1]) * \
                          torch.sin(boxes2[..., dim:dim + 1])

        boxes1 = torch.cat([boxes1[..., :dim], rad_pred_encoding,
                            boxes1[..., dim + 1:]], dim=-1)
        boxes2 = torch.cat([boxes2[..., :dim], rad_tg_encoding,
                            boxes2[..., dim + 1:]], dim=-1)
        return boxes1, boxes2


    def logging(self, epoch, batch_id, batch_len, writer = None, suffix=""):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        flow_loss = self.loss_dict.get('flow_loss', 0)
        state_loss = self.loss_dict.get('state_loss', 0)
        total_loss = self.loss_dict.get('total_loss', 0)
        reg_loss = self.loss_dict.get('reg_loss', 0)
        conf_loss = self.loss_dict.get('conf_loss', 0)
            
        print_msg = ("[epoch %d][%d/%d]%s, || Loss: %.4f || Conf Loss: %.4f"
                    " || Loc Loss: %.4f || Flow Loss: %.4f || State Loss: %.4f" % (
                        epoch, batch_id + 1, batch_len, suffix, 
                        total_loss, conf_loss, reg_loss, flow_loss, state_loss))
        
        if self.use_dir:
            dir_loss = self.loss_dict['dir_loss']
            print_msg += " || Dir Loss: %.4f" % dir_loss.item()

        print(print_msg)

        if not writer is None:
            writer.add_scalar('Regression_loss'+suffix, reg_loss,
                            epoch*batch_len + batch_id)
            writer.add_scalar('Confidence_loss'+suffix, conf_loss,
                            epoch*batch_len + batch_id)
                            
            if self.use_dir:
                writer.add_scalar('dir_loss'+suffix, dir_loss,
                            epoch*batch_len + batch_id)

def softmax_cross_entropy_with_logits(logits, labels):
    param = list(range(len(logits.shape)))
    transpose_param = [0] + [param[-1]] + param[1:-1]
    logits = logits.permute(*transpose_param)
    loss_ftor = torch.nn.CrossEntropyLoss(reduction="none")
    loss = loss_ftor(logits, labels.max(dim=-1)[1])
    return loss

opencood/loss/point_pillar_tc_loss_syncnet.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu


from tabnanny import verbose
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import d3d.mathh as mathh
from opencood.utils.common_utils import limit_period
from functools import partial

class WeightedSmoothL1Loss(nn.Module):
    """
    Code-wise Weighted Smooth L1 Loss modified based on fvcore.nn.smooth_l1_loss
    https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/smooth_l1_loss.py
                  | 0.5 * x ** 2 / beta   if abs(x) < beta
    smoothl1(x) = |
                  | abs(x) - 0.5 * beta   otherwise,
    where x = input - target.
    """
    def __init__(self, beta: float = 1.0 / 9.0, code_weights: list = None):
        """
        Args:
            beta: Scalar float.
                L1 to L2 change point.
                For beta values < 1e-5, L1 loss is computed.
            code_weights: (#codes) float list if not None.
                Code-wise weights.
        """
        super(WeightedSmoothL1Loss, self).__init__()
        self.beta = beta
        if code_weights is not None:
            self.code_weights = np.array(code_weights, dtype=np.float32)
            self.code_weights = torch.from_numpy(self.code_weights).cuda()

    @staticmethod
    def smooth_l1_loss(diff, beta):
        if beta < 1e-5:
            loss = torch.abs(diff)
        else:
            n = torch.abs(diff)
            loss = torch.where(n < beta, 0.5 * n ** 2 / beta, n - 0.5 * beta)

        return loss

    def forward(self, input: torch.Tensor,
                target: torch.Tensor, weights: torch.Tensor = None):
        """
        Args:
            input: (B, #anchors, #codes) float tensor.
                Ecoded predicted locations of objects.
            target: (B, #anchors, #codes) float tensor.
                Regression targets.
            weights: (B, #anchors) float tensor if not None.

        Returns:
            loss: (B, #anchors) float tensor.
                Weighted smooth l1 loss without reduction.
        """
        target = torch.where(torch.isnan(target), input, target)  # ignore nan targets

        diff = input - target
        loss = self.smooth_l1_loss(diff, self.beta)

        # anchor-wise weighting
        if weights is not None:
            assert weights.shape[0] == loss.shape[0] and weights.shape[1] == loss.shape[1]
            loss = loss * weights.unsqueeze(-1)

        return loss




class KLLoss(nn.Module):
    def __init__(self, args):
        super(KLLoss, self).__init__()

        self.angle_weight = args['angle_weight']
        self.uncertainty_dim = args['uncertainty_dim']
        if args['xy_loss_type'] == "l2":
            self.xy_loss = self.kl_loss_l2
        elif args['xy_loss_type'] == "l1":
            self.xy_loss = self.kl_loss_l1
        else:
            raise "not implemented"

        if args['angle_loss_type'] == "l2":
            self.angle_loss = self.kl_loss_l2
        elif args['angle_loss_type'] == "von":
            lambda_V = args['lambda_V']
            s0 = args['s0']
            limit_period = args['limit_period']
            self.angle_loss = partial(self.kl_loss_angular, lambda_V=lambda_V, s0=s0, limit_period=limit_period)
        else:
            raise "not implemented"




    @staticmethod
    def kl_loss_l2(diff, s):
        """
        Args:
            diff: [B, 2]
            s:    [B, 2]
        Returns:
            loss: [B, 2]
        """
        loss = 0.5*(torch.exp(-s) * (diff**2) + s)
        return loss
    
    @staticmethod
    def kl_loss_l1(diff, s):
        """
        Args:
            diff: [B, 2]
            s:    [B, 2]
        Returns:
            loss: [B, 2]
        """
        loss = 0.5*torch.exp(-s) * torch.abs(diff) + s
        return loss
    
    @staticmethod
    def kl_loss_angular(diff, s, lambda_V=1, s0=1, limit_period=False):
        """
        Args:
            diff: [B, 1]
            s:    [B, 1]
            if limit_period, 
                diff + 180 ~ diff. 
        Returns:
            loss: [B, 1]
        """
        exp_minus_s = torch.exp(-s)
        if limit_period:
            cos_abs = torch.abs(torch.cos(diff))
            loss = loss = torch.log(mathh.i0e_cuda(exp_minus_s)*torch.exp(exp_minus_s)) - exp_minus_s * cos_abs.detach() + lambda_V * F.elu(s-s0)
        else:
            loss = torch.log(mathh.i0e_cuda(exp_minus_s)*torch.exp(exp_minus_s)) - exp_minus_s * torch.cos(diff) + lambda_V * F.elu(s-s0)

        return loss

    
    def forward(self, input: torch.Tensor,
                      target: torch.Tensor, 
                      sm: torch.Tensor, 
                      weights: torch.Tensor = None):
        target = torch.where(torch.isnan(target), input, target)  # ignore nan targets
        
        

        if self.uncertainty_dim == 3:
            xy_diff = input[...,:2] - target[...,:2]
            loss1 = self.xy_loss(xy_diff, sm[...,:2])
            
            theta_diff = input[...,7:8] - target[...,7:8]

            loss2 = self.angle_weight * self.angle_loss(theta_diff, sm[...,2:3])

            loss = torch.cat((loss1, loss2), dim=-1)
            
        elif self.uncertainty_dim == 7:
            ## is this right?
            other_diff = input[...,:6] - target[...,:6]
            theta_diff = input[...,7:8] - target[...,7:8]

            diff = torch.cat((other_diff, theta_diff), dim=-1)
            loss = self.xy_loss(diff, sm)

        elif self.uncertainty_dim == 2:
            xy_diff = input[...,:2] - target[...,:2]
            loss = self.xy_loss(xy_diff, sm[...,:2])
        else:
            raise "not implemented"

        # anchor-wise weighting
        if weights is not None:
            assert weights.shape[0] == loss.shape[0] and weights.shape[1] == loss.shape[1]
            loss = loss * weights.unsqueeze(-1)
        
        return loss



class PointPillarUncertaintyLoss(nn.Module):
    def __init__(self, args):
        super(PointPillarUncertaintyLoss, self).__init__()
        self.reg_loss_func = WeightedSmoothL1Loss()
        self.alpha = 0.25
        self.gamma = 2.0

        self.cls_weight = args['cls_weight']
        self.kl_weight = args['kl_weight']
        self.reg_coe = args['reg']
        self.uncertainty_dim = args['kl_args']['uncertainty_dim']

        if 'dir_args' in args.keys():
            self.use_dir = True
            self.dir_weight = args['dir_args']['dir_weight']
            self.dir_offset = args['dir_args']['args']['dir_offset']
            self.num_bins = args['dir_args']['args']['num_bins']
            anchor_yaw = np.deg2rad(np.array(args['dir_args']['anchor_yaw']))  # for direction classification
            self.anchor_yaw_map = torch.from_numpy(anchor_yaw).view(1,-1,1)  # [1,2,1]
            self.anchor_num = self.anchor_yaw_map.shape[1]

        else:
            self.use_dir =False


        self.kl_loss_func = KLLoss(args['kl_args'])

        self.loss_dict = {}

    def forward(self, output_dict, target_dict):
        """
        Parameters
        ----------
        output_dict : dict
        target_dict : dict
        """
        rm = output_dict['rm']  # [B, 14, 50, 176]
        psm = output_dict['psm'] # [B, 2, 50, 176]
        sm = output_dict['sm']  # log of sigma^2 / scale [B, 6, 50 176]
        targets = target_dict['targets']

        cls_preds = psm.permute(0, 2, 3, 1).contiguous() # N, C, H, W -> N, H, W, C

        box_cls_labels = target_dict['pos_equal_one']  # [B, 50, 176, 2] 
        box_cls_labels = box_cls_labels.view(psm.shape[0], -1).contiguous() # -> [B, 50*176*2], two types of anchor

        positives = box_cls_labels > 0
        negatives = box_cls_labels == 0
        negative_cls_weights = negatives * 1.0
        cls_weights = (negative_cls_weights + 1.0 * positives).float() # all 1
        reg_weights = positives.float()

        pos_normalizer = positives.sum(1, keepdim=True).float() # positive number per sample
        reg_weights /= torch.clamp(pos_normalizer, min=1.0)
        cls_weights /= torch.clamp(pos_normalizer, min=1.0)
        cls_targets = box_cls_labels
        cls_targets = cls_targets.unsqueeze(dim=-1)

        cls_targets = cls_targets.squeeze(dim=-1)
        one_hot_targets = torch.zeros(
            *list(cls_targets.shape), 2,
            dtype=cls_preds.dtype, device=cls_targets.device
        )
        one_hot_targets.scatter_(-1, cls_targets.unsqueeze(dim=-1).long(), 1.0)
        cls_preds = cls_preds.view(psm.shape[0], -1, 1)
        one_hot_targets = one_hot_targets[..., 1:]

        cls_loss_src = self.cls_loss_func(cls_preds,
                                          one_hot_targets,
                                          weights=cls_weights)  # [N, M]
        cls_loss = cls_loss_src.sum() / psm.shape[0]
        conf_loss = cls_loss * self.cls_weight

        ########## regression ##########
        rm = rm.permute(0, 2, 3, 1).contiguous()
        rm = rm.view(rm.size(0), -1, 7)
        targets = targets.view(targets.size(0), -1, 7)

        box_preds_sin, reg_targets_sin = self.add_sin_difference_dim(rm,
                                                                 targets)
        loc_loss_src =\
            self.reg_loss_func(box_preds_sin[...,:7],
                               reg_targets_sin[...,:7],
                               weights=reg_weights)
        reg_loss = loc_loss_src.sum() / rm.shape[0]
        reg_loss *= self.reg_coe


        ######## direction ##########
        if self.use_dir:
            dir_targets = self.get_direction_target(targets)
            N =  output_dict["dm"].shape[0]
            dir_logits = output_dict["dm"].permute(0, 2, 3, 1).contiguous().view(N, -1, 2) # [N, H*W*#anchor, 2]


            dir_loss = softmax_cross_entropy_with_logits(dir_logits.view(-1, self.anchor_num), dir_targets.view(-1, self.anchor_num)) 

            dir_loss = dir_loss.view(dir_logits.shape[:2]) * reg_weights # [N, H*W*anchor_num]

            dir_loss = dir_loss.sum() * self.dir_weight / N

        ######## kl #########
        sm = sm.permute(0, 2, 3, 1).contiguous() # [N, H, W, #anchor_num * 3]
        sm = sm.view(sm.size(0), -1, self.uncertainty_dim)

        kl_loss_src = \
            self.kl_loss_func(box_preds_sin,
                              reg_targets_sin,
                              sm,
                              reg_weights)

        kl_loss = kl_loss_src.sum() / sm.shape[0]
        kl_loss *= self.kl_weight

        total_loss = reg_loss + conf_loss + kl_loss

        self.loss_dict.update({'total_loss': total_loss,
                               'reg_loss': reg_loss,
                               'conf_loss': conf_loss,
                               'kl_loss': kl_loss})
        
        if self.use_dir:
            total_loss += dir_loss
            self.loss_dict.update({'dir_loss': dir_loss})


        return total_loss

    def get_direction_target(self, reg_targets):
        """
        Args:
            reg_targets:  [N, H * W * #anchor_num, 7]
                The last term is (theta_gt - theta_a)
        
        Returns:
            dir_targets:
                theta_gt: [N, H * W * #anchor_num, NUM_BIN] 
                NUM_BIN = 2
        """
        # (1, 2, 1)
        H_times_W_times_anchor_num = reg_targets.shape[1]
        anchor_map = self.anchor_yaw_map.repeat(1, H_times_W_times_anchor_num//self.anchor_num, 1).to(reg_targets.device) # [1, H * W * #anchor_num, 1]
        rot_gt = reg_targets[..., -1] + anchor_map[..., -1] # [N, H*W*anchornum]
        offset_rot = limit_period(rot_gt - self.dir_offset, 0, 2 * np.pi)
        dir_cls_targets = torch.floor(offset_rot / (2 * np.pi / self.num_bins)).long()  # [N, H*W*anchornum]
        dir_cls_targets = torch.clamp(dir_cls_targets, min=0, max=self.num_bins - 1)
        # one_hot:
        # if rot_gt > 0, then the label is 1, then the regression target is [0, 1]
        dir_cls_targets = one_hot_f(dir_cls_targets, self.num_bins)
        return dir_cls_targets



    def cls_loss_func(self, input: torch.Tensor,
                      target: torch.Tensor,
                      weights: torch.Tensor):
        """
        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets
            weights: (B, #anchors) float tensor.
                Anchor-wise weights.

        Returns:
            weighted_loss: (B, #anchors, #classes) float tensor after weighting.
        """
        pred_sigmoid = torch.sigmoid(input)
        alpha_weight = target * self.alpha + (1 - target) * (1 - self.alpha)
        pt = target * (1.0 - pred_sigmoid) + (1.0 - target) * pred_sigmoid
        focal_weight = alpha_weight * torch.pow(pt, self.gamma)

        bce_loss = self.sigmoid_cross_entropy_with_logits(input, target)

        loss = focal_weight * bce_loss

        if weights.shape.__len__() == 2 or \
                (weights.shape.__len__() == 1 and target.shape.__len__() == 2):
            weights = weights.unsqueeze(-1)

        assert weights.shape.__len__() == loss.shape.__len__()

        return loss * weights

    @staticmethod
    def sigmoid_cross_entropy_with_logits(input: torch.Tensor, target: torch.Tensor):
        """ PyTorch Implementation for tf.nn.sigmoid_cross_entropy_with_logits:
            max(x, 0) - x * z + log(1 + exp(-abs(x))) in
            https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits

        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets

        Returns:
            loss: (B, #anchors, #classes) float tensor.
                Sigmoid cross entropy loss without reduction
        """
        loss = torch.clamp(input, min=0) - input * target + \
               torch.log1p(torch.exp(-torch.abs(input)))
        return loss

    @staticmethod
    def add_sin_difference_dim(boxes1, boxes2, dim=6):
        """
        This is different with other loss function.
        Here we especially retain the angel

            Add sin difference ?
            Replace sin difference !

        Returns:
            [B, H*W, 7] -> [B, H*W, 8]
        """
        assert dim != -1

        # sin(theta1 - theta2) = sin(theta1)*cos(theta2) - cos(theta1)*sin(theta2) 

        rad_pred_encoding = torch.sin(boxes1[..., dim:dim + 1]) * \
                            torch.cos(boxes2[..., dim:dim + 1])

        rad_tg_encoding = torch.cos(boxes1[..., dim: dim + 1]) * \
                          torch.sin(boxes2[..., dim: dim + 1])

        # boxes1 = torch.cat([boxes1[..., :dim], rad_pred_encoding,
        #                     boxes1[..., dim + 1:]], dim=-1)
        # boxes2 = torch.cat([boxes2[..., :dim], rad_tg_encoding,
        #                     boxes2[..., dim + 1:]], dim=-1)
        
        boxes1_encoded = torch.cat([boxes1[..., :dim], rad_pred_encoding,
                            boxes1[..., dim:]], dim=-1)
        boxes2_encoded = torch.cat([boxes2[..., :dim], rad_tg_encoding,
                            boxes2[..., dim:]], dim=-1)

        return boxes1_encoded, boxes2_encoded


    def logging(self, epoch, batch_id, batch_len, writer = None):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        total_loss = self.loss_dict['total_loss']
        reg_loss = self.loss_dict['reg_loss']
        conf_loss = self.loss_dict['conf_loss']
        kl_loss = self.loss_dict['kl_loss']
        

        print_msg = ("[epoch %d][%d/%d], || Loss: %.4f || Conf Loss: %.4f"
                    " || Loc Loss: %.4f || KL Loss: %.4f" % (
                        epoch, batch_id + 1, batch_len,
                        total_loss.item(), conf_loss.item(), reg_loss.item(),  kl_loss.item()))
        
        if self.use_dir:
            dir_loss = self.loss_dict['dir_loss']
            print_msg += " || Dir Loss: %.4f" % dir_loss.item()

        print(print_msg)

        if not writer is None:
            writer.add_scalar('Regression_loss', reg_loss.item(),
                            epoch*batch_len + batch_id)
            writer.add_scalar('Confidence_loss', conf_loss.item(),
                            epoch*batch_len + batch_id)
            writer.add_scalar('kl_loss', kl_loss.item(),
                            epoch*batch_len + batch_id)
            if self.use_dir:
                writer.add_scalar('dir_loss', dir_loss.item(),
                            epoch*batch_len + batch_id)

def one_hot_f(tensor, depth, dim=-1, on_value=1.0, dtype=torch.float32):
    tensor_onehot = torch.zeros(*list(tensor.shape), depth, dtype=dtype, device=tensor.device) # [4, 70400, 2]
    tensor_onehot.scatter_(dim, tensor.unsqueeze(dim).long(), on_value)                        # [4, 70400, 2]
    return tensor_onehot

def softmax_cross_entropy_with_logits(logits, labels):
    param = list(range(len(logits.shape)))
    transpose_param = [0] + [param[-1]] + param[1:-1]
    logits = logits.permute(*transpose_param)
    loss_ftor = torch.nn.CrossEntropyLoss(reduction="none")
    loss = loss_ftor(logits, labels.max(dim=-1)[1])
    return loss

opencood/loss/point_pillar_uncertainty_loss.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu

from icecream import ic
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from opencood.models.sub_modules.v2v_robust_module import regroup
from opencood.utils.transformation_utils import tfm_to_pose, tfm_to_pose_torch
torch.set_printoptions(precision=3, sci_mode=False)

class WeightedSmoothL1Loss(nn.Module):
    """
    Code-wise Weighted Smooth L1 Loss modified based on fvcore.nn.smooth_l1_loss
    https://github.com/facebookresearch/fvcore/blob/master/fvcore/nn/smooth_l1_loss.py
                  | 0.5 * x ** 2 / beta   if abs(x) < beta
    smoothl1(x) = |
                  | abs(x) - 0.5 * beta   otherwise,
    where x = input - target.
    """
    def __init__(self, beta: float = 1.0 / 9.0, code_weights: list = None):
        """
        Args:
            beta: Scalar float.
                L1 to L2 change point.
                For beta values < 1e-5, L1 loss is computed.
            code_weights: (#codes) float list if not None.
                Code-wise weights.
        """
        super(WeightedSmoothL1Loss, self).__init__()
        self.beta = beta
        if code_weights is not None:
            self.code_weights = np.array(code_weights, dtype=np.float32)
            self.code_weights = torch.from_numpy(self.code_weights).cuda()

    @staticmethod
    def smooth_l1_loss(diff, beta):
        if beta < 1e-5:
            loss = torch.abs(diff)
        else:
            n = torch.abs(diff)
            loss = torch.where(n < beta, 0.5 * n ** 2 / beta, n - 0.5 * beta)

        return loss

    def forward(self, input: torch.Tensor,
                target: torch.Tensor, weights: torch.Tensor = None):
        """
        Args:
            input: (B, #anchors, #codes) float tensor.
                Ecoded predicted locations of objects.
            target: (B, #anchors, #codes) float tensor.
                Regression targets.
            weights: (B, #anchors) float tensor if not None.

        Returns:
            loss: (B, #anchors) float tensor.
                Weighted smooth l1 loss without reduction.
        """
        target = torch.where(torch.isnan(target), input, target)  # ignore nan targets

        diff = input - target
        loss = self.smooth_l1_loss(diff, self.beta)

        # anchor-wise weighting
        if weights is not None:
            assert weights.shape[0] == loss.shape[0] and weights.shape[1] == loss.shape[1]
            loss = loss * weights.unsqueeze(-1)

        return loss


class PointPillarV2VRobustLoss(nn.Module):
    def __init__(self, args):
        super(PointPillarV2VRobustLoss, self).__init__()
        self.reg_loss_func = WeightedSmoothL1Loss()
        self.score_loss_func = nn.BCELoss(reduce=True, reduction="mean")
        self.pose_loss_func = nn.SmoothL1Loss(reduce=True, reduction="mean", beta=1.0/9)
        self.alpha = 0.25
        self.gamma = 2.0

        self.cls_weight = args['cls_weight']
        self.reg_coe = args['reg']
        self.score_weight = args['score_weight']
        self.pose_weight = args['pose_weight']
        self.loss_dict = {}

    def forward(self, output_dict, target_dict):
        """
        Parameters
        ----------
        output_dict : dict
        target_dict : dict
        """
        stage = output_dict['stage']

        if stage == 0 or stage == 2:
            rm = output_dict['rm']  # [B, 14, 50, 176]
            psm = output_dict['psm'] # [B, 2, 50, 176]
            targets = target_dict['targets']

            cls_preds = psm.permute(0, 2, 3, 1).contiguous() # N, C, H, W -> N, H, W, C

            box_cls_labels = target_dict['pos_equal_one']  # [B, 50, 176, 2]
            box_cls_labels = box_cls_labels.view(psm.shape[0], -1).contiguous()

            positives = box_cls_labels > 0
            negatives = box_cls_labels == 0
            negative_cls_weights = negatives * 1.0
            cls_weights = (negative_cls_weights + 1.0 * positives).float()
            reg_weights = positives.float()

            pos_normalizer = positives.sum(1, keepdim=True).float()
            reg_weights /= torch.clamp(pos_normalizer, min=1.0)
            cls_weights /= torch.clamp(pos_normalizer, min=1.0)
            cls_targets = box_cls_labels
            cls_targets = cls_targets.unsqueeze(dim=-1)

            cls_targets = cls_targets.squeeze(dim=-1)
            one_hot_targets = torch.zeros(
                *list(cls_targets.shape), 2,
                dtype=cls_preds.dtype, device=cls_targets.device
            )
            one_hot_targets.scatter_(-1, cls_targets.unsqueeze(dim=-1).long(), 1.0)
            cls_preds = cls_preds.view(psm.shape[0], -1, 1)
            one_hot_targets = one_hot_targets[..., 1:]

            cls_loss_src = self.cls_loss_func(cls_preds,
                                            one_hot_targets,
                                            weights=cls_weights)  # [N, M]
            cls_loss = cls_loss_src.sum() / psm.shape[0]
            conf_loss = cls_loss * self.cls_weight

            # regression
            rm = rm.permute(0, 2, 3, 1).contiguous()
            rm = rm.view(rm.size(0), -1, 7)
            targets = targets.view(targets.size(0), -1, 7)
            box_preds_sin, reg_targets_sin = self.add_sin_difference(rm,
                                                                    targets)
            loc_loss_src =\
                self.reg_loss_func(box_preds_sin,
                                reg_targets_sin,
                                weights=reg_weights)
            reg_loss = loc_loss_src.sum() / rm.shape[0]
            reg_loss *= self.reg_coe

            total_loss = reg_loss + conf_loss

            self.loss_dict.update({'total_loss': total_loss,
                                'reg_loss': reg_loss,
                                'conf_loss': conf_loss})
        else:
            total_loss = 0

        

        # robust v2vnet part
        record_len = target_dict['record_len'] # we can also put this in output_dict
        if stage == 0:
            scores = output_dict['scores']
            choice = output_dict['choice']
            
            score_loss = self.attention_loss(scores, choice, record_len)
            total_loss += self.score_weight * score_loss
            self.loss_dict.update({'total_loss': total_loss,
                                    'score_loss': score_loss})
                                    
        elif stage == 1 or stage == 2:
            pairwise_corr = output_dict['pairwise_corr']
            pairwise_t_matrix = output_dict['pairwise_t_matrix']
            pairwise_t_matrix_gt = target_dict['pairwise_t_matrix']

            pose_loss = self.pose_loss(pairwise_corr, pairwise_t_matrix, pairwise_t_matrix_gt, record_len)
            total_loss += self.pose_weight * pose_loss
            self.loss_dict.update({'total_loss': total_loss,
                                    'pose_loss': pose_loss})
        

        return total_loss

    def attention_loss(self, scores, choices, record_len):
        """
        Args:
            scores: (B, L, L)
                scores[b,i,i] is already 0. 
            choices: (sum(N_cav), 1)
                0 is strong noise, 1 is weak noise
            record_len: 
                list, shape [B]
        """
        # first build gt label from choice
        B = scores.shape[0]
        choice_split = regroup(choices, record_len)
        label = torch.zeros_like(scores, device=scores.device)
        mask = torch.zeros_like(scores, device=scores.device)
        for b in range(B):
            N = record_len[b]
            choice = choice_split[b].float() # [N, 1]
            choice = choice @ choice.T # [N, N]
            
            gamma = 0.85
            label[b,:N,:N] = choice * gamma + (1-choice) * (1-gamma) # [N, N]

            mask[b,:N,:N] = 1
            mask[b,range(N),range(N)] = 0

        mask = mask.bool()

        input = torch.masked_select(scores, mask) 
        target = torch.masked_select(label, mask)
        print("input:", input)
        print("target:", target)

        return self.score_loss_func(input, target)

    def pose_loss(self, pairwise_corr,  pairwise_t_matrix, pairwise_t_matrix_gt, record_len):
        """
        Args:
            pairwise_corr: [B, L, L, 3]
            pairwise_t_matrix/pairwise_t_matrix_gt: [B,L,L,4,4]
            record_len: list, shape [B]
        """

        pairwise_t_matrix_gt = pairwise_t_matrix_gt.float()
        B, L = pairwise_t_matrix.shape[:2]
        mask = torch.zeros((B, L, L), device = pairwise_t_matrix.device)

        for b in range(B):
            N = record_len[b]
            mask[b,:N,:N] = 1
            mask[b,range(N), range(N)] = 0
        
        pair_corr_gt = torch.linalg.solve(pairwise_t_matrix.transpose(-2,-1), pairwise_t_matrix_gt.transpose(-2,-1)).transpose(-2,-1)
 
        yaw = pairwise_corr[..., 2] # [B,L,L]
        yaw_gt = torch.rad2deg(torch.atan2(pair_corr_gt[...,1,0], pair_corr_gt[...,0,0])) # [B,L,L]

        x = pairwise_corr[..., 0] # [B,L,L]
        x_gt = pair_corr_gt[..., 0,3]

        y = pairwise_corr[..., 1] # [B,L,L]
        y_gt = pair_corr_gt[..., 1,3]

        mask = mask.bool()
        mask = mask.view(B,L,L) # [B, L, L, ]

        input_x = torch.masked_select(x, mask)
        target_x = torch.masked_select(x_gt, mask)

        input_y = torch.masked_select(y, mask)
        target_y = torch.masked_select(y_gt, mask)

        input_yaw = torch.masked_select(yaw, mask)
        target_yaw = torch.masked_select(yaw_gt, mask)

        loss_x = self.pose_loss_func(input_x, target_x)
        loss_y = self.pose_loss_func(input_y, target_y)
        loss_yaw = self.pose_loss_func(input_yaw, target_yaw)

        lambda_trans = 2/3
        lambda_rot = 1/3
        print("x:\n", input_x, "\n", target_x)
        print("y:\n", input_y, "\n", target_y)
        print("yaw:\n", input_yaw, "\n", target_yaw)

        return lambda_trans * (loss_x + loss_y) + lambda_rot * loss_yaw



    def cls_loss_func(self, input: torch.Tensor,
                      target: torch.Tensor,
                      weights: torch.Tensor):
        """
        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets
            weights: (B, #anchors) float tensor.
                Anchor-wise weights.

        Returns:
            weighted_loss: (B, #anchors, #classes) float tensor after weighting.
        """
        pred_sigmoid = torch.sigmoid(input)
        alpha_weight = target * self.alpha + (1 - target) * (1 - self.alpha)
        pt = target * (1.0 - pred_sigmoid) + (1.0 - target) * pred_sigmoid
        focal_weight = alpha_weight * torch.pow(pt, self.gamma)

        bce_loss = self.sigmoid_cross_entropy_with_logits(input, target)

        loss = focal_weight * bce_loss

        if weights.shape.__len__() == 2 or \
                (weights.shape.__len__() == 1 and target.shape.__len__() == 2):
            weights = weights.unsqueeze(-1)

        assert weights.shape.__len__() == loss.shape.__len__()

        return loss * weights

    @staticmethod
    def sigmoid_cross_entropy_with_logits(input: torch.Tensor, target: torch.Tensor):
        """ PyTorch Implementation for tf.nn.sigmoid_cross_entropy_with_logits:
            max(x, 0) - x * z + log(1 + exp(-abs(x))) in
            https://www.tensorflow.org/api_docs/python/tf/nn/sigmoid_cross_entropy_with_logits

        Args:
            input: (B, #anchors, #classes) float tensor.
                Predicted logits for each class
            target: (B, #anchors, #classes) float tensor.
                One-hot encoded classification targets

        Returns:
            loss: (B, #anchors, #classes) float tensor.
                Sigmoid cross entropy loss without reduction
        """
        loss = torch.clamp(input, min=0) - input * target + \
               torch.log1p(torch.exp(-torch.abs(input)))
        return loss

    @staticmethod
    def add_sin_difference(boxes1, boxes2, dim=6):
        assert dim != -1
        rad_pred_encoding = torch.sin(boxes1[..., dim:dim + 1]) * \
                            torch.cos(boxes2[..., dim:dim + 1])
        rad_tg_encoding = torch.cos(boxes1[..., dim:dim + 1]) * \
                          torch.sin(boxes2[..., dim:dim + 1])

        boxes1 = torch.cat([boxes1[..., :dim], rad_pred_encoding,
                            boxes1[..., dim + 1:]], dim=-1)
        boxes2 = torch.cat([boxes2[..., :dim], rad_tg_encoding,
                            boxes2[..., dim + 1:]], dim=-1)
        return boxes1, boxes2


    def logging(self, epoch, batch_id, batch_len, writer = None):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        total_loss = self.loss_dict['total_loss'].item()
        if 'reg_loss' in self.loss_dict:
            reg_loss = self.loss_dict['reg_loss'].item()
        else:
            reg_loss = 0
        if 'conf_loss' in self.loss_dict:
            conf_loss = self.loss_dict['conf_loss'].item()
        else:
            conf_loss = 0
        if "score_loss" in self.loss_dict:
            score_loss = self.loss_dict['score_loss']
        else:
            score_loss = 0
        if "pose_loss" in self.loss_dict:
            pose_loss = self.loss_dict['pose_loss']
        else:
            pose_loss = 0

        print("[epoch %d][%d/%d], || Loss: %.4f || Conf Loss: %.4f"
              " || Loc Loss: %.4f || Score Loss: %.4f || Pose Loss: %.4f" % (
                  epoch, batch_id + 1, batch_len,
                  total_loss, conf_loss, reg_loss, score_loss, pose_loss))

        if not writer is None:
            writer.add_scalar('Regression_loss', reg_loss,
                            epoch*batch_len + batch_id)
            writer.add_scalar('Confidence_loss', conf_loss,
                            epoch*batch_len + batch_id)
            writer.add_scalar('Score_loss', score_loss,
                            epoch*batch_len + batch_id)
            writer.add_scalar('Pose_loss', pose_loss,
                            epoch*batch_len + batch_id)
opencood/loss/point_pillar_v2v_robust_loss.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib



import torch
import torch.nn as nn
import torch.nn.functional as F


class VoxelNetLoss(nn.Module):
    def __init__(self, args):
        super(VoxelNetLoss, self).__init__()
        self.smoothl1loss = nn.SmoothL1Loss(size_average=False)
        self.alpha = args['alpha']
        self.beta = args['beta']
        self.reg_coe = args['reg']
        self.loss_dict = {}

    def forward(self, output_dict, target_dict):
        """
        Parameters
        ----------
        output_dict : dict
        target_dict : dict
        """
        rm = output_dict['rm']
        psm = output_dict['psm']

        pos_equal_one = target_dict['pos_equal_one']
        neg_equal_one = target_dict['neg_equal_one']
        targets = target_dict['targets']

        p_pos = F.sigmoid(psm.permute(0, 2, 3, 1))
        rm = rm.permute(0, 2, 3, 1).contiguous()
        rm = rm.view(rm.size(0), rm.size(1), rm.size(2), -1, 7)
        targets = targets.view(targets.size(0), targets.size(1),
                               targets.size(2), -1, 7)
        pos_equal_one_for_reg = pos_equal_one.unsqueeze(
            pos_equal_one.dim()).expand(-1, -1, -1, -1, 7)

        rm_pos = rm * pos_equal_one_for_reg
        targets_pos = targets * pos_equal_one_for_reg

        cls_pos_loss = -pos_equal_one * torch.log(p_pos + 1e-6)
        cls_pos_loss = cls_pos_loss.sum() / (pos_equal_one.sum() + 1e-6)

        cls_neg_loss = -neg_equal_one * torch.log(1 - p_pos + 1e-6)
        cls_neg_loss = cls_neg_loss.sum() / (neg_equal_one.sum() + 1e-6)

        reg_loss = self.smoothl1loss(rm_pos, targets_pos)
        reg_loss = reg_loss / (pos_equal_one.sum() + 1e-6)
        conf_loss = self.alpha * cls_pos_loss + self.beta * cls_neg_loss

        total_loss = self.reg_coe * reg_loss + conf_loss

        self.loss_dict.update({'total_loss': total_loss,
                               'reg_loss': reg_loss,
                               'conf_loss': conf_loss})

        return total_loss

    def logging(self, epoch, batch_id, batch_len, writer = None):
        """
        Print out  the loss function for current iteration.

        Parameters
        ----------
        epoch : int
            Current epoch for training.
        batch_id : int
            The current batch.
        batch_len : int
            Total batch length in one iteration of training,
        writer : SummaryWriter
            Used to visualize on tensorboard
        """
        total_loss = self.loss_dict['total_loss']
        reg_loss = self.loss_dict['reg_loss']
        conf_loss = self.loss_dict['conf_loss']

        print("[epoch %d][%d/%d], || Loss: %.4f || Conf Loss: %.4f"
              " || Loc Loss: %.4f" % (
                  epoch, batch_id + 1, batch_len,
                  total_loss.item(), conf_loss.item(), reg_loss.item()))

        if not writer is None:
            writer.add_scalar('Regression_loss', reg_loss.item(),
                            epoch*batch_len + batch_id)
            writer.add_scalar('Confidence_loss', conf_loss.item(),
                            epoch*batch_len + batch_id)

opencood/loss/voxel_net_loss.py
MediaBrain-SJTU/CoBEVFlow
import torch
from torch import nn
import numpy as np

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
from opencood.models.sub_modules.cia_ssd_utils import SSFA, Head


class CIASSD(nn.Module):
    def __init__(self, args):
        super(CIASSD, self).__init__()
        lidar_range = np.array(args['lidar_range'])
        grid_size = np.round((lidar_range[3:6] - lidar_range[:3]) /
                             np.array(args['voxel_size'])).astype(np.int64)
        self.vfe = MeanVFE(args['mean_vfe'], args['mean_vfe']['num_point_features'])
        self.spconv_block = VoxelBackBone8x(args['spconv'],
                                            input_channels=args['spconv']['num_features_in'],
                                            grid_size=grid_size)
        self.map_to_bev = HeightCompression(args['map2bev'])
        self.ssfa = SSFA(args['ssfa'])
        self.head = Head(**args['head'])

    def forward(self, batch_dict):
        voxel_features = batch_dict['processed_lidar']['voxel_features']
        voxel_coords = batch_dict['processed_lidar']['voxel_coords']
        voxel_num_points = batch_dict['processed_lidar']['voxel_num_points']

        # save memory
        batch_dict.pop('processed_lidar')
        batch_dict.update({'voxel_features': voxel_features,
                           'voxel_coords': voxel_coords,
                           'voxel_num_points': voxel_num_points})

        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['spatial_features'])
        out = self.head(out)
        batch_dict['preds_dict_stage1'] = out

        return batch_dict



if __name__=="__main__":
    model = SSFA(None)
    print(model)
opencood/models/ciassd.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yue Hu <phyllis1sjtu@outlook.com>
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import numpy as np

class Communication(nn.Module):
    def __init__(self, args):
        super(Communication, self).__init__()
        
        self.smooth = False
        self.thre = args['thre']
        if 'gaussian_smooth' in args:
            # Gaussian Smooth
            self.smooth = True
            kernel_size = args['gaussian_smooth']['k_size']
            c_sigma = args['gaussian_smooth']['c_sigma']
            self.gaussian_filter = nn.Conv2d(1, 1, kernel_size=kernel_size, stride=1, padding=(kernel_size-1)//2)
            self.init_gaussian_filter(kernel_size, c_sigma)
            self.gaussian_filter.requires_grad = False
        
    def init_gaussian_filter(self, k_size=5, sigma=1):
        def _gen_gaussian_kernel(k_size=5, sigma=1):
            center = k_size // 2
            x, y = np.mgrid[0 - center : k_size - center, 0 - center : k_size - center]
            g = 1 / (2 * np.pi * sigma) * np.exp(-(np.square(x) + np.square(y)) / (2 * np.square(sigma)))
            return g
        gaussian_kernel = _gen_gaussian_kernel(k_size, sigma)
        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 forward(self, batch_confidence_maps, record_len, pairwise_t_matrix):
        # batch_confidence_maps:[(L1, H, W), (L2, H, W), ...]
        # pairwise_t_matrix: (B,L,L,2,3)
        # thre: threshold of objectiveness
        # a_ji = (1 - q_i)*q_ji
        B, L, _, _, _ = pairwise_t_matrix.shape
        _, _, H, W = batch_confidence_maps[0].shape
        
        communication_masks = []
        communication_rates = []
        batch_communication_maps = []
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            # t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            ori_communication_maps = batch_confidence_maps[b].sigmoid().max(dim=1)[0].unsqueeze(1) # dim1=2 represents the confidence of two anchors
            
            if self.smooth:
                communication_maps = self.gaussian_filter(ori_communication_maps)
            else:
                communication_maps = ori_communication_maps

            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.thre, ones_mask, zeros_mask)# TODO: thre=0.01, 所以这里面mask全1 

            communication_rate = communication_mask[0].sum()/(H*W)

            # communication_mask = warp_affine_simple(communication_mask,
            #                                 t_matrix[0, :, :, :],
            #                                 (H, W))
            
            communication_mask_nodiag = communication_mask.clone() #TODO: question, 这里面把 0, 2, 4, ... 车的mask置1，是为什么呀
            ones_mask = torch.ones_like(communication_mask).to(communication_mask.device)
            communication_mask_nodiag[::2] = ones_mask[::2]

            communication_masks.append(communication_mask_nodiag)
            communication_rates.append(communication_rate) #TODO: question, 这里面先计算 rate，然后再把 mask 修改了 
            batch_communication_maps.append(ori_communication_maps*communication_mask_nodiag) # (confidence > thre 的位置以及偶数id的cav 是1，其余是0) * confidence
        communication_rates = sum(communication_rates)/B
        communication_masks = torch.concat(communication_masks, dim=0) # confidence > thre 的位置以及偶数id的cav 是1，其余是0；然后concat到一起
        return batch_communication_maps, communication_masks, communication_rates 
opencood/models/comm_modules/where2comm.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yue Hu <phyllis1sjtu@outlook.com>, Sizhe Wei <sizhewei@sjtu.edu.cn>
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import numpy as np

class Communication(nn.Module):
    def __init__(self, args):
        super(Communication, self).__init__()
        
        self.smooth = False
        self.thre = args['thre']
        self.num_blocks_thres = 1
        if 'num_blocks_thres' in args:
            self.num_blocks_thres = args['num_blocks_thres']
        if 'gaussian_smooth' in args:
            # Gaussian Smooth
            self.smooth = True
            kernel_size = args['gaussian_smooth']['k_size']
            c_sigma = args['gaussian_smooth']['c_sigma']
            self.gaussian_filter = nn.Conv2d(1, 1, kernel_size=kernel_size, stride=1, padding=(kernel_size-1)//2)
            self.init_gaussian_filter(kernel_size, c_sigma)
            self.gaussian_filter.requires_grad = False
        
    def init_gaussian_filter(self, k_size=5, sigma=1):
        def _gen_gaussian_kernel(k_size=5, sigma=1):
            center = k_size // 2
            x, y = np.mgrid[0 - center : k_size - center, 0 - center : k_size - center]
            g = 1 / (2 * np.pi * sigma) * np.exp(-(np.square(x) + np.square(y)) / (2 * np.square(sigma)))
            return g
        gaussian_kernel = _gen_gaussian_kernel(k_size, sigma)
        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 forward(self, batch_confidence_maps, record_len, pairwise_t_matrix):
        """ 

        Parameters:
        -----------
        batch_confidence_maps:[(L1xk, 2, H, W), (L2xk, 2, H, W), ..., (L_B x k, 2, H, W)]
        confidence_maps (without regroup func): (B, 2, H, W)  B
        pairwise_t_matrix: (B,L,k,2,3)

        Returns:
        --------
        batch_communication_maps: list
            orignal confidence map x mask

        communication_masks: list / torch.Tensor(修改后是 list, 方便下游处理任务)
            mask, shape:(B, 1, H, W). 0 or 1

        communication_rates: float

        Memos:
        ------
        self.thre: threshold of objectiveness
        a_ji = (1 - q_i)*q_ji
        """
        
        '''
        B: batch_size
        L: max cav num
        k: num of past frames
        H, W: feature map height and weight
        '''
        B, L, k, _, _ = pairwise_t_matrix.shape
        _, _, H, W = batch_confidence_maps[0].shape
        
        communication_masks = []
        communication_rates = []
        batch_communication_maps = []
        all_blocks = H*W
        num_valid_thres = int(all_blocks * self.num_blocks_thres)
        for b in range(B):

            ori_communication_maps = batch_confidence_maps[b].sigmoid().max(dim=1)[0].unsqueeze(1) # dim1=2 represents the confidence of two anchors, (num_cav * k, 1, H, W)
            
            if self.smooth:
                communication_maps = self.gaussian_filter(ori_communication_maps)
            else:
                communication_maps = ori_communication_maps

            if self.num_blocks_thres<1 :
                curr_batch = communication_maps.shape[0] # num_cav * k
                communication_mask = (communication_maps>self.thre).to(torch.int)
                valid_nums = torch.sum(communication_mask, dim=(-1, -2, -3))

                final_comm_mask = torch.zeros_like(communication_mask)
                for i in range(curr_batch):
                    if valid_nums[i] > num_valid_thres :
                        # 选择 comm_maps[i] 里面 top num_valid 的位置
                        tmp = communication_maps[i].reshape((1, -1))
                        _, idx = torch.topk(tmp, num_valid_thres)
                        # print(idx)
                        curr_mask = torch.zeros_like(tmp)
                        curr_mask[0, idx[0]] = 1
                        final_comm_mask[i] = curr_mask.reshape((1, H, W))
                    else:
                        final_comm_mask[i] = communication_mask[i]
                communication_mask = final_comm_mask
            else:
                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.thre, ones_mask, zeros_mask)

            communication_rate = communication_mask[0].sum()/(H*W)

            communication_masks.append(communication_mask)
            communication_rates.append(communication_rate) 
            batch_communication_maps.append(ori_communication_maps*communication_mask) # (confidence > thre 的位置以及偶数id的cav 是1，其余是0) * confidence
        communication_rates = sum(communication_rates)/B
        # communication_masks = torch.concat(communication_masks, dim=0) # confidence > thre 的位置以及偶数id的cav 是1，其余是0；然后concat到一起
        return batch_communication_maps, communication_masks, communication_rates 
opencood/models/comm_modules/where2comm_multisweep.py
MediaBrain-SJTU/CoBEVFlow
import random, os

import torch
from torch import nn
import numpy as np

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
from opencood.models.sub_modules.cia_ssd_utils import SSFA, Head
from opencood.models.sub_modules.vsa import VoxelSetAbstraction
from opencood.models.sub_modules.roi_head import RoIHead
from opencood.models.sub_modules.matcher import Matcher
from opencood.data_utils.post_processor.fpvrcnn_postprocessor import \
    FpvrcnnPostprocessor

class FPVRCNN(nn.Module):
    def __init__(self, args):
        super(FPVRCNN, self).__init__()
        lidar_range = np.array(args['lidar_range'])
        grid_size = np.round((lidar_range[3:6] - lidar_range[:3]) /
                             np.array(args['voxel_size'])).astype(np.int64)
        self.vfe = MeanVFE(args['mean_vfe'],
                           args['mean_vfe']['num_point_features'])
        self.spconv_block = VoxelBackBone8x(args['spconv'],
                                            input_channels=args['spconv'][
                                                'num_features_in'],
                                            grid_size=grid_size)
        self.map_to_bev = HeightCompression(args['map2bev'])
        self.ssfa = SSFA(args['ssfa'])
        self.head = Head(**args['head'])
        self.post_processor = FpvrcnnPostprocessor(args['post_processer'],
                                                   train=True)
        self.vsa = VoxelSetAbstraction(args['vsa'], args['voxel_size'],
                                       args['lidar_range'],
                                       num_bev_features=128,
                                       num_rawpoint_features=3)
        self.matcher = Matcher(args['matcher'], args['lidar_range'])
        self.roi_head = RoIHead(args['roi_head'])
        self.train_stage2 = args['activate_stage2']
        self.discrete_ratio = args['voxel_size'][0]
        self.proj_first = args['proj_first']

    def forward(self, batch_dict):
        voxel_features = batch_dict['processed_lidar']['voxel_features']
        voxel_coords = batch_dict['processed_lidar']['voxel_coords']
        voxel_num_points = batch_dict['processed_lidar']['voxel_num_points']

        # save memory
        batch_dict.pop('processed_lidar')
        batch_dict.update({'voxel_features': voxel_features,
                           'voxel_coords': voxel_coords,
                           'voxel_num_points': voxel_num_points,
                           'batch_size': int(batch_dict['record_len'].sum()),
                           'proj_first': self.proj_first,
                           'lidar_pose': batch_dict['lidar_pose']})

        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['spatial_features'])
        batch_dict['preds_dict_stage1'] = self.head(out) 

        data_dict, output_dict = {}, {}
        data_dict['ego'], output_dict['ego'] = batch_dict, batch_dict

        pred_box3d_list, scores_list = \
            self.post_processor.post_process(data_dict, output_dict,
                                             stage1=True)

        # if proj_first is False
        # the boxes are predicted in each coordinate              
        batch_dict['det_boxes'] = pred_box3d_list
        batch_dict['det_scores'] = scores_list

        if pred_box3d_list is not None and self.train_stage2:
            batch_dict = self.vsa(batch_dict)
            batch_dict = self.matcher(batch_dict)
            batch_dict = self.roi_head(batch_dict)

        return batch_dict





if __name__ == "__main__":
    model = SSFA(None)
    print(model)
opencood/models/fpvrcnn.py
MediaBrain-SJTU/CoBEVFlow

import numpy as np
import torch
import functools
import torch.nn as nn
import torch.nn.functional as F
from timm.models.layers import DropPath, trunc_normal_
import ipdb
import copy
import matplotlib.pyplot as plt
import os

def initialize_weights(net_l, scale=1):
    if not isinstance(net_l, list):
        net_l = [net_l]
    for net in net_l:
        for m in net.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
                m.weight.data *= scale  # for residual block
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.Linear):
                nn.init.kaiming_normal_(m.weight, a=0, mode='fan_in')
                m.weight.data *= scale
                if m.bias is not None:
                    m.bias.data.zero_()
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias.data, 0.0)


def make_layer(block, n_layers):
    layers = []
    for _ in range(n_layers):
        layers.append(block())
    return nn.Sequential(*layers)


class ResidualBlock_noBN(nn.Module):
    '''Residual block w/o BN
    ---Conv-ReLU-Conv-+-
     |________________|
    '''

    def __init__(self, nf=64):
        super(ResidualBlock_noBN, self).__init__()
        self.conv1 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)
        self.conv2 = nn.Conv2d(nf, nf, 3, 1, 1, bias=True)

        # initialization
        initialize_weights([self.conv1, self.conv2], 0.1)

    def forward(self, x):
        identity = x
        out = F.relu(self.conv1(x), inplace=True)
        out = self.conv2(out)
        return identity + out

def window_partition(x, window_size):
    """Partition the input video sequences into several windows along spatial 
    dimensions.
    Args:
        x (torch.Tensor): (B, D, H, W, C)
        window_size (tuple[int]): Window size
    Returns: 
        windows: (B*nW, D, Wh, Ww, C)
    """
    B, D, H, W, C = x.shape
    # B, D, num_Hwin, Wh, num_Wwin, Ww, C
    x = x.view(B, D, H // window_size[0], window_size[0], W // window_size[1], window_size[1], C) 
    windows = x.permute(0, 2, 4, 1, 3, 5, 6).contiguous().view(-1, D, window_size[0], window_size[1], C)
    return windows

def window_reverse(windows, window_size, B, D, H, W):
    """Reverse window partition.
    Args:
        windows (torch.Tensor): (B*nW, D, Wh, Ww, C)
        window_size (tuple[int]): Window size
        B (int): Number of batches
        D (int): Number of frames
        H (int): Height of image
        W (int): Width of image
    Returns:
        x: (B, D, H, W, C)
    """
    x = windows.view(B, H // window_size[0], W // window_size[1], D, window_size[0], window_size[1], -1)
    x = x.permute(0, 3, 1, 4, 2, 5, 6).contiguous().view(B, D, H, W, -1)
    return x

def get_window_size(x_size, window_size, shift_size=None):
    """Adjust window size and shift size based on the size of the input.
    Args:
        x_size (tuple[int]): The shape of x.
        window_size (tuple[int]): Window size.
        shift_size (tuple[int], optional): Shift size. Defaults to None.
    Returns:
        use_window_size: Window size for use.
        use_shift_size: Shift size for use.
    """
    use_window_size = list(window_size)
    if shift_size is not None:
        use_shift_size = list(shift_size)
    for i in range(len(x_size)):
        if x_size[i] <= window_size[i]:
            use_window_size[i] = x_size[i]
            if shift_size is not None:
                use_shift_size[i] = 0

    if shift_size is None:
        return tuple(use_window_size)
    else:
        return tuple(use_window_size), tuple(use_shift_size)

class Mlp(nn.Module):
    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        self.fc1 = nn.Linear(in_features, hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(hidden_features, out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x

class WindowAttention3D(nn.Module):
    """Window based multi-head self/cross attention (W-MSA/W-MCA) module with relative 
    position bias. 
    It supports both of shifted and non-shifted window.
    """
    def __init__(self, dim, num_frames_q, num_frames_kv, window_size, num_heads, 
                 qkv_bias=True, qk_scale=None, attn_drop=0., proj_drop=0.):
        """Initialization function.
        Args:
            dim (int): Number of input channels.
            num_frames (int): Number of input frames.
            window_size (tuple[int]): The size of the window.
            num_heads (int): Number of attention heads.
            qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Defaults to True.
            qk_scale (float, optional): Override default qk scale of head_dim ** -0.5 if set. Defaults to None.
            attn_drop (float, optional): Dropout ratio of attention weight. Defaults to 0.0
            proj_drop (float, optional): Dropout ratio of output. Defaults to 0.0
        """
        super().__init__()
        self.dim = dim
        self.num_frames_q = num_frames_q # D1
        self.num_frames_kv = num_frames_kv # D2
        self.window_size = window_size  # Wh, Ww
        self.num_heads = num_heads # nH
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim ** -0.5

        # Define a parameter table of relative position bias
        self.relative_position_bias_table = nn.Parameter(
            torch.zeros((2 * num_frames_q - 1) * (2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)
        )  # 2*D-1 * 2*Wh-1 * 2*Ww-1, nH

        # Get pair-wise relative position index for each token inside the window
        coords_d_q = torch.arange(self.num_frames_q)
        # print(self.num_frames_q,self.num_frames_kv)
        if int((self.num_frames_q + 1) // self.num_frames_kv) != 0:
            coords_d_kv = torch.arange(0, self.num_frames_q, int((self.num_frames_q + 1) // self.num_frames_kv))
        else:
            coords_d_kv = torch.tensor([0]*self.num_frames_kv)
        # print(coords_d_kv,'coords_d_kv')
        # coords_d_kv = torch.tensor([0])
        coords_h = torch.arange(self.window_size[0])
        coords_w = torch.arange(self.window_size[1])
        coords_q = torch.stack(torch.meshgrid([coords_d_q, coords_h, coords_w]))  # 3, D2, Wh, Ww
        coords_kv = torch.stack(torch.meshgrid([coords_d_kv, coords_h, coords_w]))  # 3, D1, Wh, Ww
        # print('coords_q',coords_q.shape,'coords_kv',coords_kv.shape)
        coords_q_flatten = torch.flatten(coords_q, 1)  # 3, D1*Wh*Ww
        coords_kv_flatten = torch.flatten(coords_kv, 1)  # 3, D2*Wh*Ww
        relative_coords = coords_q_flatten[:, :, None] - coords_kv_flatten[:, None, :]  # 3, D1*Wh*Ww, D2*Wh*Ww
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()  # D1*Wh*Ww, D2*Wh*Ww, 3
        relative_coords[:, :, 0] += self.num_frames_q - 1  # shift to start from 0
        relative_coords[:, :, 1] += self.window_size[0] - 1
        relative_coords[:, :, 2] += self.window_size[1] - 1
        relative_coords[:, :, 0] *= (2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1)
        relative_coords[:, :, 1] *= 2 * self.window_size[1] - 1
        relative_position_index = relative_coords.sum(-1)  # D1*Wh*Ww, D2*Wh*Ww
        self.register_buffer("relative_position_index", relative_position_index)

        self.q = nn.Linear(dim, dim, bias=qkv_bias)
        self.kv = nn.Linear(dim, dim * 2, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        trunc_normal_(self.relative_position_bias_table, std=.02)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, q, kv=None, mask=None):
        """Forward function.
        Args:
            q (torch.Tensor): (B*nW, D1*Wh*Ww, C)
            kv (torch.Tensor): (B*nW, D2*Wh*Ww, C). Defaults to None.
            mask (torch.Tensor, optional): Mask for shifted window attention (nW, D1*Wh*Ww, D2*Wh*Ww). Defaults to None.
        Returns:
            torch.Tensor: (B*nW, D1*Wh*Ww, C)
        """
        # ipdb.set_trace()
        kv = q if kv is None else kv
        B_, N1, C = q.shape # N1 = D1*Wh*Ww, B_ = B*nW
        B_, N2, C = kv.shape # N2 = D2*Wh*Ww, B_ = B*nW
        # print('N1',N1,'N2',N2)
        q = self.q(q).reshape(B_, N1, 1, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        kv = self.kv(kv).reshape(B_, N2, 2, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
        q, k, v = q[0], kv[0], kv[1] # B_, nH, N1(2), C
        q = q * self.scale
        attn = (q @ k.transpose(-2, -1)) # B_, nH, N1, N2

        relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)].view(
            N1, N2, -1)  # D1*Wh*Ww, D2*Wh*Ww, nH
        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, D1*Wh*Ww, D2*Wh*Ww
        attn = attn + relative_position_bias.unsqueeze(0) # B_, nH, D1*Wh*Ww, D2*Wh*Ww

        if mask is not None:
            nW = mask.shape[0]
            attn = attn.view(B_ // nW, nW, self.num_heads, N1, N2) + mask.unsqueeze(1).unsqueeze(0) # B, nW, nH, D1*Wh*Ww, D2*Wh*Ww
            attn = attn.view(-1, self.num_heads, N1, N2)
            attn = self.softmax(attn)
        else:
            attn = self.softmax(attn)

        attn = self.attn_drop(attn)

        x = (attn @ v).transpose(1, 2).reshape(B_, N1, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x, attn

class VSTSREncoderTransformerBlock(nn.Module):
    """Video spatial-temporal super-resolution encoder transformer block.
    """
    def __init__(self, dim, num_heads, num_frames=4, window_size=(8, 8), 
                 shift_size=(0, 0), mlp_ratio=4., qkv_bias=True, qk_scale=None,
                 drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, 
                 norm_layer=nn.LayerNorm,encoder_id = 0, layer_id = 0):
        """Initialization function.
        Args:
            dim (int): Number of input channels. 
            num_heads (int): Number of attention heads.
            num_frames (int): Number of input frames.
            window_size (tuple[int], optional): Window size. Defaults to 8.
            shift_size (tuple[int], optional): Shift size. Defaults to 0.
            mlp_ratio (int, optional): Ratio of mlp hidden dim to embedding dim. Defaults to 4.
            qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Defaults to True.
            qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. Defaults to None.
            drop (float, optional): Dropout rate. Defaults to 0.
            attn_drop (float, optional): Attention dropout rate. Defaults to 0.
            drop_path (float, optional):  Stochastic depth rate. Defaults to 0.
            act_layer (nn.Module, optional): Activation layer. Defaults to nn.GELU.
            norm_layer (nn.Module, optional): Normalization layer. Defaults to nn.LayerNorm.
        """
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.num_frames = num_frames
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        assert 0 <= self.shift_size[0] < self.window_size[0], "shift_size must in 0-win_size"
        assert 0 <= self.shift_size[1] < self.window_size[1], "shift_size must in 0-win_size"

        self.norm1 = norm_layer(dim)
        self.attn = WindowAttention3D(
            dim, num_frames_q=self.num_frames, num_frames_kv=self.num_frames,
            window_size=self.window_size, num_heads=num_heads, 
            qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, 
            proj_drop=drop
        )

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        self.encoder_id = encoder_id
        self.layer_id = layer_id

    def forward(self, x, mask_matrix, data_id = 0):
        """Forward function.
        Args:
            x (torch.Tensor): (B, D, H, W, C)
            mask_matrix (torch.Tensor): (nW*B, D*Wh*Ww, D*Wh*Ww)
        Returns:
            torch.Tensor: (B, D, H, W, C)
        """
        # ipdb.set_trace()
        B, D, H, W, C = x.shape
        window_size, shift_size = get_window_size((H, W), self.window_size, self.shift_size)

        shortcut = x
        x = self.norm1(x)

        # Padding
        pad_b = (window_size[0] - H % window_size[0]) % window_size[0]
        pad_r = (window_size[1] - W % window_size[1]) % window_size[1]
        x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b, 0, 0))
        _, _, Hp, Wp, _ = x.shape

        # cyclic shift
        if any(i > 0 for i in shift_size):
            shifted_x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(2, 3))
            attn_mask = mask_matrix
        else:
            shifted_x = x
            attn_mask = None

        # partition windows
        x_windows = window_partition(shifted_x, window_size)  # nW*B, D, window_size, window_size, C
        x_windows = x_windows.view(-1, D * window_size[0] * window_size[1], C)  # nW*B, D*window_size*window_size, C

        # W-MSA/SW-MSA
        attn_windows = self.attn(x_windows, mask=attn_mask)[0]  # nW*B, D*window_size*window_size, C

        # merge windows
        attn_windows = attn_windows.view(-1, D, window_size[0], window_size[1], C)
        shifted_x = window_reverse(attn_windows, window_size, B, D, Hp, Wp)  # B, D, H, W, C
        # save_dir = '/dssg/home/acct-eezy/eezy-user1/wz/multi-cam-pnp_11/vis_stpt/{}/'.format(data_id)
        # if not os.path.exists(save_dir):
        #     os.mkdir(save_dir)
        
        # attention_matrix = copy.deepcopy(shifted_x).cpu()
        # for t in range(3):
        #     feature_map = attention_matrix[0, t]
        #     max_feature_map = torch.max(feature_map, -1)[0]
            
        #     plt.matshow(-max_feature_map.numpy(), cmap=plt.cm.Blues)
        #     plt.grid(False)
        #     # plt.imshow(-max_feature_map.numpy())
        #     plt.savefig(save_dir + 'encoder_{}_depth_{}_t_{}.png'.format(self.encoder_id, self.layer_id, t), bbox_inches='tight')
        # torch.save(shifted_x, save_dir + 'encoder_{}_depth_{}.pt'.format(self.encoder_id, self.layer_id))
        # ipdb.set_trace()

        # reverse cyclic shift
        if any(i > 0 for i in shift_size):
            x = torch.roll(shifted_x, shifts=(shift_size[0], shift_size[1]), dims=(2, 3))
        else:
            x = shifted_x

        if pad_r > 0 or pad_b > 0:
            x = x[:, :, :H, :W, :].contiguous()

        # FFN
        x = shortcut + self.drop_path(x)
        x = x + self.drop_path(self.mlp(self.norm2(x)))

        return x

class VSTSRDecoderTransformerBlock(nn.Module):
    """Video spatial-temporal super-resolution decoder transformer block.
    """
    def __init__(self, dim, num_heads, num_frames=4, num_out_frames=5, window_size=(8, 8), 
                 shift_size=(0, 0), mlp_ratio=4., qkv_bias=True, qk_scale=None,
                 drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, 
                 norm_layer=nn.LayerNorm):
        """Initialization function.
        Args:
            dim (int): Number of input channels. 
            num_heads (int): Number of attention heads.
            num_frames (int): Number of input frames.
            window_size (tuple[int], optional): Window size. Defaults to 8.
            shift_size (tuple[int], optional): Shift size. Defaults to 0.
            mlp_ratio (int, optional): Ratio of mlp hidden dim to embedding dim. Defaults to 4.
            qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Defaults to True.
            qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. Defaults to None.
            drop (float, optional): Dropout rate. Defaults to 0.
            attn_drop (float, optional): Attention dropout rate. Defaults to 0.
            drop_path (float, optional):  Stochastic depth rate. Defaults to 0.
            act_layer (nn.Module, optional): Activation layer. Defaults to nn.GELU.
            norm_layer (nn.Module, optional): Normalization layer. Defaults to nn.LayerNorm.
        """
        super().__init__()
        self.dim = dim
        self.num_heads = num_heads
        self.num_frames = num_frames
        # self.num_out_frames = 2 * num_frames - 1
        self.num_out_frames = num_out_frames
        self.window_size = window_size
        self.shift_size = shift_size
        self.mlp_ratio = mlp_ratio
        assert 0 <= self.shift_size[0] < self.window_size[0], "shift_size must in 0-win_size"
        assert 0 <= self.shift_size[1] < self.window_size[1], "shift_size must in 0-win_size"

        self.norm1 = norm_layer(dim)
        self.attn1 = WindowAttention3D(
            dim, num_frames_q=self.num_out_frames, 
            num_frames_kv=self.num_out_frames, window_size=self.window_size, 
            num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, 
            attn_drop=attn_drop, proj_drop=drop
        )
        self.attn2 = WindowAttention3D(
            dim, num_frames_q=self.num_out_frames, 
            num_frames_kv=self.num_frames, window_size=self.window_size, 
            num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, 
            attn_drop=attn_drop, proj_drop=drop
        )

        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
        self.norm2 = norm_layer(dim)
        self.norm3 = norm_layer(dim)
        self.norm_kv = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
        
       

    def forward(self, x, attn_kv, mask_matrix_q, mask_matrix_qkv):
        """Forward function.
        Args:
            x (torch.Tensor): (B, D1, H, W, C)
            attn_kv (torch.Tensor): (B, D2, H, W, C)
            mask_matrix_q (torch.Tensor): (nW*B, D1*Wh*Ww, D1*Wh*Ww)
            mask_matrix_qkv (torch.Tensor): (nW*B, D1*Wh*Ww, D2*Wh*Ww)
        Returns:
            torch.Tensor: (B, D1, H, W, C)
        """
        B, D1, H, W, C = x.shape
        B, D2, H, W, C = attn_kv.shape
        window_size, shift_size = get_window_size((H, W), self.window_size, self.shift_size)

        # shortcut = x
        # x = self.norm1(x)

        # Padding
        pad_b = (window_size[0] - H % window_size[0]) % window_size[0]
        pad_r = (window_size[1] - W % window_size[1]) % window_size[1]
        # x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b, 0, 0))
        # _, _, Hp, Wp, _ = x.shape

        # # cyclic shift
        # if any(i > 0 for i in shift_size):
        #     shifted_x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(2, 3))
        #     attn_mask_q = mask_matrix_q
        #     attn_mask_qkv = mask_matrix_qkv
        # else:
        #     shifted_x = x
        #     attn_mask_q = None
        #     attn_mask_qkv = None

        # # partition windows
        # x_windows = window_partition(shifted_x, window_size)  # nW*B, D1, window_size, window_size, C
        # x_windows = x_windows.view(-1, D1 * window_size[0] * window_size[1], C)  # nW*B, D1*window_size*window_size, C

        # # W-MSA/SW-MSA for query
        # attn_windows = self.attn1(x_windows, mask=attn_mask_q)[0] # nW*B, D1*window_size*window_size, C
        # attn_windows = attn_windows.view(-1, D1, window_size[0], window_size[1], C)
        # shifted_x = window_reverse(attn_windows, window_size, B, D1, Hp, Wp) # B, D1, Hp, Wp, C

        # # reverse cyclic shift
        # if any(i > 0 for i in shift_size):
        #     x = torch.roll(shifted_x, shifts=(shift_size[0], shift_size[1]), dims=(2, 3))
        # else:
        #     x = shifted_x
        
        # if pad_r > 0 or pad_b > 0:
        #     x = x[:, :, :H, :W, :].contiguous()

        # x = shortcut + self.drop_path(x)

        shortcut = x
        x = self.norm2(x)
        attn_kv = self.norm_kv(attn_kv)
        x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b, 0, 0))
        _, _, Hp, Wp, _ = x.shape
        attn_kv = F.pad(attn_kv, (0, 0, 0, pad_r, 0, pad_b, 0, 0))

        # cyclic shift
        if any(i > 0 for i in shift_size):
            shifted_x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(2, 3))
            shifted_attn_kv = torch.roll(attn_kv, shifts=(-shift_size[0], -shift_size[1]), dims=(2, 3))
            attn_mask_q = mask_matrix_q
            attn_mask_qkv = mask_matrix_qkv
        else:
            shifted_x = x
            shifted_attn_kv = attn_kv
            attn_mask_q = None
            attn_mask_qkv = None

        # partition windows
        x_windows = window_partition(shifted_x, window_size)  # nW*B, D1, window_size, window_size, C
        attn_kv_windows = window_partition(shifted_attn_kv, window_size)  # nW*B, D2, window_size, window_size, C
        x_windows = x_windows.view(-1, D1 * window_size[0] * window_size[1], C)  # nW*B, D1*window_size*window_size, C
        attn_kv_windows = attn_kv_windows.view(-1, D2 * window_size[0] * window_size[1], C)  # nW*B, D2*window_size*window_size, C

        # W-MSA/SW-MSA
        attn_windows = self.attn2(x_windows, attn_kv_windows, mask=attn_mask_qkv)[0]  # nW*B, D1*window_size*window_size, C

        # merge windows
        attn_windows = attn_windows.view(-1, D1, window_size[0], window_size[1], C)
        shifted_x = window_reverse(attn_windows, window_size, B, D1, Hp, Wp)  # B, D1, H, W, C

        # reverse cyclic shift
        if any(i > 0 for i in shift_size):
            x = torch.roll(shifted_x, shifts=(shift_size[0], shift_size[1]), dims=(2, 3))
        else:
            x = shifted_x

        if pad_r > 0 or pad_b > 0:
            x = x[:, :, :H, :W, :].contiguous()

        x = shortcut + self.drop_path(x)

        # FFN
        x = x + self.drop_path(self.mlp(self.norm3(x)))

        return x

# class VSTSRDecoderTransformerBlock(nn.Module):
#     """Video spatial-temporal super-resolution decoder transformer block.
#     """
#     def __init__(self, dim, num_heads, num_frames=4, num_out_frames=5, window_size=(8, 8), 
#                  shift_size=(0, 0), mlp_ratio=4., qkv_bias=True, qk_scale=None,
#                  drop=0., attn_drop=0., drop_path=0., act_layer=nn.GELU, 
#                  norm_layer=nn.LayerNorm):
#         """Initialization function.

#         Args:
#             dim (int): Number of input channels. 
#             num_heads (int): Number of attention heads.
#             num_frames (int): Number of input frames.
#             window_size (tuple[int], optional): Window size. Defaults to 8.
#             shift_size (tuple[int], optional): Shift size. Defaults to 0.
#             mlp_ratio (int, optional): Ratio of mlp hidden dim to embedding dim. Defaults to 4.
#             qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Defaults to True.
#             qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. Defaults to None.
#             drop (float, optional): Dropout rate. Defaults to 0.
#             attn_drop (float, optional): Attention dropout rate. Defaults to 0.
#             drop_path (float, optional):  Stochastic depth rate. Defaults to 0.
#             act_layer (nn.Module, optional): Activation layer. Defaults to nn.GELU.
#             norm_layer (nn.Module, optional): Normalization layer. Defaults to nn.LayerNorm.
#         """
#         super().__init__()
#         self.dim = dim
#         self.num_heads = num_heads
#         self.num_frames = num_frames
#         # self.num_out_frames = 2 * num_frames - 1
#         self.num_out_frames = num_out_frames
#         self.window_size = window_size
#         self.shift_size = shift_size
#         self.mlp_ratio = mlp_ratio
#         assert 0 <= self.shift_size[0] < self.window_size[0], "shift_size must in 0-win_size"
#         assert 0 <= self.shift_size[1] < self.window_size[1], "shift_size must in 0-win_size"

#         self.norm1 = norm_layer(dim)
#         self.attn1 = WindowAttention3D(
#             dim, num_frames_q=self.num_out_frames, 
#             num_frames_kv=self.num_out_frames, window_size=self.window_size, 
#             num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, 
#             attn_drop=attn_drop, proj_drop=drop
#         )
#         self.attn2 = WindowAttention3D(
#             dim, num_frames_q=self.num_out_frames, 
#             num_frames_kv=self.num_frames, window_size=self.window_size, 
#             num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, 
#             attn_drop=attn_drop, proj_drop=drop
#         )

#         self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
#         self.norm2 = norm_layer(dim)
#         self.norm3 = norm_layer(dim)
#         self.norm_kv = norm_layer(dim)
#         mlp_hidden_dim = int(dim * mlp_ratio)
#         self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
       

#     def forward(self, x, attn_kv, mask_matrix_q, mask_matrix_qkv):
#         """Forward function.

#         Args:
#             x (torch.Tensor): (B, D1, H, W, C)
#             attn_kv (torch.Tensor): (B, D2, H, W, C)
#             mask_matrix_q (torch.Tensor): (nW*B, D1*Wh*Ww, D1*Wh*Ww)
#             mask_matrix_qkv (torch.Tensor): (nW*B, D1*Wh*Ww, D2*Wh*Ww)

#         Returns:
#             torch.Tensor: (B, D1, H, W, C)
#         """
#         B, D1, H, W, C = x.shape
#         B, D2, H, W, C = attn_kv.shape
#         window_size, shift_size = get_window_size((H, W), self.window_size, self.shift_size)

#         shortcut = x
#         x = self.norm1(x)

#         # Padding
#         pad_b = (window_size[0] - H % window_size[0]) % window_size[0]
#         pad_r = (window_size[1] - W % window_size[1]) % window_size[1]
#         x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b, 0, 0))
#         _, _, Hp, Wp, _ = x.shape

#         # cyclic shift
#         if any(i > 0 for i in shift_size):
#             shifted_x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(2, 3))
#             attn_mask_q = mask_matrix_q
#             attn_mask_qkv = mask_matrix_qkv
#         else:
#             shifted_x = x
#             attn_mask_q = None
#             attn_mask_qkv = None

#         # partition windows
#         x_windows = window_partition(shifted_x, window_size)  # nW*B, D1, window_size, window_size, C
#         x_windows = x_windows.view(-1, D1 * window_size[0] * window_size[1], C)  # nW*B, D1*window_size*window_size, C

#         # W-MSA/SW-MSA for query
#         attn_windows = self.attn1(x_windows, mask=attn_mask_q)[0] # nW*B, D1*window_size*window_size, C
#         attn_windows = attn_windows.view(-1, D1, window_size[0], window_size[1], C)
#         shifted_x = window_reverse(attn_windows, window_size, B, D1, Hp, Wp) # B, D1, Hp, Wp, C

#         # reverse cyclic shift
#         if any(i > 0 for i in shift_size):
#             x = torch.roll(shifted_x, shifts=(shift_size[0], shift_size[1]), dims=(2, 3))
#         else:
#             x = shifted_x
        
#         if pad_r > 0 or pad_b > 0:
#             x = x[:, :, :H, :W, :].contiguous()

#         x = shortcut + self.drop_path(x)

#         shortcut = x
#         x = self.norm2(x)
#         attn_kv = self.norm_kv(attn_kv)
#         x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b, 0, 0))
#         attn_kv = F.pad(attn_kv, (0, 0, 0, pad_r, 0, pad_b, 0, 0))

#         # cyclic shift
#         if any(i > 0 for i in shift_size):
#             shifted_x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(2, 3))
#             shifted_attn_kv = torch.roll(attn_kv, shifts=(-shift_size[0], -shift_size[1]), dims=(2, 3))
#             attn_mask_q = mask_matrix_q
#             attn_mask_qkv = mask_matrix_qkv
#         else:
#             shifted_x = x
#             shifted_attn_kv = attn_kv
#             attn_mask_q = None
#             attn_mask_qkv = None

#         # partition windows
#         x_windows = window_partition(shifted_x, window_size)  # nW*B, D1, window_size, window_size, C
#         attn_kv_windows = window_partition(shifted_attn_kv, window_size)  # nW*B, D2, window_size, window_size, C
#         x_windows = x_windows.view(-1, D1 * window_size[0] * window_size[1], C)  # nW*B, D1*window_size*window_size, C
#         attn_kv_windows = attn_kv_windows.view(-1, D2 * window_size[0] * window_size[1], C)  # nW*B, D2*window_size*window_size, C

#         # W-MSA/SW-MSA
#         attn_windows = self.attn2(x_windows, attn_kv_windows, mask=attn_mask_qkv)[0]  # nW*B, D1*window_size*window_size, C

#         # merge windows
#         attn_windows = attn_windows.view(-1, D1, window_size[0], window_size[1], C)
#         shifted_x = window_reverse(attn_windows, window_size, B, D1, Hp, Wp)  # B, D1, H, W, C

#         # reverse cyclic shift
#         if any(i > 0 for i in shift_size):
#             x = torch.roll(shifted_x, shifts=(shift_size[0], shift_size[1]), dims=(2, 3))
#         else:
#             x = shifted_x

#         if pad_r > 0 or pad_b > 0:
#             x = x[:, :, :H, :W, :].contiguous()

#         x = shortcut + self.drop_path(x)

#         # FFN
#         x = x + self.drop_path(self.mlp(self.norm3(x)))

#         return x

class EncoderLayer(nn.Module):
    def __init__(self, dim, depth, num_heads, num_frames, window_size=(8, 8), 
                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., 
                 drop_path=0., norm_layer=nn.LayerNorm, encoder_id = 0):
        """Encoder layer
        Args:
            dim (int): Number of feature channels
            depth (int): Depths of this stage.
            num_heads (int): Number of attention head.
            num_frames (int]): Number of input frames.
            window_size (tuple[int], optional): Window size. Defaults to (8, 8).
            mlp_ratio (int, optional): Ratio of mlp hidden dim to embedding dim. Defaults to 4.
            qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Defaults to True.
            qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. Defaults to None.
            drop (float, optional): Dropout rate. Defaults to 0.
            attn_drop (float, optional): Attention dropout rate. Defaults to 0.
            drop_path (float, optional): Stochastic depth rate. Defaults to 0.
            norm_layer (nn.Module, optional): Normalization layer. Defaults to nn.LayerNorm.
        """
        super().__init__()
        self.window_size = window_size
        self.shift_size = tuple(i // 2 for i in window_size)
        self.depth = depth
        self.encoder_id = encoder_id
        # Build blocks
        self.blocks = nn.ModuleList([
            VSTSREncoderTransformerBlock(dim=dim, num_heads=num_heads,
            num_frames=num_frames,window_size=window_size, 
            shift_size=(0, 0) if (i % 2 == 0) else self.shift_size,
            mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
            drop=drop, attn_drop=attn_drop,
            drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
            norm_layer=norm_layer,
            encoder_id = self.encoder_id,
            layer_id = i)
        for i in range(depth)])


    def forward(self, x, data_id = 0):
        """Forward function.
        Args:
            x (torch.Tensor): (B, D, C, H, W)
        Returns:
            torch.Tensor: (B, D, C, H, W)
        """
        # ipdb.set_trace()
        B, D, C, H, W = x.shape
        x = x.permute(0, 1, 3, 4, 2) # B, D, H, W, C

        window_size, shift_size = get_window_size((H, W), self.window_size, self.shift_size)

        Hp = int(np.ceil(H / window_size[0])) * window_size[0]
        Wp = int(np.ceil(W / window_size[1])) * window_size[1]

        img_mask = torch.zeros((1, D, Hp, Wp, 1), device=x.device) # 1, D, H, W, 1
        h_slices = (slice(0, -window_size[0]),
                    slice(-window_size[0], -shift_size[0]),
                    slice(-shift_size[0], None))
        w_slices = (slice(0, -window_size[1]),
                    slice(-window_size[1], -shift_size[1]),
                    slice(-shift_size[1], None))
        cnt = 0
        for h in h_slices:
            for w in w_slices:
                img_mask[:, :, h, w, :] = cnt
                cnt += 1

        mask_windows = window_partition(img_mask, window_size) # nW, D, Wh, Ww, 1
        mask_windows = mask_windows.view(-1, D * window_size[0] * window_size[1]) # nW, D*Wh*Ww
        attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) # nW, D*Wh*Ww, D*Wh*Ww
        attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))

        for blk in self.blocks:
            x = blk(x, attn_mask, data_id)

        x = x.permute(0, 1, 4, 2, 3) # B, D, C, H, W

        return x

class DecoderLayer(nn.Module):
    def __init__(self, dim, depth, num_heads, num_frames, num_out_frames, window_size=(8, 8), 
                 mlp_ratio=4., qkv_bias=True, qk_scale=None, drop=0., attn_drop=0., 
                 drop_path=0., norm_layer=nn.LayerNorm):
        """Decoder layer
        Args:
            dim (int): Number of feature channels
            depth (int): Depths of this stage.
            num_heads (int): Number of attention head.
            num_frames (int]): Number of input frames.
            window_size (tuple[int], optional): Window size. Defaults to (8, 8).
            mlp_ratio (int, optional): Ratio of mlp hidden dim to embedding dim. Defaults to 4.
            qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Defaults to True.
            qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. Defaults to None.
            drop (float, optional): Dropout rate. Defaults to 0.
            attn_drop (float, optional): Attention dropout rate. Defaults to 0.
            drop_path (float, optional): Stochastic depth rate. Defaults to 0.
            norm_layer (nn.Module, optional): Normalization layer. Defaults to nn.LayerNorm.
        """
        super().__init__()
        self.window_size = window_size
        self.shift_size = tuple(i // 2 for i in window_size)
        self.depth = depth

        # Build blocks
        self.blocks = nn.ModuleList([
            VSTSRDecoderTransformerBlock(dim=dim, num_heads=num_heads,
            num_frames=num_frames,num_out_frames=num_out_frames, window_size=window_size, 
            shift_size=(0, 0) if (i % 2 == 0) else self.shift_size,
            mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
            drop=drop, attn_drop=attn_drop,
            drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
            norm_layer=norm_layer)
        for i in range(depth)])

    def forward(self, x, attn_kv):
        """Forward function.
        Args:
            x (torch.Tensor): (B, D1, C, H, W)
            attn_kv (torch.Tensor): (B, D2, C, H, W)
        Returns:
            torch.Tensor: (B, D1, C, H, W)
        """
        B, D1, C, H, W = x. shape
        _, D2, C, _, _ = attn_kv.shape
        x = x.permute(0, 1, 3, 4, 2) # B, D1, H, W, C
        attn_kv = attn_kv.permute(0, 1, 3, 4, 2) # B, D2, H, W, C

        window_size, shift_size = get_window_size((H, W), self.window_size, self.shift_size)

        Hp = int(np.ceil(H / window_size[0])) * window_size[0]
        Wp = int(np.ceil(W / window_size[1])) * window_size[1]

        img_mask_q = torch.zeros((1, D1, Hp, Wp, 1), device=x.device) # 1, D1, H, W, 1
        img_mask_kv = torch.zeros((1, D2, Hp, Wp, 1), device=x.device) # 1, D2, H, W, 1
        h_slices = (slice(0, -window_size[0]),
                    slice(-window_size[0], -shift_size[0]),
                    slice(-shift_size[0], None))
        w_slices = (slice(0, -window_size[1]),
                    slice(-window_size[1], -shift_size[1]),
                    slice(-shift_size[1], None))
        cnt = 0
        for h in h_slices:
            for w in w_slices:
                img_mask_q[:, :, h, w, :] = cnt
                img_mask_kv[:, :, h, w, :] = cnt
                cnt += 1

        mask_windows_q = window_partition(img_mask_q, window_size) # nW, D1, Wh, Ww, 1
        mask_windows_kv = window_partition(img_mask_kv, window_size) # nW, D2, Wh, Ww, 1
        mask_windows_q = mask_windows_q.view(-1, D1 * window_size[0] * window_size[1]) # nW, D1*Wh*Ww
        mask_windows_kv = mask_windows_kv.view(-1, D2 * window_size[0] * window_size[1]) # nW, D2*Wh*Ww
        attn_mask_q = mask_windows_q.unsqueeze(1) - mask_windows_q.unsqueeze(2) # nW, D1*Wh*Ww, D1*Wh*Ww
        attn_mask_qkv = mask_windows_kv.unsqueeze(1) - mask_windows_q.unsqueeze(2) # nW, D1*Wh*Ww, D2*Wh*Ww
        attn_mask_q = attn_mask_q.masked_fill(attn_mask_q != 0, float(-100.0)).masked_fill(attn_mask_q == 0, float(0.0))
        attn_mask_qkv = attn_mask_qkv.masked_fill(attn_mask_qkv != 0, float(-100.0)).masked_fill(attn_mask_qkv == 0, float(0.0))

        for blk in self.blocks:
            x = blk(x, attn_kv, attn_mask_q, attn_mask_qkv)

        x = x.permute(0, 1, 4, 2, 3) # B, D, C, H, W

        return x


class InputProj(nn.Module):
    """Video input projection
    Args:
        in_channels (int): Number of input image channels. Default: 3.
        embed_dim (int): Number of output channels. Default: 32.
        kernel_size (int): Size of the convolution kernel. Default: 3
        stride (int): Stride of the convolution. Default: 1
        norm_layer (nn.Module, optional): Normalization layer. Default: None.
        act_layer (nn.Module): Activation layer. Default: nn.LeakyReLU.
    """
    def __init__(self, in_channels=3, embed_dim=32, kernel_size=3, stride=1, 
                 norm_layer=None, act_layer=nn.LeakyReLU):
        super().__init__()

        self.embed_dim = embed_dim

        self.proj = nn.Sequential(
            nn.Conv2d(in_channels, embed_dim, kernel_size=kernel_size, 
                      stride=stride, padding=kernel_size//2),
            act_layer(inplace=True)
        )

        if norm_layer is not None:
            self.norm = norm_layer(embed_dim)
        else:
            self.norm = None

    def forward(self, x):
        """
        Args:
            x (torch.Tensor): (B, D, C, H, W)
        Returns:
            torch.Tensor: (B, D, C, H, W)
        """
        B, D, C, H, W = x.shape
        x = x.reshape(-1, C, H, W) # B*D, C, H, W
        x = self.proj(x).reshape(B, D, -1, H, W) # B, D, C, H, W
        if self.norm is not None:
            x = x.permute(0, 1, 3, 4, 2) # B, D, H, W, C
            x = self.norm(x)
            x = x.permute(0, 1, 4, 2, 3) # B, D, C, H, W
        return x

class Downsample(nn.Module):
    def __init__(self, in_chans, out_chans):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_chans, out_chans, kernel_size=4, stride=2, padding=1),
        )

    def forward(self, x):
        """Forward function.
        Args:
            x (torch.Tensor): (B, D, C, H, W)
        Returns:
            torch.Tensor: (B, D, C, H, W)
        """
        B, D, C, H, W = x.shape
        x = x.view(-1, C, H, W) # B*D, C, H, W
        out = self.conv(x).view(B, D, -1, H // 2, W // 2)  # B, D, C, H, W
        return out

class Upsample(nn.Module):
    def __init__(self, in_chans, out_chans):
        super().__init__()
        self.deconv = nn.Sequential(
            nn.ConvTranspose2d(in_chans, out_chans, kernel_size=2, stride=2),
        )

    def forward(self, x):
        """Forward function.
        Args:
            x (torch.Tensor): (B, D, C, H, W)
        Returns:
            torch.Tensor: (B, D, C, H, W)
        """
        B, D, C, H, W = x. shape
        x = x.view(-1, C, H, W) # B*D, C, H, W
        out = self.deconv(x).view(B, D, -1, H * 2, W * 2) # B, D, C, H, W
        return out


class HDmapConv(nn.Module):
    def __init__(self, in_channels=2, out_channels=144):
        super().__init__()
        self.hdmap_encoder=nn.Sequential(
            nn.Conv2d(in_channels, 32, kernel_size=5, stride=2, padding=2),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, out_channels, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, hdmap_x):
        """
        hdmap_x: b,2,200,200
        output: b,out_c,25,25
        """
        return self.hdmap_encoder(hdmap_x)


class HDmapfeatConv(nn.Module):
    def __init__(self, in_channels=256, out_channels=144):
        super().__init__()
        self.hdmap_feat_encoder = nn.Sequential(
            nn.Conv2d(in_channels, 256, kernel_size=5, stride=2, padding=2),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.Conv2d(256, out_channels, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
            nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True),
        )

    def forward(self, hdmap_x):
        """
        hdmap_x: b,256,200,200
        output: b,out_c,25,25
        """
        return self.hdmap_feat_encoder(hdmap_x)


class STPT(nn.Module):
    # def __init__(self, in_chans=3, embed_dim=96, 
    #              depths=[8, 8, 8, 8, 8, 8, 8, 8], 
    #              num_heads=[2, 4, 8, 16, 16, 8, 4, 2], num_frames=4,
    #              window_sizes=[(4,4), (4,4), (4,4), (4,4), (4,4), (4,4), (4,4), (4,4)], 
    #              mlp_ratio=2., qkv_bias=True, qk_scale=None,
    #              drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
    #              norm_layer=nn.LayerNorm, patch_norm=True, 
    #              back_RBs=0):
    def __init__(self, in_chans=3, embed_dim=96, 
                 depths=[4, 4, 4, 4, 4, 4, 4, 4], 
                 num_heads=[2, 4, 8, 16, 16, 8, 4, 2], num_frames=4, num_out_frames=5,
                 window_sizes=[(2,2), (2,2), (2,2), (2,2), (2,2), (2,2), (2,2), (2,2)], 
                 mlp_ratio=2., qkv_bias=True, qk_scale=None,
                 drop_rate=0., attn_drop_rate=0., drop_path_rate=0.1,
                 norm_layer=nn.LayerNorm, patch_norm=True, 
                 back_RBs=0, use_hdmap=False, use_hdmap_feat=False):

    
        """
        Args:
            in_chans (int, optional): Number of input image channels. Defaults to 3.
            embed_dim (int, optional): Number of projection output channels. Defaults to 32.
            depths (list[int], optional): Depths of each Transformer stage. Defaults to [2, 2, 2, 2, 2, 2, 2, 2].
            num_heads (list[int], optional): Number of attention head of each stage. Defaults to [2, 4, 8, 16, 16, 8, 4, 2].
            num_frames (int, optional): Number of input frames. Defaults to 4.
            window_size (tuple[int], optional): Window size. Defaults to (8, 8).
            mlp_ratio (int, optional): Ratio of mlp hidden dim to embedding dim. Defaults to 4..
            qkv_bias (bool, optional): If True, add a learnable bias to query, key, value. Defaults to True.
            qk_scale (float | None, optional): Override default qk scale of head_dim ** -0.5 if set. Defaults to None.
            drop_rate (float, optional): Dropout rate. Defaults to 0.
            attn_drop_rate (float, optional): Attention dropout rate. Defaults to 0.
            drop_path_rate (float, optional): Stochastic depth rate. Defaults to 0.1.
            norm_layer (nn.Module, optional): Normalization layer. Defaults to nn.LayerNorm.
            patch_norm (bool, optional): If True, add normalization after patch embedding. Defaults to True.
            back_RBs (int, optional): Number of residual blocks for super resolution. Defaults to 10.
            use_hdmap: use hdmap for bev query
        """
        super().__init__()
        print(depths)
        print(num_heads)
        self.num_layers = len(depths)
        self.num_enc_layers = self.num_layers // 2
        self.num_dec_layers = self.num_layers // 2
        self.scale = 2 ** (self.num_enc_layers - 1)
        dec_depths = depths[self.num_enc_layers:]
        self.embed_dim = embed_dim
        self.patch_norm = patch_norm
        self.num_in_frames = num_frames
        # self.num_out_frames = 2 * num_frames - 1
        self.num_out_frames = num_out_frames

        self.pos_drop = nn.Dropout(p=drop_rate)

        # Stochastic depth 
        enc_dpr= [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths[:self.num_enc_layers]))] 
        dec_dpr = enc_dpr[::-1]

        self.input_proj = InputProj(in_channels=in_chans, embed_dim=embed_dim,
                                    kernel_size=3, stride=1, act_layer=nn.LeakyReLU)

        # Encoder
        self.encoder_layers = nn.ModuleList()
        self.downsample = nn.ModuleList()
        for i_layer in range(self.num_enc_layers):
            encoder_layer = EncoderLayer(
                    dim=embed_dim, 
                    depth=depths[i_layer], num_heads=num_heads[i_layer], 
                    num_frames=num_frames, window_size=window_sizes[i_layer], mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias, qk_scale=qk_scale, drop=drop_rate, 
                    attn_drop=attn_drop_rate,
                    drop_path=enc_dpr[sum(depths[:i_layer]):sum(depths[:i_layer + 1])],
                    norm_layer=norm_layer,
                    encoder_id = i_layer
            )
            downsample = Downsample(embed_dim, embed_dim)
            self.encoder_layers.append(encoder_layer)
            self.downsample.append(downsample)


        # Decoder
        self.decoder_layers = nn.ModuleList()
        self.upsample = nn.ModuleList()
        for i_layer in range(self.num_dec_layers):
            decoder_layer = DecoderLayer(
                    dim=embed_dim, 
                    depth=depths[i_layer + self.num_enc_layers], 
                    num_heads=num_heads[i_layer + self.num_enc_layers], 
                    num_frames=num_frames, num_out_frames = num_out_frames, window_size=window_sizes[i_layer + self.num_enc_layers], mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias, qk_scale=qk_scale, drop=drop_rate, 
                    attn_drop=attn_drop_rate,
                    drop_path=dec_dpr[sum(dec_depths[:i_layer]):sum(dec_depths[:i_layer + 1])],
                    norm_layer=norm_layer
            )
            self.decoder_layers.append(decoder_layer)
            if i_layer != self.num_dec_layers - 1:
                upsample = Upsample(embed_dim, embed_dim)
                self.upsample.append(upsample)

        # Reconstruction block
        # ResidualBlock_noBN_f = functools.partial(ResidualBlock_noBN, nf=embed_dim)
        # self.recon_trunk = make_layer(ResidualBlock_noBN_f, back_RBs)
        # Upsampling
        # self.upconv1 = nn.Conv2d(embed_dim, embed_dim * 4, 3, 1, 1, bias=True)
        # self.upconv2 = nn.Conv2d(embed_dim, 64 * 4, 3, 1, 1, bias=True)
        # self.pixel_shuffle = nn.PixelShuffle(2)
        # self.HRconv = nn.Conv2d(64, 64, 3, 1, 1, bias=True)
        # self.conv_last = nn.Conv2d(64, 3, 3, 1, 1, bias=True)

        # hdmap query encoder
        self.use_hdmap = use_hdmap
        self.use_hdmap_feat = use_hdmap_feat
        if self.use_hdmap and not self.use_hdmap_feat:
            self.hdmap_conv = HDmapConv(out_channels=self.embed_dim)
        elif self.use_hdmap_feat:
            self.hdmap_feat_conv = HDmapfeatConv(out_channels=self.embed_dim)
        
        self.future_query = nn.Parameter(torch.Tensor(1, self.num_out_frames, self.embed_dim, 100 // self.scale, 352 // self.scale), requires_grad=True)
        # self.future_query = nn.Parameter(torch.Tensor(1, 1, self.embed_dim, 25, 25), requires_grad=True)
        nn.init.xavier_uniform_(self.future_query)
        
        self.col_embed = nn.Embedding(25, self.embed_dim // 3)
        self.row_embed = nn.Embedding(25, self.embed_dim // 3)
        self.t_embed = nn.Embedding(self.num_out_frames, self.embed_dim // 3)
        
        self.my_conv = nn.Conv2d(embed_dim, in_chans, 3, 1, 1)

        # Activation function
        self.lrelu = nn.LeakyReLU(negative_slope=0.1, inplace=True)

        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def forward(self, x, hdmap=None, data_id = 0):
        """
        hdmap: if hdmap not none, use hdmap feat to bev query
        """
        B, D, C, H, W = x.size()  # D input video frames
        # x = x.permute(0, 2, 1, 3, 4)
        # upsample_x = F.interpolate(x, (2*D-1, H*4, W*4), mode='trilinear', align_corners=False)
        # x = x.permute(0, 2, 1, 3, 4)
        # ipdb.set_trace()
        x = self.input_proj(x) # B, D, C, H, W

        Hp = int(np.ceil(H / self.scale)) * self.scale
        Wp = int(np.ceil(W / self.scale)) * self.scale
        x = F.pad(x, (0, Wp - W, 0, Hp - H))

        encoder_features = []
        for i_layer in range(self.num_enc_layers):
            x = self.encoder_layers[i_layer](x, data_id)
            encoder_features.append(x)
            if i_layer != self.num_enc_layers - 1:
                x = self.downsample[i_layer](x)
        # ipdb.set_trace()
        _, _, C, h, w = x.size()
        # print(x.shape)
        #####################################################################################
        # TODO: Use interpolation for queries
        # y = torch.zeros((B, self.num_out_frames, C, h, w), device=x.device)
        # for i in range(self.num_out_frames):
        #     if i % 2 == 0:
        #         y[:, i, :, :, :] = x[:, i//2]
        #     else:
        #         y[:, i, :, :, :] = (x[:, i//2] + x[:, i//2 + 1]) / 2
        #####################################################################################
        
        
        x_grid = torch.arange(25, device = x.device)
        y_grid = torch.arange(25, device = x.device)
        t_grid = torch.arange(self.num_out_frames, device = x.device)
        
        x_embed = self.col_embed(x_grid)
        y_embed = self.row_embed(y_grid)
        t_embed = self.t_embed(t_grid)
        
        query_embed = torch.cat((
                                x_embed.unsqueeze(0).unsqueeze(0).repeat(self.num_out_frames, 25, 1, 1),
                                y_embed.unsqueeze(1).unsqueeze(0).repeat(self.num_out_frames, 1, 25, 1),
                                t_embed.unsqueeze(1).unsqueeze(1).repeat(1, 25, 25, 1)),
                                dim = -1
                                ).permute(0, 3, 1, 2).repeat(B, 1, 1, 1, 1)
        
        # print("query_embed", query_embed.shape)  # 1,5,144,25,25
        # print("future_query", self.future_query.shape)  # 1,5,144,25,25
        
        # y = self.future_query.repeat(B, self.num_out_frames, 1, 1, 1)
        y = self.future_query.repeat(B, 1, 1, 1, 1)
        # print('y', y.shape)  # 1,5,144,25,25

        # y = y + query_embed

        if hdmap is not None:
            if not self.use_hdmap_feat:
                hdmap_feat = self.hdmap_conv(hdmap)  # b, 144, 25, 25
                hdmap_feat = hdmap_feat.unsqueeze(dim=1).repeat(1, self.num_out_frames, 1, 1, 1)
                y = y + hdmap_feat
            else:
                hdmap_feat = self.hdmap_feat_conv(hdmap)
                hdmap_feat = hdmap_feat.unsqueeze(dim=1).repeat(1, self.num_out_frames, 1, 1, 1)
                y = y + hdmap_feat
        
        for i_layer in range(self.num_dec_layers):
            y = self.decoder_layers[i_layer](y, encoder_features[-i_layer - 1])
            if i_layer != self.num_dec_layers - 1:
                y = self.upsample[i_layer](y)
        # ipdb.set_trace()
        y = y[:, :, :, :H, :W].contiguous()
        # Super-resolution
        B, D, C, H, W = y.size()
        y = y.view(B*D, C, H, W)
        ############################################################################
        # out = self.recon_trunk(y)
        # out = self.lrelu(self.pixel_shuffle(self.upconv1(out)))
        # out = self.lrelu(self.pixel_shuffle(self.upconv2(out)))

        # out = self.lrelu(self.HRconv(out))
        # out = self.conv_last(out)
        ############################################################################
        
        out = self.my_conv(y)
        _, _, H, W = out.size()
        outs = out.view(B, self.num_out_frames, -1, H, W)
        outs = outs# + upsample_x.permute(0, 2, 1, 3, 4)
        return outs
    
    
    
if __name__ == '__main__':
    device = torch.device('cuda',7)
    a = torch.randn((4, 3, 128, 100, 352)).to(device)
    model = STPT(in_chans=128,
                depths=[4, 4], 
                num_heads=[2, 2], 
                window_sizes=[(4,4), (4,4)],
                num_frames=3,
                num_out_frames=1,
                embed_dim=48,
                use_hdmap=False,
                use_hdmap_feat=False).to(device)
    
    output = model(a)
    print(output.shape)     #[1, 5, 32, 200, 200])
opencood/models/fuse_modules/STPT.py
MediaBrain-SJTU/CoBEVFlow
import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SyncLSTM(nn.Module):
    def __init__(self, channel_size = 256, h = 32, w = 32, k = 3, TM_Flag = False, compressed_size = 32):
        super(SyncLSTM, self).__init__()
        self.k = k

        self.channel_size = channel_size
        self.compressed_size = compressed_size
        self.time_weight = Time_Modulation(input_channel = 2*self.compressed_size)
        self.lstmcell = MotionLSTM(h,w, self.compressed_size)
        self.init_c = nn.parameter.Parameter(torch.rand(self.compressed_size, h, w))
        self.TM_Flag = TM_Flag

        self.ratio = int(math.sqrt(channel_size / compressed_size))
        self.conv_pre_1 = nn.Conv2d(self.channel_size, self.ratio * self.compressed_size, kernel_size=3, stride=1, padding=1)
        self.conv_pre_2 = nn.Conv2d(self.ratio * self.compressed_size, self.compressed_size, kernel_size=3, stride=1, padding=1)
        self.bn_pre_1 = nn.BatchNorm2d(self.ratio * self.compressed_size)
        self.bn_pre_2 = nn.BatchNorm2d(self.compressed_size)
        self.conv_after_1 = nn.Conv2d(self.compressed_size, self.ratio * self.compressed_size, kernel_size=3, stride=1, padding=1)
        self.conv_after_2 = nn.Conv2d(self.ratio * self.compressed_size, self.channel_size, kernel_size=3, stride=1, padding=1)
        self.bn_after_1 = nn.BatchNorm2d(self.ratio * self.compressed_size)
        self.bn_after_2 = nn.BatchNorm2d(self.channel_size)

    def forward(self, x_raw, delta_t):
        batch, seq, channel, h, w = x_raw.shape
        if self.compressed_size != self.channel_size:
            x = F.relu(self.bn_pre_1(self.conv_pre_1(x_raw.reshape(-1,channel,h,w))))
            x = F.relu(self.bn_pre_2(self.conv_pre_2(x)))
            x = x.view(batch, seq, self.compressed_size, h, w)
        else:
            x = x_raw

        if delta_t[0] > 0:
            self.delta_t = delta_t[0]
            h = x[:,0]
            c = self.init_c
            for i in range(1, self.k):
                h,c = self.lstmcell(x[:,i], (h,c))
            for t in range(int(self.delta_t - 1)):
                print('ok')
                h,c = self.lstmcell(h, (h,c))
            if self.TM_Flag:
                w = self.time_weight(torch.cat([x[:,-1], h],1), delta_t[0])
                w = torch.tanh(0.1 * int(delta_t[0] - 1) * w)
                res = w * h + (1-w) * x[:,-1]
            else:
                res = h
        else:
            res = x[:,-1]
        if self.compressed_size != self.channel_size:
            res = F.relu(self.bn_after_1(self.conv_after_1(res)))
            res = F.relu(self.bn_after_2(self.conv_after_2(res)))
            # res = res.view(batch,channel,h,w)
        else:
            res = res
        return res.unsqueeze(1)

class MotionLSTM(nn.Module):
    def __init__(self, h,w, input_channel_size, hidden_size = 0):
        super().__init__()
        self.input_channel_size = input_channel_size  # channel size
        self.hidden_size = hidden_size
        self.h = h
        self.w = w

        #i_t 
        # self.U_i = nn.Parameter(torch.Tensor(input_channel_size, hidden_size)) 
        # self.V_i = nn.Parameter(torch.Tensor(hidden_size, hidden_size)) 

        self.U_i = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.V_i = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.b_i = nn.Parameter(torch.Tensor(1, self.input_channel_size, h, w))

        # #f_t 
        # self.U_f = nn.Parameter(torch.Tensor(input_channel_size, hidden_size)) 
        # self.V_f = nn.Parameter(torch.Tensor(hidden_size, hidden_size)) 
        # self.b_f = nn.Parameter(torch.Tensor(hidden_size)) 
        self.U_f = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.V_f = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.b_f = nn.Parameter(torch.Tensor(1, self.input_channel_size, h, w))

        # #c_t 
        # self.U_c = nn.Parameter(torch.Tensor(input_channel_size, hidden_size)) 
        # self.V_c = nn.Parameter(torch.Tensor(hidden_size, hidden_size)) 
        # self.b_c = nn.Parameter(torch.Tensor(hidden_size)) 
        self.U_c = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.V_c = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.b_c = nn.Parameter(torch.Tensor(1, self.input_channel_size, h, w))

        # #o_t 
        # self.U_o = nn.Parameter(torch.Tensor(input_channel_size, hidden_size)) 
        # self.V_o = nn.Parameter(torch.Tensor(hidden_size, hidden_size)) 
        # self.b_o = nn.Parameter(torch.Tensor(hidden_size)) 
        self.U_o = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.V_o = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.b_o = nn.Parameter(torch.Tensor(1, self.input_channel_size, h, w))

        # self.init_weights()

    # def init_weights(self):
    #     stdv = 1.0 / math.sqrt(self.hidden_size)
    #     for weight in self.parameters():
    #         weight.data.uniform_(-stdv, stdv)

    def forward(self,x,init_states=None): 
        """ 
        assumes x.shape represents (batch_size, sequence_size, input_channel_size) 
        """ 
        h, c = init_states 
        i = torch.sigmoid(self.U_i(x) + self.V_i(h) + self.b_i) 
        f = torch.sigmoid(self.U_f(x) + self.V_f(h) + self.b_f) 
        g = torch.tanh(self.U_c(x) + self.V_c(h) + self.b_c) 
        o = torch.sigmoid(self.U_o(x) + self.V_o(x) + self.b_o) 
        c_out = f * c + i * g 
        h_out = o *  torch.tanh(c_out) 

        # hidden_seq.append(h_t.unsqueeze(0)) 

        # #reshape hidden_seq p/ retornar 
        # hidden_seq = torch.cat(hidden_seq, dim=0) 
        # hidden_seq = hidden_seq.transpose(0, 1).contiguous() 
        return (h_out, c_out)



class STPN_MotionLSTM(nn.Module):
    def __init__(self, height_feat_size = 16):
        super(STPN_MotionLSTM, self).__init__()

        # self.conv3d_1 = Conv3D(4, 8, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
        # self.conv3d_2 = Conv3D(8, 8, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
        # self.conv3d_1 = Conv3D(64, 64, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
        # self.conv3d_2 = Conv3D(128, 128, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))

        self.conv1_1 = nn.Conv2d(height_feat_size, 2*height_feat_size, kernel_size=3, stride=2, padding=1)
        self.conv1_2 = nn.Conv2d(2*height_feat_size, 2*height_feat_size, kernel_size=3, stride=1, padding=1)

        self.conv2_1 = nn.Conv2d(2*height_feat_size, 4*height_feat_size, kernel_size=3, stride=2, padding=1)
        self.conv2_2 = nn.Conv2d(4*height_feat_size, 4*height_feat_size, kernel_size=3, stride=1, padding=1)

        self.conv7_1 = nn.Conv2d(6*height_feat_size, 2*height_feat_size, kernel_size=3, stride=1, padding=1)
        self.conv7_2 = nn.Conv2d(2*height_feat_size, 2*height_feat_size, kernel_size=3, stride=1, padding=1)

        self.conv8_1 = nn.Conv2d(3*height_feat_size , height_feat_size, kernel_size=3, stride=1, padding=1)
        self.conv8_2 = nn.Conv2d(height_feat_size, height_feat_size, kernel_size=3, stride=1, padding=1)

        self.bn1_1 = nn.BatchNorm2d(2*height_feat_size)
        self.bn1_2 = nn.BatchNorm2d(2*height_feat_size)

        self.bn2_1 = nn.BatchNorm2d(4*height_feat_size)
        self.bn2_2 = nn.BatchNorm2d(4*height_feat_size)

        self.bn7_1 = nn.BatchNorm2d(2*height_feat_size)
        self.bn7_2 = nn.BatchNorm2d(2*height_feat_size)

        self.bn8_1 = nn.BatchNorm2d(1*height_feat_size)
        self.bn8_2 = nn.BatchNorm2d(1*height_feat_size)

    def forward(self, x):

        x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
        x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
        x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
        x_2 = F.relu(self.bn2_1(self.conv2_1(x_1)))
        x_2 = F.relu(self.bn2_2(self.conv2_2(x_2)))
        x_7 = F.relu(self.bn7_1(self.conv7_1(torch.cat((F.interpolate(x_2, scale_factor=(2, 2)), x_1), dim=1))))
        x_7 = F.relu(self.bn7_2(self.conv7_2(x_7)))
        x_8 = F.relu(self.bn8_1(self.conv8_1(torch.cat((F.interpolate(x_7, scale_factor=(2, 2)), x), dim=1))))
        res_x = F.relu(self.bn8_2(self.conv8_2(x_8)))

        return res_x


class Time_Modulation(nn.Module):
    def __init__(self, input_channel = 128):
        super(Time_Modulation, self).__init__()
        self.input_channel = input_channel
        self.conv1_channel = int(self.input_channel / 2)
        self.conv2_channel = int(self.conv1_channel / 2)
        # self.ratio = math.sqrt()
        self.conv1 = nn.Conv2d(self.input_channel, self.conv1_channel, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(self.conv1_channel, self.conv2_channel, kernel_size=3, stride=1, padding=1)
        self.convl1 = nn.Conv2d(1, 8, kernel_size=3, stride=1, padding=1)
        self.convl2 = nn.Conv2d(8,self.conv2_channel, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(int(2 * self.conv2_channel), 8, kernel_size = 3, stride=1, padding = 1) 
        # self.linear3 = nn.Linear(16,8)
        self.conv4 = nn.Conv2d(8, 1, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(self.conv1_channel)
        self.bn2 = nn.BatchNorm2d(self.conv2_channel)
        self.bn3 = nn.BatchNorm2d(8)
        self.bn4 = nn.BatchNorm2d(1)
        self.bnl1 = nn.BatchNorm2d(8)
        self.bnl2 = nn.BatchNorm2d(self.conv2_channel)
        self.linear3 = nn.Linear(int(2 * self.conv2_channel), 8)
        self.linear4 = nn.Linear(8, 1)
    def forward(self, x, delta_t):
        a,b,c,d = x.size()
        y = delta_t * torch.ones((a,1,c,d)).to(x.device)
        t_y = F.relu(self.bnl1(self.convl1(y)))
        t_y = F.relu(self.bnl2(self.convl2(t_y)))
        t_x = F.relu(self.bn1(self.conv1(x)))
        t_x = F.relu(self.bn2(self.conv2(t_x)))
        t_xy = torch.cat([t_x, t_y], 1)
        t_xy = F.relu(self.bn3(self.conv3(t_xy)))
        t_xy = torch.sigmoid(self.bn4(self.conv4(t_xy)))
        # t_xy = F.relu(self.bn3(self.linear3(t_xy)))
        # t_xy = torch.sigmoid(self.bn4(self.linear4(t_xy)))
        return t_xy
opencood/models/fuse_modules/SyncNet.py
MediaBrain-SJTU/CoBEVFlow
"""
Implementation of deformable fusion

The design is: for ego agent f_0 and collaborative agent f_1.

f_0[x0,y0] may not correspond to f_1[x0,y0]

So it will learn an offset (delta_x and delta_y) for this pixel position.
Then f_0[x0,y0] will fuse with f_1[x0+delta_x, y0+delta_y]
"""

from this import d
import torch
import torch.nn as nn
import torch.nn.functional as F
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
from icecream import ic

def regroup(x, record_len):
    cum_sum_len = torch.cumsum(record_len, dim=0)
    split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
    return split_x


class RigidOffset(nn.Module):
    """ Learn a rigid transformation grid for the whole feature map
    """

    def __init__(self, in_ch, hidden_ch=32):
        super(RigidOffset, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_ch, hidden_ch, kernel_size=(3, 3), padding=1),
            nn.LeakyReLU(negative_slope=0.01),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(hidden_ch, hidden_ch, kernel_size=(3, 3), padding=1),
            nn.LeakyReLU(negative_slope=0.01),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(hidden_ch, hidden_ch, kernel_size=(3, 3), padding=1),
            nn.LeakyReLU(negative_slope=0.01),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.AdaptiveAvgPool2d(output_size=1),
            nn.Flatten(),
            nn.Linear(in_features=hidden_ch, out_features=hidden_ch, bias=True),
            nn.LeakyReLU(negative_slope=0.01),
            nn.Linear(in_features=hidden_ch, out_features=3, bias=True),
        )

    def forward(self, x, return_M=False):
        """
        Args:
            x.shape:(sum(record_len_minus1), 2C, H, W)
        Returns:
            out.shape: (sum(record_len_minus1), H, W, 2)
        """
        N, _, H, W = x.shape
        xytheta = self.model(x)  # [sum(record_len_minus1), 3], 3 corresponds to x, y, theta


        cos = torch.cos(xytheta[:, 2])
        sin = torch.sin(xytheta[:, 2])

        M = torch.zeros((N, 2, 3), device=x.device)
        M[:, 0, 0] = cos
        M[:, 0, 1] = sin
        M[:, 1, 0] = -sin
        M[:, 1, 1] = cos
        M[:, 0, 2] = xytheta[:, 0]
        M[:, 1, 2] = xytheta[:, 1]

        grid = F.affine_grid(M, size=x.shape)

        if return_M:
            return grid, M

        return grid


class ArbitraryOffset(nn.Module):
    """ Learn a offset/residual grid for each pixel
    """

    def __init__(self, in_ch, out_ch=2, hidden_ch=32):
        """
        Args:
            in_ch: is 2 times feature channel, since they concat together
        """
        super(ArbitraryOffset, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_ch, hidden_ch, 3, 1, 1),
            nn.InstanceNorm2d(hidden_ch),
            nn.LeakyReLU(negative_slope=0.01),
            nn.Conv2d(hidden_ch, hidden_ch // 2, 3, 1, 1),
            nn.InstanceNorm2d(hidden_ch // 2),
            nn.LeakyReLU(negative_slope=0.01),
            nn.Conv2d(hidden_ch // 2, hidden_ch // 4, 1, 1, 0),
            nn.InstanceNorm2d(hidden_ch // 4),
            nn.LeakyReLU(negative_slope=0.01),
            nn.Conv2d(hidden_ch // 4, 2, 1, 1, 0)
        )

    def forward(self, x):
        """
        Args:
            x.shape:(sum(record_len_minus1), 2C, H, W)
        Returns:
            out.shape: (sum(record_len_minus1), H, W, 2)
        """
        N, _, H, W = x.shape

        x = self.model(x)

        grid_residual = x.reshape(N, H, W, 2)

        M_origin = torch.Tensor([[[1, 0, 0], [0, 1, 0]]])
        grid_origin = F.affine_grid(M_origin, size=(1, 1, H, W)).to(x.device)

        grid = grid_residual + grid_origin
        return grid


class DeformFusion(nn.Module):
    """ deformable fusion for multiscale feature map
        For each pixel in ego agent's feature map,
        it will learn a offset to fuse the feature.
    """

    def __init__(self, in_ch, deform_method, cycle_consist_loss=False):
        """
        Args:
            in_ch: channels num of one agent's feature map.
        """
        super(DeformFusion, self).__init__()
        self.cycle_consistency_loss = cycle_consist_loss

        if deform_method == "rigid":
            self.grid_net = RigidOffset(in_ch * 2)
        elif deform_method == "arbitrary":
            self.grid_net = ArbitraryOffset(in_ch * 2)
        

    def forward(self, features, record_len, pairwise_t_matrix, lidar_pose=None):
        """
        Args:
            features: List[torch.Tensor]
                multiscale features. features[i] is (sum(cav), C, H, W), different i, different C, H, W
            record_len: torch.tensor
                record cav number
            pairwise_t_matrix: torch.Tensor,
                already normalized. shape [B, N_max, N_max, 2, 3]
            lidar_pose: torch.Tensor
                shape [(sum(cav), 6)], this is only used to calculate intersection. If proj_first=False, then equal to pairwise_t_matrix
        """

        ##### first align them to ego coordinate, espeically when proj_first = False.
        device = features[0].device
        record_len_minus1 = record_len - 1

        if(torch.sum(record_len_minus1)==0):
            return features

        ms_split_x = [regroup(features[i], record_len) for i in range(len(features))]
        ms_split_x_warp = []

        for split_x in ms_split_x:  # different scale
            split_x_warp = []
            H, W = split_x[0].shape[2:]
            for b, xx in enumerate(split_x):  # different samples
                N = xx.shape[0]
                t_matrix = pairwise_t_matrix[b][:N, :N, :, :]
                i = 0  # ego
                split_x_warp.append(warp_affine_simple(xx, t_matrix[i, :, :, :], (H, W)))  # [N_,C,H,W], N_ varies
            ms_split_x_warp.append(split_x_warp)


        ##### we caculate the grid by scale=1 feature, and share it with all scales.
        split_x = ms_split_x_warp[0]  # first scale
        H, W = split_x[0].shape[:2]

        cat_features = []
        for b, xx in enumerate(split_x):
            N = xx.shape[0]
            cat_feature = torch.cat([xx[0:1].expand(N - 1, -1, -1, -1), xx[1:]], dim=1)  # (N-1, 2C, H, W)
            cat_features.append(cat_feature)

        cat_feature = torch.cat(cat_features, dim=0)  # (sum(record_len_minus1), 2C, H, W)

        grid_offset = self.grid_net(cat_feature)  # (sum(record_len_minus1), H, W, 2)

        grid = grid_offset  # (sum(record_len_minus1),H,W,2)
        ms_grid = [grid[:,::2**i,::2**i,:] for i in range(len(features))]

        ms_split_grid = [regroup(grid, record_len_minus1) for grid in ms_grid]  # [[N1-1,H,W,2], [N2-1,H,W,2],...], shared for all scales.

        #####  fusion
        ms_fused_features = []
        for scale, split_x in enumerate(ms_split_x_warp):
            fused_features = []
            for b, xx in enumerate(split_x):
                if xx.shape[0] == 1:
                    fused_features.append(xx[0])
                else:
                    neighbor_feature_deform = torch.cat([F.grid_sample(xx[1:], ms_split_grid[scale][b]), xx[0:1]], dim=0) # (N-1, C, H, W)
                    fuesd_feature = torch.max(neighbor_feature_deform, dim=0)[0]  
                    fused_features.append(fuesd_feature)
            ms_fused_features.append(torch.stack(fused_features))


        if self.cycle_consistency_loss:
            split_x = ms_split_x[0]  # before warping to the ego agent, scale = 1
            H, W = split_x[0].shape[2:]

            cat_features = []
            for b, xx in enumerate(split_x):
                N = xx.shape[0]
                t_matrix = pairwise_t_matrix[b][:N, :N, :, :]
                """
                    [agent0, agent1] -> M_0
                    [agent1, agent2] -> M_1
                    ...
                    [agentN-1,agent0] -> M_N-1
                    M_0@M_1@...@M_N-1 = I

                    The latter should align to the former agent.
                """
                latter_agent = torch.cat([xx[1:],xx[:1]], dim=0) # [agent1,agent2,..., agent0]
                t_matrix_adj = torch.stack([t_matrix[i,(i+1)%N] for i in range(N)])
                latter_agent_warp = warp_affine_simple(latter_agent, t_matrix_adj, dsize=(H,W))
                cat_feature = torch.cat([xx, latter_agent_warp], dim=1)
                cat_features.append(cat_feature)

            cat_feature = torch.cat(cat_features, dim=0)  # (sum(record_len), 2C, H, W)
            _, M = self.grid_net(cat_feature, return_M=True)  # (sum(record_len)*H*W, 2)

            M_homo = F.pad(M, (0, 0, 0, 1), "constant", 0)  # pad 2nd to last by (0, 1)
            M_homo[:, 2, 2] = 1

            split_M = regroup(M_homo, record_len)

        return ms_fused_features


if __name__ == "__main__":
    features = [torch.randn(4,64,200,704), torch.randn(4,128,100,352), torch.randn(4,256,50,176)]
    record_len = torch.tensor([1,3])
    pairwise_t_matirx = torch.eye(4).view(1,1,1,4,4).expand(2,5,5,4,4)

    model = DeformFusion(in_ch=64, deform_method='rigid', cycle_consist_loss=True)

    out = model(features, record_len, pairwise_t_matirx)
    for xx in out:
        print(xx.shape)
opencood/models/fuse_modules/deform_fuse.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
Implementation of transformer encoder fusion.
It is only a method to fuse features
Not rely on specific backbone.
"""

import torch
import torch.nn as nn

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
import torch.nn.functional as F
from icecream import ic
from mmcv.ops import DeformConv2dPack as DCN2d

class MultiheadAttBlock(nn.Module):
    def __init__(self, channels, n_head=8, dropout=0):
        super(MultiheadAttBlock, self).__init__()
        self.attn = nn.MultiheadAttention(channels, n_head, dropout)

    def forward(self, q, k, v):
        """
        order (seq, batch, feature)
        Args:
            q: (1, H*W, C)
            k: (N, H*W, C)
            v: (N, H*W, C)
        Returns:
            outputs: ()
        """
        context, weight = self.attn(q,k,v) # (1, H*W, C)

        return context

class TransformerBlock(nn.Module):
    def __init__(self, channels, n_head=8, dropout=0):
        super(TransformerBlock, self).__init__()
        self.attn = nn.MultiheadAttention(channels, n_head, dropout)
        self.linear1 = nn.Linear(channels, channels)
        self.linear2 = nn.Linear(channels, channels)

        self.norm1 = nn.LayerNorm(channels)
        self.norm2 = nn.LayerNorm(channels)

        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.relu = nn.ReLU()

    def forward(self, q, k, v):
        """
        order (seq, batch, feature)
        Args:
            q: (1, H*W, C)
            k: (N, H*W, C)
            v: (N, H*W, C)
        Returns:
            outputs: ()
        """
        residual = q
        context, weight = self.attn(q,k,v) # (1, H*W, C)
        context = self.dropout1(context)
        output1 = self.norm1(residual + context)

        # feed forward net
        residual = output1 # (1, H*W, C)
        context = self.linear2(self.relu(self.linear1(output1)))
        context = self.dropout2(context)
        output2 = self.norm2(residual + context)

        return output2



class DeformTransformerFusion(nn.Module):
    def __init__(self, args):
        super(DeformTransformerFusion, self).__init__()
        
        self.channels = args['in_channels']
        self.n_head = args['n_head']
        self.dropout = args['dropout_rate']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    0.16m one pixel? I think it's 0.4m per pixel, according to [200, 704]
        self.downsample_rate = args['downsample_rate']  # 4, downsample rate from original feature map [200, 704]

        self.deform_conv1 = DCN2d(self.channels, self.channels, kernel_size=1, stride=1, padding=0)
        self.deform_conv2 = DCN2d(self.channels, self.channels, kernel_size=3, stride=1, padding=1)
        if args['only_attention']:
            self.transformer_block = MultiheadAttBlock(self.channels, self.n_head, self.dropout)
        else:
            self.transformer_block = TransformerBlock(self.channels, self.n_head, self.dropout)

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
            
        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        split_x = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L,1,H,W)).to(x)
                roi_mask[b,i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :],(H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        out = []
        # iterate each batch
        for b in range(B):

            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            updated_node_features = []

            # update each node i
            i = 0 # ego
            # (N,1,H,W)
            mask = roi_mask[b, i, :N, ...]

            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                            t_matrix[i, :, :, :],
                                            (H, W))
            
            neighbor_feature_deform1_flat = self.deform_conv1(neighbor_feature).view(N,C,H*W)
            neighbor_feature_deform2_flat = self.deform_conv2(neighbor_feature).view(N,C,H*W)

            neighbor_feature_flat = neighbor_feature.view(N,C,H*W)  # (N, C, H*W)


            
            query = neighbor_feature_flat[0:1,...].permute(0,2,1)  # (1, H*W, C)

            key0 = neighbor_feature_flat.permute(0,2,1)  # (N, H*W, C)
            key1 = neighbor_feature_deform1_flat.permute(0,2,1) # (N, H*W, C)
            key2 = neighbor_feature_deform2_flat.permute(0,2,1) # (N, H*W, C)

            key = torch.cat((key0,key1,key2), dim=0)  # (3N, H*W, C)

            value = key

            fusion_result = self.transformer_block(query, key, value)  # (1, H*W, C)
            fusion_result = fusion_result.permute(0,2,1).reshape(1, C, H, W)[0]

            out.append(fusion_result)
            
        out = torch.stack(out)
        
        return out









opencood/models/fuse_modules/deform_transformer_fuse.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
Implementation of V2VNet Fusion
"""

import torch
import torch.nn as nn

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
from opencood.models.sub_modules.convgru import ConvGRU
from icecream import ic
from matplotlib import pyplot as plt
from icecream import ic

class MaxFusion(nn.Module):
    def __init__(self, args):
        super(MaxFusion, self).__init__()

        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]


    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
            
        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        split_x = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        # (B*L,L,1,H,W)
        # roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        # for b in range(B):
        #     N = record_len[b]
        #     for i in range(N):
        #         one_tensor = torch.ones((L,1,H,W)).to(x)
        #         roi_mask[b,i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :],(H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        out = []
        # iterate each batch
        for b in range(B):

            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :] # [N, N, 2, 3]

            updated_node_features = []

            # update each node i
            i = 0 # ego
            # (N,1,H,W)
            # mask = roi_mask[b, i, :N, ...]
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                            t_matrix[i, :, :, :],
                                            (H, W))
            out.append(torch.max(neighbor_feature, dim=0)[0])

            # # TODO: debug use
            # viz_path = '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out'
            # torch.save(neighbor_feature, viz_path+'/neighbor_feature.pt')

        out = torch.stack(out)
        
        return out


    def forward_debug(self, x, origin_x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding
        Used for debug and visualization

        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        origin_x: torch.Tensor
            pillars (sum(n_cav), C, H * downsample_rate, W * downsample_rate)
            
        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 
            
        Returns
        -------
        Fused feature.
        """
        from matplotlib import pyplot as plt

        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        split_x = self.regroup(x, record_len)
        split_origin_x = self.regroup(origin_x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L,1,H,W)).to(x)
                roi_mask[b,i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :],(H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        # visualize warped feature map
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0 # ego
            mask = roi_mask[b, i, :N, ...]
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                            t_matrix[i, :, :, :],
                                            (H, W))
            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx],0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/feature_{b}_{idx}")
                plt.clf()
                plt.imshow(mask[idx][0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/mask_feature_{b}_{idx}")
                plt.clf()


        
        # visualize origin pillar feature 
        origin_node_features = split_origin_x

        for b in range(B):
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            i = 0 # ego
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(origin_node_features[b],
                                            t_matrix[i, :, :, :],
                                            (H*self.downsample_rate, W*self.downsample_rate))

            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx],0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/origin_{b}_{idx}")
                plt.clf()
opencood/models/fuse_modules/max_fuse.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
Implementation of V2VNet Fusion
"""

import torch
import torch.nn as nn

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
from opencood.models.sub_modules.convgru import ConvGRU
from icecream import ic
from matplotlib import pyplot as plt
from icecream import ic

class MeanFusion(nn.Module):
    def __init__(self, args):
        super(MeanFusion, self).__init__()

        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    0.16m one pixel? I think it's 0.4m per pixel, according to [200, 704]
        self.downsample_rate = args['downsample_rate']  # 4, downsample rate from original feature map [200, 704]

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
            
        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        split_x = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L,1,H,W)).to(x)
                roi_mask[b,i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :],(H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        out = []
        # iterate each batch
        for b in range(B):

            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0 # ego
            # (N,1,H,W)
            mask = roi_mask[b, i, :N, ...]
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                            t_matrix[i, :, :, :],
                                            (H, W))
            
            out.append(torch.mean(neighbor_feature, dim=0))
        out = torch.stack(out)
        
        return out
















# from matplotlib import pyplot as plt
# for idx in range(3):
#     plt.imshow(torch.max(neighbor_feature[idx],0)[0].detach().cpu().numpy())
#     plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/in_mean/agent{idx}")
#     plt.clf()
#     plt.imshow(mask[idx][0].detach().cpu().numpy())
#     plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/in_mean/mask{idx}")
#     plt.clf()

# raise
opencood/models/fuse_modules/mean_fuse.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# Date: 2022/12/06
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

import torch.nn.functional as F
from turtle import update
# import ipdb
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
from opencood.models.comm_modules.where2comm_multisweep import Communication

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class PositionalEncoding_spatial(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model_feature, d_model_hidden, dropout, max_len=300,mode='cat'):
        super(PositionalEncoding_spatial, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        d_model = int(d_model_hidden/2)
        # Compute the positional encodings once in log space.
        pe = torch.zeros((max_len, d_model))
        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe_ = pe.unsqueeze(0)#batch,len,feature
        pe_h = pe_.unsqueeze(2).repeat(1,1,max_len,1)
        pe_w = pe_.unsqueeze(1).repeat(1,max_len,1,1)
        pe = torch.cat((pe_h,pe_w),dim=-1)

        self.register_buffer('pe', pe)

        self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x):
        if self.mode == 'cat':
            tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
            tmp = tmp.repeat(x.shape[0], 1, 1, 1)
            x = self.linear(torch.cat((x, tmp),dim=-1))

        elif self.mode == 'init':
            x = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)

        else:
            tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
            tmp = tmp.repeat(x.shape[0], 1, 1, 1)
            x = x + tmp

        return self.dropout(x)

class PositionalEncoding_spatial_custom(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model_feature, d_model_hidden, dropout,mode='cat'):
        super(PositionalEncoding_spatial_custom, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = int(d_model_hidden/2)
        # Compute the positional encodings once in log space.
        
        self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x, nonego_idx):#x(batch,num,feature)
   
        pe_h = torch.zeros((x.shape[1], self.d_model)).to(x.device)
        pe_w = torch.zeros((x.shape[1], self.d_model)).to(x.device)
        position_h, position_w = nonego_idx[0],nonego_idx[1]
        position_h = position_h.unsqueeze(1)
        position_w = position_w.unsqueeze(1)

        div_term = torch.exp(torch.arange(0, self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(position_h.device)
        pe_h[:, 0::2] = torch.sin(position_h * div_term)
        pe_h[:, 1::2] = torch.cos(position_h * div_term)

        pe_w[:, 0::2] = torch.sin(position_w * div_term)
        pe_w[:, 1::2] = torch.cos(position_w * div_term)
       
        pe = torch.cat((pe_h,pe_w),dim=-1).unsqueeze(0).repeat(x.shape[0],1,1)

        if self.mode == 'cat':
            x = self.linear(torch.cat((x, pe),dim=-1))

        else:
            x = x + Variable(pe,requires_grad=False)

        return self.dropout(x)

class PositionalEncoding_sensor_dist(nn.Module):

    def __init__(self, d_model, dropout):
        super(PositionalEncoding_sensor_dist, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = d_model
        self.linear = nn.Linear(2*d_model,d_model)
        self.mode = 'cat'
         
    def forward(self, x, sensor_dist):#x.shape=[batch,H,W,F],senfor_dist.shape=[batch,H,W]

        x = x.reshape(-1,x.shape[-1])
        sensor_dist = sensor_dist.reshape(-1)

        pe = torch.zeros(x.shape[0], self.d_model).to(x.device)
        position = sensor_dist.unsqueeze(1) 
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(x.device)
         
        pe[:, 0::2] = torch.sin(position * div_term)  
        pe[:, 1::2] = torch.cos(position * div_term)   

        if self.mode == 'cat':
            x = self.linear(torch.cat((x,Variable(pe,require_grad=False)),dim=-1))

        else:
            x = x + Variable(pe,requires_grad=False)
        
        return self.dropout(x)

class PositionalEncoding_irregular_time(nn.Module):

    def __init__(self, d_model_feature, d_model_hidden, dropout):
        super(PositionalEncoding_irregular_time, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = d_model_feature
        self.linear = nn.Linear(d_model_feature+d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x, time):#x.shape=[batch,height,width,feature],time.shape=[batch]

        pe = torch.zeros(x.shape[0],self.d_model).to(x.device)
        position = time.unsqueeze(1) 
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(x.device)
         
        pe[:, 0::2] = torch.sin(position * div_term)  
        pe[:, 1::2] = torch.cos(position * div_term)   

        if self.mode == 'cat':
            if len(x.shape) == 4:
                pe = pe.unsqueeze(1).unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
            else:
                pe = pe.unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], 1)
            x = self.linear(torch.cat((x, tmp),dim=-1))

        else:
            if len(x.shape) == 4:
                pe = pe.unsqueeze(1).unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
            else:
                pe = pe.unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], 1)
            x = x + tmp

        
        return self.dropout(x)

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    
    # ipdb.set_trace()
    
    # print('score',scores.shape,'mask',mask.shape)

    if mask is not None:
        # mask = mask.transpose(1,2).to(query.device)
        mask = mask.to(query.device)
        scores = scores.masked_fill(mask == 0, -1e10) 

    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)

    return torch.matmul(p_attn, value), p_attn

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, q_dim, k_dim,v_dim, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h  
        self.h = h 
        self.q_dim = q_dim
        self.k_dim = k_dim
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.linear_q = nn.Linear(q_dim,d_model)
        self.linear_k = nn.Linear(k_dim,d_model)
        self.linear_v = nn.Linear(v_dim,d_model)
        self.linear_out = nn.Linear(d_model,q_dim)
        self.attn = None  
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k 

        query = self.linear_q(query).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
        key = self.linear_k(key).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
        value = self.linear_v(value).view(nbatches,-1,self.h,self.d_k).transpose(1,2)

        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(query, key, value, mask=mask,
                                 dropout=self.dropout)

        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(1, 2).contiguous() \
            .view(nbatches ,-1, self.h * self.d_k)

        return self.linear_out(x)

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
 
    def forward(self, x):
        "Norm"
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, x, sublayer):
        # print('sub',x.shape,self.dropout(sublayer(self.norm(x))).shape)
        return x + self.dropout(sublayer(self.norm(x))) #残差连接

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model,d_out, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_out)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

class EncoderLayer_temporal(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer_temporal, self).__init__()
        self.self_attn = self_attn #sublayer 1 
        self.feed_forward = feed_forward #sublayer 2 
        self.sublayer = clones(SublayerConnection(size, dropout), 2) #拷贝两个SublayerConnection，一个为了attention，一个是独自作为简单的神经网络
        self.size = size

    def forward(self, q, k, v, mask):
        "Follow Figure 1 (left) for connections."
        # print('x',x.shape,'attn',self.self_attn(x,x,x,mask).shape)
        x = self.sublayer[0](q, lambda q: self.self_attn(q, k, v, mask)) #attention层 对应层1
        return self.sublayer[1](x, self.feed_forward) # 对应 层2

class Spatial_conv(nn.Module):
    def __init__(self, input_channels, output_channels, kernel_size):
        super(Spatial_conv,self).__init__()
        self.conv = nn.Conv2d(in_channels = input_channels,
                              out_channels = output_channels,
                              kernel_size = (kernel_size,kernel_size),
                              stride = (1,1),
                              padding = (int((kernel_size-1)/2),int((kernel_size-1)/2))
        )

        self.relu = nn.ReLU()
        # self.bn = nn.BatchNorm2d(num_features=output_channels)

    def forward(self, x):
        init_x = x.clone()
        x = self.conv(x)
        x = self.relu(x)
        # x = self.bn(x)
        return x+init_x

class SpatialAttention_mtf(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention_mtf, self).__init__()

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, curr, prev):#####(batch, channel, h, w)
        avg_out = torch.mean(curr, dim=1, keepdim=True)
        max_out, _ = torch.max(curr, dim=1, keepdim=True)
        y = torch.cat([avg_out, max_out], dim=1)
        y = self.conv1(y)
        return self.sigmoid(y) * prev ####(batch-1, channel, h, w)

class ScaledDotProductAttention(nn.Module):
    """
    Scaled Dot-Product Attention proposed in "Attention Is All You Need"
    Compute the dot products of the query with all keys, divide each by sqrt(dim),
    and apply a softmax function to obtain the weights on the values
    Args: dim, mask
        dim (int): dimention of attention
        mask (torch.Tensor): tensor containing indices to be masked
    Inputs: query, key, value, mask
        - **query** (batch, q_len, d_model): tensor containing projection
          vector for decoder.
        - **key** (batch, k_len, d_model): tensor containing projection
          vector for encoder.
        - **value** (batch, v_len, d_model): tensor containing features of the
          encoded input sequence.
        - **mask** (-): tensor containing indices to be masked
    Returns: context, attn
        - **context**: tensor containing the context vector from
          attention mechanism.
        - **attn**: tensor containing the attention (alignment) from the
          encoder outputs.
    """

    def __init__(self, dim):
        super(ScaledDotProductAttention, self).__init__()
        self.sqrt_dim = np.sqrt(dim)

    def forward(self, query, key, value, mask):
        score = torch.bmm(query, key.transpose(1, 2)) / self.sqrt_dim
        if mask != None:
            score = score.masked_fill(mask == 0, -1e10) 
        attn = F.softmax(score, -1)
        context = torch.bmm(attn, value)
        return context

class AttentionFusion(nn.Module):
    def __init__(self, settings, input_dim, hidden_dim, num_heads, num_layers, neighbor_range, height, width, dropout, max_len=300):
        super(AttentionFusion, self).__init__()
        self.sweep_length = settings['sweep_length']
        self.if_spatial_encoding = settings['if_spatial_encoding']
        self.if_ego_time_encoding = settings['if_ego_time_encoding']
        self.if_nonego_time_encoding = settings['if_nonego_time_encoding']
        self.if_sensor_encoding = settings['if_sensor_encoding']
        self.if_time_attn_aggre = settings['if_time_attn_aggre']
        self.if_spatial_conv = settings['if_spatial_conv']
        self.if_dotproductattn = settings['if_dotproductattn']
        self.use_mask = settings['use_mask']
        self.confidence_fetch = settings['confidence_fetch']  
        self.input_dim = input_dim 
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.neighbor_range = neighbor_range
        self.height = height
        self.width = width
        if self.if_spatial_encoding:
            self.spatial_encoding = PositionalEncoding_spatial(self.input_dim, hidden_dim, dropout, max_len)
            self.spatial_encoding_nonego = PositionalEncoding_spatial_custom(self.input_dim, hidden_dim, dropout)
        if self.if_ego_time_encoding or self.if_nonego_time_encoding:
            self.time_encoding = PositionalEncoding_irregular_time(self.input_dim, hidden_dim, dropout)
        if self.if_sensor_encoding:
            self.sensor_dist_encoding = PositionalEncoding_sensor_dist(hidden_dim, dropout)

        c = copy.deepcopy
        temporal_attn = MultiHeadedAttention(self.num_heads, self.input_dim, self.input_dim, self.input_dim, self.hidden_dim, dropout)

        ff = PositionwiseFeedForward(self.input_dim, self.input_dim, self.hidden_dim, dropout)

        if self.if_time_attn_aggre:
            self.layers_temporal = clones(EncoderLayer_temporal(self.input_dim, c(temporal_attn), c(ff), dropout),num_layers)
        # self.temporal_aggregator = MultiHeadedAttention(self.num_heads, hidden_dim, self.input_dim, self.input_dim, self.hidden_dim, dropout)
        if self.if_spatial_conv:  
            self.layers_spatial = clones(Spatial_conv(self.input_dim,self.input_dim,kernel_size=5),num_layers)
        # self.layers_spatial = Spatial_conv(self.input_dim,self.input_dim,kernel_size=5)
        if self.if_dotproductattn:
            self.dotproductattn = ScaledDotProductAttention(self.input_dim)

        print("================Rain attention Init================")

    def forward(self, features, sensor_dist, time_stamps, valid_mask):
        """
        Parameters:
        -----------
        features: torch.Tensor, shape (B, C, H, W)
            neighbor features in ego view
            B = sum_{i}^{num_cav} k_i, num_cav 表示cav数量, k_i 表示对于cav_i 有多少观测帧
            C : feature 维度，一般为 256
            H : height, 一般为 100
            W : width, 一般为 252

        sensor_dist: torch.Tensor, shape (B, H, W)
            distance between each block and the sensor
            
        time_stamps: torch.Tensor, shape (B)
            表示每个位置的时间信息

        valid_mask: (B,H,W)
            为0/1,表示是否在该位置有效

        Returns:
        --------
        output_feature: torch.Tensor, shape (C, H, W)
            ego cav's fused feature at current frame.
        """

        return torch.max(features, dim=0)[0]

        # ego_features = features[:1].clone()
        # return ego_features[0]
        # if self.if_time_attn_aggre != True and self.if_dotproductattn != True:
        #     features = torch.max(features, dim=0)[0]
        #     if self.if_spatial_conv:
        #         features = features.unsqueeze(0)
        #         features = self.layers_spatial[0](features)
        #     return features.squeeze(0)

        #### BEGIN ST atten ####
        #### Spatial
        cav_num, C, H, W = features.shape
        # features = features.view(cav_num, C, -1).permute(2, 0, 1) #  (H*W, cav_num, C), perform self attention on each pixel.
        # masks = valid_mask.view(cav_num, H*W).permute(1,0)
        # masks = masks.unsqueeze(-1).repeat(1,1,cav_num)
        
        # features = self.dotproductattn(features, features, features, masks)
        # features = features.permute(1, 0, 2).view(cav_num, H, W, C)

        #### Temporal
        features = features.permute(0, 2, 3, 1)  #(B, H, W, C)
        features = self.time_encoding(features,time_stamps) #(B, H, W, C)
        features = features.view(cav_num, -1, C).permute(1, 0, 2) #  (H*W, B, C), perform self attention on each time_stamps.
        masks = valid_mask.view(cav_num, H*W).permute(1, 0).unsqueeze(-1).repeat(1, 1, cav_num)
        features = self.dotproductattn(features, features, features, masks) #(HW, B, C) 
        features = features.permute(1, 2, 0).view(cav_num, C, H, W)[0]

        return features

        #### END ST atten ####

        if self.confidence_fetch == False:
            cav_num, C, H, W = features.shape
            features = features.view(cav_num, C, -1).permute(2, 0, 1) #  (H*W, cav_num, C), perform self attention on each pixel.
            masks = valid_mask.view(cav_num, H*W).permute(1,0)
            masks = masks.unsqueeze(-1).repeat(1,1,cav_num)
            
            features = self.dotproductattn(features, features, features, masks)
            features = features.permute(1, 2, 0).view(cav_num, C, H, W)[0] 
            return features

        valid_mask = valid_mask.squeeze(1)
        
        features = features.permute(0,2,3,1)
        
        features = torch.cat((features[0:1],features[self.sweep_length:]),dim=0)
        

        ##########################################  三种encoding
        if self.if_spatial_encoding:
            features[0:1] = self.spatial_encoding(features[0:1])# features.shape = [batch_all,height,width,feature_dim]
        if self.if_sensor_encoding and sensor_dist!=-1:
            sensor_dist = torch.cat((sensor_dist[0:1],sensor_dist[self.sweep_length:]))
            features = self.sensor_dist_encoding(features,sensor_dist)
        time_stamps = torch.cat((time_stamps[0:1],time_stamps[self.sweep_length:]),dim=0)
        valid_mask = torch.cat((valid_mask[0:1],valid_mask[self.sweep_length:]),dim=0)
        valid_mask_sum = valid_mask.sum(0)
        nonego_idx = torch.where(valid_mask_sum>1)
 
        features_to_aggre = features[:,nonego_idx[0],nonego_idx[1],:] 
        if self.if_ego_time_encoding:
            features[0:1] = self.time_encoding(features[0:1],time_stamps[0:1])

        if features_to_aggre.shape[1] != 0:
            if self.if_spatial_encoding:
                features_to_aggre = self.spatial_encoding_nonego(features_to_aggre,nonego_idx)
            if self.if_nonego_time_encoding:
                features_to_aggre = self.time_encoding(features_to_aggre,time_stamps)
        

        ###########################################

        ################################################### 
        # ipdb.set_trace()
        masks = valid_mask[:,nonego_idx[0],nonego_idx[1]]#(valid_num,len)
        batch_len = masks.shape[0]
        masks = masks.permute(1,0).unsqueeze(-1).repeat(1,1,batch_len)

        features_to_aggre = features_to_aggre.permute(1,0,2)
        for l in range(self.num_layers):
                
            if masks.shape[0] != 0:
                if self.if_dotproductattn:
                    # if self.use_mask == False:
                    #     masks =  None
                    # features_to_aggre = self.dotproductattn(features_to_aggre,features_to_aggre,features_to_aggre,masks)
                    features_to_aggre[:,0] = torch.max(features_to_aggre,dim=1)[0]
                else:
                    features_to_aggre = self.layers_temporal[l](features_to_aggre,features_to_aggre,features_to_aggre,masks)
                features[:,nonego_idx[0],nonego_idx[1],:] = features_to_aggre.permute(1,0,2)
           

            if self.if_spatial_conv:
                features = features.permute(0,3,1,2)
                features_backup = features.clone()
                features[0:1] = self.layers_spatial[l](features_backup[0:1])
                features = features.permute(0,2,3,1)

                features_to_aggre[:,0:1] = features[0:1,nonego_idx[0],nonego_idx[1],:].permute(1,0,2)

        return features[0].permute(2,0,1)

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_fuse(nn.Module):
    def __init__(self, args):
        super(raindrop_fuse, self).__init__()
        
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.agg_mode == 'RAIN':
            d_ff = args['d_ff']
            h = args['h']
            dropout = args['dropout']
            input_dim = args['input_dim']
            num_layers = args['num_layers']
            neighbor_range = args['neighbor_range']
            max_len = args['max_len']
            height = 100
            width = 252

            if self.multi_scale:
                layer_nums = args['layer_nums']
                num_filters = args['num_filters']
                self.num_levels = len(layer_nums)
                self.fuse_modules = nn.ModuleList()
                for idx in range(self.num_levels):
                    fuse_network = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len)
                    for p in fuse_network.parameters():
                        if p.dim() > 1:
                            nn.init.xavier_uniform_(p)

                    height = int(height/2)
                    width = int(width/2)
                    input_dim = input_dim * 2
                    d_ff = d_ff * 2
                    self.fuse_modules.append(fuse_network)
            else:
                self.fuse_modules = AttentionFusion(input_dim, d_ff, h, num_layers, neighbor_range, 100, 252, dropout, max_len)
                # This was important from their code. 
                # Initialize parameters with Glorot / fan_avg.
                for p in self.fuse_modules.parameters():
                    if p.dim() > 1:
                        nn.init.xavier_uniform_(p)

        else:
            if self.multi_scale:
                layer_nums = args['layer_nums']
                num_filters = args['num_filters']
                self.num_levels = len(layer_nums)
                self.fuse_modules = nn.ModuleList()
                for idx in range(self.num_levels):
                    if self.agg_mode == 'MAX':
                        fuse_network = MaxFusion()
                    else:
                        # self.agg_mode == 'ATTEN':
                        fuse_network = AttenFusion(num_filters[idx])
                    self.fuse_modules.append(fuse_network)
            else: 
                if self.agg_mode == 'MAX': # max fusion, debug use
                    self.fuse_modules = MaxFusion()

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None, noise_pairwise_t_matrix=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
            
        Returns
        -------
        Fused feature.
        """

        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        if noise_pairwise_t_matrix is not None:
            noise_pairwise_t_matrix = noise_pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
            noise_pairwise_t_matrix[...,0,1] = noise_pairwise_t_matrix[...,0,1] * H / W
            noise_pairwise_t_matrix[...,1,0] = noise_pairwise_t_matrix[...,1,0] * W / H
            noise_pairwise_t_matrix[...,0,2] = noise_pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
            noise_pairwise_t_matrix[...,1,2] = noise_pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        if self.multi_scale:
            ups = []
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(x)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
            for i in range(self.num_levels):
                x = feats[i] if with_resnet else backbone.blocks[i](x)

                ############ 1. Communication (Mask the features) #########
                if i==0:
                    if self.communication:
                        batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                        batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                        _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                        communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                        x = x * communication_masks_tensor
                    else:
                        communication_rates = torch.tensor(0).to(x.device)
                else:
                    if self.communication:
                        communication_masks_tensor = F.max_pool2d(communication_masks_tensor.to(torch.float32), kernel_size=2)
                        # TODO: scale = 1, 2 不加 mask
                        x = x * communication_masks_tensor
                
                ############ 2. Split the confidence map #######################
                # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                batch_node_features = self.regroup(x, record_len, K)

                ############ 3. Fusion ####################################
                x_fuse = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    if noise_pairwise_t_matrix is not None:
                        t_matrix = noise_pairwise_t_matrix[b][:N, 0, :, :] #(N, 2, 3)
                    else:
                        t_matrix = pairwise_t_matrix[b][:N, 0, :, :] #(N, 2, 3)
                    # t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))
                    record_frames = np.ones((N))*K
                    # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                    
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))  

                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])
                    else: # ATTEN, MAX, Transformer
                        # x_fuse.append(node_features[0])
                        # print('warning, neighbor losses') # TODO: uncomment
                        x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])

                x_fuse = torch.stack(x_fuse)

                ############ 4. Deconv ####################################
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[i](x_fuse))
                else:
                    ups.append(x_fuse)
                
            if len(ups) > 1:
                x_fuse = torch.cat(ups, dim=1)
            elif len(ups) == 1:
                x_fuse = ups[0]
            
            if len(backbone.deblocks) > self.num_levels:
                x_fuse = backbone.deblocks[-1](x_fuse)
                
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_confidence_maps = self.regroup(rm, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                # _, (B, 1, H, W), float
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix,
                                                (H, W))
                record_frames = np.ones((N))*K
                # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                if self.agg_mode == "RAIN":
                    # for sensor embedding
                    sensor_dist = -1# (B, H, W)
                    x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)

            # self.fuse_modsules(x_fuse, record_len)
        
        return x_fuse, communication_rates, {}

# # -*- coding: utf-8 -*-
# # Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# # Date: 2022/12/06
# # License: TDG-Attribution-NonCommercial-NoDistrib

# import torch
# import torch.nn as nn
# import easydict
# import copy
# import math
# from torch.autograd import Variable
# import torch.nn.functional as F
# import random
# # from dgl.nn.pytorch.factory import KNNGraph
# # import dgl
# import numpy as np
# import os

# import torch.nn.functional as F
# from turtle import update
# # import ipdb
# from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
# from opencood.models.comm_modules.where2comm_multisweep import Communication

# def clones(module, N):
#     "Produce N identical layers."
#     return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

# class PositionalEncoding_spatial(nn.Module):
#     "Implement the PE function."

#     def __init__(self, d_model_feature, d_model_hidden, dropout, max_len=300,mode='cat'):
#         super(PositionalEncoding_spatial, self).__init__()
#         self.dropout = nn.Dropout(p=dropout)
#         d_model = int(d_model_hidden/2)
#         # Compute the positional encodings once in log space.
#         pe = torch.zeros((max_len, d_model))
#         position = torch.arange(0, max_len).unsqueeze(1)

#         div_term = torch.exp(torch.arange(0, d_model, 2) *
#                              -(math.log(10000.0) / d_model))
#         pe[:, 0::2] = torch.sin(position * div_term)
#         pe[:, 1::2] = torch.cos(position * div_term)
#         pe_ = pe.unsqueeze(0)#batch,len,feature
#         pe_h = pe_.unsqueeze(2).repeat(1,1,max_len,1)
#         pe_w = pe_.unsqueeze(1).repeat(1,max_len,1,1)
#         pe = torch.cat((pe_h,pe_w),dim=-1)

#         self.register_buffer('pe', pe)

#         self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
#         self.mode = 'plus'

#     def forward(self, x):
#         if self.mode == 'cat':
#             tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
#             tmp = tmp.repeat(x.shape[0], 1, 1, 1)
#             x = self.linear(torch.cat((x, tmp),dim=-1))

#         elif self.mode == 'init':
#             x = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)

#         else:
#             tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
#             tmp = tmp.repeat(x.shape[0], 1, 1, 1)
#             x = x + tmp

#         return self.dropout(x)

# class PositionalEncoding_spatial_custom(nn.Module):
#     "Implement the PE function."

#     def __init__(self, d_model_feature, d_model_hidden, dropout,mode='cat'):
#         super(PositionalEncoding_spatial_custom, self).__init__()
#         self.dropout = nn.Dropout(p=dropout)
#         self.d_model = int(d_model_hidden/2)
#         # Compute the positional encodings once in log space.
        
#         self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
#         self.mode = 'plus'

#     def forward(self, x, nonego_idx):#x(batch,num,feature)
   
#         pe_h = torch.zeros((x.shape[1], self.d_model)).to(x.device)
#         pe_w = torch.zeros((x.shape[1], self.d_model)).to(x.device)
#         position_h, position_w = nonego_idx[0],nonego_idx[1]
#         position_h = position_h.unsqueeze(1)
#         position_w = position_w.unsqueeze(1)

#         div_term = torch.exp(torch.arange(0, self.d_model, 2) *
#                              -(math.log(10000.0) / self.d_model)).to(position_h.device)
#         pe_h[:, 0::2] = torch.sin(position_h * div_term)
#         pe_h[:, 1::2] = torch.cos(position_h * div_term)

#         pe_w[:, 0::2] = torch.sin(position_w * div_term)
#         pe_w[:, 1::2] = torch.cos(position_w * div_term)
       
#         pe = torch.cat((pe_h,pe_w),dim=-1).unsqueeze(0).repeat(x.shape[0],1,1)

#         if self.mode == 'cat':
#             x = self.linear(torch.cat((x, pe),dim=-1))

#         else:
#             x = x + Variable(pe,requires_grad=False)

#         return self.dropout(x)

# class PositionalEncoding_sensor_dist(nn.Module):

#     def __init__(self, d_model, dropout):
#         super(PositionalEncoding_sensor_dist, self).__init__()
#         self.dropout = nn.Dropout(p=dropout)
#         self.d_model = d_model
#         self.linear = nn.Linear(2*d_model,d_model)
#         self.mode = 'cat'
         
#     def forward(self, x, sensor_dist):#x.shape=[batch,H,W,F],senfor_dist.shape=[batch,H,W]

#         x = x.reshape(-1,x.shape[-1])
#         sensor_dist = sensor_dist.reshape(-1)

#         pe = torch.zeros(x.shape[0], self.d_model).to(x.device)
#         position = sensor_dist.unsqueeze(1) 
#         div_term = torch.exp(torch.arange(0., self.d_model, 2) *
#                              -(math.log(10000.0) / self.d_model)).to(x.device)
         
#         pe[:, 0::2] = torch.sin(position * div_term)  
#         pe[:, 1::2] = torch.cos(position * div_term)   

#         if self.mode == 'cat':
#             x = self.linear(torch.cat((x,Variable(pe,require_grad=False)),dim=-1))

#         else:
#             x = x + Variable(pe,requires_grad=False)
        
#         return self.dropout(x)

# class PositionalEncoding_irregular_time(nn.Module):

#     def __init__(self, d_model_feature, d_model_hidden, dropout):
#         super(PositionalEncoding_irregular_time, self).__init__()
#         self.dropout = nn.Dropout(p=dropout)
#         self.d_model = d_model_feature
#         self.linear = nn.Linear(d_model_feature+d_model_hidden, d_model_feature)
#         self.mode = 'plus'

#     def forward(self, x, time):#x.shape=[batch,height,width,feature],time.shape=[batch]

#         pe = torch.zeros(x.shape[0],self.d_model).to(x.device)
#         position = time.unsqueeze(1) 
#         div_term = torch.exp(torch.arange(0., self.d_model, 2) *
#                              -(math.log(10000.0) / self.d_model)).to(x.device)
         
#         pe[:, 0::2] = torch.sin(position * div_term)  
#         pe[:, 1::2] = torch.cos(position * div_term)   

#         if self.mode == 'cat':
#             if len(x.shape) == 4:
#                 pe = pe.unsqueeze(1).unsqueeze(1)
#                 tmp = Variable(pe, requires_grad=False)
#                 tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
#             else:
#                 pe = pe.unsqueeze(1)
#                 tmp = Variable(pe, requires_grad=False)
#                 tmp = tmp.repeat(1, x.shape[1], 1)
#             x = self.linear(torch.cat((x, tmp),dim=-1))

#         else:
#             if len(x.shape) == 4:
#                 pe = pe.unsqueeze(1).unsqueeze(1)
#                 tmp = Variable(pe, requires_grad=False)
#                 tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
#             else:
#                 pe = pe.unsqueeze(1)
#                 tmp = Variable(pe, requires_grad=False)
#                 tmp = tmp.repeat(1, x.shape[1], 1)
#             x = x + tmp

        
#         return self.dropout(x)

# def attention(query, key, value, mask=None, dropout=None):
#     "Compute 'Scaled Dot Product Attention'"
#     d_k = query.size(-1)
#     scores = torch.matmul(query, key.transpose(-2, -1)) \
#              / math.sqrt(d_k)
    
#     # ipdb.set_trace()
    
#     # print('score',scores.shape,'mask',mask.shape)

#     if mask is not None:
#         # mask = mask.transpose(1,2).to(query.device)
#         mask = mask.to(query.device)
#         scores = scores.masked_fill(mask == 0, -1e10) 

#     p_attn = F.softmax(scores, dim = -1)
#     if dropout is not None:
#         p_attn = dropout(p_attn)

#     return torch.matmul(p_attn, value), p_attn

# class MultiHeadedAttention(nn.Module):
#     def __init__(self, h, q_dim, k_dim,v_dim, d_model, dropout=0.1):
#         "Take in model size and number of heads."
#         super(MultiHeadedAttention, self).__init__()
#         assert d_model % h == 0
#         # We assume d_v always equals d_k
#         self.d_k = d_model // h  
#         self.h = h 
#         self.q_dim = q_dim
#         self.k_dim = k_dim
#         self.linears = clones(nn.Linear(d_model, d_model), 4)
#         self.linear_q = nn.Linear(q_dim,d_model)
#         self.linear_k = nn.Linear(k_dim,d_model)
#         self.linear_v = nn.Linear(v_dim,d_model)
#         self.linear_out = nn.Linear(d_model,q_dim)
#         self.attn = None  
#         self.dropout = nn.Dropout(p=dropout)

#     def forward(self, query, key, value, mask=None):
#         "Implements Figure 2"
#         if mask is not None:
#             # Same mask applied to all h heads.
#             mask = mask.unsqueeze(1)
#         nbatches = query.size(0)

#         # 1) Do all the linear projections in batch from d_model => h x d_k 

#         query = self.linear_q(query).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
#         key = self.linear_k(key).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
#         value = self.linear_v(value).view(nbatches,-1,self.h,self.d_k).transpose(1,2)

#         # 2) Apply attention on all the projected vectors in batch. 
#         x, self.attn = attention(query, key, value, mask=mask,
#                                  dropout=self.dropout)

#         # 3) "Concat" using a view and apply a final linear. 
#         x = x.transpose(1, 2).contiguous() \
#             .view(nbatches ,-1, self.h * self.d_k)

#         return self.linear_out(x)

# class LayerNorm(nn.Module):
#     def __init__(self, features, eps=1e-6):
#         super(LayerNorm, self).__init__()
#         self.a_2 = nn.Parameter(torch.ones(features))
#         self.b_2 = nn.Parameter(torch.zeros(features))
#         self.eps = eps
 
#     def forward(self, x):
#         "Norm"
#         mean = x.mean(-1, keepdim=True)
#         std = x.std(-1, keepdim=True)
#         return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

# class SublayerConnection(nn.Module):
#     def __init__(self, size, dropout):
#         super(SublayerConnection, self).__init__()
#         self.norm = LayerNorm(size)
#         self.dropout = nn.Dropout(dropout)
 
#     def forward(self, x, sublayer):
#         # print('sub',x.shape,self.dropout(sublayer(self.norm(x))).shape)
#         return x + self.dropout(sublayer(self.norm(x))) #残差连接

# class PositionwiseFeedForward(nn.Module):
#     "Implements FFN equation."
#     def __init__(self, d_model,d_out, d_ff, dropout=0.1):
#         super(PositionwiseFeedForward, self).__init__()
#         self.w_1 = nn.Linear(d_model, d_ff)
#         self.w_2 = nn.Linear(d_ff, d_out)
#         self.dropout = nn.Dropout(dropout)

#     def forward(self, x):
#         return self.w_2(self.dropout(F.relu(self.w_1(x))))

# class EncoderLayer_temporal(nn.Module):
#     "Encoder is made up of self-attn and feed forward (defined below)"
#     def __init__(self, size, self_attn, feed_forward, dropout):
#         super(EncoderLayer_temporal, self).__init__()
#         self.self_attn = self_attn #sublayer 1 
#         self.feed_forward = feed_forward #sublayer 2 
#         self.sublayer = clones(SublayerConnection(size, dropout), 2) #拷贝两个SublayerConnection，一个为了attention，一个是独自作为简单的神经网络
#         self.size = size

#     def forward(self, q, k, v, mask):
#         "Follow Figure 1 (left) for connections."
#         # print('x',x.shape,'attn',self.self_attn(x,x,x,mask).shape)
#         x = self.sublayer[0](q, lambda q: self.self_attn(q, k, v, mask)) #attention层 对应层1
#         return self.sublayer[1](x, self.feed_forward) # 对应 层2

# class Spatial_conv(nn.Module):
#     def __init__(self, input_channels, output_channels, kernel_size):
#         super(Spatial_conv,self).__init__()
#         self.conv = nn.Conv2d(in_channels = input_channels,
#                               out_channels = output_channels,
#                               kernel_size = (kernel_size,kernel_size),
#                               stride = (1,1),
#                               padding = (int((kernel_size-1)/2),int((kernel_size-1)/2))
#         )

#         self.relu = nn.ReLU()
#         # self.bn = nn.BatchNorm2d(num_features=output_channels)

#     def forward(self, x):
#         init_x = x.clone()
#         x = self.conv(x)
#         x = self.relu(x)
#         # x = self.bn(x)
#         return x+init_x

# class SpatialAttention_mtf(nn.Module):
#     def __init__(self, kernel_size=7):
#         super(SpatialAttention_mtf, self).__init__()

#         self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
#         self.sigmoid = nn.Sigmoid()

#     def forward(self, curr, prev):#####(batch, channel, h, w)
#         avg_out = torch.mean(curr, dim=1, keepdim=True)
#         max_out, _ = torch.max(curr, dim=1, keepdim=True)
#         y = torch.cat([avg_out, max_out], dim=1)
#         y = self.conv1(y)
#         return self.sigmoid(y) * prev ####(batch-1, channel, h, w)

# class ScaledDotProductAttention(nn.Module):
#     """
#     Scaled Dot-Product Attention proposed in "Attention Is All You Need"
#     Compute the dot products of the query with all keys, divide each by sqrt(dim),
#     and apply a softmax function to obtain the weights on the values
#     Args: dim, mask
#         dim (int): dimention of attention
#         mask (torch.Tensor): tensor containing indices to be masked
#     Inputs: query, key, value, mask
#         - **query** (batch, q_len, d_model): tensor containing projection
#           vector for decoder.
#         - **key** (batch, k_len, d_model): tensor containing projection
#           vector for encoder.
#         - **value** (batch, v_len, d_model): tensor containing features of the
#           encoded input sequence.
#         - **mask** (-): tensor containing indices to be masked
#     Returns: context, attn
#         - **context**: tensor containing the context vector from
#           attention mechanism.
#         - **attn**: tensor containing the attention (alignment) from the
#           encoder outputs.
#     """

#     def __init__(self, dim):
#         super(ScaledDotProductAttention, self).__init__()
#         self.sqrt_dim = np.sqrt(dim)

#     def forward(self, query, key, value, mask):
#         score = torch.bmm(query, key.transpose(1, 2)) / self.sqrt_dim
#         if mask != None:
#             score = score.masked_fill(mask == 0, -1e10) 
#         attn = F.softmax(score, -1)
#         context = torch.bmm(attn, value)
#         return context

# class AttentionFusion(nn.Module):
#     def __init__(self, settings, input_dim, hidden_dim, num_heads, num_layers, neighbor_range, height, width, dropout, max_len=300,compensation=True):
#         super(AttentionFusion, self).__init__()
#         self.sweep_length = settings['sweep_length']
#         self.if_max_aggre = settings['if_max_aggre']
#         self.if_spatial_encoding = settings['if_spatial_encoding']
#         self.if_ego_time_encoding = settings['if_ego_time_encoding']
#         self.if_nonego_time_encoding = settings['if_nonego_time_encoding']
#         self.if_sensor_encoding = settings['if_sensor_encoding']
#         self.if_time_attn_aggre = settings['if_time_attn_aggre']
#         self.if_spatial_conv = settings['if_spatial_conv']
#         self.if_dotproductattn = settings['if_dotproductattn']
#         self.use_mask = settings['use_mask']
#         self.confidence_fetch = settings['confidence_fetch']
#         self.if_conv_aggre = settings['if_conv_aggre']
#         self.sup_individual = settings['sup_individual']
#         self.if_swin_trans = settings['if_swin_trans']
#         self.input_dim = input_dim 
#         self.hidden_dim = hidden_dim
#         self.num_heads = num_heads
#         self.num_layers = num_layers
#         self.neighbor_range = neighbor_range
#         self.height = height
#         self.width = width
#         self.compensation=compensation
#         if self.compensation:
#             if self.if_spatial_encoding:
#                 self.spatial_encoding = PositionalEncoding_spatial(self.input_dim, hidden_dim, dropout, max_len)
#                 self.spatial_encoding_nonego = PositionalEncoding_spatial_custom(self.input_dim, hidden_dim, dropout)
#             if self.if_ego_time_encoding or self.if_nonego_time_encoding:
#                 self.time_encoding = PositionalEncoding_irregular_time(self.input_dim, hidden_dim, dropout)
#             if self.if_sensor_encoding:
#                 self.sensor_dist_encoding = PositionalEncoding_sensor_dist(hidden_dim, dropout)

#             c = copy.deepcopy
#             temporal_attn = MultiHeadedAttention(self.num_heads, self.input_dim, self.input_dim, self.input_dim, self.hidden_dim, dropout)

#             ff = PositionwiseFeedForward(self.input_dim, self.input_dim, self.hidden_dim, dropout)

#             if self.if_time_attn_aggre:
#                 self.layers_temporal = clones(EncoderLayer_temporal(self.input_dim, c(temporal_attn), c(ff), dropout),num_layers)
#             # self.temporal_aggregator = MultiHeadedAttention(self.num_heads, hidden_dim, self.input_dim, self.input_dim, self.hidden_dim, dropout)
#             if self.if_spatial_conv:  
#                 self.layers_spatial = clones(Spatial_conv(self.input_dim,self.input_dim,kernel_size=5),num_layers)
#             if self.if_swin_trans:
#                 self.swin_trans = STPT(in_chans=self.input_dim,
#                     depths=[4, 4], 
#                     num_heads=[2, 2], 
#                     window_sizes=[(4,4), (4,4)],
#                     num_frames=3,
#                     num_out_frames=1,
#                     embed_dim=48,
#                     use_hdmap=False,
#                     use_hdmap_feat=False)
            
#             # self.layers_spatial = Spatial_conv(self.input_dim,self.input_dim,kernel_size=5)
#             if self.if_dotproductattn:
#                 self.dotproductattn = ScaledDotProductAttention(self.input_dim)
#             if self.if_conv_aggre:
#                 self.conv_aggre = SpatialAttention_mtf()
#         print("================Rain attention Init================")




#     def forward(self, features, sensor_dist, time_stamps, valid_mask):
#         """
#         Parameters:
#         -----------
#         features: torch.Tensor, shape (B, C, H, W)
#             neighbor features in ego view
#             B : 按照顺序排列: ego[curr], cav_1[curr], cav_1[past_1], ..., cav_2[past_k], ....
#             C : feature 维度，一般为 256
#             H : height, 一般为 100
#             W : width, 一般为 252

#         sensor_dist: torch.Tensor, shape (B, H, W)
#             distance between each block and the sensor
            
#         time_stamps: torch.Tensor, shape (B)
#             表示每个位置的时间信息

#         valid_mask: (B,1, H,W)
#             为0/1,表示是否在该位置有效

#         Returns:
#         --------
#         output_feature: torch.Tensor, shape (C, H, W)
#             ego cav's fused feature at current frame.
#         """
#         return torch.max(features, dim=0)[0]
        
#         if self.if_max_aggre:
#             features = torch.max(features, dim=0)[0]
#             if self.if_spatial_conv:
#                 features = features.unsqueeze(0)
#                 features = self.layers_spatial[0](features)
#             return features.squeeze(0)
#         if self.confidence_fetch == False:
#             if self.sup_individual:
#                 curr_features = features[1::self.sweep_length+1].clone()
#                 latency_features = features[2::self.sweep_length+1].clone()
#                 if self.compensation==False:
                    
#                     final_fuse_feature = torch.max(torch.cat((features[:1],curr_features),dim=0),dim=0)[0]
#                     final_fuse_feature_latency = torch.max(torch.cat((features[:1],latency_features),dim=0),dim=0)[0]

#                     return final_fuse_feature_latency,final_fuse_feature,final_fuse_feature_latency,0


#                 cav_num, C, H, W = features.shape
#                 individual_recon_loss = 0
                
#                 num_nonego = int((cav_num-1)/(self.sweep_length+1)) 
#                 nonego_features = features[1:].clone()
#                 nonego_timestamps = time_stamps[1:]
#                 nonego_features = nonego_features.permute(0,2,3,1)
#                 nonego_features = self.time_encoding(nonego_features,nonego_timestamps)
#                 nonego_features = nonego_features.permute(0,3,1,2)

#                 nonego_splits = nonego_features.split(self.sweep_length+1,0)
#                 nonego_splits = torch.stack(nonego_splits,dim=0)  ###[nonego_cars, T, C,H,W]
#                 nonego_history = nonego_splits[:,1:]
#                 # nonego_splits = nonego_splits.permute(0,3,4,1,2).reshape(num_nonego*H*W, self.sweep_length, C)
#                 # time_attned_feature = self.layers_temporal[0](nonego_splits, nonego_splits, nonego_splits, None)
#                 # time_attned_feature = nonego_splits
#                 # time_attned_feature = time_attned_feature.reshape(num_nonego, H, W, self.sweep_length*C).permute(0,3,1,2) 
                
#                 hidden = self.swin_trans(nonego_history)##############[non_ego_cars, 1,C,H,W]
#                 # query = self.predict_query.repeat(num_nonego,1,1)
#                 # final_estimations = self.predict_attention(query, hidden, hidden, None).reshape(num_nonego, H, W, -1).permute(0,3,1,2)
#                 final_estimations = hidden.squeeze(1)##############[non_ego_cars,C,H,W]
                
#                 individual_recon_loss = torch.torch.nn.functional.smooth_l1_loss(final_estimations,nonego_splits[:,0],reduction='mean')
#                 ###############skimage.dilation##########################################
#                 final_fuse_estimation = torch.max(torch.cat((features[:1],final_estimations),dim=0), dim=0)[0]
#                 final_fuse_feature = torch.max(torch.cat((features[:1],curr_features),dim=0),dim=0)[0]
#                 fuse_recon_loss = torch.torch.nn.functional.smooth_l1_loss(final_fuse_estimation,final_fuse_feature,reduction='mean')

#                 final_fuse_feature_latency = torch.max(torch.cat((features[:1],latency_features),dim=0),dim=0)[0]
#                 ipdb.set_trace()
#                 # print('#######################individual:',torch.torch.nn.functional.smooth_l1_loss(latest_features[1:],curr_features[1:]),individual_recon_loss)
#                 # latest_features = features[self.sweep_length-1::self.sweep_length].clone()
#                 # latest_features[0] = features[0]

#                 # final1 = torch.max(torch.cat((features[:1],latest_features[1:]),dim=0), dim=0)[0]
#                 # cnt = 5
#                 # torch.save(curr_features,'/GPFS/data/yuxiwei/yuxiwei/OpenCOOD/vis/current'+str(cnt)+'.pt')
#                 # torch.save(latest_features,'/GPFS/data/yuxiwei/yuxiwei/OpenCOOD/vis/latest'+str(cnt)+'.pt')
#                 # cnt += 1
#                 # ipdb.set_trace()
#                 # print('#######################:fuse',torch.torch.nn.functional.smooth_l1_loss(final1,final_fuse_feature),fuse_recon_loss)
             
#                 return final_fuse_estimation,final_fuse_feature,final_fuse_feature_latency, individual_recon_loss+fuse_recon_loss
#                 # return final_fuse_feature, 0

            
#             num_cav = int(features.shape[0]/self.sweep_length)
#             curr_features = features[:self.sweep_length:]
#             features = torch.cat((features[:1],features[num_cav:]),dim=0)
#             time_stamps = torch.cat((time_stamps[:1],time_stamps[num_cav:]),dim=0)
#             cav_num, C, H, W = features.shape
#             ################################
#             # features = features[::2]
#             # valid_mask = valid_mask[::2]
#             # cav_num = int(cav_num/2)
#             ################################


#             if self.if_conv_aggre:
#                 curr = features[0:1].clone()
#                 prev = features[1:].clone()
#                 prev = self.conv_aggre(curr,prev)
#                 features = torch.cat((curr,prev),dim=0)
#                 features = features.permute(0,2,3,1)
#                 time_stamps = time_stamps*-1
#                 features = self.time_encoding(features,time_stamps)
#                 features = features.permute(0,3,1,2)
#                 if self.if_spatial_conv:
#                     features = self.layers_spatial[0](features)
#                 return torch.max(features,dim=0)[0]
            
#             if self.if_ego_time_encoding:
#                 # s1 = time.time()
#                 features = features.permute(0,2,3,1)
#                 time_stamps = time_stamps*-1
#                 features = self.time_encoding(features,time_stamps)
#                 features = features.permute(0,3,1,2)
#                 # s2 = time.time()
#                 # print('TE',s2-s1)
            
            
#             features = features.view(cav_num, C, -1).permute(2, 0, 1) #  (H*W, cav_num, C), perform self attention on each pixel.
#             # masks = valid_mask.view(cav_num, H*W).permute(1,0)
#             # masks = masks.unsqueeze(-1).repeat(1,1,cav_num)
            


#             if self.if_dotproductattn:
#                 features = self.dotproductattn(features, features, features, masks)
#             if self.if_time_attn_aggre:
#                 # s1 = time.time()
#                 features = self.layers_temporal[0](features, features, features, None)
#                 # s2 = time.time()
#                 # print('trans',s2-s1)
            
           

#             features = features.permute(1, 2, 0).view(cav_num, C, H, W)[:1] 
#             if self.if_spatial_conv:
#                 # s1 = time.time()
#                 features = self.layers_spatial[0](features)[0]
#                 # s2 = time.time()
#                 # print('conv',s2-s1)
#             if self.sup_individual:
#                 features = torch.max(final_estimations, dim=0)[0]
#             curr_fuse_features = torch.max(curr_features, dim=0)[0]
#             recon_loss = torch.torch.nn.functional.smooth_l1_loss(features,curr_fuse_features)

#             return features, recon_loss


#         valid_mask = valid_mask.squeeze(1)
        
#         features = features.permute(0,2,3,1)
        
#         features = torch.cat((features[0:1],features[self.sweep_length:]),dim=0)
        

#         ##########################################  三种encoding
#         if self.if_spatial_encoding:
#             features[0:1] = self.spatial_encoding(features[0:1])# features.shape = [batch_all,height,width,feature_dim]
#         if self.if_sensor_encoding and sensor_dist!=-1:
#             sensor_dist = torch.cat((sensor_dist[0:1],sensor_dist[self.sweep_length:]))
#             features = self.sensor_dist_encoding(features,sensor_dist)
#         time_stamps = torch.cat((time_stamps[0:1],time_stamps[self.sweep_length:]),dim=0)
#         valid_mask = torch.cat((valid_mask[0:1],valid_mask[self.sweep_length:]),dim=0)
#         valid_mask_sum = valid_mask.sum(0)
#         nonego_idx = torch.where(valid_mask_sum>1)
 
#         features_to_aggre = features[:,nonego_idx[0],nonego_idx[1],:] 
#         if self.if_ego_time_encoding:
#             features[0:1] = self.time_encoding(features[0:1],time_stamps[0:1])

#         if features_to_aggre.shape[1] != 0:
#             if self.if_spatial_encoding:
#                 features_to_aggre = self.spatial_encoding_nonego(features_to_aggre,nonego_idx)
#             if self.if_nonego_time_encoding:
#                 features_to_aggre = self.time_encoding(features_to_aggre,time_stamps)
        

#         ###########################################

#         ################################################### 
#         # ipdb.set_trace()
#         masks = valid_mask[:,nonego_idx[0],nonego_idx[1]]#(valid_num,len)
#         batch_len = masks.shape[0]
#         masks = masks.permute(1,0).unsqueeze(-1).repeat(1,1,batch_len)

#         features_to_aggre = features_to_aggre.permute(1,0,2)
#         for l in range(self.num_layers):
                
#             if masks.shape[0] != 0:
#                 if self.if_dotproductattn:
#                     # if self.use_mask == False:
#                     #     masks =  None
#                     # features_to_aggre = self.dotproductattn(features_to_aggre,features_to_aggre,features_to_aggre,masks)
#                     features_to_aggre[:,0] = torch.max(features_to_aggre,dim=1)[0]
#                 else:
#                     features_to_aggre = self.layers_temporal[l](features_to_aggre,features_to_aggre,features_to_aggre,masks)
#                 features[:,nonego_idx[0],nonego_idx[1],:] = features_to_aggre.permute(1,0,2)
           

#             if self.if_spatial_conv:
#                 features = features.permute(0,3,1,2)
#                 features_backup = features.clone()
#                 features[0:1] = self.layers_spatial[l](features_backup[0:1])
#                 features = features.permute(0,2,3,1)

#                 features_to_aggre[:,0:1] = features[0:1,nonego_idx[0],nonego_idx[1],:].permute(1,0,2)

#         return features[0].permute(2,0,1)

# class MaxFusion(nn.Module):
#     def __init__(self):
#         super(MaxFusion, self).__init__()

#     def forward(self, x):
#         return torch.max(x, dim=0)[0]
        
# class raindrop_fuse(nn.Module):
#     def __init__(self, args):
#         super(raindrop_fuse, self).__init__()
        
#         self.communication = False
#         self.round = 1
#         if 'communication' in args:
#             self.communication = True
#             self.naive_communication = Communication(args['communication'])
#             if 'round' in args['communication']:
#                 self.round = args['communication']['round']
#         self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
#         self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
#         self.agg_mode = args['agg_operator']['mode']
#         self.multi_scale = args['multi_scale']

#         #################################################
#         if self.agg_mode == 'RAIN':
#             d_ff = args['d_ff']
#             h = args['h']
#             dropout = args['dropout']
#             input_dim = args['input_dim']
#             num_layers = args['num_layers']
#             neighbor_range = args['neighbor_range']
#             max_len = args['max_len']
#             height = 100
#             width = 352

#             if self.multi_scale:
#                 layer_nums = args['layer_nums']
#                 num_filters = args['num_filters']
#                 self.num_levels = len(layer_nums)
#                 self.fuse_modules = nn.ModuleList()
#                 for idx in range(self.num_levels):
#                     fuse_network = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len)
#                     for p in fuse_network.parameters():
#                         if p.dim() > 1:
#                             nn.init.xavier_uniform_(p)

#                     height = int(height/2)
#                     width = int(width/2)
#                     input_dim = input_dim * 2
#                     d_ff = d_ff * 2
#                     self.fuse_modules.append(fuse_network)
#             else:
#                 self.fuse_modules = AttentionFusion(input_dim, d_ff, h, num_layers, neighbor_range, 100, 252, dropout, max_len)
#                 # This was important from their code. 
#                 # Initialize parameters with Glorot / fan_avg.
#                 for p in self.fuse_modules.parameters():
#                     if p.dim() > 1:
#                         nn.init.xavier_uniform_(p)

#         else:
#             if self.multi_scale:
#                 layer_nums = args['layer_nums']
#                 num_filters = args['num_filters']
#                 self.num_levels = len(layer_nums)
#                 self.fuse_modules = nn.ModuleList()
#                 for idx in range(self.num_levels):
#                     if self.agg_mode == 'MAX':
#                         fuse_network = MaxFusion()
#                     else:
#                         # self.agg_mode == 'ATTEN':
#                         fuse_network = AttenFusion(num_filters[idx])
#                     self.fuse_modules.append(fuse_network)
#             else: 
#                 if self.agg_mode == 'MAX': # max fusion, debug use
#                     self.fuse_modules = MaxFusion()

#     def regroup(self, x, len, k=0):
#         """
#         split into different batch and time k

#         Parameters
#         ----------
#         x : torch.Tensor
#             input data, (B, ...)

#         len : list 
#             cav num in differnt batch, eg [3, 3]

#         k : torch.Tensor
#             num of past frames
        
#         Returns:
#         --------
#         split_x : list
#             different cav's lidar feature
#             for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
#         """
#         cum_sum_len = torch.cumsum(len*k, dim=0)
#         split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
#         return split_x

#     def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None):
#         """
#         Fusion forwarding.
        
#         Parameters
#         ----------
#         x : torch.Tensor
#             input data, (sum(n_cav), C, H, W)
        
#         rm : torch.Tensor
#             confidence map, (sum(n_cav), 2, H, W)

#         record_len : list
#             shape: (B)
            
#         pairwise_t_matrix : torch.Tensor
#             The transformation matrix from each cav to ego, 
#             shape: (B, L, L, 4, 4) 

#         time_diffs : torch.Tensor
#             time interval, (sum(n_cav), )
            
#         Returns
#         -------
#         Fused feature.
#         """
#         _, C, H, W = x.shape
#         B, L, K = pairwise_t_matrix.shape[:3]

#         # (B,L,k,2,3)
#         pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
#         # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
#         pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
#         pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
#         pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
#         pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

#         if self.multi_scale:
#             ups = []
#             # backbone.__dict__()
#             with_resnet = True if hasattr(backbone, 'resnet') else False
#             if with_resnet:
#                 feats = backbone.resnet(x)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
#             for i in range(self.num_levels):
#                 x = feats[i] if with_resnet else backbone.blocks[i](x)

#                 ############ 1. Communication (Mask the features) #########
#                 if i==0:
#                     if self.communication:
#                         batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
#                         batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
#                         _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
#                         communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
#                         # TODO scale = 0 不加 mask
#                         x = x * communication_masks_tensor
#                     else:
#                         communication_rates = torch.tensor(0).to(x.device)
#                 else:
#                     if self.communication:
#                         communication_masks_tensor = F.max_pool2d(communication_masks_tensor, kernel_size=2)
#                         # TODO scale = 1, 2 不加 mask
#                         # x = x * communication_masks_tensor
                
#                 ############ 2. Split the confidence map #######################
#                 # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
#                 # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
#                 # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
#                 batch_node_features = self.regroup(x, record_len, K)

#                 ############ 3. Fusion ####################################
#                 x_fuse = []
#                 for b in range(B):
#                     # number of valid agent
#                     N = record_len[b]
#                     # t_matrix[i, j]-> from i to j
#                     t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
#                     node_features = batch_node_features[b]
#                     C, H, W = node_features.shape[1:]
#                     # if node_features.shape[0] != t_matrix.shape[0]:
#                     neighbor_feature = warp_affine_simple(node_features,
#                                                     t_matrix,
#                                                     (H, W))
#                     record_frames = np.ones((N))*K
                    
#                     if self.agg_mode == 'RAIN':
#                         # for sensor embedding
#                         sensor_dist = -1# (B, H, W)
#                         x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
#                         # # TODO for scale debug
#                         # if i==self.num_levels-1:
#                         #     x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
#                         # else:
#                         #     x_fuse.append(neighbor_feature[0])
#                     else: # ATTEN, MAX, Transformer
#                         x_fuse.append(self.fuse_modules[i](neighbor_feature))
#                         # x_fuse.append(neighbor_feature[0]) # TODO: for single
#                         # print("=== single feature! ===")
#                         # # TODO for scale debug
#                         # if i==self.num_levels-1:
#                         #     x_fuse.append(self.fuse_modules[i](neighbor_feature))
#                         # else:
#                         #     x_fuse.append(neighbor_feature[0])

#                 x_fuse = torch.stack(x_fuse)

#                 ############ 4. Deconv ####################################
#                 if len(backbone.deblocks) > 0:
#                     ups.append(backbone.deblocks[i](x_fuse))
#                 else:
#                     ups.append(x_fuse)
                
#             if len(ups) > 1:
#                 x_fuse = torch.cat(ups, dim=1)
#             elif len(ups) == 1:
#                 x_fuse = ups[0]
            
#             if len(backbone.deblocks) > self.num_levels:
#                 x_fuse = backbone.deblocks[-1](x_fuse)
                
#         else:
#             ############ 1. Split the features #######################
#             # split x:[(L1, C, H, W), (L2, C, H, W), ...]
#             # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
#             batch_node_features = self.regroup(x, record_len, K)
#             batch_confidence_maps = self.regroup(rm, record_len, K)
#             batch_time_intervals = self.regroup(time_diffs, record_len, K)

#             ############ 2. Communication (Mask the features) #########
#             if self.communication:
#                 # _, (B, 1, H, W), float
#                 _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
#             else:
#                 communication_rates = torch.tensor(0).to(x.device)
            
#             ############ 3. Fusion ####################################
#             x_fuse = []
#             for b in range(B):
#                 # number of valid agent
#                 N = record_len[b]
#                 # (N,N,4,4)
#                 # t_matrix[i, j]-> from i to j
#                 t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
#                 node_features = batch_node_features[b]
#                 if self.communication:
#                     node_features = node_features * communication_masks[b]
#                 neighbor_feature = warp_affine_simple(node_features,
#                                                 t_matrix,
#                                                 (H, W))
#                 record_frames = np.ones((N))*K
#                 # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
#                 if self.agg_mode == "RAIN":
#                     # for sensor embedding
#                     sensor_dist = -1# (B, H, W)
#                     x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
#                 else: # ATTEN, MAX, Transformer
#                     x_fuse.append(self.fuse_modules(neighbor_feature))
#             x_fuse = torch.stack(x_fuse)

#             # self.fuse_modsules(x_fuse, record_len)
        
#         return x_fuse, communication_rates, {}


opencood/models/fuse_modules/raindrop_attn.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# Date: 2022/12/06
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

import torch.nn.functional as F
from turtle import update
# import ipdb
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
from opencood.models.comm_modules.where2comm_multisweep import Communication

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class PositionalEncoding_spatial(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model_feature, d_model_hidden, dropout, max_len=300,mode='cat'):
        super(PositionalEncoding_spatial, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        d_model = int(d_model_hidden/2)
        # Compute the positional encodings once in log space.
        pe = torch.zeros((max_len, d_model))
        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe_ = pe.unsqueeze(0)#batch,len,feature
        pe_h = pe_.unsqueeze(2).repeat(1,1,max_len,1)
        pe_w = pe_.unsqueeze(1).repeat(1,max_len,1,1)
        pe = torch.cat((pe_h,pe_w),dim=-1)

        self.register_buffer('pe', pe)

        self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x):
        if self.mode == 'cat':
            tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
            tmp = tmp.repeat(x.shape[0], 1, 1, 1)
            x = self.linear(torch.cat((x, tmp),dim=-1))

        elif self.mode == 'init':
            x = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)

        else:
            tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
            tmp = tmp.repeat(x.shape[0], 1, 1, 1)
            x = x + tmp

        return self.dropout(x)

class PositionalEncoding_spatial_custom(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model_feature, d_model_hidden, dropout,mode='cat'):
        super(PositionalEncoding_spatial_custom, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = int(d_model_hidden/2)
        # Compute the positional encodings once in log space.
        
        self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x, nonego_idx):#x(batch,num,feature)
   
        pe_h = torch.zeros((x.shape[1], self.d_model)).to(x.device)
        pe_w = torch.zeros((x.shape[1], self.d_model)).to(x.device)
        position_h, position_w = nonego_idx[0],nonego_idx[1]
        position_h = position_h.unsqueeze(1)
        position_w = position_w.unsqueeze(1)

        div_term = torch.exp(torch.arange(0, self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(position_h.device)
        pe_h[:, 0::2] = torch.sin(position_h * div_term)
        pe_h[:, 1::2] = torch.cos(position_h * div_term)

        pe_w[:, 0::2] = torch.sin(position_w * div_term)
        pe_w[:, 1::2] = torch.cos(position_w * div_term)
       
        pe = torch.cat((pe_h,pe_w),dim=-1).unsqueeze(0).repeat(x.shape[0],1,1)

        if self.mode == 'cat':
            x = self.linear(torch.cat((x, pe),dim=-1))

        else:
            x = x + Variable(pe,requires_grad=False)

        return self.dropout(x)

class PositionalEncoding_sensor_dist(nn.Module):

    def __init__(self, d_model, dropout):
        super(PositionalEncoding_sensor_dist, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = d_model
        self.linear = nn.Linear(2*d_model,d_model)
        self.mode = 'cat'
         
    def forward(self, x, sensor_dist):#x.shape=[batch,H,W,F],senfor_dist.shape=[batch,H,W]

        x = x.reshape(-1,x.shape[-1])
        sensor_dist = sensor_dist.reshape(-1)

        pe = torch.zeros(x.shape[0], self.d_model).to(x.device)
        position = sensor_dist.unsqueeze(1) 
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(x.device)
         
        pe[:, 0::2] = torch.sin(position * div_term)  
        pe[:, 1::2] = torch.cos(position * div_term)   

        if self.mode == 'cat':
            x = self.linear(torch.cat((x,Variable(pe,require_grad=False)),dim=-1))

        else:
            x = x + Variable(pe,requires_grad=False)
        
        return self.dropout(x)

class PositionalEncoding_irregular_time(nn.Module):

    def __init__(self, d_model_feature, d_model_hidden, dropout):
        super(PositionalEncoding_irregular_time, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = d_model_feature
        self.linear = nn.Linear(d_model_feature+d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x, time):#x.shape=[batch,height,width,feature],time.shape=[batch]

        pe = torch.zeros(x.shape[0],self.d_model).to(x.device)
        position = time.unsqueeze(1) 
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(x.device)
         
        pe[:, 0::2] = torch.sin(position * div_term)  
        pe[:, 1::2] = torch.cos(position * div_term)   

        if self.mode == 'cat':
            if len(x.shape) == 4:
                pe = pe.unsqueeze(1).unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
            else:
                pe = pe.unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], 1)
            x = self.linear(torch.cat((x, tmp),dim=-1))

        else:
            if len(x.shape) == 4:
                pe = pe.unsqueeze(1).unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
            else:
                pe = pe.unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], 1)
            x = x + tmp

        
        return self.dropout(x)

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    
    # ipdb.set_trace()
    
    # print('score',scores.shape,'mask',mask.shape)

    if mask is not None:
        # mask = mask.transpose(1,2).to(query.device)
        mask = mask.to(query.device)
        scores = scores.masked_fill(mask == 0, -1e10) 

    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)

    return torch.matmul(p_attn, value), p_attn

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, q_dim, k_dim,v_dim, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h  
        self.h = h 
        self.q_dim = q_dim
        self.k_dim = k_dim
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.linear_q = nn.Linear(q_dim,d_model)
        self.linear_k = nn.Linear(k_dim,d_model)
        self.linear_v = nn.Linear(v_dim,d_model)
        self.linear_out = nn.Linear(d_model,q_dim)
        self.attn = None  
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k 

        query = self.linear_q(query).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
        key = self.linear_k(key).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
        value = self.linear_v(value).view(nbatches,-1,self.h,self.d_k).transpose(1,2)

        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(query, key, value, mask=mask,
                                 dropout=self.dropout)

        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(1, 2).contiguous() \
            .view(nbatches ,-1, self.h * self.d_k)

        return self.linear_out(x)

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
 
    def forward(self, x):
        "Norm"
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, x, sublayer):
        # print('sub',x.shape,self.dropout(sublayer(self.norm(x))).shape)
        return x + self.dropout(sublayer(self.norm(x))) #残差连接

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model,d_out, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_out)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

class EncoderLayer_temporal(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer_temporal, self).__init__()
        self.self_attn = self_attn #sublayer 1 
        self.feed_forward = feed_forward #sublayer 2 
        self.sublayer = clones(SublayerConnection(size, dropout), 2) #拷贝两个SublayerConnection，一个为了attention，一个是独自作为简单的神经网络
        self.size = size

    def forward(self, q, k, v, mask):
        "Follow Figure 1 (left) for connections."
        # print('x',x.shape,'attn',self.self_attn(x,x,x,mask).shape)
        x = self.sublayer[0](q, lambda q: self.self_attn(q, k, v, mask)) #attention层 对应层1
        return self.sublayer[1](x, self.feed_forward) # 对应 层2

class Spatial_conv(nn.Module):
    def __init__(self, input_channels, output_channels, kernel_size):
        super(Spatial_conv,self).__init__()
        self.conv = nn.Conv2d(in_channels = input_channels,
                              out_channels = output_channels,
                              kernel_size = (kernel_size,kernel_size),
                              stride = (1,1),
                              padding = (int((kernel_size-1)/2),int((kernel_size-1)/2))
        )

        self.relu = nn.ReLU()
        # self.bn = nn.BatchNorm2d(num_features=output_channels)

    def forward(self, x):
        init_x = x.clone()
        x = self.conv(x)
        x = self.relu(x)
        # x = self.bn(x)
        return x+init_x

class SpatialAttention_mtf(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention_mtf, self).__init__()

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, curr, prev):#####(batch, channel, h, w)
        avg_out = torch.mean(curr, dim=1, keepdim=True)
        max_out, _ = torch.max(curr, dim=1, keepdim=True)
        y = torch.cat([avg_out, max_out], dim=1)
        y = self.conv1(y)
        return self.sigmoid(y) * prev ####(batch-1, channel, h, w)

class ScaledDotProductAttention(nn.Module):
    """
    Scaled Dot-Product Attention proposed in "Attention Is All You Need"
    Compute the dot products of the query with all keys, divide each by sqrt(dim),
    and apply a softmax function to obtain the weights on the values
    Args: dim, mask
        dim (int): dimention of attention
        mask (torch.Tensor): tensor containing indices to be masked
    Inputs: query, key, value, mask
        - **query** (batch, q_len, d_model): tensor containing projection
          vector for decoder.
        - **key** (batch, k_len, d_model): tensor containing projection
          vector for encoder.
        - **value** (batch, v_len, d_model): tensor containing features of the
          encoded input sequence.
        - **mask** (-): tensor containing indices to be masked
    Returns: context, attn
        - **context**: tensor containing the context vector from
          attention mechanism.
        - **attn**: tensor containing the attention (alignment) from the
          encoder outputs.
    """

    def __init__(self, dim):
        super(ScaledDotProductAttention, self).__init__()
        self.sqrt_dim = np.sqrt(dim)

    def forward(self, query, key, value, mask):
        score = torch.bmm(query, key.transpose(1, 2)) / self.sqrt_dim
        if mask != None:
            score = score.masked_fill(mask == 0, -1e10) 
        attn = F.softmax(score, -1)
        context = torch.bmm(attn, value)
        return context

class AttentionFusion(nn.Module):
    def __init__(self, settings, input_dim, hidden_dim, num_heads, num_layers, neighbor_range, height, width, dropout, max_len=300):
        super(AttentionFusion, self).__init__()
        self.sweep_length = settings['sweep_length']
        self.if_spatial_encoding = settings['if_spatial_encoding']
        self.if_ego_time_encoding = settings['if_ego_time_encoding']
        self.if_nonego_time_encoding = settings['if_nonego_time_encoding']
        self.if_sensor_encoding = settings['if_sensor_encoding']
        self.if_time_attn_aggre = settings['if_time_attn_aggre']
        self.if_spatial_conv = settings['if_spatial_conv']
        self.if_dotproductattn = settings['if_dotproductattn']
        self.use_mask = settings['use_mask']
        self.confidence_fetch = settings['confidence_fetch']  
        self.input_dim = input_dim 
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.neighbor_range = neighbor_range
        self.height = height
        self.width = width
        if self.if_spatial_encoding:
            self.spatial_encoding = PositionalEncoding_spatial(self.input_dim, hidden_dim, dropout, max_len)
            self.spatial_encoding_nonego = PositionalEncoding_spatial_custom(self.input_dim, hidden_dim, dropout)
        if self.if_ego_time_encoding or self.if_nonego_time_encoding:
            self.time_encoding = PositionalEncoding_irregular_time(self.input_dim, hidden_dim, dropout)
        if self.if_sensor_encoding:
            self.sensor_dist_encoding = PositionalEncoding_sensor_dist(hidden_dim, dropout)

        c = copy.deepcopy
        temporal_attn = MultiHeadedAttention(self.num_heads, self.input_dim, self.input_dim, self.input_dim, self.hidden_dim, dropout)

        ff = PositionwiseFeedForward(self.input_dim, self.input_dim, self.hidden_dim, dropout)

        if self.if_time_attn_aggre:
            self.layers_temporal = clones(EncoderLayer_temporal(self.input_dim, c(temporal_attn), c(ff), dropout),num_layers)
        # self.temporal_aggregator = MultiHeadedAttention(self.num_heads, hidden_dim, self.input_dim, self.input_dim, self.hidden_dim, dropout)
        if self.if_spatial_conv:  
            self.layers_spatial = clones(Spatial_conv(self.input_dim,self.input_dim,kernel_size=5),num_layers)
        # self.layers_spatial = Spatial_conv(self.input_dim,self.input_dim,kernel_size=5)
        if self.if_dotproductattn:
            self.dotproductattn = ScaledDotProductAttention(self.input_dim)

        print("================Rain attention Init================")

    


    def forward(self, features, sensor_dist, time_stamps, valid_mask):
        """
        Parameters:
        -----------
        features: torch.Tensor, shape (B, C, H, W)
            neighbor features in ego view
            B = sum_{i}^{num_cav} k_i, num_cav 表示cav数量, k_i 表示对于cav_i 有多少观测帧
            C : feature 维度，一般为 256
            H : height, 一般为 100
            W : width, 一般为 252

        sensor_dist: torch.Tensor, shape (B, H, W)
            distance between each block and the sensor
            
        time_stamps: torch.Tensor, shape (B)
            表示每个位置的时间信息

        valid_mask: (B,H,W)
            为0/1,表示是否在该位置有效

        Returns:
        --------
        output_feature: torch.Tensor, shape (C, H, W)
            ego cav's fused feature at current frame.
        """
        
        ego_features = features[:1].clone()
        return ego_features[0]
        # if self.if_time_attn_aggre != True and self.if_dotproductattn != True:
        #     features = torch.max(features, dim=0)[0]
        #     if self.if_spatial_conv:
        #         features = features.unsqueeze(0)
        #         features = self.layers_spatial[0](features)
        #     return features.squeeze(0)

        #### BEGIN ST atten ####
        #### Spatial
        cav_num, C, H, W = features.shape
        # features = features.view(cav_num, C, -1).permute(2, 0, 1) #  (H*W, cav_num, C), perform self attention on each pixel.
        # masks = valid_mask.view(cav_num, H*W).permute(1,0)
        # masks = masks.unsqueeze(-1).repeat(1,1,cav_num)
        
        # features = self.dotproductattn(features, features, features, masks)
        # features = features.permute(1, 0, 2).view(cav_num, H, W, C)

        #### Temporal
        features = features.permute(0, 2, 3, 1)  #(B, H, W, C)
        features = self.time_encoding(features,time_stamps) #(B, H, W, C)
        features = features.view(cav_num, -1, C).permute(1, 0, 2) #  (H*W, B, C), perform self attention on each time_stamps.
        masks = valid_mask.view(cav_num, H*W).permute(1, 0).unsqueeze(-1).repeat(1, 1, cav_num)
        features = self.dotproductattn(features, features, features, masks) #(HW, B, C) 
        features = features.permute(1, 2, 0).view(cav_num, C, H, W)[0]

        return features

        #### END ST atten ####

        if self.confidence_fetch == False:
            cav_num, C, H, W = features.shape
            features = features.view(cav_num, C, -1).permute(2, 0, 1) #  (H*W, cav_num, C), perform self attention on each pixel.
            masks = valid_mask.view(cav_num, H*W).permute(1,0)
            masks = masks.unsqueeze(-1).repeat(1,1,cav_num)
            
            features = self.dotproductattn(features, features, features, masks)
            features = features.permute(1, 2, 0).view(cav_num, C, H, W)[0] 
            return features

        valid_mask = valid_mask.squeeze(1)
        
        features = features.permute(0,2,3,1)
        
        features = torch.cat((features[0:1],features[self.sweep_length:]),dim=0)
        

        ##########################################  三种encoding
        if self.if_spatial_encoding:
            features[0:1] = self.spatial_encoding(features[0:1])# features.shape = [batch_all,height,width,feature_dim]
        if self.if_sensor_encoding and sensor_dist!=-1:
            sensor_dist = torch.cat((sensor_dist[0:1],sensor_dist[self.sweep_length:]))
            features = self.sensor_dist_encoding(features,sensor_dist)
        time_stamps = torch.cat((time_stamps[0:1],time_stamps[self.sweep_length:]),dim=0)
        valid_mask = torch.cat((valid_mask[0:1],valid_mask[self.sweep_length:]),dim=0)
        valid_mask_sum = valid_mask.sum(0)
        nonego_idx = torch.where(valid_mask_sum>1)
 
        features_to_aggre = features[:,nonego_idx[0],nonego_idx[1],:] 
        if self.if_ego_time_encoding:
            features[0:1] = self.time_encoding(features[0:1],time_stamps[0:1])

        if features_to_aggre.shape[1] != 0:
            if self.if_spatial_encoding:
                features_to_aggre = self.spatial_encoding_nonego(features_to_aggre,nonego_idx)
            if self.if_nonego_time_encoding:
                features_to_aggre = self.time_encoding(features_to_aggre,time_stamps)
        

        ###########################################

        ################################################### 
        # ipdb.set_trace()
        masks = valid_mask[:,nonego_idx[0],nonego_idx[1]]#(valid_num,len)
        batch_len = masks.shape[0]
        masks = masks.permute(1,0).unsqueeze(-1).repeat(1,1,batch_len)

        features_to_aggre = features_to_aggre.permute(1,0,2)
        for l in range(self.num_layers):
                
            if masks.shape[0] != 0:
                if self.if_dotproductattn:
                    # if self.use_mask == False:
                    #     masks =  None
                    # features_to_aggre = self.dotproductattn(features_to_aggre,features_to_aggre,features_to_aggre,masks)
                    features_to_aggre[:,0] = torch.max(features_to_aggre,dim=1)[0]
                else:
                    features_to_aggre = self.layers_temporal[l](features_to_aggre,features_to_aggre,features_to_aggre,masks)
                features[:,nonego_idx[0],nonego_idx[1],:] = features_to_aggre.permute(1,0,2)
           

            if self.if_spatial_conv:
                features = features.permute(0,3,1,2)
                features_backup = features.clone()
                features[0:1] = self.layers_spatial[l](features_backup[0:1])
                features = features.permute(0,2,3,1)

                features_to_aggre[:,0:1] = features[0:1,nonego_idx[0],nonego_idx[1],:].permute(1,0,2)

        return features[0].permute(2,0,1)

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_fuse(nn.Module):
    def __init__(self, args):
        super(raindrop_fuse, self).__init__()
        
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.agg_mode == 'RAIN':
            d_ff = args['d_ff']
            h = args['h']
            dropout = args['dropout']
            input_dim = args['input_dim']
            num_layers = args['num_layers']
            neighbor_range = args['neighbor_range']
            max_len = args['max_len']
            height = 100
            width = 252

            if self.multi_scale:
                self.fuse_modules = nn.ModuleList()
                for idx in range(self.num_levels):
                    fuse_network = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len)
                    for p in fuse_network.parameters():
                        if p.dim() > 1:
                            nn.init.xavier_uniform_(p)

                    height = int(height/2)
                    width = int(width/2)
                    input_dim = input_dim * 2
                    d_ff = d_ff * 2
                    self.fuse_modules.append(fuse_network)
            else:
                self.fuse_modules = AttentionFusion(input_dim, d_ff, h, num_layers, neighbor_range, 100, 252, dropout, max_len)
                # This was important from their code. 
                # Initialize parameters with Glorot / fan_avg.
                for p in self.fuse_modules.parameters():
                    if p.dim() > 1:
                        nn.init.xavier_uniform_(p)

        else:
            if self.multi_scale:
                layer_nums = args['layer_nums']
                num_filters = args['num_filters']
                self.num_levels = len(layer_nums)
                self.fuse_modules = nn.ModuleList()
                for idx in range(self.num_levels):
                    if self.agg_mode == 'MAX':
                        fuse_network = MaxFusion()
                    else:
                        # self.agg_mode == 'ATTEN':
                        fuse_network = AttenFusion(num_filters[idx])
                    self.fuse_modules.append(fuse_network)
            else: 
                if self.agg_mode == 'MAX': # max fusion, debug use
                    self.fuse_modules = MaxFusion()

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        if self.multi_scale:
            ups = []
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(x)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
            for i in range(self.num_levels):
                x = feats[i] if with_resnet else backbone.blocks[i](x)

                ############ 1. Communication (Mask the features) #########
                if i==0:
                    if self.communication:
                        batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                        batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                        _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                        communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                        x = x * communication_masks_tensor
                    else:
                        communication_rates = torch.tensor(0).to(x.device)
                else:
                    if self.communication:
                        communication_masks_tensor = F.max_pool2d(communication_masks_tensor, kernel_size=2)
                        # TODO: scale = 1, 2 不加 mask
                        # x = x * communication_masks_tensor
                
                ############ 2. Split the confidence map #######################
                # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                batch_node_features = self.regroup(x, record_len, K)

                ############ 3. Fusion ####################################
                x_fuse = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))
                    record_frames = np.ones((N))*K
                    # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                    
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])
                    else: # ATTEN, MAX, Transformer
                        x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])

                x_fuse = torch.stack(x_fuse)

                ############ 4. Deconv ####################################
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[i](x_fuse))
                else:
                    ups.append(x_fuse)
                
            if len(ups) > 1:
                x_fuse = torch.cat(ups, dim=1)
            elif len(ups) == 1:
                x_fuse = ups[0]
            
            if len(backbone.deblocks) > self.num_levels:
                x_fuse = backbone.deblocks[-1](x_fuse)
                
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_confidence_maps = self.regroup(rm, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                # _, (B, 1, H, W), float
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix,
                                                (H, W))
                record_frames = np.ones((N))*K
                # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                if self.agg_mode == "RAIN":
                    # for sensor embedding
                    sensor_dist = -1# (B, H, W)
                    x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)

            # self.fuse_modsules(x_fuse, record_len)
        
        return x_fuse, communication_rates, {}
opencood/models/fuse_modules/raindrop_attn_modifying.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# Date: 2022/12/06
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

import torch.nn.functional as F
from turtle import update
# import ipdb
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
from opencood.models.comm_modules.where2comm_multisweep import Communication

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_fuse(nn.Module):
    def __init__(self, args):
        super(raindrop_fuse, self).__init__()
        
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.agg_mode == 'MAX': # max fusion, debug use
            self.fuse_modules = MaxFusion()

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        ############ 1. Split the features #######################
        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        batch_node_features = self.regroup(x, record_len, K)
        batch_confidence_maps = self.regroup(rm, record_len, K)
        batch_time_intervals = self.regroup(time_diffs, record_len, K)

        ############ 2. Communication (Mask the features) #########
        if self.communication:
            # _, (B, 1, H, W), float
            _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
        else:
            communication_rates = torch.tensor(0).to(x.device)
        
        ############ 3. Fusion ####################################
        x_fuse = []
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
            node_features = batch_node_features[b]
            if self.communication:
                node_features = node_features * communication_masks[b]
            neighbor_feature = warp_affine_simple(node_features,
                                            t_matrix,
                                            (H, W))
            record_frames = np.ones((N))*K
            # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
            if self.agg_mode == "RAIN":
                # for sensor embedding
                sensor_dist = -1# (B, H, W)
                x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
            else: # ATTEN, MAX, Transformer
                x_fuse.append(self.fuse_modules(neighbor_feature))
        x_fuse = torch.stack(x_fuse)

        # self.fuse_modsules(x_fuse, record_len)
        
        return x_fuse, communication_rates, {}


opencood/models/fuse_modules/raindrop_attn_single_2_fused.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
from opencood.models.sub_modules.convgru import ConvGRU
from icecream import ic
from matplotlib import pyplot as plt
# from opencood.models.sub_modules.SyncLSTM import SyncLSTM
from opencood.models.sub_modules.MotionNet import STPN, MotionPrediction, StateEstimation, FlowUncPrediction
# from opencood.models.sub_modules.dcn_net import DCNNet
from opencood.models.comm_modules.where2comm_multisweep import Communication
from opencood.models.fuse_modules.SyncNet import SyncLSTM

import torch.nn as nn

class FineTuneFlow(nn.Module):
    def __init__(self):
        super(FineTuneFlow, self).__init__()
        self.conv1 = nn.Conv2d(2, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 2, kernel_size=3, padding=1)

    def forward(self, x):
        x = x.permute(0, 3, 1, 2)
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.relu(self.conv2(x))
        x = self.conv3(x)
        x = x.permute(0, 2, 3, 1)
        return x

class FineTuneFlow_2(nn.Module):
    def __init__(self):
        super(FineTuneFlow_2, self).__init__()
        self.conv1 = nn.Conv2d(4, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 2, kernel_size=3, padding=1)

    def forward(self, x_1, x_2):
        x = torch.cat([x_1, x_2], dim=-1)
        x = x.permute(0, 3, 1, 2)
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.relu(self.conv2(x))
        x = self.conv3(x)
        x = x.permute(0, 2, 3, 1)
        return x


def get_warped_feature_mask(flow, updated_grid):
    ''' 
    get the mask of the warped feature map, where the value is 1 means the location is valid, and 0 means the location is invalid.
    ----------
    Args:
        flow: (2, H, W)
        updated_grid: (2, H, W)

    Returns:
        mask: (H, W)
    '''
    def get_large_flow_indices(flow, threshold=0):
        '''
        find the indices of the flow map where the flow is larger than the threshold, which means these locations have been moved.
        ----------
        Args:
            flow: (2, H, W)
            
        '''
        max_values, max_indices = torch.max(torch.abs(flow[:2]), dim=0)
        large_indices = torch.nonzero(max_values > threshold, as_tuple=False)
        return large_indices

    def remove_duplicate_points(points):
        unique_points, inverse_indices = torch.unique(points, sorted=True, return_inverse=True, dim=0)
        return unique_points, inverse_indices
    
    def get_nonzero_idx(flow, idx):
        flow_values = flow[:, idx[:, 0], idx[:, 1]]
        nonzero_idx = torch.nonzero(torch.abs(flow_values).sum(dim=0) == 0, as_tuple=False).squeeze()
        return idx[nonzero_idx]
    
    flow_idx = get_large_flow_indices(flow)

    mask = torch.ones(flow.shape[-2], flow.shape[-1])
    if flow_idx.shape[0] == 0:
        return mask

    # print(flow_idx)
    updated_grid_points_tmp = updated_grid[:, flow_idx[:,0], flow_idx[:,1]].to(torch.int64).T
    # change the order of dim 1
    updated_grid_points = torch.zeros_like(updated_grid_points_tmp)
    updated_grid_points[:, 0] = updated_grid_points_tmp[:, 1]
    updated_grid_points[:, 1] = updated_grid_points_tmp[:, 0]
    # print(updated_grid_points)

    unique_points_idx, _ = remove_duplicate_points(updated_grid_points)
    # print(unique_points_idx)

    nonzero_idx = get_nonzero_idx(flow, unique_points_idx)
    # print(nonzero_idx)
    
    # mask out the outlier idx

    if len(nonzero_idx.shape) > 1:
        mask[nonzero_idx[:, 0], nonzero_idx[:, 1]] = 0
    else: 
        mask[nonzero_idx[0], nonzero_idx[1]] = 0

    return mask

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_fuse(nn.Module):
    def __init__(self, args, design_mode=0):
        super(raindrop_fuse, self).__init__()
        
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.multi_scale:
            layer_nums = args['layer_nums']
            num_filters = args['num_filters']
            self.num_levels = len(layer_nums)
            self.fuse_modules = nn.ModuleList()
            for idx in range(self.num_levels):
                if self.agg_mode == 'MAX':
                    fuse_network = MaxFusion()
                else:
                    # self.agg_mode == 'ATTEN':
                    fuse_network = AttenFusion(num_filters[idx])
                self.fuse_modules.append(fuse_network)
        else: 
            if self.agg_mode == 'MAX': # max fusion, debug use
                self.fuse_modules = MaxFusion()

        self.design = design_mode
        if self.design == 1:
            self.fine_conv = FineTuneFlow()
        elif self.design == 2:
            self.stpn = STPN(args['channel_size'])
            self.motion_pred = MotionPrediction(seq_len=1)
            self.state_classify = StateEstimation(motion_category_num=1)
            self.fine_conv = FineTuneFlow()
            self.fine_conv_2 = FineTuneFlow_2()
        elif self.design == 3:
            self.stpn = STPN(args['channel_size'])
            self.motion_pred = MotionPrediction(seq_len=1)
            self.state_classify = StateEstimation(motion_category_num=1)
        elif self.design == 4: # SyncNet
            design4_h = 200; design4_w = 704; design4_k = 3; design4_TM_Flag = False
            self.syncnet = SyncLSTM(channel_size = 64, h = design4_h, w = design4_w, k = design4_k, TM_Flag = design4_TM_Flag)
            print(f"*** SyncNet init finished! channel_size = 64, h = {design4_h}, w = {design4_w}, k = {design4_k}, TM_Flag = {design4_TM_Flag} ***")

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def generateFlow(self, feats, pairwise_t_matrix, record_len, flow_gt=None):
        '''
        1. generate flow from feature sequence
        2. then update the feature
        Note:
            1. ego feature does not need to be updated

        params:
            feats: (sum(B,N,K), C, H, W)
            pairwise_t_matrix: (B, L, K, 2, 3)
            record_len: (B)

        return:
            flow: 
        '''
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]
        batch_node_features = self.regroup(feats, record_len, k=K)

        # debug use
        batch_flow_gt = self.regroup(flow_gt, record_len, 2)

        updated_features_list = []
        flow_list = []
        state_class_pred_list = []
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            # t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
            node_features = batch_node_features[b]
            # neighbor_feature = warp_affine_simple(node_features, t_matrix, (H, W))
            node_features = node_features.view(-1, K, C, H, W) # (N, k, C, H, W)
            bevs = self.stpn(node_features)
            flow = self.motion_pred(bevs) # [N, 2, H, W]
            flow_list.append(flow)

            # TODO: debug use
            flow = batch_flow_gt[b].view(-1, 2, H, W)
            # print('=== gt flow used! ===')
            
            # Motion State Classification head
            state_class_pred = self.state_classify(bevs)
            state_class_pred_list.append(state_class_pred)

            # Given disp shift feature
            x_coord = torch.arange(W).float()   # [0, ..., W]
            y_coord = torch.arange(H).float()   # [0, ..., H]
            y, x = torch.meshgrid(y_coord, x_coord)  # [H, W], [H, W]
            grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(flow.shape[0], -1, -1, -1).to(flow.device)  # N, 2, H, W
            # updated_grid = grid + flow * (state_class_pred.sigmoid() > self.flow_thre)
            updated_grid = grid + flow # TODO: 这个位置为什么是减号
            
            # TODO: debug use, do not use gt flow, keep the original delay feature as the fusion input. 
            # updated_grid = grid
            
            # generate the mask for filtering out the pixels which moved to other locations but not filled by others
            mask_list = []
            for i in range(flow.shape[0]):
                mask_list.append(get_warped_feature_mask(flow[i], updated_grid[i]))
            mask = torch.stack(mask_list, dim=0)
            mask = mask.unsqueeze(1).repeat(1, C, 1, 1).to(flow.device)
            ##################################
            
            updated_grid[:, 0, :, :] = updated_grid[:, 0, :, :] / (W / 2.0) - 1.0
            updated_grid[:, 1, :, :] = updated_grid[:, 1, :, :] / (H / 2.0) - 1.0
            latest_node_features = node_features[:, 0, :, :, :] # (N, C, H, W)
            updated_features = F.grid_sample(latest_node_features, grid=updated_grid.permute(0, 2, 3, 1), mode='bilinear', align_corners=False)
            # mask the features
            updated_features = mask * updated_features
            # ego feature use the latest feature (no delay)
            updated_features[0, :, :, :] = latest_node_features[0, :, :, :]

            updated_features_list.append(updated_features)
        
        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)

        # TODO: for debug use
        # debug_path = '/remote-home/share/sizhewei/logs/where2comm_flow_debug/viz_flow/'
        # torch.save(feats, debug_path + 'feats.pt')
        # torch.save(updated_features_all, debug_path + 'updated_features_all.pt')

        flow_all = torch.cat(flow_list, dim=0)  # (sum(B,N), 2, H, W)
        state_class_pred_all = torch.cat(state_class_pred_list, dim=0)  # (sum(B,N), 2, H, W)
        
        return updated_features_all, flow_all, state_class_pred_all

    def update_features_boxflow(self, feats, pairwise_t_matrix, record_len, flow_map, reserved_mask, flow_gt):
        """
        Update features with box flow.
        
        Parameters
        ----------
        feats : torch.Tensor
            input data, 
            shape: (sum(n_cav), C, H, W)
        
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        record_len : list
            shape: (B)
            
        flow_map: torch.Tensor
            flow generate by past 2 frames detection boxes
            shape: (sum(N_b), H, W, 2)

        reserved_mask: torch.Tensor
            shape: (sum(N_b), C, H, W)

        gt_flow: torch.Tensor
            shape: (sum(N_b), H, W, 2)
            
        Returns
        -------
        Warped feature.
        """
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]
        batch_node_features = self.regroup(feats, record_len, k=K)

        batch_flow_map = self.regroup(flow_map, record_len, k=1)
        batch_reserved_mask = self.regroup(reserved_mask, record_len, k=1)

        updated_features_list = []
        flow_list = []
        state_class_pred_list = []
        gt_flow_map_list = []
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            node_features = batch_node_features[b]
            node_features = node_features.view(-1, K, C, H, W) # (N, k, C, H, W)
            flow = batch_flow_map[b] # .view(-1, 2, H, W) # [N, H, W, 2]
            # flow = self.fine_conv(flow) # [N, H, W, 2] 
            # flow_list.append(flow)
            
            latest_node_features = node_features[:, 0, :, :, :] # past_0 features, (N, C, H, W) 
            updated_features = F.grid_sample(latest_node_features, grid=flow, mode='bilinear', align_corners=False)

            updated_features = updated_features*batch_reserved_mask[b] # (N, C, H, W)
            # updated_features[0, ...] = node_features[0, 0, :, :, :] # (N, C, H, W)
            updated_features_list.append(updated_features)

        # flow_pred = torch.cat(flow_list, dim=0) # (sum(N_b), H, W, 2)
        # gt_flow_norm = torch.cat(gt_flow_map_list, dim=0) # (sum(N_b), H, W, 2)
        # # compute the flow map loss:
        # loss = F.smooth_l1_loss(flow_pred, gt_flow_norm)

        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)

        return updated_features_all
    
    def update_features_boxflow_design_1(self, feats, pairwise_t_matrix, record_len, flow_map, reserved_mask, flow_gt):
        """
        Update features with box flow.
        
        Parameters
        ----------
        feats : torch.Tensor
            input data, 
            shape: (sum(n_cav), C, H, W)
        
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        record_len : list
            shape: (B)
            
        flow_map: torch.Tensor
            flow generate by past 2 frames detection boxes
            shape: (sum(N_b), H, W, 2)

        reserved_mask: torch.Tensor
            shape: (sum(N_b), C, H, W)

        gt_flow: torch.Tensor
            shape: (sum(N_b), H, W, 2)
            
        Returns
        -------
        Warped feature.
        """
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]
        batch_node_features = self.regroup(feats, record_len, k=K)

        batch_flow_map = self.regroup(flow_map, record_len, k=1)
        batch_reserved_mask = self.regroup(reserved_mask, record_len, k=1)

        # debug use
        batch_flow_gt = self.regroup(flow_gt, record_len, k=2)

        updated_features_list = []
        flow_list = []
        state_class_pred_list = []
        gt_flow_map_list = []
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            node_features = batch_node_features[b]
            node_features = node_features.view(-1, K, C, H, W) # (N, k, C, H, W)
            flow = batch_flow_map[b] # .view(-1, 2, H, W) # [N, H, W, 2]
            # flow = self.fine_conv(flow) # [N, H, W, 2] # TODO: conv
            flow_list.append(flow)
            # # motion net flow
            # bevs = self.stpn(node_features)
            # flow_pred = self.motion_pred(bevs) # [N, 2, H, W]
            # # Motion State Classification head
            # state_class_pred = self.state_classify(bevs)
            # state_class_pred_list.append(state_class_pred)
            
            latest_node_features = node_features[:, 0, :, :, :] # past_0 features, (N, C, H, W) 
            updated_features = F.grid_sample(latest_node_features, grid=flow, mode='nearest', align_corners=False)

            updated_features = updated_features*batch_reserved_mask[b] # (N, C, H, W)
            updated_features_list.append(updated_features)

            # normalizing GT flow: gt_flow_map [N, H, W, 2]
            # Given disp shift feature
            x_coord = torch.arange(W).float()   # [0, ..., W]
            y_coord = torch.arange(H).float()   # [0, ..., H]
            y, x = torch.meshgrid(y_coord, x_coord)  # [H, W], [H, W]
            grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(N, -1, -1, -1).to(flow.device)
            gt_flow_delta = batch_flow_gt[b].view(-1, 2, H, W)
            gt_flow_map = grid - gt_flow_delta
            gt_flow_map[:, 0, :, :] = gt_flow_map[:, 0, :, :] / (W / 2.0) - 1.0
            gt_flow_map[:, 1, :, :] = gt_flow_map[:, 1, :, :] / (H / 2.0) - 1.0
            gt_flow_map = gt_flow_map.permute(0, 2, 3, 1) 
            gt_flow_map_list.append(gt_flow_map)

        flow_pred = torch.cat(flow_list, dim=0) # (sum(N_b), H, W, 2)
        gt_flow_norm = torch.cat(gt_flow_map_list, dim=0) # (sum(N_b), H, W, 2)
        # compute the flow map loss:
        loss = F.smooth_l1_loss(flow_pred, gt_flow_norm)

        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)

        return updated_features_all, loss
    
    def update_features_boxflow_design_2(self, feats, pairwise_t_matrix, record_len, flow_map, reserved_mask, flow_gt):
        """
        Update features with box flow.
        
        Parameters
        ----------
        feats : torch.Tensor
            input data, 
            shape: (sum(n_cav), C, H, W)
        
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        record_len : list
            shape: (B)
            
        flow_map: torch.Tensor
            flow generate by past 2 frames detection boxes
            shape: (sum(N_b), H, W, 2)

        reserved_mask: torch.Tensor
            shape: (sum(N_b), C, H, W)

        gt_flow: torch.Tensor
            shape: (sum(N_b), H, W, 2)
            
        Returns
        -------
        Warped feature.
        """
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]
        batch_node_features = self.regroup(feats, record_len, k=K)

        batch_flow_map = self.regroup(flow_map, record_len, k=1)
        batch_reserved_mask = self.regroup(reserved_mask, record_len, k=1)

        # debug use
        batch_flow_gt = self.regroup(flow_gt, record_len, k=2)

        updated_features_list = []
        flow_list = []
        state_class_pred_list = []
        gt_flow_map_list = []
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            node_features = batch_node_features[b]
            node_features = node_features.view(-1, K, C, H, W) # (N, k, C, H, W)
            flow_box = batch_flow_map[b] # .view(-1, 2, H, W) # [N, H, W, 2]
            # flow_box = self.fine_conv(flow_box) # [N, H, W, 2]
            
            # motion net flow
            bevs = self.stpn(node_features)
            flow_delta_pred = self.motion_pred(bevs) # [N, 2, H, W]
            # Motion State Classification head
            state_class_pred = self.state_classify(bevs)
            state_class_pred_list.append(state_class_pred)
            # Given disp shift feature
            x_coord = torch.arange(W).float()   # [0, ..., W]
            y_coord = torch.arange(H).float()   # [0, ..., H]
            y, x = torch.meshgrid(y_coord, x_coord)  # [H, W], [H, W]
            grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(N, -1, -1, -1).to(flow_box.device)
            # flow_pred = grid - flow_delta_pred
            # # normalizing flow: flow_pred [N, H, W, 2]
            # flow_pred[:, 0, :, :] = flow_pred[:, 0, :, :] / (W / 2.0) - 1.0
            # flow_pred[:, 1, :, :] = flow_pred[:, 1, :, :] / (H / 2.0) - 1.0
            # flow_pred = flow_pred.permute(0, 2, 3, 1) # [N, H, W, 2]

            # normalizing flow_delta_pred 
            flow_delta_pred[:, 0, :, :] = flow_delta_pred[:, 0, :, :] / (W / 2.0) - 1.0
            flow_delta_pred[:, 1, :, :] = flow_delta_pred[:, 1, :, :] / (H / 2.0) - 1.0
            flow_delta_pred = flow_delta_pred.permute(0, 2, 3, 1) # [N, H, W, 2]
            flow = self.fine_conv(flow_delta_pred + flow_box) # [N, H, W, 2]

            # flow = self.fine_conv_2(flow_box, flow_pred) # [N, H, W, 2]
            flow_list.append(flow)
            
            latest_node_features = node_features[:, 0, :, :, :] # past_0 features, (N, C, H, W) 
            updated_features = F.grid_sample(latest_node_features, grid=flow, mode='nearest', align_corners=False)

            updated_features = updated_features*batch_reserved_mask[b] # (N, C, H, W)
            updated_features_list.append(updated_features)

            # normalizing GT flow: gt_flow_map [N, H, W, 2]
            gt_flow_delta = batch_flow_gt[b].view(-1, 2, H, W)
            gt_flow_map = grid - gt_flow_delta
            gt_flow_map[:, 0, :, :] = gt_flow_map[:, 0, :, :] / (W / 2.0) - 1.0
            gt_flow_map[:, 1, :, :] = gt_flow_map[:, 1, :, :] / (H / 2.0) - 1.0
            gt_flow_map = gt_flow_map.permute(0, 2, 3, 1) 
            gt_flow_map_list.append(gt_flow_map)

        flow_all = torch.cat(flow_list, dim=0) # (sum(N_b), H, W, 2)
        state_class_pred_all = torch.cat(state_class_pred_list, dim=0)  # (sum(B,N), 2, H, W)
        gt_flow_norm = torch.cat(gt_flow_map_list, dim=0) # (sum(N_b), H, W, 2)
        # compute the flow map loss:
        loss = F.smooth_l1_loss(flow_all, gt_flow_norm)

        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)

        return updated_features_all, loss, flow_all, state_class_pred_all

    # syncnet
    def generate_estimated_feats(self, feats, pairwise_t_matrix, record_len):
        '''
        update the feature by SyncNet
        Note:
            1. ego feature does not need to be updated

        params:
            feats: (sum(B,N,K), C, H, W)
            pairwise_t_matrix: (B, L, K, 2, 3)
            record_len: (B)

        return:
            updated feature: [N, C, H, W]
        '''
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]
        batch_node_features = self.regroup(feats, record_len, k=K)

        updated_features_list = []
        flow_list = []
        # all_recon_loss = 0
        for b in range(B):
            all_feats = batch_node_features[b]
            non_ego_feats = all_feats[K:].clone()
            non_ego_feats = non_ego_feats.reshape(-1, K, C, H, W)
            non_ego_feats = torch.flip(non_ego_feats, [1])
            estimated_non_ego_feats = self.syncnet(non_ego_feats, [1]) # [N-1, 1, C, H, W]
            # estimated_non_ego_feats = non_ego_feats[:, 1:2] # TODO:
            updated_features_list.append(torch.cat([all_feats[0].unsqueeze(0), estimated_non_ego_feats[:,0]], dim=0))

        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)
        
        return updated_features_all

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None, flow_gt=None, box_flow=None, reserved_mask=None, viz_bbx_flag=False, noise_pairwise_t_matrix=None, num_roi_thres=-1):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
        
        flow_gt: ground truth flow, generate by object center and id

        box_flow: flow generate by past 2 frames detection boxes
            
        Returns
        -------
        Fused feature
        flow_map loss
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        if noise_pairwise_t_matrix is not None:
            noise_pairwise_t_matrix = noise_pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
            noise_pairwise_t_matrix[...,0,1] = noise_pairwise_t_matrix[...,0,1] * H / W
            noise_pairwise_t_matrix[...,1,0] = noise_pairwise_t_matrix[...,1,0] * W / H
            noise_pairwise_t_matrix[...,0,2] = noise_pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
            noise_pairwise_t_matrix[...,1,2] = noise_pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        # single frame TODO: 
        # batch_confidence_maps = self.regroup(psm_single, record_len, 1) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
        self.num_roi_thres = num_roi_thres
        if self.num_roi_thres > 0:
            self.thre = 0.01
            _, _, H, W = rm.shape
            num_valid_thres = self.num_roi_thres * 10
            communication_maps = rm.sigmoid().max(dim=1)[0].unsqueeze(1)
            communication_masks_list = []
            curr_batch = communication_maps.shape[0] # num_cav * k
            communication_mask = (communication_maps>self.thre).to(torch.int)
            valid_nums = torch.sum(communication_mask, dim=(-1, -2, -3))

            final_comm_mask = torch.zeros_like(communication_mask)
            for i in range(curr_batch):
                if valid_nums[i] > num_valid_thres :
                    # 选择 comm_maps[i] 里面 top num_valid 的位置
                    tmp = communication_maps[i].reshape((1, -1))
                    _, idx = torch.topk(tmp, num_valid_thres)
                    # print(idx)
                    curr_mask = torch.zeros_like(tmp)
                    curr_mask[0, idx[0]] = 1
                    final_comm_mask[i] = curr_mask.reshape((1, H, W))
                else:
                    final_comm_mask[i] = communication_mask[i]

            #     communication_mask = final_comm_mask
            #     communication_masks_list.append(communication_mask)
            # # _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            # communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
            # print(final_comm_mask.shape)
            # final_comm_mask scale x 2 [B, 1, H, W] -> [B, 1, 2H, 2W]

            # 将张量插值到目标大小 (B, 1, 2H, 2W)
            interpolated_tensor = F.interpolate(final_comm_mask.to(torch.float32).to(x.device), scale_factor=2, mode='nearest')
            interpolated_tensor[:K, ...] = 1

            x = x*interpolated_tensor
        
        # 2. feature compensation with flow
        # 2.1 generate flow, 在同一个坐标系内，计算每个cav的flow
        # 2.2 compensation
        # x: (BxNxK, C, H, W) -> (BxN, C, H, W)
        if self.design == 0 or self.design == 5:
            updated_features = self.update_features_boxflow(x, pairwise_t_matrix, record_len, box_flow, reserved_mask, flow_gt)
        elif self.design == 1:
            updated_features, flow_recon_loss = self.update_features_boxflow_design_1(x, pairwise_t_matrix, record_len, box_flow, reserved_mask, flow_gt)
        elif self.design == 2:
            updated_features, flow_recon_loss, flow, state_preds = self.update_features_boxflow_design_2(x, pairwise_t_matrix, record_len, box_flow, reserved_mask, flow_gt)
        elif self.design == 3:
            updated_features, flow, state_preds = self.generateFlow(x, pairwise_t_matrix, record_len, flow_gt) # (BxN, C, H, W)
        elif self.design == 4: # SyncNet
            updated_features = self.generate_estimated_feats(x, pairwise_t_matrix, record_len) # (BxN, C, H, W)
        
        # 3. feature fusion
        if self.multi_scale:
            ups = []
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(updated_features)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
            for i in range(self.num_levels):
                x = feats[i] if with_resnet else backbone.blocks[i](x)  # (BxN, C', H, W)

                ############ 1. Communication (Mask the features) #########
                if i==0:
                    if self.communication:
                        batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                        batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                        _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                        communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                        # x = x * communication_masks_tensor
                    else:
                        communication_rates = torch.tensor(0).to(x.device)
                else:
                    if self.communication:
                        communication_masks_tensor = F.max_pool2d(communication_masks_tensor, kernel_size=2)
                        # TODO: scale = 1, 2 不加 mask
                        # x = x * communication_masks_tensor
                
                ############ 2. Split the confidence map #######################
                # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                batch_node_features = self.regroup(x, record_len, 1)

                ############ 3. Fusion ####################################
                x_fuse = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    if noise_pairwise_t_matrix is not None:
                        t_matrix = noise_pairwise_t_matrix[b][:N, 0, :, :] #(N, 2, 3)
                    else:
                        t_matrix = pairwise_t_matrix[b][:N, 0, :, :] #(N, 2, 3)
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))
                    record_frames = np.ones((N))*K
                    # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                    
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])
                    else: # ATTEN, MAX, Transformer
                        x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])

                x_fuse = torch.stack(x_fuse)

                ############ 4. Deconv ####################################
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[i](x_fuse))
                else:
                    ups.append(x_fuse)
                
            if len(ups) > 1:
                x_fuse = torch.cat(ups, dim=1)
            elif len(ups) == 1:
                x_fuse = ups[0]
            
            if len(backbone.deblocks) > self.num_levels:
                x_fuse = backbone.deblocks[-1](x_fuse)
                
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_confidence_maps = self.regroup(rm, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                # _, (B, 1, H, W), float
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix,
                                                (H, W))
                record_frames = np.ones((N))*K
                # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                if self.agg_mode == "RAIN":
                    # for sensor embedding
                    sensor_dist = -1# (B, H, W)
                    x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)

            # self.fuse_modsules(x_fuse, record_len)
        if self.design == 0 or self.design == 5:
            if viz_bbx_flag:
                return  x_fuse, communication_rates, {}, updated_features
            return x_fuse, communication_rates, {}
        elif self.design == 1:
            return x_fuse, communication_rates, {}, flow_recon_loss
        elif self.design == 2:
            return x_fuse, communication_rates, {}, flow_recon_loss #, flow, state_preds
        elif self.design == 3:
            return x_fuse, communication_rates, {}, flow, state_preds
        elif self.design == 4:
            return x_fuse, communication_rates, {}

class TemporalFusion(nn.Module):
    def __init__(self, args):
        super(TemporalFusion, self).__init__()
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        # channel_size = args['channel_size']
        # spatial_size = args['spatial_size']
        # TM_Flag = args['TM_Flag']
        # compressed_size = args['compressed_size']
        # self.compensation_net = SyncLSTM(channel_size, spatial_size, TM_Flag, compressed_size)
        self.stpn = STPN(height_feat_size=args['channel_size'])
        self.flow_thre = args['flow_thre']
        self.motion_pred = MotionPrediction(seq_len=1)
        self.state_classify = StateEstimation(motion_category_num=1)

        self.flow_unc_flag = False
        if 'flow_unc_flag' in args:
            self.flow_unc_flag = True
            self.flow_unc_pred = FlowUncPrediction(seq_len=1)

        self.dcn = False
        if 'dcn' in args:
            self.dcn = True
            self.dcn_net = DCNNet(args['dcn'])

        if 'FlowPredictionFix' in args.keys() and args['FlowPredictionFix']:
            self.FlowPredictionFix()

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def FlowPredictionFix(self):
        print('------ Flow Prediction Fix, Only train DCN ------')
        for p in self.stpn.parameters():
            p.requires_grad = False
        for p in self.motion_pred.parameters():
            p.requires_grad = False
        for p in self.state_classify.parameters():
            p.requires_grad = False

    def forward(self, x, record_len, pairwise_t_matrix, data_dict=None):
        """
        Fusion forwarding.

        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        feat_seqs = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # iteratively warp feature to current timestamp
        batch_feat_seqs = []
        for b in range(B):
            # number of valid timestamps
            K = record_len[b]
            t_matrix = pairwise_t_matrix[b][:K, :K, :, :]
            curr_feat_seq = warp_affine_simple(feat_seqs[b],
                                               t_matrix[0, :, :, :],
                                               (H, W))
            batch_feat_seqs.append(curr_feat_seq[None, ...])
        batch_feat_seqs = torch.cat(batch_feat_seqs, dim=0)  # b, K, c, h, w
        batch_hist_feat_seqs = batch_feat_seqs[:, 1:].flip(1)

        # Backbone network
        bevs = self.stpn(batch_hist_feat_seqs)  # b, K, c, h, w

        # Motion Displacement prediction
        flow = self.motion_pred(bevs)
        flow = flow.view(-1, 2, bevs.size(-2), bevs.size(-1))

        flow_unc = None
        if self.flow_unc_flag:
            flow_unc = self.flow_unc_pred(bevs)
            flow_unc = flow_unc.view(-1, 2, bevs.size(-2), bevs.size(-1))

        # flow = data_dict['flow_gt']

        # Motion State Classification head
        state_class_pred = self.state_classify(bevs)

        # Given disp shift feature
        x_coord = torch.arange(bevs.size(-1)).float()
        y_coord = torch.arange(bevs.size(-2)).float()
        y, x = torch.meshgrid(y_coord, x_coord)
        grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(flow.shape[0], -1, -1, -1).to(
            flow.device)
        # updated_grid = grid + flow * (state_class_pred.sigmoid() > self.flow_thre)
        updated_grid = grid - flow
        updated_grid[:, 0, :, :] = updated_grid[:, 0, :, :] / (bevs.size(-1) / 2.0) - 1.0
        updated_grid[:, 1, :, :] = updated_grid[:, 1, :, :] / (bevs.size(-2) / 2.0) - 1.0
        out = F.grid_sample(batch_feat_seqs[:, 1], grid=updated_grid.permute(0, 2, 3, 1), mode='bilinear')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=updated_grid.permute(0,2,3,1), mode='nearest')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=grid.permute(0,2,3,1), mode='bilinear')
        # out = batch_feat_seqs[:,1]
        if self.dcn:
            out = self.dcn_net(out)

        flow_dict = {'flow':flow, 'flow_unc': flow_unc}
        return flow_dict, state_class_pred, out

    def forward_debug(self, x, origin_x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding
        Used for debug and visualization


        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        origin_x: torch.Tensor
            pillars (sum(n_cav), C, H * downsample_rate, W * downsample_rate)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        from matplotlib import pyplot as plt

        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        split_x = self.regroup(x, record_len)
        split_origin_x = self.regroup(origin_x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L, 1, H, W)).to(x)
                roi_mask[b, i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :], (H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        # visualize warped feature map
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0  # ego
            mask = roi_mask[b, i, :N, ...]
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H, W))
            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/feature_{b}_{idx}")
                plt.clf()
                plt.imshow(mask[idx][0].detach().cpu().numpy())
                plt.savefig(
                    f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/mask_feature_{b}_{idx}")
                plt.clf()

        # visualize origin pillar feature
        origin_node_features = split_origin_x

        for b in range(B):
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            i = 0  # ego
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(origin_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H * self.downsample_rate, W * self.downsample_rate))

            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/origin_{b}_{idx}")
                plt.clf()
opencood/models/fuse_modules/raindrop_flow.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
from opencood.models.sub_modules.convgru import ConvGRU
from icecream import ic
from matplotlib import pyplot as plt
# from opencood.models.sub_modules.SyncLSTM import SyncLSTM
from opencood.models.sub_modules.MotionNet import STPN, MotionPrediction, StateEstimation, FlowUncPrediction
# from opencood.models.sub_modules.dcn_net import DCNNet
from opencood.models.comm_modules.where2comm_multisweep import Communication

import torch.nn as nn

class FineTuneFlow(nn.Module):
    def __init__(self):
        super(FineTuneFlow, self).__init__()
        self.conv1 = nn.Conv2d(2, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 2, kernel_size=3, padding=1)

    def forward(self, x):
        x = x.permute(0, 3, 1, 2)
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.relu(self.conv2(x))
        x = self.conv3(x)
        x = x.permute(0, 2, 3, 1)
        return x

class FineTuneFlow_2(nn.Module):
    def __init__(self):
        super(FineTuneFlow_2, self).__init__()
        self.conv1 = nn.Conv2d(4, 16, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(32, 2, kernel_size=3, padding=1)

    def forward(self, x_1, x_2):
        x = torch.cat([x_1, x_2], dim=-1)
        x = x.permute(0, 3, 1, 2)
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.relu(self.conv2(x))
        x = self.conv3(x)
        x = x.permute(0, 2, 3, 1)
        return x


def get_warped_feature_mask(flow, updated_grid):
    ''' 
    get the mask of the warped feature map, where the value is 1 means the location is valid, and 0 means the location is invalid.
    ----------
    Args:
        flow: (2, H, W)
        updated_grid: (2, H, W)

    Returns:
        mask: (H, W)
    '''
    def get_large_flow_indices(flow, threshold=0):
        '''
        find the indices of the flow map where the flow is larger than the threshold, which means these locations have been moved.
        ----------
        Args:
            flow: (2, H, W)
            
        '''
        max_values, max_indices = torch.max(torch.abs(flow[:2]), dim=0)
        large_indices = torch.nonzero(max_values > threshold, as_tuple=False)
        return large_indices

    def remove_duplicate_points(points):
        unique_points, inverse_indices = torch.unique(points, sorted=True, return_inverse=True, dim=0)
        return unique_points, inverse_indices
    
    def get_nonzero_idx(flow, idx):
        # print(idx[:, 0].max())
        # print(idx[:, 1].max())
        flow_values = flow[:, idx[:, 0], idx[:, 1]]
        nonzero_idx = torch.nonzero(torch.abs(flow_values).sum(dim=0) == 0, as_tuple=False).squeeze()
        return idx[nonzero_idx]
    
    def filter_tensor(A, H, W):
        mask = (A[:, 0] < H) & (A[:, 1] < W)
        filtered_A = A[mask]
        return filtered_A
    
    flow_idx = get_large_flow_indices(flow)

    _, H, W = flow.shape

    mask = torch.ones(H, W)
    if flow_idx.shape[0] == 0:
        return mask

    # print(flow_idx)
    updated_grid_points_tmp = updated_grid[:, flow_idx[:,0], flow_idx[:,1]].to(torch.int64).T
    # change the order of dim 1
    updated_grid_points = torch.zeros_like(updated_grid_points_tmp)
    updated_grid_points[:, 0] = updated_grid_points_tmp[:, 1]
    updated_grid_points[:, 1] = updated_grid_points_tmp[:, 0]
    # print(updated_grid_points)

    unique_points_idx, _ = remove_duplicate_points(updated_grid_points)
    # print(unique_points_idx)
    
    # mast out the out of range indices
    unique_points_idx = filter_tensor(unique_points_idx, H, W)
    
    nonzero_idx = get_nonzero_idx(flow, unique_points_idx)
    # print(nonzero_idx)

    if len(nonzero_idx.shape) > 1:
        mask[nonzero_idx[:, 0], nonzero_idx[:, 1]] = 0
    else: 
        mask[nonzero_idx[0], nonzero_idx[1]] = 0

    return mask

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_fuse(nn.Module):
    def __init__(self, args, design_mode=0):
        super(raindrop_fuse, self).__init__()
        
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.multi_scale:
            layer_nums = args['layer_nums']
            num_filters = args['num_filters']
            self.num_levels = len(layer_nums)
            self.fuse_modules = nn.ModuleList()
            for idx in range(self.num_levels):
                if self.agg_mode == 'MAX':
                    fuse_network = MaxFusion()
                else:
                    # self.agg_mode == 'ATTEN':
                    fuse_network = AttenFusion(num_filters[idx])
                self.fuse_modules.append(fuse_network)
        else: 
            if self.agg_mode == 'MAX': # max fusion, debug use
                self.fuse_modules = MaxFusion()

        self.design = design_mode
        if self.design == 1:
            self.fine_conv = FineTuneFlow()
        elif self.design == 2:
            self.stpn = STPN(args['channel_size'])
            self.motion_pred = MotionPrediction(seq_len=1)
            self.state_classify = StateEstimation(motion_category_num=1)
            self.fine_conv = FineTuneFlow()
            self.fine_conv_2 = FineTuneFlow_2()
        else:
            self.stpn = STPN(args['channel_size'])
            self.motion_pred = MotionPrediction(seq_len=1)
            self.state_classify = StateEstimation(motion_category_num=1)

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def generateFlow(self, feats, pairwise_t_matrix, pastk_2_past0_tr_mats, record_len, flow_gt=None, warp_mask=None):
        '''
        1. generate flow from feature sequence
        2. then update the feature
        Note:
            1. ego feature does not need to be updated

        params:
            feats: (sum(B,N,K), C, H, W)
            pairwise_t_matrix: (B, L, K, 2, 3)
            pastk_2_past0_tr_mats: (sum(n_cav), K, 2, 3)
            record_len: (B)

        return:
            flow: 
        '''
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]
        batch_node_features = self.regroup(feats, record_len, k=K)
        batch_pastk_2_past0_tr_mats = self.regroup(pastk_2_past0_tr_mats, record_len, k=1)

        # debug use
        batch_flow_gt = self.regroup(flow_gt, record_len, 2)
        batch_warp_mask = self.regroup(warp_mask, record_len, C)

        updated_features_list = []
        flow_list = []
        state_class_pred_list = []
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            # t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
            node_features = batch_node_features[b]
            t_matrix = batch_pastk_2_past0_tr_mats[b].view(-1, 2, 3) # (Nxk, 2, 3)
            node_features = warp_affine_simple(node_features, t_matrix, (H, W))
            node_features = node_features.view(-1, K, C, H, W) # (N, k, C, H, W)

            # debug_path = '/remote-home/share/sizhewei/logs/where2comm_motionnet_flow_regular_new/debug_path'
            # torch.save(node_features, debug_path + '/node_features.pt')
            
            # MotionNet generating flow
            bevs = self.stpn(node_features)
            flow = self.motion_pred(bevs) # [N, 2, H, W]
            flow_list.append(flow)
            
            # Motion State Classification head
            state_class_pred = self.state_classify(bevs)
            state_class_pred_list.append(state_class_pred)

            # # TODO: use gt flow for debug use
            # flow = batch_flow_gt[b].view(-1, 2, H, W)  # (N, 2, H, W)
            # mask = batch_warp_mask[b].view(-1, C, H, W)  # (N, 1, H, W)
            # flow_list.append(flow)
            # state_class_pred = torch.ones(1, H, W)
            # state_class_pred_list.append(state_class_pred)
            # # flow_old = flow.clone()
            # ##########################################

            # Original flow warp code
            # Given disp shift feature
            x_coord = torch.arange(W).float()   # [0, ..., W]
            y_coord = torch.arange(H).float()   # [0, ..., H]
            y, x = torch.meshgrid(y_coord, x_coord)  # [H, W], [H, W]
            grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(flow.shape[0], -1, -1, -1).to(flow.device)  # N, 2, H, W
            # updated_grid = grid + flow * (state_class_pred.sigmoid() > self.flow_thre)
            updated_grid = grid + flow
            
            # # TODO: debug use, do not use gt flow, keep the original delay feature as the fusion input. 
            # # updated_grid = grid
            
            # # generate the mask for filtering out the pixels which moved to other locations but not filled by others
            mask_list = []
            for i in range(flow.shape[0]):
                mask_list.append(get_warped_feature_mask(flow[i], updated_grid[i]))
            mask = torch.stack(mask_list, dim=0)
            mask = mask.unsqueeze(1).repeat(1, C, 1, 1).to(flow.device)
            # ################################
            
            updated_grid[:, 0, :, :] = updated_grid[:, 0, :, :] / (W / 2.0) - 1.0
            updated_grid[:, 1, :, :] = updated_grid[:, 1, :, :] / (H / 2.0) - 1.0
            latest_node_features = node_features[:, 0, :, :, :] # (N, C, H, W)
            updated_features = F.grid_sample(latest_node_features, grid=updated_grid.permute(0, 2, 3, 1), mode='bilinear', align_corners=False)
            
            # use szwei flow generator
            # latest_node_features = node_features[:, 0, :, :, :] # (N, C, H, W)
            # updated_features = F.grid_sample(latest_node_features, grid=flow.permute(0, 2, 3, 1), mode='bilinear', align_corners=False)
            
            # mask the features
            updated_features = mask * updated_features
            # ego feature use the latest feature (no delay)
            updated_features[0, :, :, :] = latest_node_features[0, :, :, :]

            updated_features_list.append(updated_features)
        
        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)

        # TODO: for debug use
        # debug_path = '/remote-home/share/sizhewei/logs/where2comm_flow_debug/viz_flow/'
        # torch.save(feats, debug_path + 'feats.pt')
        # torch.save(updated_features_all, debug_path + 'updated_features_all.pt')

        flow_all = torch.cat(flow_list, dim=0)  # (sum(B,N), 2, H, W)
        state_class_pred_all = torch.cat(state_class_pred_list, dim=0)  # (sum(B,N), 2, H, W)
        
        return updated_features_all, flow_all, state_class_pred_all

    def update_features_boxflow(self, feats, pairwise_t_matrix, record_len, flow_map, reserved_mask, flow_gt):
        """
        Update features with box flow.
        
        Parameters
        ----------
        feats : torch.Tensor
            input data, 
            shape: (sum(n_cav), C, H, W)
        
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        record_len : list
            shape: (B)
            
        flow_map: torch.Tensor
            flow generate by past 2 frames detection boxes
            shape: (sum(N_b), H, W, 2)

        reserved_mask: torch.Tensor
            shape: (sum(N_b), C, H, W)

        gt_flow: torch.Tensor
            shape: (sum(N_b), H, W, 2)
            
        Returns
        -------
        Warped feature.
        """
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]
        batch_node_features = self.regroup(feats, record_len, k=K)

        batch_flow_map = self.regroup(flow_map, record_len, k=1)
        batch_reserved_mask = self.regroup(reserved_mask, record_len, k=1)

        # debug use
        batch_flow_gt = self.regroup(flow_gt, record_len, k=2)

        updated_features_list = []
        flow_list = []
        state_class_pred_list = []
        gt_flow_map_list = []
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            node_features = batch_node_features[b]
            node_features = node_features.view(-1, K, C, H, W) # (N, k, C, H, W)
            flow = batch_flow_map[b] # .view(-1, 2, H, W) # [N, H, W, 2]
            flow = self.fine_conv(flow) # [N, H, W, 2]
            flow_list.append(flow)
            # # motion net flow
            # bevs = self.stpn(node_features)
            # flow_pred = self.motion_pred(bevs) # [N, 2, H, W]
            # # Motion State Classification head
            # state_class_pred = self.state_classify(bevs)
            # state_class_pred_list.append(state_class_pred)
            
            latest_node_features = node_features[:, 0, :, :, :] # past_0 features, (N, C, H, W) 
            updated_features = F.grid_sample(latest_node_features, grid=flow, mode='nearest', align_corners=False)

            updated_features = updated_features*batch_reserved_mask[b] # (N, C, H, W)
            updated_features_list.append(updated_features)

            # normalizing GT flow: gt_flow_map [N, H, W, 2]
            # Given disp shift feature
            x_coord = torch.arange(W).float()   # [0, ..., W]
            y_coord = torch.arange(H).float()   # [0, ..., H]
            y, x = torch.meshgrid(y_coord, x_coord)  # [H, W], [H, W]
            grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(N, -1, -1, -1).to(flow.device)
            gt_flow_delta = batch_flow_gt[b].view(-1, 2, H, W)
            gt_flow_map = grid - gt_flow_delta
            gt_flow_map[:, 0, :, :] = gt_flow_map[:, 0, :, :] / (W / 2.0) - 1.0
            gt_flow_map[:, 1, :, :] = gt_flow_map[:, 1, :, :] / (H / 2.0) - 1.0
            gt_flow_map = gt_flow_map.permute(0, 2, 3, 1) 
            gt_flow_map_list.append(gt_flow_map)

        flow_pred = torch.cat(flow_list, dim=0) # (sum(N_b), H, W, 2)
        gt_flow_norm = torch.cat(gt_flow_map_list, dim=0) # (sum(N_b), H, W, 2)
        # compute the flow map loss:
        loss = F.smooth_l1_loss(flow_pred, gt_flow_norm)

        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)

        return updated_features_all, loss
    
    def update_features_boxflow_design_2(self, feats, pairwise_t_matrix, record_len, flow_map, reserved_mask, flow_gt):
        """
        Update features with box flow.
        
        Parameters
        ----------
        feats : torch.Tensor
            input data, 
            shape: (sum(n_cav), C, H, W)
        
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        record_len : list
            shape: (B)
            
        flow_map: torch.Tensor
            flow generate by past 2 frames detection boxes
            shape: (sum(N_b), H, W, 2)

        reserved_mask: torch.Tensor
            shape: (sum(N_b), C, H, W)

        gt_flow: torch.Tensor
            shape: (sum(N_b), H, W, 2)
            
        Returns
        -------
        Warped feature.
        """
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]
        batch_node_features = self.regroup(feats, record_len, k=K)

        batch_flow_map = self.regroup(flow_map, record_len, k=1)
        batch_reserved_mask = self.regroup(reserved_mask, record_len, k=1)

        # debug use
        batch_flow_gt = self.regroup(flow_gt, record_len, k=2)

        updated_features_list = []
        flow_list = []
        state_class_pred_list = []
        gt_flow_map_list = []
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            node_features = batch_node_features[b]
            node_features = node_features.view(-1, K, C, H, W) # (N, k, C, H, W)
            flow_box = batch_flow_map[b] # .view(-1, 2, H, W) # [N, H, W, 2]
            # flow_box = self.fine_conv(flow_box) # [N, H, W, 2]
            
            # motion net flow
            bevs = self.stpn(node_features)
            flow_delta_pred = self.motion_pred(bevs) # [N, 2, H, W]
            # Motion State Classification head
            state_class_pred = self.state_classify(bevs)
            state_class_pred_list.append(state_class_pred)
            # Given disp shift feature
            x_coord = torch.arange(W).float()   # [0, ..., W]
            y_coord = torch.arange(H).float()   # [0, ..., H]
            y, x = torch.meshgrid(y_coord, x_coord)  # [H, W], [H, W]
            grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(N, -1, -1, -1).to(flow_box.device)
            # flow_pred = grid - flow_delta_pred
            # # normalizing flow: flow_pred [N, H, W, 2]
            # flow_pred[:, 0, :, :] = flow_pred[:, 0, :, :] / (W / 2.0) - 1.0
            # flow_pred[:, 1, :, :] = flow_pred[:, 1, :, :] / (H / 2.0) - 1.0
            # flow_pred = flow_pred.permute(0, 2, 3, 1) # [N, H, W, 2]

            # normalizing flow_delta_pred 
            flow_delta_pred[:, 0, :, :] = flow_delta_pred[:, 0, :, :] / (W / 2.0) - 1.0
            flow_delta_pred[:, 1, :, :] = flow_delta_pred[:, 1, :, :] / (H / 2.0) - 1.0
            flow_delta_pred = flow_delta_pred.permute(0, 2, 3, 1) # [N, H, W, 2]
            flow = self.fine_conv(flow_delta_pred + flow_box) # [N, H, W, 2]

            # flow = self.fine_conv_2(flow_box, flow_pred) # [N, H, W, 2]
            flow_list.append(flow)
            
            latest_node_features = node_features[:, 0, :, :, :] # past_0 features, (N, C, H, W) 
            updated_features = F.grid_sample(latest_node_features, grid=flow, mode='nearest', align_corners=False)

            updated_features = updated_features*batch_reserved_mask[b] # (N, C, H, W)
            updated_features_list.append(updated_features)

            # normalizing GT flow: gt_flow_map [N, H, W, 2]
            gt_flow_delta = batch_flow_gt[b].view(-1, 2, H, W)
            gt_flow_map = grid - gt_flow_delta
            gt_flow_map[:, 0, :, :] = gt_flow_map[:, 0, :, :] / (W / 2.0) - 1.0
            gt_flow_map[:, 1, :, :] = gt_flow_map[:, 1, :, :] / (H / 2.0) - 1.0
            gt_flow_map = gt_flow_map.permute(0, 2, 3, 1) 
            gt_flow_map_list.append(gt_flow_map)

        flow_all = torch.cat(flow_list, dim=0) # (sum(N_b), H, W, 2)
        state_class_pred_all = torch.cat(state_class_pred_list, dim=0)  # (sum(B,N), 2, H, W)
        gt_flow_norm = torch.cat(gt_flow_map_list, dim=0) # (sum(N_b), H, W, 2)
        # compute the flow map loss:
        loss = F.smooth_l1_loss(flow_all, gt_flow_norm)

        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)

        return updated_features_all, loss, flow_all, state_class_pred_all

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None, flow_gt=None, box_flow=None, reserved_mask=None, warp_mask=None, pastk_2_past0_tr_mats=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
        
        flow_gt: ground truth flow, generate by object center and id

        box_flow: flow generate by past 2 frames detection boxes
            
        Returns
        -------
        Fused feature
        flow_map loss
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        # (sum(n_cav), k, 4, 4) -> (sum(n_cav), k, 2, 3)
        pastk_2_past0_tr_mats = pastk_2_past0_tr_mats[:,:,[0, 1],:][:,:,:,[0, 1, 3]]
        pastk_2_past0_tr_mats[...,0,1] = pastk_2_past0_tr_mats[...,0,1] * H / W
        pastk_2_past0_tr_mats[...,1,0] = pastk_2_past0_tr_mats[...,1,0] * W / H
        pastk_2_past0_tr_mats[...,0,2] = pastk_2_past0_tr_mats[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pastk_2_past0_tr_mats[...,1,2] = pastk_2_past0_tr_mats[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        # 2. feature compensation with flow
        # 2.1 generate flow, 在同一个坐标系内，计算每个cav的flow
        # 2.2 compensation
        # x: (BxNxK, C, H, W) -> (BxN, C, H, W)
        if self.design == 1:
            updated_features, flow_recon_loss = self.update_features_boxflow(x, pairwise_t_matrix, record_len, box_flow, reserved_mask, flow_gt)
        elif self.design == 2:
            updated_features, flow_recon_loss, flow, state_preds = self.update_features_boxflow_design_2(x, pairwise_t_matrix, record_len, box_flow, reserved_mask, flow_gt)
        else:
            updated_features, flow, state_preds = self.generateFlow(x, pairwise_t_matrix, pastk_2_past0_tr_mats, record_len, flow_gt, warp_mask) # (BxN, C, H, W)
        
        # 3. feature fusion
        if self.multi_scale:
            ups = []
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(updated_features)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
            for i in range(self.num_levels):
                x = feats[i] if with_resnet else backbone.blocks[i](x)  # (BxN, C', H, W)

                ############ 1. Communication (Mask the features) #########
                if i==0:
                    if self.communication:
                        batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                        batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                        _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                        communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                        # x = x * communication_masks_tensor
                    else:
                        communication_rates = torch.tensor(0).to(x.device)
                else:
                    if self.communication:
                        communication_masks_tensor = F.max_pool2d(communication_masks_tensor, kernel_size=2)
                        # TODO: scale = 1, 2 不加 mask
                        # x = x * communication_masks_tensor
                
                ############ 2. Split the confidence map #######################
                # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                batch_node_features = self.regroup(x, record_len, 1)

                ############ 3. Fusion ####################################
                x_fuse = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    t_matrix = pairwise_t_matrix[b][:N, 0, :, :] #(N, 2, 3)
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))
                    record_frames = np.ones((N))*K
                    # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                    
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])
                    else: # ATTEN, MAX, Transformer
                        x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])

                x_fuse = torch.stack(x_fuse)

                ############ 4. Deconv ####################################
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[i](x_fuse))
                else:
                    ups.append(x_fuse)
                
            if len(ups) > 1:
                x_fuse = torch.cat(ups, dim=1)
            elif len(ups) == 1:
                x_fuse = ups[0]
            
            if len(backbone.deblocks) > self.num_levels:
                x_fuse = backbone.deblocks[-1](x_fuse)
                
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_confidence_maps = self.regroup(rm, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                # _, (B, 1, H, W), float
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix,
                                                (H, W))
                record_frames = np.ones((N))*K
                # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                if self.agg_mode == "RAIN":
                    # for sensor embedding
                    sensor_dist = -1# (B, H, W)
                    x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)

            # self.fuse_modsules(x_fuse, record_len)
        
        if self.design == 1:
            return x_fuse, communication_rates, {}, flow_recon_loss
        elif self.design == 2:
            return x_fuse, communication_rates, {}, flow_recon_loss #, flow, state_preds
        else:
            return x_fuse, communication_rates, {}, flow, state_preds

class TemporalFusion(nn.Module):
    def __init__(self, args):
        super(TemporalFusion, self).__init__()
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        # channel_size = args['channel_size']
        # spatial_size = args['spatial_size']
        # TM_Flag = args['TM_Flag']
        # compressed_size = args['compressed_size']
        # self.compensation_net = SyncLSTM(channel_size, spatial_size, TM_Flag, compressed_size)
        self.stpn = STPN(height_feat_size=args['channel_size'])
        self.flow_thre = args['flow_thre']
        self.motion_pred = MotionPrediction(seq_len=1)
        self.state_classify = StateEstimation(motion_category_num=1)

        self.flow_unc_flag = False
        if 'flow_unc_flag' in args:
            self.flow_unc_flag = True
            self.flow_unc_pred = FlowUncPrediction(seq_len=1)

        self.dcn = False
        if 'dcn' in args:
            self.dcn = True
            self.dcn_net = DCNNet(args['dcn'])

        if 'FlowPredictionFix' in args.keys() and args['FlowPredictionFix']:
            self.FlowPredictionFix()

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def FlowPredictionFix(self):
        print('------ Flow Prediction Fix, Only train DCN ------')
        for p in self.stpn.parameters():
            p.requires_grad = False
        for p in self.motion_pred.parameters():
            p.requires_grad = False
        for p in self.state_classify.parameters():
            p.requires_grad = False

    def forward(self, x, record_len, pairwise_t_matrix, data_dict=None):
        """
        Fusion forwarding.

        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        feat_seqs = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # iteratively warp feature to current timestamp
        batch_feat_seqs = []
        for b in range(B):
            # number of valid timestamps
            K = record_len[b]
            t_matrix = pairwise_t_matrix[b][:K, :K, :, :]
            curr_feat_seq = warp_affine_simple(feat_seqs[b],
                                               t_matrix[0, :, :, :],
                                               (H, W))
            batch_feat_seqs.append(curr_feat_seq[None, ...])
        batch_feat_seqs = torch.cat(batch_feat_seqs, dim=0)  # b, K, c, h, w
        batch_hist_feat_seqs = batch_feat_seqs[:, 1:].flip(1)

        # Backbone network
        bevs = self.stpn(batch_hist_feat_seqs)  # b, K, c, h, w

        # Motion Displacement prediction
        flow = self.motion_pred(bevs)
        flow = flow.view(-1, 2, bevs.size(-2), bevs.size(-1))

        flow_unc = None
        if self.flow_unc_flag:
            flow_unc = self.flow_unc_pred(bevs)
            flow_unc = flow_unc.view(-1, 2, bevs.size(-2), bevs.size(-1))

        # flow = data_dict['flow_gt']

        # Motion State Classification head
        state_class_pred = self.state_classify(bevs)

        # Given disp shift feature
        x_coord = torch.arange(bevs.size(-1)).float()
        y_coord = torch.arange(bevs.size(-2)).float()
        y, x = torch.meshgrid(y_coord, x_coord)
        grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(flow.shape[0], -1, -1, -1).to(
            flow.device)
        # updated_grid = grid + flow * (state_class_pred.sigmoid() > self.flow_thre)
        updated_grid = grid - flow
        updated_grid[:, 0, :, :] = updated_grid[:, 0, :, :] / (bevs.size(-1) / 2.0) - 1.0
        updated_grid[:, 1, :, :] = updated_grid[:, 1, :, :] / (bevs.size(-2) / 2.0) - 1.0
        out = F.grid_sample(batch_feat_seqs[:, 1], grid=updated_grid.permute(0, 2, 3, 1), mode='bilinear')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=updated_grid.permute(0,2,3,1), mode='nearest')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=grid.permute(0,2,3,1), mode='bilinear')
        # out = batch_feat_seqs[:,1]
        if self.dcn:
            out = self.dcn_net(out)

        flow_dict = {'flow':flow, 'flow_unc': flow_unc}
        return flow_dict, state_class_pred, out

    def forward_debug(self, x, origin_x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding
        Used for debug and visualization


        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        origin_x: torch.Tensor
            pillars (sum(n_cav), C, H * downsample_rate, W * downsample_rate)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        from matplotlib import pyplot as plt

        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        split_x = self.regroup(x, record_len)
        split_origin_x = self.regroup(origin_x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L, 1, H, W)).to(x)
                roi_mask[b, i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :], (H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        # visualize warped feature map
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0  # ego
            mask = roi_mask[b, i, :N, ...]
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H, W))
            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/feature_{b}_{idx}")
                plt.clf()
                plt.imshow(mask[idx][0].detach().cpu().numpy())
                plt.savefig(
                    f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/mask_feature_{b}_{idx}")
                plt.clf()

        # visualize origin pillar feature
        origin_node_features = split_origin_x

        for b in range(B):
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            i = 0  # ego
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(origin_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H * self.downsample_rate, W * self.downsample_rate))

            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/origin_{b}_{idx}")
                plt.clf()
opencood/models/fuse_modules/raindrop_flow_regular.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# Date: 2022/12/06
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

import torch.nn.functional as F
from turtle import update
# import ipdb
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
from opencood.models.comm_modules.where2comm_multisweep import Communication
from opencood.models.fuse_modules.raindrop_attn import raindrop_fuse

from opencood.models.fuse_modules.STPT import STPT
# from opencood.models.fuse_modules.SyncNet import SyncLSTM

if_save_pt = False  # TODO: for debug use, save pt file

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class PositionalEncoding_spatial(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model_feature, d_model_hidden, dropout, max_len=300,mode='cat'):
        super(PositionalEncoding_spatial, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        d_model = int(d_model_hidden/2)
        # Compute the positional encodings once in log space.
        pe = torch.zeros((max_len, d_model))
        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe_ = pe.unsqueeze(0)#batch,len,feature
        pe_h = pe_.unsqueeze(2).repeat(1,1,max_len,1)
        pe_w = pe_.unsqueeze(1).repeat(1,max_len,1,1)
        pe = torch.cat((pe_h,pe_w),dim=-1)

        self.register_buffer('pe', pe)

        self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x):
        if self.mode == 'cat':
            tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
            tmp = tmp.repeat(x.shape[0], 1, 1, 1)
            x = self.linear(torch.cat((x, tmp),dim=-1))

        elif self.mode == 'init':
            x = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)

        else:
            tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
            tmp = tmp.repeat(x.shape[0], 1, 1, 1)
            x = x + tmp

        return self.dropout(x)

class PositionalEncoding_spatial_custom(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model_feature, d_model_hidden, dropout,mode='cat'):
        super(PositionalEncoding_spatial_custom, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = int(d_model_hidden/2)
        # Compute the positional encodings once in log space.
        
        self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x, nonego_idx):#x(batch,num,feature)
   
        pe_h = torch.zeros((x.shape[1], self.d_model)).to(x.device)
        pe_w = torch.zeros((x.shape[1], self.d_model)).to(x.device)
        position_h, position_w = nonego_idx[0],nonego_idx[1]
        position_h = position_h.unsqueeze(1)
        position_w = position_w.unsqueeze(1)

        div_term = torch.exp(torch.arange(0, self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(position_h.device)
        pe_h[:, 0::2] = torch.sin(position_h * div_term)
        pe_h[:, 1::2] = torch.cos(position_h * div_term)

        pe_w[:, 0::2] = torch.sin(position_w * div_term)
        pe_w[:, 1::2] = torch.cos(position_w * div_term)
       
        pe = torch.cat((pe_h,pe_w),dim=-1).unsqueeze(0).repeat(x.shape[0],1,1)

        if self.mode == 'cat':
            x = self.linear(torch.cat((x, pe),dim=-1))

        else:
            x = x + Variable(pe,requires_grad=False)

        return self.dropout(x)


class PositionalEncoding_sensor_dist(nn.Module):

    def __init__(self, d_model, dropout):
        super(PositionalEncoding_sensor_dist, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = d_model
        self.linear = nn.Linear(2*d_model,d_model)
        self.mode = 'cat'
         
    def forward(self, x, sensor_dist):#x.shape=[batch,H,W,F],senfor_dist.shape=[batch,H,W]

        x = x.reshape(-1,x.shape[-1])
        sensor_dist = sensor_dist.reshape(-1)

        pe = torch.zeros(x.shape[0], self.d_model).to(x.device)
        position = sensor_dist.unsqueeze(1) 
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(x.device)
         
        pe[:, 0::2] = torch.sin(position * div_term)  
        pe[:, 1::2] = torch.cos(position * div_term)   

        if self.mode == 'cat':
            x = self.linear(torch.cat((x,Variable(pe,require_grad=False)),dim=-1))

        else:
            x = x + Variable(pe,requires_grad=False)
        
        return self.dropout(x)

class PositionalEncoding_irregular_time(nn.Module):

    def __init__(self, d_model_feature, d_model_hidden, dropout):
        super(PositionalEncoding_irregular_time, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = d_model_feature
        self.linear = nn.Linear(d_model_feature+d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x, time):#x.shape=[batch,height,width,feature],time.shape=[batch]

        pe = torch.zeros(x.shape[0],self.d_model).to(x.device)
        position = time.unsqueeze(1) 
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(x.device)
         
        pe[:, 0::2] = torch.sin(position * div_term)  
        pe[:, 1::2] = torch.cos(position * div_term)   

        if self.mode == 'cat':
            if len(x.shape) == 4:
                pe = pe.unsqueeze(1).unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
            else:
                pe = pe.unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], 1)
            x = self.linear(torch.cat((x, tmp),dim=-1))

        else:
            if len(x.shape) == 4:
                pe = pe.unsqueeze(1).unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
            else:
                pe = pe.unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], 1)
            x = x + tmp

        return self.dropout(x)

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    
    # ipdb.set_trace()
    
    # print('score',scores.shape,'mask',mask.shape)

    if mask is not None:
        # mask = mask.transpose(1,2).to(query.device)
        mask = mask.to(query.device)
        scores = scores.masked_fill(mask == 0, -1e10) 

    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)

    return torch.matmul(p_attn, value), p_attn


class MultiHeadedAttention(nn.Module):
    def __init__(self, h, q_dim, k_dim,v_dim, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h  
        self.h = h 
        self.q_dim = q_dim
        self.k_dim = k_dim
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.linear_q = nn.Linear(q_dim,d_model)
        self.linear_k = nn.Linear(k_dim,d_model)
        self.linear_v = nn.Linear(v_dim,d_model)
        self.linear_out = nn.Linear(d_model,q_dim)
        self.attn = None  
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k 

        query = self.linear_q(query).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
        key = self.linear_k(key).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
        value = self.linear_v(value).view(nbatches,-1,self.h,self.d_k).transpose(1,2)

        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(query, key, value, mask=mask,
                                 dropout=self.dropout)

        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(1, 2).contiguous() \
            .view(nbatches ,-1, self.h * self.d_k)

        return self.linear_out(x)

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
 
    def forward(self, x):
        "Norm"
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, x, sublayer):
        # print('sub',x.shape,self.dropout(sublayer(self.norm(x))).shape)
        return x + self.dropout(sublayer(self.norm(x))) #残差连接

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model,d_out, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_out)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

class EncoderLayer_temporal(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer_temporal, self).__init__()
        self.self_attn = self_attn #sublayer 1 
        self.feed_forward = feed_forward #sublayer 2 
        self.sublayer = clones(SublayerConnection(size, dropout), 2) #拷贝两个SublayerConnection，一个为了attention，一个是独自作为简单的神经网络
        self.size = size

    def forward(self, q, k, v, mask):
        "Follow Figure 1 (left) for connections."
        # print('x',x.shape,'attn',self.self_attn(x,x,x,mask).shape)
        x = self.sublayer[0](q, lambda q: self.self_attn(q, k, v, mask)) #attention层 对应层1
        return self.sublayer[1](x, self.feed_forward) # 对应 层2

class Spatial_conv(nn.Module):
    def __init__(self, input_channels, output_channels, kernel_size):
        super(Spatial_conv,self).__init__()
        self.conv = nn.Conv2d(in_channels = input_channels,
                              out_channels = output_channels,
                              kernel_size = (kernel_size,kernel_size),
                              stride = (1,1),
                              padding = (int((kernel_size-1)/2),int((kernel_size-1)/2))
        )

        self.relu = nn.ReLU()
        # self.bn = nn.BatchNorm2d(num_features=output_channels)

    def forward(self, x):
        init_x = x.clone()
        x = self.conv(x)
        x = self.relu(x)
        # x = self.bn(x)
        return x+init_x

class SpatialAttention_mtf(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention_mtf, self).__init__()

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, curr, prev):#####(batch, channel, h, w)
        avg_out = torch.mean(curr, dim=1, keepdim=True)
        max_out, _ = torch.max(curr, dim=1, keepdim=True)
        y = torch.cat([avg_out, max_out], dim=1)
        y = self.conv1(y)
        return self.sigmoid(y) * prev ####(batch-1, channel, h, w)

class ScaledDotProductAttention(nn.Module):
    """
    Scaled Dot-Product Attention proposed in "Attention Is All You Need"
    Compute the dot products of the query with all keys, divide each by sqrt(dim),
    and apply a softmax function to obtain the weights on the values
    Args: dim, mask
        dim (int): dimention of attention
        mask (torch.Tensor): tensor containing indices to be masked
    Inputs: query, key, value, mask
        - **query** (batch, q_len, d_model): tensor containing projection
          vector for decoder.
        - **key** (batch, k_len, d_model): tensor containing projection
          vector for encoder.
        - **value** (batch, v_len, d_model): tensor containing features of the
          encoded input sequence.
        - **mask** (-): tensor containing indices to be masked
    Returns: context, attn
        - **context**: tensor containing the context vector from
          attention mechanism.
        - **attn**: tensor containing the attention (alignment) from the
          encoder outputs.
    """

    def __init__(self, dim):
        super(ScaledDotProductAttention, self).__init__()
        self.sqrt_dim = np.sqrt(dim)

    def forward(self, query, key, value, mask):
        score = torch.bmm(query, key.transpose(1, 2)) / self.sqrt_dim
        if mask != None:
            score = score.masked_fill(mask == 0, -1e10) 
        attn = F.softmax(score, -1)
        context = torch.bmm(attn, value)
        return context

class AttentionFusion(nn.Module):
    def __init__(self, settings, input_dim, hidden_dim, num_heads, num_layers, neighbor_range, height, width, dropout, max_len=300,compensation=True):
        super(AttentionFusion, self).__init__()
        self.sweep_length = settings['sweep_length']
        self.if_max_aggre = settings['if_max_aggre']
        self.if_spatial_encoding = settings['if_spatial_encoding']
        self.if_ego_time_encoding = settings['if_ego_time_encoding']
        self.if_nonego_time_encoding = settings['if_nonego_time_encoding']
        self.if_sensor_encoding = settings['if_sensor_encoding']
        self.if_time_attn_aggre = settings['if_time_attn_aggre']
        self.if_spatial_conv = settings['if_spatial_conv']
        self.if_dotproductattn = settings['if_dotproductattn']
        self.use_mask = settings['use_mask']
        self.confidence_fetch = settings['confidence_fetch']
        self.if_conv_aggre = settings['if_conv_aggre']
        self.sup_individual = settings['sup_individual']
        
        # self.if_swin_trans = False
        # self.if_syncnet = False
        # if settings['compensation_mode'] == 'SWIN':
        self.if_swin_trans = True
        # elif settings['compensation_mode'] == 'SYNCNET':
            # self.if_syncnet = True

        # assert self.if_swin_trans != self.if_syncnet  # 两种补偿模块只能二选一

        self.input_dim = input_dim 
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.neighbor_range = neighbor_range
        self.height = height
        self.width = width
        self.compensation=compensation
        if self.compensation:
            if 'dilation_size' in settings:
                self.dilation_size = settings['dilation_size']
            else:
                self.dilation_size = 1
            self.dilate_mask = nn.MaxPool2d(kernel_size=self.dilation_size, stride=1, padding=int((self.dilation_size-1)/2))
            if self.if_spatial_encoding:
                self.spatial_encoding = PositionalEncoding_spatial(self.input_dim, hidden_dim, dropout, max_len)
                self.spatial_encoding_nonego = PositionalEncoding_spatial_custom(self.input_dim, hidden_dim, dropout)
            if self.if_ego_time_encoding or self.if_nonego_time_encoding:
                self.time_encoding = PositionalEncoding_irregular_time(self.input_dim, hidden_dim, dropout)
            if self.if_sensor_encoding:
                self.sensor_dist_encoding = PositionalEncoding_sensor_dist(hidden_dim, dropout)

            c = copy.deepcopy
            temporal_attn = MultiHeadedAttention(self.num_heads, self.input_dim, self.input_dim, self.input_dim, self.hidden_dim, dropout)

            ff = PositionwiseFeedForward(self.input_dim, self.input_dim, self.hidden_dim, dropout)

            if self.if_time_attn_aggre:
                self.layers_temporal = clones(EncoderLayer_temporal(self.input_dim, c(temporal_attn), c(ff), dropout),num_layers)
            # self.temporal_aggregator = MultiHeadedAttention(self.num_heads, hidden_dim, self.input_dim, self.input_dim, self.hidden_dim, dropout)
            if self.if_spatial_conv:  
                self.layers_spatial = clones(Spatial_conv(self.input_dim,self.input_dim,kernel_size=5),num_layers)
            if self.if_swin_trans:
                self.swin_trans = STPT(in_chans=self.input_dim,
                    depths=[4, 4], 
                    num_heads=[2, 2], 
                    window_sizes=[(4,4), (4,4)],
                    num_frames=3,
                    num_out_frames=1,
                    embed_dim=48,
                    use_hdmap=False,
                    use_hdmap_feat=False)
            # elif self.if_syncnet:
            #     self.syncnet = SyncLSTM(channel_size=self.input_dim,
            #                             h=self.height,
            #                             w=self.width, 
            #                             k=self.sweep_length,
            #                             TM_Flag=True)
            # self.layers_spatial = Spatial_conv(self.input_dim,self.input_dim,kernel_size=5)
            if self.if_dotproductattn:
                self.dotproductattn = ScaledDotProductAttention(self.input_dim)
            if self.if_conv_aggre:
                self.conv_aggre = SpatialAttention_mtf()
        
        print("================Rain attention Init================")


    

    def forward(self, features, sensor_dist, time_stamps, valid_mask):
        """
        Parameters:
        -----------
        features: torch.Tensor, shape (B, C, H, W)
            neighbor features in ego view
            B : 按照顺序排列: ego[curr], cav_1[curr], cav_1[past_1], ..., cav_2[past_k], ....
            C : feature 维度，一般为 256
            H : height, 一般为 100
            W : width, 一般为 252

        sensor_dist: torch.Tensor, shape (B, H, W)
            distance between each block and the sensor
            
        time_stamps: torch.Tensor, shape (B)
            表示每个位置的时间信息

        valid_mask: (B, 1, H, W)
            为0/1,表示是否在该位置有效

        Returns:
        --------
        output_feature: torch.Tensor, shape (C, H, W)
            ego cav's fused feature at current frame.
        """
        if self.if_max_aggre:
            features = torch.max(features, dim=0)[0]
            if self.if_spatial_conv:
                features = features.unsqueeze(0)
                features = self.layers_spatial[0](features)
            return features.squeeze(0)
        if self.confidence_fetch == False:
            if self.sup_individual:
                curr_features = features[1::self.sweep_length+1].clone()    # [n, C, H, W]
                latency_features = features[2::self.sweep_length+1].clone()  # [n, C, H, W]
                if self.compensation==False:
                    
                    final_fuse_feature = torch.max(torch.cat((features[:1],curr_features),dim=0),dim=0)[0]
                    final_fuse_feature_latency = torch.max(torch.cat((features[:1],latency_features),dim=0),dim=0)[0]

                    return final_fuse_feature_latency,final_fuse_feature,final_fuse_feature_latency,0,0


                cav_num, C, H, W = features.shape
                individual_recon_loss = 0
                latency_recon_loss = 0
                # debug use, set true to save pt for viz
                
                if if_save_pt:
                    torch.save(time_stamps, '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'time_stamps'+'.pt')
                    torch.save(features[1:], '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'features'+'.pt')
                
                num_nonego = int((cav_num-1)/(self.sweep_length+1)) 
                nonego_features = features[1:].clone()                  # [nonego_cavs(1+k), C, H, W]
                
                nonego_splits = nonego_features.split(self.sweep_length+1,0) # ([1+k, C, H, W], [], ..., [])
                nonego_splits = torch.stack(nonego_splits,dim=0)        # [nonego_cars, 1+k, C, H, W]
                # torch.save(nonego_splits, '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'nonego_splits'+'.pt')
                nonego_history = nonego_splits[:,1:]                    # [nonego_cars, k, C, H, W]

                # if self.if_swin_trans:
                nonego_timestamps = time_stamps[1:]                     # [nonego_cavs(1+k), ]
                nonego_timestamps_splits = nonego_timestamps.split(self.sweep_length+1, 0)
                nonego_timestamps_splits = torch.stack(nonego_timestamps_splits,dim=0)   # [n, 1+k, ]
                nonego_timestamps_history = nonego_timestamps_splits[:, 1:]


                nonego_history = nonego_history.reshape(-1, C, H, W).permute(0,2,3,1)
                nonego_timestamps_history = nonego_timestamps_history.reshape(-1)
                # nonego_features = nonego_features.permute(0,2,3,1)      # 
                nonego_history = self.time_encoding(nonego_history,nonego_timestamps_history)
                nonego_history = nonego_history.permute(0,3,1,2).reshape(num_nonego, self.sweep_length, C, H, W)      # [nonego_cavs, k, C, H, W]

                # mask after TE, valid_mask: [B, H, W]
                non_ego_mask = valid_mask[1:]                           # [nonego_cavs*(k+1), 1, H, W]
                non_ego_mask_split = non_ego_mask.split(self.sweep_length+1,0) # ([1+k, 1, H, W], ... ) nonego_cavs 
                non_ego_mask = torch.stack(non_ego_mask_split, dim=0)   # nonego_cavs, 1+k, H, W
                if if_save_pt:
                    torch.save(non_ego_mask, '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'nonego_mask_splits'+'.pt')

                non_ego_mask_past = non_ego_mask[:, 1:]                 # nonego_cavs, k, H, W
                nonego_history = nonego_history * non_ego_mask_past        # [nonego_cavs, k, C, H, W]
            
                hidden = self.swin_trans(nonego_history)                # [non_ego_cars, 1, C, H, W]

                non_ego_mask_past_union = (torch.sum(non_ego_mask_past, dim=1)>0).to(torch.float32) # nonego_cavs, 1, H, W
                non_ego_mask_past_union = self.dilate_mask(non_ego_mask_past_union)
                if if_save_pt:
                    torch.save(non_ego_mask_past_union,'/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ f'non_ego_mask_past_union_dilated_{self.dilation_size}' + '.pt')  

                # elif self.if_syncnet:
                #     hidden = self.syncnet(nonego_history,[1])
                #     non_ego_mask_past_union = torch.ones((num_nonego, C, H, W)).to(hidden)

                # query = self.predict_query.repeat(num_nonego,1,1)
                # final_estimations = self.predict_attention(query, hidden, hidden, None).reshape(num_nonego, H, W, -1).permute(0,3,1,2)
                final_estimations = hidden.squeeze(1)                   # [non_ego_cars, C, H, W]
                if if_save_pt:
                    torch.save(final_estimations, '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'nonego_estimations'+'.pt')

                final_estimations = final_estimations * non_ego_mask_past_union

                # debug use: uncomment to 用带有delay的 feature 直接作为feature 输出
                # final_estimations = nonego_splits[:,1]

                # debug use, uncomment to retrieve current mask
                # ego_mask_curr = valid_mask[0]
                # nonego_mask_curr = non_ego_mask[:, 0] # [nonego_cav, 1, H, W]
                # torch.save(nonego_mask_curr,'/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ f'nonego_mask_curr_{self.dilation_size}' + '.pt')

                individual_recon_loss_pos = torch.torch.nn.functional.smooth_l1_loss(final_estimations,nonego_splits[:,0],reduction='sum')/C
                if non_ego_mask_past_union.sum() != 0:
                    individual_recon_loss_pos /= (non_ego_mask_past_union.sum())

                # individual_recon_loss_neg = torch.torch.nn.functional.smooth_l1_loss((1-non_ego_mask_past_union)*final_estimations,nonego_splits[:,0],reduction='sum')
                # if (1-non_ego_mask_past_union).sum() != 0:
                    # individual_recon_loss_neg /= ((1-non_ego_mask_past_union).sum())

                # individual_recon_loss = individual_recon_loss_pos + individual_recon_loss_neg
                individual_recon_loss = individual_recon_loss_pos

                latency_recon_loss = torch.torch.nn.functional.smooth_l1_loss(latency_features, nonego_splits[:,0],reduction='sum')/C
                if non_ego_mask_past_union.sum() != 0:
                    latency_recon_loss /= (non_ego_mask_past_union.sum())


                ###############skimage.dilation##########################################
                final_fuse_estimation = torch.max(torch.cat((features[:1],final_estimations),dim=0), dim=0)[0]
                final_fuse_feature = torch.max(torch.cat((features[:1],curr_features),dim=0),dim=0)[0]
                # fuse_recon_loss = torch.torch.nn.functional.smooth_l1_loss(final_fuse_estimation,final_fuse_feature,reduction='mean')

                final_fuse_feature_latency = torch.max(torch.cat((features[:1],latency_features),dim=0),dim=0)[0]
                # ipdb.set_trace()
                # print('#######################individual:',torch.torch.nn.functional.smooth_l1_loss(latest_features[1:],curr_features[1:]),individual_recon_loss)
                # latest_features = features[self.sweep_length-1::self.sweep_length].clone()
                # latest_features[0] = features[0]

                # final1 = torch.max(torch.cat((features[:1],latest_features[1:]),dim=0), dim=0)[0]
                # cnt = 5
                # torch.save(curr_features,'/GPFS/data/yuxiwei/yuxiwei/OpenCOOD/vis/current'+str(cnt)+'.pt')
                # torch.save(latest_features,'/GPFS/data/yuxiwei/yuxiwei/OpenCOOD/vis/latest'+str(cnt)+'.pt')
                # cnt += 1
                # ipdb.set_trace()
                # print('#######################:fuse',torch.torch.nn.functional.smooth_l1_loss(final1,final_fuse_feature),fuse_recon_loss)
             
                return final_fuse_estimation,final_fuse_feature,final_fuse_feature_latency, individual_recon_loss, latency_recon_loss # +fuse_recon_loss
                # return final_fuse_feature, 0

            
            num_cav = int(features.shape[0]/self.sweep_length)
            curr_features = features[:self.sweep_length:]
            features = torch.cat((features[:1],features[num_cav:]),dim=0)
            time_stamps = torch.cat((time_stamps[:1],time_stamps[num_cav:]),dim=0)
            cav_num, C, H, W = features.shape
            ################################
            # features = features[::2]
            # valid_mask = valid_mask[::2]
            # cav_num = int(cav_num/2)
            ################################


            if self.if_conv_aggre:
                curr = features[0:1].clone()
                prev = features[1:].clone()
                prev = self.conv_aggre(curr,prev)
                features = torch.cat((curr,prev),dim=0)
                features = features.permute(0,2,3,1)
                time_stamps = time_stamps*-1
                features = self.time_encoding(features,time_stamps)
                features = features.permute(0,3,1,2)
                if self.if_spatial_conv:
                    features = self.layers_spatial[0](features)
                return torch.max(features,dim=0)[0]
            
            if self.if_ego_time_encoding:
                # s1 = time.time()
                features = features.permute(0,2,3,1)
                time_stamps = time_stamps*-1
                features = self.time_encoding(features,time_stamps)
                features = features.permute(0,3,1,2)
                # s2 = time.time()
                # print('TE',s2-s1)
            
            
            features = features.view(cav_num, C, -1).permute(2, 0, 1) #  (H*W, cav_num, C), perform self attention on each pixel.
            # masks = valid_mask.view(cav_num, H*W).permute(1,0)
            # masks = masks.unsqueeze(-1).repeat(1,1,cav_num)
            


            if self.if_dotproductattn:
                features = self.dotproductattn(features, features, features, masks)
            if self.if_time_attn_aggre:
                # s1 = time.time()
                features = self.layers_temporal[0](features, features, features, None)
                # s2 = time.time()
                # print('trans',s2-s1)
            
           

            features = features.permute(1, 2, 0).view(cav_num, C, H, W)[:1] 
            if self.if_spatial_conv:
                # s1 = time.time()
                features = self.layers_spatial[0](features)[0]
                # s2 = time.time()
                # print('conv',s2-s1)
            if self.sup_individual:
                features = torch.max(final_estimations, dim=0)[0]
            curr_fuse_features = torch.max(curr_features, dim=0)[0]
            recon_loss = torch.torch.nn.functional.smooth_l1_loss(features,curr_fuse_features)

            return features, recon_loss


        valid_mask = valid_mask.squeeze(1)
        
        features = features.permute(0,2,3,1)
        
        features = torch.cat((features[0:1],features[self.sweep_length:]),dim=0)
        

        ##########################################  三种encoding
        if self.if_spatial_encoding:
            features[0:1] = self.spatial_encoding(features[0:1])# features.shape = [batch_all,height,width,feature_dim]
        if self.if_sensor_encoding and sensor_dist!=-1:
            sensor_dist = torch.cat((sensor_dist[0:1],sensor_dist[self.sweep_length:]))
            features = self.sensor_dist_encoding(features,sensor_dist)
        time_stamps = torch.cat((time_stamps[0:1],time_stamps[self.sweep_length:]),dim=0)
        valid_mask = torch.cat((valid_mask[0:1],valid_mask[self.sweep_length:]),dim=0)
        valid_mask_sum = valid_mask.sum(0)
        nonego_idx = torch.where(valid_mask_sum>1)
 
        features_to_aggre = features[:,nonego_idx[0],nonego_idx[1],:] 
        if self.if_ego_time_encoding:
            features[0:1] = self.time_encoding(features[0:1],time_stamps[0:1])

        if features_to_aggre.shape[1] != 0:
            if self.if_spatial_encoding:
                features_to_aggre = self.spatial_encoding_nonego(features_to_aggre,nonego_idx)
            if self.if_nonego_time_encoding:
                features_to_aggre = self.time_encoding(features_to_aggre,time_stamps)
        ###########################################

        ################################################### 
        # ipdb.set_trace()
        masks = valid_mask[:,nonego_idx[0],nonego_idx[1]]#(valid_num,len)
        batch_len = masks.shape[0]
        masks = masks.permute(1,0).unsqueeze(-1).repeat(1,1,batch_len)

        features_to_aggre = features_to_aggre.permute(1,0,2)
        for l in range(self.num_layers):
                
            if masks.shape[0] != 0:
                if self.if_dotproductattn:
                    # if self.use_mask == False:
                    #     masks =  None
                    # features_to_aggre = self.dotproductattn(features_to_aggre,features_to_aggre,features_to_aggre,masks)
                    features_to_aggre[:,0] = torch.max(features_to_aggre,dim=1)[0]
                else:
                    features_to_aggre = self.layers_temporal[l](features_to_aggre,features_to_aggre,features_to_aggre,masks)
                features[:,nonego_idx[0],nonego_idx[1],:] = features_to_aggre.permute(1,0,2)


            if self.if_spatial_conv:
                features = features.permute(0,3,1,2)
                features_backup = features.clone()
                features[0:1] = self.layers_spatial[l](features_backup[0:1])
                features = features.permute(0,2,3,1)

                features_to_aggre[:,0:1] = features[0:1,nonego_idx[0],nonego_idx[1],:].permute(1,0,2)

        return features[0].permute(2,0,1)

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_swin(nn.Module):
    def __init__(self, args):
        super(raindrop_swin, self).__init__()
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.agg_mode == 'RAIN':
            d_ff = args['d_ff']
            h = args['h']
            dropout = args['dropout']
            input_dim = args['input_dim']
            num_layers = args['num_layers']
            neighbor_range = args['neighbor_range']
            max_len = args['max_len']
            height = 100
            width = 352

            self.fuse_layer1 = False
            if self.multi_scale:
                layer_nums = args['layer_nums']
                num_filters = args['num_filters']
                self.num_levels = len(layer_nums)
                if self.fuse_layer1:
                    self.fuse_modules = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len)
                    for p in self.fuse_modules.parameters():
                        if p.dim() > 1:
                            nn.init.xavier_uniform_(p)
                else:
                    self.fuse_modules = nn.ModuleList()
                    for idx in range(self.num_levels):
                        if idx==0:
                            fuse_network = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len, compensation=True)
                        else:
                            fuse_network = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len, compensation=False)
                        for p in fuse_network.parameters():
                            if p.dim() > 1:
                                nn.init.xavier_uniform_(p)

                        height = int(height/2)
                        width = int(width/2)
                        input_dim = input_dim * 2
                        d_ff = d_ff * 2
                        self.fuse_modules.append(fuse_network)
            else:
                self.fuse_modules = AttentionFusion(input_dim, d_ff, h, num_layers, neighbor_range, 100, 252, dropout, max_len)
                # This was important from their code. 
                # Initialize parameters with Glorot / fan_avg.
                for p in self.fuse_modules.parameters():
                    if p.dim() > 1:
                        nn.init.xavier_uniform_(p)

        else:
            if self.multi_scale:
                layer_nums = args['layer_nums']
                num_filters = args['num_filters']
                self.num_levels = len(layer_nums)
                self.fuse_modules = nn.ModuleList()
                for idx in range(self.num_levels):
                    if self.agg_mode == 'MAX':
                        fuse_network = MaxFusion()
                    else:
                        # self.agg_mode == 'ATTEN':
                        fuse_network = AttenFusion(num_filters[idx])
                    self.fuse_modules.append(fuse_network)
            else: 
                if self.agg_mode == 'MAX': # max fusion, debug use
                    self.fuse_modules = MaxFusion()

        print("##### Raindrop + ** SWIN Trans ** w.o. single supervision Init! #####")

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k-(k-1), dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        
        return split_x

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]
        # print('num',x.shape[0],pairwise_t_matrix.shape)
        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2
        
        if self.multi_scale:
            
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(x)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
            if self.fuse_layer1:
                all_recon_loss = 0
                layer1_feats = backbone.resnet.forward_layer1(x)  # [B, C, H, W]
                if self.communication:
                    batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                    batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                    _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                    communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                    # debug use, to save original feature:
                    # debug_feats = layer1_feats.
                    # TODO scale = 0 不加 mask
                    layer1_feats = layer1_feats * communication_masks_tensor

                batch_node_features = self.regroup(layer1_feats, record_len, K)
                x_fuse = []
                x_fuse_curr = []
                x_fuse_latency = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    # print('node and t mat',node_features.shape,t_matrix.shape,record_len)
                    t_matrix = t_matrix[K-1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))
                    
                    # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                    record_frames = np.ones((N))*K; record_frames[0]-=(K-1)
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        # x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        features,features_curr,features_latency, recon_loss = self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b])
                        all_recon_loss += recon_loss # TODO: 是否需要添加 /= B ? 
                        x_fuse.append(features)
                        x_fuse_curr.append(features_curr)
                        x_fuse_latency.append(features_latency)
                x_fuse = torch.stack(x_fuse)  ###(batch, channel, h, w)
                x_fuse_curr = torch.stack(x_fuse_curr)
                x_fuse_latency = torch.stack(x_fuse_latency)

                layer2_feat, layer3_feat = backbone.resnet.forward_from_layer1(x_fuse)
                ups = []
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[0](x_fuse))
                    ups.append(backbone.deblocks[1](layer2_feat))
                    ups.append(backbone.deblocks[2](layer3_feat))
                
                if len(ups) > 1:
                    x_fuse = torch.cat(ups, dim=1)
                elif len(ups) == 1:
                    x_fuse = ups[0]
                
                if len(backbone.deblocks) > self.num_levels:
                    x_fuse = backbone.deblocks[-1](x_fuse)

                
                layer2_feat_curr, layer3_feat_curr = backbone.resnet.forward_from_layer1(x_fuse_curr)
                ups = []
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[0](x_fuse_curr))
                    ups.append(backbone.deblocks[1](layer2_feat_curr))
                    ups.append(backbone.deblocks[2](layer3_feat_curr))
                    
                if len(ups) > 1:
                    x_fuse_curr = torch.cat(ups, dim=1)
                elif len(ups) == 1:
                    x_fuse_curr = ups[0]
                
                if len(backbone.deblocks) > self.num_levels:
                    x_fuse_curr = backbone.deblocks[-1](x_fuse_curr)

                
                layer2_feat_latency, layer3_feat_latency = backbone.resnet.forward_from_layer1(x_fuse_latency)
                ups = []
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[0](x_fuse_latency))
                    ups.append(backbone.deblocks[1](layer2_feat_latency))
                    ups.append(backbone.deblocks[2](layer3_feat_latency))
                    
                if len(ups) > 1:
                    x_fuse_latency = torch.cat(ups, dim=1)
                elif len(ups) == 1:
                    x_fuse_latency = ups[0]
                
                if len(backbone.deblocks) > self.num_levels:
                    x_fuse_latency = backbone.deblocks[-1](x_fuse_latency)

            else:
                ups = []
                ups_curr = []
                ups_latency = []
                all_recon_loss = 0
                all_latency_recon_loss = 0
                for i in range(self.num_levels):
                    x = feats[i] if with_resnet else backbone.blocks[i](x)
                    ############ 1. Communication (Mask the features) #########
                    if i==0:
                        if self.communication:
                            batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                            batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                            _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                            communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                            # for debug use
                            if if_save_pt:
                                x_wo_mask = x.clone()
                            
                            # if self.if_syncnet: # 如果使用 SyncNet 作为补偿模块，则不使用mask
                            #     communication_masks_tensor = torch.ones_like(communication_masks_tensor)
                            
                            x = x * communication_masks_tensor
                        else:
                            communication_rates = torch.tensor(0).to(x.device)
                    else:
                        if self.communication:
                            communication_masks_tensor = F.max_pool2d(communication_masks_tensor, kernel_size=2)
                            # scale = 1, 2 不加 mask
                            # x = x * communication_masks_tensor
                    
                    ############ 2. Split the confidence map #######################
                    # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                    # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                    # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                    # print('x',x.shape)
                    batch_node_features = self.regroup(x, record_len, K)
                    batch_comm_masks = self.regroup(communication_masks_tensor, record_len, K)
                    
                    if if_save_pt:
                        batch_x_wo_mask = self.regroup(x_wo_mask, record_len, K)
                    
                    ############ 3. Fusion ####################################
                    x_fuse = []
                    x_fuse_curr = []
                    x_fuse_latency = []
                    for b in range(B):
                        # number of valid agent
                        N = record_len[b]
                        # t_matrix[i, j]-> from i to j
                        t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) # [Nx(k+1), 2, 3]
                        node_features = batch_node_features[b]
                        node_masks = batch_comm_masks[b]
                        C, H, W = node_features.shape[1:]
                        # print('node and t mat',node_features.shape,t_matrix.shape,record_len)
                        t_matrix = t_matrix[K-1:]
                        neighbor_feature = warp_affine_simple(node_features,
                                                        t_matrix,
                                                        (H, W))
                        neighbor_mask = warp_affine_simple(node_masks, t_matrix, (H, W))

                        if if_save_pt:
                            neighbor_wo_mask = warp_affine_simple(batch_x_wo_mask[b], t_matrix, (H, W))
                            torch.save(neighbor_wo_mask[1:], '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'feature_wo_mask'+'.pt')

                        # TODO: 这个位置
                        record_frames = np.ones((N))*K
                        # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                        
                        if self.agg_mode == 'RAIN':
                            # for sensor embedding
                            sensor_dist = -1# (B, H, W)
                            # x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                            # features,features_curr,features_latency, recon_loss = self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b])
                            features, features_curr, features_latency, recon_loss, latency_recon_loss = self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], neighbor_mask)
                            x_fuse.append(features)
                            x_fuse_curr.append(features_curr)
                            x_fuse_latency.append(features_latency)
                            all_recon_loss += recon_loss
                            all_latency_recon_loss += latency_recon_loss
                            # # TODO for scale debug
                            # if i==self.num_levels-1:
                            #     x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                            # else:
                            #     x_fuse.append(neighbor_feature[0])
                        else: # ATTEN, MAX, Transformer
                            x_fuse.append(self.fuse_modules[i](neighbor_feature))
                            # # TODO for scale debug
                            # if i==self.num_levels-1:
                            #     x_fuse.append(self.fuse_modules[i](neighbor_feature))
                            # else:
                            #     x_fuse.append(neighbor_feature[0])

                    x_fuse = torch.stack(x_fuse)
                    x_fuse_curr = torch.stack(x_fuse_curr)
                    x_fuse_latency = torch.stack(x_fuse_latency)
                    ############ 4. Deconv ####################################
                    if len(backbone.deblocks) > 0:
                        ups.append(backbone.deblocks[i](x_fuse))
                        ups_curr.append(backbone.deblocks[i](x_fuse_curr))
                        ups_latency.append(backbone.deblocks[i](x_fuse_latency))
                    else:
                        ups.append(x_fuse)
                        ups_curr.append(x_fuse_curr)
                        ups_latency.append(x_fuse_latency)
                    
                if len(ups) > 1:
                    x_fuse = torch.cat(ups, dim=1)
                    x_fuse_curr = torch.cat(ups_curr, dim=1)
                    x_fuse_latency = torch.cat(ups_latency, dim=1)
                elif len(ups) == 1:
                    x_fuse = ups[0]
                    x_fuse_curr = ups_curr[0]
                    x_fuse_latency = ups_latency[0]
                
                if len(backbone.deblocks) > self.num_levels:
                    x_fuse = backbone.deblocks[-1](x_fuse)
                    x_fuse_curr = backbone.deblocks[-1](x_fuse_curr)
                    x_fuse_latency = backbone.deblocks[-1](x_fuse_latency)
                
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_confidence_maps = self.regroup(rm, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                # _, (B, 1, H, W), float
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix,
                                                (H, W))
                record_frames = np.ones((N))*K
                # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                if self.agg_mode == "RAIN":
                    # for sensor embedding
                    sensor_dist = -1# (B, H, W)
                    x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)

            # self.fuse_modsules(x_fuse, record_len)
        
        return x_fuse,x_fuse_curr,x_fuse_latency, communication_rates, all_recon_loss, all_latency_recon_loss, {}
opencood/models/fuse_modules/raindrop_swin.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# Date: 2022/12/06
# License: TDG-Attribution-NonCommercial-NoDistrib

"""
where2comm + swin transformer + single supervisation.
"""

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

import torch.nn.functional as F
from turtle import update
# import ipdb
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
from opencood.models.comm_modules.where2comm_multisweep import Communication

from opencood.models.fuse_modules.STPT import STPT

if_save_pt = False  # TODO: for debug use, save pt file
pt_save_path = '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out'

class PixelWeightedFusionSoftmax(nn.Module):
    def __init__(self, channel):
        super(PixelWeightedFusionSoftmax, self).__init__()

        self.conv1_1 = nn.Conv2d(channel * 2, 128, kernel_size=1, stride=1, padding=0)
        self.bn1_1 = nn.BatchNorm2d(128)

        self.conv1_2 = nn.Conv2d(128, 32, kernel_size=1, stride=1, padding=0)
        self.bn1_2 = nn.BatchNorm2d(32)

        self.conv1_3 = nn.Conv2d(32, 8, kernel_size=1, stride=1, padding=0)
        self.bn1_3 = nn.BatchNorm2d(8)

        self.conv1_4 = nn.Conv2d(8, 1, kernel_size=1, stride=1, padding=0)
        # self.bn1_4 = nn.BatchNorm2d(1)

    def forward(self, x):
        x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
        x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
        x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
        x_1 = F.relu(self.bn1_3(self.conv1_3(x_1)))
        x_1 = F.relu(self.conv1_4(x_1))

        return x_1


class mask_fuse(nn.Module):
    def __init__(self, channel):
        super(mask_fuse, self).__init__()
        self.weight_conv = PixelWeightedFusionSoftmax(channel)
        self.soft_max = nn.Softmax(dim=0)
    def forward(self, ego_feature, nonego_feature):#####[cars, C, H, W]
        ''' 
        ego_feature [1, C, H, W]
        nonego_feature [n, C, H, W]
        '''
        all_feature = torch.cat((ego_feature,nonego_feature),dim=0)
        ego_feature_repeat = ego_feature.repeat(nonego_feature.shape[0]+1,1,1,1)

        cat_feature = torch.cat((ego_feature_repeat,all_feature),dim=1)

        mask_map = self.weight_conv(cat_feature)
        mask_weight = self.soft_max(mask_map)

        final_feature = torch.sum(mask_weight*all_feature, dim=0)

        return final_feature

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class PositionalEncoding_spatial(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model_feature, d_model_hidden, dropout, max_len=300,mode='cat'):
        super(PositionalEncoding_spatial, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        d_model = int(d_model_hidden/2)
        # Compute the positional encodings once in log space.
        pe = torch.zeros((max_len, d_model))
        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe_ = pe.unsqueeze(0)#batch,len,feature
        pe_h = pe_.unsqueeze(2).repeat(1,1,max_len,1)
        pe_w = pe_.unsqueeze(1).repeat(1,max_len,1,1)
        pe = torch.cat((pe_h,pe_w),dim=-1)

        self.register_buffer('pe', pe)

        self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x):
        if self.mode == 'cat':
            tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
            tmp = tmp.repeat(x.shape[0], 1, 1, 1)
            x = self.linear(torch.cat((x, tmp),dim=-1))

        elif self.mode == 'init':
            x = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)

        else:
            tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
            tmp = tmp.repeat(x.shape[0], 1, 1, 1)
            x = x + tmp

        return self.dropout(x)

class PositionalEncoding_spatial_custom(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model_feature, d_model_hidden, dropout,mode='cat'):
        super(PositionalEncoding_spatial_custom, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = int(d_model_hidden/2)
        # Compute the positional encodings once in log space.
        
        self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x, nonego_idx):#x(batch,num,feature)
   
        pe_h = torch.zeros((x.shape[1], self.d_model)).to(x.device)
        pe_w = torch.zeros((x.shape[1], self.d_model)).to(x.device)
        position_h, position_w = nonego_idx[0],nonego_idx[1]
        position_h = position_h.unsqueeze(1)
        position_w = position_w.unsqueeze(1)

        div_term = torch.exp(torch.arange(0, self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(position_h.device)
        pe_h[:, 0::2] = torch.sin(position_h * div_term)
        pe_h[:, 1::2] = torch.cos(position_h * div_term)

        pe_w[:, 0::2] = torch.sin(position_w * div_term)
        pe_w[:, 1::2] = torch.cos(position_w * div_term)
       
        pe = torch.cat((pe_h,pe_w),dim=-1).unsqueeze(0).repeat(x.shape[0],1,1)

        if self.mode == 'cat':
            x = self.linear(torch.cat((x, pe),dim=-1))

        else:
            x = x + Variable(pe,requires_grad=False)

        return self.dropout(x)


class PositionalEncoding_sensor_dist(nn.Module):

    def __init__(self, d_model, dropout):
        super(PositionalEncoding_sensor_dist, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = d_model
        self.linear = nn.Linear(2*d_model,d_model)
        self.mode = 'cat'
         
    def forward(self, x, sensor_dist):#x.shape=[batch,H,W,F],senfor_dist.shape=[batch,H,W]

        x = x.reshape(-1,x.shape[-1])
        sensor_dist = sensor_dist.reshape(-1)

        pe = torch.zeros(x.shape[0], self.d_model).to(x.device)
        position = sensor_dist.unsqueeze(1) 
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(x.device)
         
        pe[:, 0::2] = torch.sin(position * div_term)  
        pe[:, 1::2] = torch.cos(position * div_term)   

        if self.mode == 'cat':
            x = self.linear(torch.cat((x,Variable(pe,require_grad=False)),dim=-1))

        else:
            x = x + Variable(pe,requires_grad=False)
        
        return self.dropout(x)

class PositionalEncoding_irregular_time(nn.Module):

    def __init__(self, d_model_feature, d_model_hidden, dropout):
        super(PositionalEncoding_irregular_time, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = d_model_feature
        self.linear = nn.Linear(d_model_feature+d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x, time):#x.shape=[batch,height,width,feature],time.shape=[batch]

        pe = torch.zeros(x.shape[0],self.d_model).to(x.device)
        position = time.unsqueeze(1) 
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(x.device)
         
        pe[:, 0::2] = torch.sin(position * div_term)  
        pe[:, 1::2] = torch.cos(position * div_term)   

        if self.mode == 'cat':
            if len(x.shape) == 4:
                pe = pe.unsqueeze(1).unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
            else:
                pe = pe.unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], 1)
            x = self.linear(torch.cat((x, tmp),dim=-1))

        else:
            if len(x.shape) == 4:
                pe = pe.unsqueeze(1).unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
            else:
                pe = pe.unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], 1)
            x = x + tmp

        return self.dropout(x)

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    
    # ipdb.set_trace()
    
    # print('score',scores.shape,'mask',mask.shape)

    if mask is not None:
        # mask = mask.transpose(1,2).to(query.device)
        mask = mask.to(query.device)
        scores = scores.masked_fill(mask == 0, -1e10) 

    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)

    return torch.matmul(p_attn, value), p_attn


class MultiHeadedAttention(nn.Module):
    def __init__(self, h, q_dim, k_dim,v_dim, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h  
        self.h = h 
        self.q_dim = q_dim
        self.k_dim = k_dim
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.linear_q = nn.Linear(q_dim,d_model)
        self.linear_k = nn.Linear(k_dim,d_model)
        self.linear_v = nn.Linear(v_dim,d_model)
        self.linear_out = nn.Linear(d_model,q_dim)
        self.attn = None  
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k 

        query = self.linear_q(query).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
        key = self.linear_k(key).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
        value = self.linear_v(value).view(nbatches,-1,self.h,self.d_k).transpose(1,2)

        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(query, key, value, mask=mask,
                                 dropout=self.dropout)

        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(1, 2).contiguous() \
            .view(nbatches ,-1, self.h * self.d_k)

        return self.linear_out(x)

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
 
    def forward(self, x):
        "Norm"
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, x, sublayer):
        # print('sub',x.shape,self.dropout(sublayer(self.norm(x))).shape)
        return x + self.dropout(sublayer(self.norm(x))) #残差连接

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model,d_out, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_out)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

class EncoderLayer_temporal(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer_temporal, self).__init__()
        self.self_attn = self_attn #sublayer 1 
        self.feed_forward = feed_forward #sublayer 2 
        self.sublayer = clones(SublayerConnection(size, dropout), 2) #拷贝两个SublayerConnection，一个为了attention，一个是独自作为简单的神经网络
        self.size = size

    def forward(self, q, k, v, mask):
        "Follow Figure 1 (left) for connections."
        # print('x',x.shape,'attn',self.self_attn(x,x,x,mask).shape)
        x = self.sublayer[0](q, lambda q: self.self_attn(q, k, v, mask)) #attention层 对应层1
        return self.sublayer[1](x, self.feed_forward) # 对应 层2

class Spatial_conv(nn.Module):
    def __init__(self, input_channels, output_channels, kernel_size):
        super(Spatial_conv,self).__init__()
        self.conv = nn.Conv2d(in_channels = input_channels,
                              out_channels = output_channels,
                              kernel_size = (kernel_size,kernel_size),
                              stride = (1,1),
                              padding = (int((kernel_size-1)/2),int((kernel_size-1)/2))
        )

        self.relu = nn.ReLU()
        # self.bn = nn.BatchNorm2d(num_features=output_channels)

    def forward(self, x):
        init_x = x.clone()
        x = self.conv(x)
        x = self.relu(x)
        # x = self.bn(x)
        return x+init_x

class SpatialAttention_mtf(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention_mtf, self).__init__()

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, curr, prev):#####(batch, channel, h, w)
        avg_out = torch.mean(curr, dim=1, keepdim=True)
        max_out, _ = torch.max(curr, dim=1, keepdim=True)
        y = torch.cat([avg_out, max_out], dim=1)
        y = self.conv1(y)
        return self.sigmoid(y) * prev ####(batch-1, channel, h, w)

class ScaledDotProductAttention(nn.Module):
    """
    Scaled Dot-Product Attention proposed in "Attention Is All You Need"
    Compute the dot products of the query with all keys, divide each by sqrt(dim),
    and apply a softmax function to obtain the weights on the values
    Args: dim, mask
        dim (int): dimention of attention
        mask (torch.Tensor): tensor containing indices to be masked
    Inputs: query, key, value, mask
        - **query** (batch, q_len, d_model): tensor containing projection
          vector for decoder.
        - **key** (batch, k_len, d_model): tensor containing projection
          vector for encoder.
        - **value** (batch, v_len, d_model): tensor containing features of the
          encoded input sequence.
        - **mask** (-): tensor containing indices to be masked
    Returns: context, attn
        - **context**: tensor containing the context vector from
          attention mechanism.
        - **attn**: tensor containing the attention (alignment) from the
          encoder outputs.
    """

    def __init__(self, dim):
        super(ScaledDotProductAttention, self).__init__()
        self.sqrt_dim = np.sqrt(dim)

    def forward(self, query, key, value, mask):
        score = torch.bmm(query, key.transpose(1, 2)) / self.sqrt_dim
        if mask != None:
            score = score.masked_fill(mask == 0, -1e10) 
        attn = F.softmax(score, -1)
        context = torch.bmm(attn, value)
        return context

class AttentionFusion(nn.Module):
    def __init__(self, settings, input_dim, hidden_dim, num_heads, num_layers, neighbor_range, height, width, dropout, max_len=300, compensation=True):
        super(AttentionFusion, self).__init__()
        self.sweep_length = settings['sweep_length']
        self.if_max_aggre = settings['if_max_aggre']
        self.if_spatial_encoding = settings['if_spatial_encoding']
        self.if_ego_time_encoding = settings['if_ego_time_encoding']
        self.if_nonego_time_encoding = settings['if_nonego_time_encoding']
        self.if_sensor_encoding = settings['if_sensor_encoding']
        self.if_time_attn_aggre = settings['if_time_attn_aggre']
        self.if_spatial_conv = settings['if_spatial_conv']
        self.if_dotproductattn = settings['if_dotproductattn']
        self.use_mask = settings['use_mask']
        self.confidence_fetch = settings['confidence_fetch']
        self.if_conv_aggre = settings['if_conv_aggre']
        
        self.sup_individual = settings['sup_individual']  # 用于单车feature 进行补偿 算recon loss
        self.if_swin_trans = True

        self.input_dim = input_dim 
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.neighbor_range = neighbor_range
        self.height = height
        self.width = width
        self.compensation=compensation
        
        self.dilation_size = 0
        if self.compensation:
            if 'dilation_size' in settings:
                self.dilation_size = settings['dilation_size']
            else:
                self.dilation_size = 1
            self.dilate_mask = nn.MaxPool2d(kernel_size=self.dilation_size, stride=1, padding=int((self.dilation_size-1)/2))
            if self.if_spatial_encoding:
                self.spatial_encoding = PositionalEncoding_spatial(self.input_dim, hidden_dim, dropout, max_len)
                self.spatial_encoding_nonego = PositionalEncoding_spatial_custom(self.input_dim, hidden_dim, dropout)
            if self.if_ego_time_encoding or self.if_nonego_time_encoding:
                self.time_encoding = PositionalEncoding_irregular_time(self.input_dim, hidden_dim, dropout)
            if self.if_sensor_encoding:
                self.sensor_dist_encoding = PositionalEncoding_sensor_dist(hidden_dim, dropout)

            c = copy.deepcopy
            temporal_attn = MultiHeadedAttention(self.num_heads, self.input_dim, self.input_dim, self.input_dim, self.hidden_dim, dropout)

            ff = PositionwiseFeedForward(self.input_dim, self.input_dim, self.hidden_dim, dropout)

            if self.if_time_attn_aggre:
                self.layers_temporal = clones(EncoderLayer_temporal(self.input_dim, c(temporal_attn), c(ff), dropout),num_layers)
            # self.temporal_aggregator = MultiHeadedAttention(self.num_heads, hidden_dim, self.input_dim, self.input_dim, self.hidden_dim, dropout)
            if self.if_spatial_conv:  
                self.layers_spatial = clones(Spatial_conv(self.input_dim,self.input_dim,kernel_size=5),num_layers)
            if self.if_swin_trans:
                self.swin_trans = STPT(in_chans=self.input_dim,
                    depths=[2, 2], 
                    num_heads=[2, 2], 
                    window_sizes=[(4,4), (4,4)],
                    num_frames=3,
                    num_out_frames=1,
                    embed_dim=36,  # TODO
                    use_hdmap=False,
                    use_hdmap_feat=False)
            
            # self.layers_spatial = Spatial_conv(self.input_dim,self.input_dim,kernel_size=5)
            if self.if_dotproductattn:
                self.dotproductattn = ScaledDotProductAttention(self.input_dim)
            if self.if_conv_aggre:
                self.conv_aggre = SpatialAttention_mtf()
        
        print(f"===== Dilation size : {self.dilation_size} =====")




    def forward(self, features, sensor_dist, time_stamps, valid_mask):
        """
        Parameters:
        -----------
        features: torch.Tensor, shape (B, C, H, W)
            neighbor features in ego view
            B : 按照顺序排列: ego[curr], cav_1[curr], cav_1[past_1], ..., cav_2[past_k], ....
            C : feature 维度，一般为 256
            H : height, 一般为 100
            W : width, 一般为 252

        sensor_dist: torch.Tensor, shape (B, H, W)
            distance between each block and the sensor
            
        time_stamps: torch.Tensor, shape (B)
            表示每个位置的时间信息

        valid_mask: (B, 1, H, W)
            为0/1,表示是否在该位置有效

        Returns:
        --------
        output_feature: torch.Tensor, shape (C, H, W)
            ego cav's fused feature at current frame.
        """
        if self.if_max_aggre:
            features = torch.max(features, dim=0)[0]
            if self.if_spatial_conv:
                features = features.unsqueeze(0)
                features = self.layers_spatial[0](features)
            return features.squeeze(0)
        if self.confidence_fetch == False:
            if self.sup_individual:
                curr_features = features[1::self.sweep_length+1].clone()    # [n, C, H, W]
                latency_features = features[2::self.sweep_length+1].clone()  # [n, C, H, W]
                if self.compensation==False:
                    all_latency_features = torch.cat((features[0].unsqueeze(0), latency_features), dim=0)
                    return all_latency_features, 0
                    
                    # final_fuse_feature = torch.max(torch.cat((features[:1],curr_features),dim=0),dim=0)[0]
                    # final_fuse_feature_latency = torch.max(torch.cat((features[:1],latency_features),dim=0),dim=0)[0]

                    # return final_fuse_feature_latency,final_fuse_feature,final_fuse_feature_latency,0,0


                cav_num, C, H, W = features.shape
                individual_recon_loss = 0
                latency_recon_loss = 0
                # debug use, set true to save pt for viz
                
                if if_save_pt:
                    torch.save(time_stamps, pt_save_path +'time_stamps'+'.pt')
                    torch.save(features[1:], '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'features'+'.pt')
                
                num_nonego = int((cav_num-1)/(self.sweep_length+1)) 
                nonego_features = features[1:].clone()                  # [nonego_cavs(1+k), C, H, W]
                
                nonego_splits = nonego_features.split(self.sweep_length+1,0) # ([1+k, C, H, W], [], ..., [])
                nonego_splits = torch.stack(nonego_splits,dim=0)        # [nonego_cars, 1+k, C, H, W]
                # torch.save(nonego_splits, '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'nonego_splits'+'.pt')
                nonego_history = nonego_splits[:,1:]                    # [nonego_cars, k, C, H, W]

                nonego_timestamps = time_stamps[1:]                     # [nonego_cavs(1+k), ]
                nonego_timestamps_splits = nonego_timestamps.split(self.sweep_length+1, 0)
                nonego_timestamps_splits = torch.stack(nonego_timestamps_splits,dim=0)   # [n, 1+k, ]
                nonego_timestamps_history = nonego_timestamps_splits[:, 1:]


                nonego_history = nonego_history.reshape(-1, C, H, W).permute(0,2,3,1)
                nonego_timestamps_history = nonego_timestamps_history.reshape(-1)
                # nonego_features = nonego_features.permute(0,2,3,1)      # 
                nonego_history = self.time_encoding(nonego_history,nonego_timestamps_history)
                nonego_history = nonego_history.permute(0,3,1,2).reshape(num_nonego, self.sweep_length, C, H, W)      # [nonego_cavs, k, C, H, W]

                # mask after TE, valid_mask: [B, H, W]
                non_ego_mask = valid_mask[1:]                           # [nonego_cavs*(k+1), 1, H, W]
                non_ego_mask_split = non_ego_mask.split(self.sweep_length+1,0) # ([1+k, 1, H, W], ... ) nonego_cavs 
                non_ego_mask = torch.stack(non_ego_mask_split, dim=0)   # nonego_cavs, 1+k, H, W
                if if_save_pt:
                    torch.save(non_ego_mask, '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'nonego_mask_splits'+'.pt')

                non_ego_mask_past = non_ego_mask[:, 1:]                 # nonego_cavs, k, H, W
                nonego_history = nonego_history * non_ego_mask_past        # [nonego_cavs, k, C, H, W]
                
                # nonego_splits = nonego_splits.permute(0,3,4,1,2).reshape(num_nonego*H*W, self.sweep_length, C)
                # time_attned_feature = self.layers_temporal[0](nonego_splits, nonego_splits, nonego_splits, None)
                # time_attned_feature = nonego_splits
                # time_attned_feature = time_attned_feature.reshape(num_nonego, H, W, self.sweep_length*C).permute(0,3,1,2) 
                
                hidden = self.swin_trans(nonego_history)                # [non_ego_cars, 1, C, H, W]
                # query = self.predict_query.repeat(num_nonego,1,1)
                # final_estimations = self.predict_attention(query, hidden, hidden, None).reshape(num_nonego, H, W, -1).permute(0,3,1,2)
                final_estimations = hidden.squeeze(1)                   # [non_ego_cars, C, H, W]
                if if_save_pt:
                    torch.save(final_estimations, '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'nonego_estimations'+'.pt')

                
                non_ego_mask_past_union = (torch.sum(non_ego_mask_past, dim=1)>0).to(torch.float32) # nonego_cavs, 1, H, W
                # torch.save(non_ego_mask_past_union,'/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ f'non_ego_mask_past_union_{self.dilation_size}' + '.pt')
                non_ego_mask_past_union = self.dilate_mask(non_ego_mask_past_union)
                if if_save_pt:
                    torch.save(non_ego_mask_past_union,'/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ f'non_ego_mask_past_union_dilated_{self.dilation_size}' + '.pt')

                final_estimations = final_estimations * non_ego_mask_past_union

                # debug use, uncomment to retrieve current mask
                ego_mask_curr = valid_mask[0]
                nonego_mask_curr = non_ego_mask[:, 0] # [nonego_cav, 1, H, W]
                # torch.save(nonego_mask_curr,'/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ f'nonego_mask_curr_{self.dilation_size}' + '.pt')

                individual_recon_loss_pos = torch.torch.nn.functional.smooth_l1_loss(final_estimations,nonego_splits[:,0],reduction='sum')/C
                if non_ego_mask_past_union.sum() != 0:
                    individual_recon_loss_pos /= (non_ego_mask_past_union.sum())

                # individual_recon_loss_neg = torch.torch.nn.functional.smooth_l1_loss((1-non_ego_mask_past_union)*final_estimations,nonego_splits[:,0],reduction='sum')
                # if (1-non_ego_mask_past_union).sum() != 0:
                    # individual_recon_loss_neg /= ((1-non_ego_mask_past_union).sum())

                # individual_recon_loss = individual_recon_loss_pos + individual_recon_loss_neg
                individual_recon_loss = individual_recon_loss_pos

                # latency_recon_loss = torch.torch.nn.functional.smooth_l1_loss(latency_features, nonego_splits[:,0],reduction='sum')/C
                # if non_ego_mask_past_union.sum() != 0:
                #     latency_recon_loss /= (non_ego_mask_past_union.sum())

                final_estimations = torch.cat((features[:1],final_estimations),dim=0)       # N, C, H, W
                # final_feature = torch.cat((features[:1],curr_features),dim=0)               # N, C, H, W
                # final_feature_latency = torch.cat((features[:1],latency_features),dim=0)    # N, C, H, W

                return final_estimations, individual_recon_loss #, latency_recon_loss

                # ###############skimage.dilation##########################################
                # final_fuse_estimation = torch.max(torch.cat((features[:1],final_estimations),dim=0), dim=0)[0]
                # final_fuse_feature = torch.max(torch.cat((features[:1],curr_features),dim=0),dim=0)[0]
                # # fuse_recon_loss = torch.torch.nn.functional.smooth_l1_loss(final_fuse_estimation,final_fuse_feature,reduction='mean')

                # final_fuse_feature_latency = torch.max(torch.cat((features[:1],latency_features),dim=0),dim=0)[0]
                # # ipdb.set_trace()
                # # print('#######################individual:',torch.torch.nn.functional.smooth_l1_loss(latest_features[1:],curr_features[1:]),individual_recon_loss)
                # # latest_features = features[self.sweep_length-1::self.sweep_length].clone()
                # # latest_features[0] = features[0]

                # # final1 = torch.max(torch.cat((features[:1],latest_features[1:]),dim=0), dim=0)[0]
                # # cnt = 5
                # # torch.save(curr_features,'/GPFS/data/yuxiwei/yuxiwei/OpenCOOD/vis/current'+str(cnt)+'.pt')
                # # torch.save(latest_features,'/GPFS/data/yuxiwei/yuxiwei/OpenCOOD/vis/latest'+str(cnt)+'.pt')
                # # cnt += 1
                # # ipdb.set_trace()
                # # print('#######################:fuse',torch.torch.nn.functional.smooth_l1_loss(final1,final_fuse_feature),fuse_recon_loss)
             
                # return final_fuse_estimation,final_fuse_feature,final_fuse_feature_latency, individual_recon_loss, latency_recon_loss # +fuse_recon_loss
                # # return final_fuse_feature, 0

            """ 
            num_cav = int(features.shape[0]/self.sweep_length)
            curr_features = features[:self.sweep_length:]
            features = torch.cat((features[:1],features[num_cav:]),dim=0)
            time_stamps = torch.cat((time_stamps[:1],time_stamps[num_cav:]),dim=0)
            cav_num, C, H, W = features.shape
            ################################
            # features = features[::2]
            # valid_mask = valid_mask[::2]
            # cav_num = int(cav_num/2)
            ################################


            if self.if_conv_aggre:
                curr = features[0:1].clone()
                prev = features[1:].clone()
                prev = self.conv_aggre(curr,prev)
                features = torch.cat((curr,prev),dim=0)
                features = features.permute(0,2,3,1)
                time_stamps = time_stamps*-1
                features = self.time_encoding(features,time_stamps)
                features = features.permute(0,3,1,2)
                if self.if_spatial_conv:
                    features = self.layers_spatial[0](features)
                return torch.max(features,dim=0)[0]
            
            if self.if_ego_time_encoding:
                # s1 = time.time()
                features = features.permute(0,2,3,1)
                time_stamps = time_stamps*-1
                features = self.time_encoding(features,time_stamps)
                features = features.permute(0,3,1,2)
                # s2 = time.time()
                # print('TE',s2-s1)
            
            
            features = features.view(cav_num, C, -1).permute(2, 0, 1) #  (H*W, cav_num, C), perform self attention on each pixel.
            # masks = valid_mask.view(cav_num, H*W).permute(1,0)
            # masks = masks.unsqueeze(-1).repeat(1,1,cav_num)
            


            if self.if_dotproductattn:
                features = self.dotproductattn(features, features, features, masks)
            if self.if_time_attn_aggre:
                # s1 = time.time()
                features = self.layers_temporal[0](features, features, features, None)
                # s2 = time.time()
                # print('trans',s2-s1)
            
           

            features = features.permute(1, 2, 0).view(cav_num, C, H, W)[:1] 
            if self.if_spatial_conv:
                # s1 = time.time()
                features = self.layers_spatial[0](features)[0]
                # s2 = time.time()
                # print('conv',s2-s1)
            if self.sup_individual:
                features = torch.max(final_estimations, dim=0)[0]
            curr_fuse_features = torch.max(curr_features, dim=0)[0]
            recon_loss = torch.torch.nn.functional.smooth_l1_loss(features,curr_fuse_features)

            return features, recon_loss
            """

        """
        valid_mask = valid_mask.squeeze(1)
        
        features = features.permute(0,2,3,1)
        
        features = torch.cat((features[0:1],features[self.sweep_length:]),dim=0)
        

        ##########################################  三种encoding
        if self.if_spatial_encoding:
            features[0:1] = self.spatial_encoding(features[0:1])# features.shape = [batch_all,height,width,feature_dim]
        if self.if_sensor_encoding and sensor_dist!=-1:
            sensor_dist = torch.cat((sensor_dist[0:1],sensor_dist[self.sweep_length:]))
            features = self.sensor_dist_encoding(features,sensor_dist)
        time_stamps = torch.cat((time_stamps[0:1],time_stamps[self.sweep_length:]),dim=0)
        valid_mask = torch.cat((valid_mask[0:1],valid_mask[self.sweep_length:]),dim=0)
        valid_mask_sum = valid_mask.sum(0)
        nonego_idx = torch.where(valid_mask_sum>1)
 
        features_to_aggre = features[:,nonego_idx[0],nonego_idx[1],:] 
        if self.if_ego_time_encoding:
            features[0:1] = self.time_encoding(features[0:1],time_stamps[0:1])

        if features_to_aggre.shape[1] != 0:
            if self.if_spatial_encoding:
                features_to_aggre = self.spatial_encoding_nonego(features_to_aggre,nonego_idx)
            if self.if_nonego_time_encoding:
                features_to_aggre = self.time_encoding(features_to_aggre,time_stamps)
        ###########################################

        ################################################### 
        # ipdb.set_trace()
        masks = valid_mask[:,nonego_idx[0],nonego_idx[1]]#(valid_num,len)
        batch_len = masks.shape[0]
        masks = masks.permute(1,0).unsqueeze(-1).repeat(1,1,batch_len)

        features_to_aggre = features_to_aggre.permute(1,0,2)
        for l in range(self.num_layers):
                
            if masks.shape[0] != 0:
                if self.if_dotproductattn:
                    # if self.use_mask == False:
                    #     masks =  None
                    # features_to_aggre = self.dotproductattn(features_to_aggre,features_to_aggre,features_to_aggre,masks)
                    features_to_aggre[:,0] = torch.max(features_to_aggre,dim=1)[0]
                else:
                    features_to_aggre = self.layers_temporal[l](features_to_aggre,features_to_aggre,features_to_aggre,masks)
                features[:,nonego_idx[0],nonego_idx[1],:] = features_to_aggre.permute(1,0,2)


            if self.if_spatial_conv:
                features = features.permute(0,3,1,2)
                features_backup = features.clone()
                features[0:1] = self.layers_spatial[l](features_backup[0:1])
                features = features.permute(0,2,3,1)

                features_to_aggre[:,0:1] = features[0:1,nonego_idx[0],nonego_idx[1],:].permute(1,0,2)

        return features[0].permute(2,0,1)
        """

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_swin_w_single(nn.Module):
    def __init__(self, args):
        super(raindrop_swin_w_single, self).__init__()
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.agg_mode == 'RAIN':
            d_ff = args['d_ff']
            h = args['h']
            dropout = args['dropout']
            input_dim = args['input_dim']
            num_layers = args['num_layers']
            neighbor_range = args['neighbor_range']
            max_len = args['max_len']
            height = 100
            width = 352
            self.sweep_length = args['fusion_setting']['sweep_length']
            self.fuse_layer1 = False
            if self.multi_scale:
                layer_nums = args['layer_nums']
                num_filters = args['num_filters']
                self.num_levels = len(layer_nums)
                if self.fuse_layer1:
                    self.fuse_modules = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len)
                    for p in self.fuse_modules.parameters():
                        if p.dim() > 1:
                            nn.init.xavier_uniform_(p)
                else:
                    self.fuse_modules = nn.ModuleList()
                    for idx in range(self.num_levels):
                        if idx==0:
                            fuse_network = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len, compensation=True)
                        else: # TODO: 换一下！
                            fuse_network = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len, compensation=False)
                        for p in fuse_network.parameters():
                            if p.dim() > 1:
                                nn.init.xavier_uniform_(p)

                        height = int(height/2)
                        width = int(width/2)
                        input_dim = input_dim * 2
                        d_ff = d_ff * 2
                        self.fuse_modules.append(fuse_network)
            else:
                input_dim = int(args['input_dim'])
                self.final_fuse = mask_fuse(input_dim)
                self.fuse_modules = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len, compensation=True)
                # This was important from their code. 
                # Initialize parameters with Glorot / fan_avg.
                for p in self.fuse_modules.parameters():
                    if p.dim() > 1:
                        nn.init.xavier_uniform_(p)

        else:
            if self.multi_scale:
                layer_nums = args['layer_nums']
                num_filters = args['num_filters']
                self.num_levels = len(layer_nums)
                self.fuse_modules = nn.ModuleList()
                for idx in range(self.num_levels):
                    if self.agg_mode == 'MAX':
                        fuse_network = MaxFusion()
                    else:
                        # self.agg_mode == 'ATTEN':
                        fuse_network = AttenFusion(num_filters[idx])
                    self.fuse_modules.append(fuse_network)
            else: 
                if self.agg_mode == 'MAX': # max fusion, debug use
                    self.fuse_modules = MaxFusion()

        print("##### Raindrop + Swin *with single supervision* Init! #####")
    
    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k-(k-1), dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        
        return split_x

    def get_t_mat_for_past_1(self, all_mat):
        """
        all_mat:        
            [N*(k+1), 2, 3]

        Return:
        -------
        mat1: [N*(k+1), 2, 3]
        mat2: [N, 2, 3]
        """
        last_row = torch.tensor([[0,0,1]]).repeat(all_mat.shape[0], 1, 1).to(all_mat)

        all_mat_with_last_row = torch.cat((all_mat,last_row), dim=1).to(torch.float32)
        nonego_past1_to_egocurr = all_mat_with_last_row[1::self.sweep_length+1]
        # print(nonego_past1_to_egocurr.shape)
        nonego_past1_to_egocurr = nonego_past1_to_egocurr.to(torch.float32)
        egocurr_to_nonego_past1 = torch.inverse(nonego_past1_to_egocurr)
        egocurr_to_nonego_past1_repeat = torch.repeat_interleave(egocurr_to_nonego_past1, self.sweep_length+1, 0)
        # print(egocurr_to_nonego_past1_repeat.shape, all_mat_with_last_row.shape)
        all_to_nonego_past1 = torch.bmm(egocurr_to_nonego_past1_repeat, all_mat_with_last_row)

        return all_to_nonego_past1[:,:2], nonego_past1_to_egocurr[:,:2]

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]
        # print('num',x.shape[0],pairwise_t_matrix.shape)
        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2
        
        if self.multi_scale:
            
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(x)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
            if self.fuse_layer1:
                all_recon_loss = 0
                layer1_feats = backbone.resnet.forward_layer1(x)  # [B, C, H, W]
                if self.communication:
                    batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                    batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                    _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                    communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                    # debug use, to save original feature:
                    # debug_feats = layer1_feats.
                    # TODO scale = 0 不加 mask
                    layer1_feats = layer1_feats * communication_masks_tensor

                batch_node_features = self.regroup(layer1_feats, record_len, K)
                x_fuse = []
                x_fuse_curr = []
                x_fuse_latency = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    # print('node and t mat',node_features.shape,t_matrix.shape,record_len)
                    t_matrix = t_matrix[K-1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))
                    
                    # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                    record_frames = np.ones((N))*K; record_frames[0]-=(K-1)
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        # x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        features,features_curr,features_latency, recon_loss = self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b])
                        all_recon_loss += recon_loss # TODO: 是否需要添加 /= B ? 
                        x_fuse.append(features)
                        x_fuse_curr.append(features_curr)
                        x_fuse_latency.append(features_latency)
                x_fuse = torch.stack(x_fuse)  ###(batch, channel, h, w)
                x_fuse_curr = torch.stack(x_fuse_curr)
                x_fuse_latency = torch.stack(x_fuse_latency)

                layer2_feat, layer3_feat = backbone.resnet.forward_from_layer1(x_fuse)
                ups = []
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[0](x_fuse))
                    ups.append(backbone.deblocks[1](layer2_feat))
                    ups.append(backbone.deblocks[2](layer3_feat))
                
                if len(ups) > 1:
                    x_fuse = torch.cat(ups, dim=1)
                elif len(ups) == 1:
                    x_fuse = ups[0]
                
                if len(backbone.deblocks) > self.num_levels:
                    x_fuse = backbone.deblocks[-1](x_fuse)

                
                layer2_feat_curr, layer3_feat_curr = backbone.resnet.forward_from_layer1(x_fuse_curr)
                ups = []
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[0](x_fuse_curr))
                    ups.append(backbone.deblocks[1](layer2_feat_curr))
                    ups.append(backbone.deblocks[2](layer3_feat_curr))
                    
                if len(ups) > 1:
                    x_fuse_curr = torch.cat(ups, dim=1)
                elif len(ups) == 1:
                    x_fuse_curr = ups[0]
                
                if len(backbone.deblocks) > self.num_levels:
                    x_fuse_curr = backbone.deblocks[-1](x_fuse_curr)

                
                layer2_feat_latency, layer3_feat_latency = backbone.resnet.forward_from_layer1(x_fuse_latency)
                ups = []
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[0](x_fuse_latency))
                    ups.append(backbone.deblocks[1](layer2_feat_latency))
                    ups.append(backbone.deblocks[2](layer3_feat_latency))
                    
                if len(ups) > 1:
                    x_fuse_latency = torch.cat(ups, dim=1)
                elif len(ups) == 1:
                    x_fuse_latency = ups[0]
                
                if len(backbone.deblocks) > self.num_levels:
                    x_fuse_latency = backbone.deblocks[-1](x_fuse_latency)

            else:
                ups = []
                ups_curr = []
                ups_latency = []
                ups_single = []
                all_recon_loss = 0
                all_latency_recon_loss = 0
                for i in range(self.num_levels):
                    x = feats[i] if with_resnet else backbone.blocks[i](x)
                    ############ 1. Communication (Mask the features) #########
                    if i==0:
                        if self.communication:
                            batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                            batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                            _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                            communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                            # for debug use
                            if if_save_pt:
                                x_wo_mask = x.clone()
                            x = x * communication_masks_tensor
                        else:
                            communication_rates = torch.tensor(0).to(x.device)
                    else:
                        if self.communication:
                            communication_masks_tensor = F.max_pool2d(communication_masks_tensor, kernel_size=2)
                            # TODO scale = 1, 2 不加 mask
                            # x = x * communication_masks_tensor
                    
                    ############ 2. Split the confidence map #######################
                    # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                    # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                    # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                    # print('x',x.shape)
                    batch_node_features = self.regroup(x, record_len, K)
                    batch_comm_masks = self.regroup(communication_masks_tensor, record_len, K)
                    
                    if if_save_pt:
                        batch_x_wo_mask = self.regroup(x_wo_mask, record_len, K)
                    
                    ############ 3. Fusion ####################################
                    x_fuse = []
                    x_single = []
                    for b in range(B):
                        # number of valid agent
                        N = record_len[b]
                        # t_matrix[i, j]-> from i to j
                        t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) # [Nx(k+1), 2, 3]

                        mat_to_past1, mat_to_ego_curr = self.get_t_mat_for_past_1(t_matrix)
                        mat_to_past1 = mat_to_past1[K-1:]
                        node_features = batch_node_features[b]              # (N-1)xk+1, C, H, W
                        node_masks = batch_comm_masks[b]                    # (N-1)xk+1, 1, H, W
                        C, H, W = node_features.shape[1:]
                        if if_save_pt:
                            torch.save(batch_x_wo_mask[b], pt_save_path + 'feature_on_cav_past_0.pt')
                            torch.save(t_matrix[K-1:], pt_save_path + 't_matrix.pt')
                            torch.save(mat_to_past1, pt_save_path + 'mat_to_past1.pt')
                            torch.save(mat_to_ego_curr, pt_save_path + 't_matrmat_to_ego_curr.pt')
                            feature_on_ego_curr = warp_affine_simple(batch_x_wo_mask[b], t_matrix[K-1:], (H, W))
                            torch.save(feature_on_ego_curr, pt_save_path + 'feature_on_ego_curr.pt')
                        if self.agg_mode == 'RAIN':
                            # for sensor embedding
                            sensor_dist = -1# (B, H, W)
                            node_features = warp_affine_simple(node_features,
                                                        mat_to_past1,
                                                        (H, W))
                            node_masks = warp_affine_simple(node_masks,
                                                        mat_to_past1,
                                                        (H, W))
                            features, recon_loss = self.fuse_modules[i](node_features, sensor_dist, batch_time_intervals[b], node_masks)
                            x_single.append(features[1:])   # B x [n*, C, H, W]
                            all_recon_loss += recon_loss
                        else: # ATTEN, MAX, Transformer
                            x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        
                        # print('node and t mat',node_features.shape,t_matrix.shape,record_len)
                        # t_matrix = t_matrix[K-1:]

                        neighbor_comp_feature = warp_affine_simple(features,
                                                        mat_to_ego_curr,
                                                        (H, W))
                        # neighbor_mask = warp_affine_simple(node_masks, mat_to_ego_curr, (H, W))

                        if if_save_pt:
                            torch.save(neighbor_comp_feature, pt_save_path+'feature_com'+'.pt')
                        # max fusion
                        neighbor_fused_feature = torch.max(neighbor_comp_feature, dim=0)[0]  # [C, H, W]
                        x_fuse.append(neighbor_fused_feature)

                        if if_save_pt:
                            neighbor_wo_mask = warp_affine_simple(batch_x_wo_mask[b], t_matrix, (H, W))
                            torch.save(neighbor_wo_mask[1:], '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'feature_wo_mask'+'.pt')

                        # TODO: 这个位置
                        record_frames = np.ones((N))*K
                        # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                        
                        
                    x_single = torch.cat(x_single, dim=0) # [BN, C, H, W]
                    x_fuse = torch.stack(x_fuse)
                    ############ 4. Deconv ####################################
                    if len(backbone.deblocks) > 0:
                        ups.append(backbone.deblocks[i](x_fuse))
                        ups_single.append(backbone.deblocks[i](x_single))

                    else:
                        ups.append(x_fuse)
                    
                if len(ups) > 1:
                    x_fuse = torch.cat(ups, dim=1)
                elif len(ups) == 1:
                    x_fuse = ups[0]

                if len(ups_single) > 1:
                    x_single = torch.cat(ups_single, dim=1)
                elif len(ups_single) == 1:
                    x_single = ups_single[0]
                
                if len(backbone.deblocks) > self.num_levels:
                    x_fuse = backbone.deblocks[-1](x_fuse)
                
        else:
            ############ 2. Communication (Mask the features) #########
            if self.communication:
                batch_confidence_maps = self.regroup(rm, record_len, K)
                # _, (B, 1, H, W), float
                _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                x = x * communication_masks_tensor
            else:
                communication_rates = torch.tensor(0).to(x.device)

            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)
            batch_comm_masks = self.regroup(communication_masks_tensor, record_len, K)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            x_single = []
            all_recon_loss = 0.0
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                mat_to_past1, mat_to_ego_curr = self.get_t_mat_for_past_1(t_matrix)
                mat_to_past1 = mat_to_past1[K-1:]
                node_features = batch_node_features[b]
                node_masks = batch_comm_masks[b]                    # (N-1)xk+1, 1, H, W
                C, H, W = node_features.shape[1:]
                if if_save_pt:
                    torch.save(node_features, pt_save_path + 'feature_on_cav_past_0.pt')
                    torch.save(t_matrix[K-1:], pt_save_path + 't_matrix.pt')
                    torch.save(mat_to_past1, pt_save_path + 'mat_to_past1.pt')
                    torch.save(mat_to_ego_curr, pt_save_path + 't_matrmat_to_ego_curr.pt')
                    feature_on_ego_curr = warp_affine_simple(node_features, t_matrix[K-1:], (H, W))
                    torch.save(feature_on_ego_curr, pt_save_path + 'feature_on_ego_curr.pt')
                if self.agg_mode == 'RAIN':
                    sensor_dist = -1# (B, H, W)
                    node_features = warp_affine_simple(node_features,
                                                    mat_to_past1,
                                                    (H, W))
                    node_masks = warp_affine_simple(node_masks,
                                                        mat_to_past1,
                                                        (H, W))
                    features, recon_loss = self.fuse_modules(node_features, sensor_dist, batch_time_intervals[b], node_masks)
                    x_single.append(features[1:])   # B x [n*, C, H, W]
                    all_recon_loss += recon_loss
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
                
                neighbor_comp_feature = warp_affine_simple(features,
                                                        mat_to_ego_curr,
                                                        (H, W))

                # max fusion # TODO: max fusion 或 其他融合方法
                # neighbor_fused_feature = torch.max(neighbor_comp_feature, dim=0)[0]  # [C, H, W]
                neighbor_fused_feature = self.final_fuse(neighbor_comp_feature[0].unsqueeze(0), neighbor_comp_feature[1:])
                x_fuse.append(neighbor_fused_feature)
            x_fuse = torch.stack(x_fuse)
            x_single = torch.cat(x_single, dim=0) # [BN, C, H, W]
        
        return x_fuse, x_single, communication_rates, all_recon_loss, {}
opencood/models/fuse_modules/raindrop_swin_w_single.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# Date: 2022/12/06
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

import torch.nn.functional as F
from turtle import update
# import ipdb
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
from opencood.models.comm_modules.where2comm_multisweep import Communication
from opencood.models.fuse_modules.SyncNet import SyncLSTM

if_save_pt = False  # TODO: for debug use, save pt file

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class PositionalEncoding_spatial(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model_feature, d_model_hidden, dropout, max_len=300,mode='cat'):
        super(PositionalEncoding_spatial, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        d_model = int(d_model_hidden/2)
        # Compute the positional encodings once in log space.
        pe = torch.zeros((max_len, d_model))
        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(math.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe_ = pe.unsqueeze(0)#batch,len,feature
        pe_h = pe_.unsqueeze(2).repeat(1,1,max_len,1)
        pe_w = pe_.unsqueeze(1).repeat(1,max_len,1,1)
        pe = torch.cat((pe_h,pe_w),dim=-1)

        self.register_buffer('pe', pe)

        self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x):
        if self.mode == 'cat':
            tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
            tmp = tmp.repeat(x.shape[0], 1, 1, 1)
            x = self.linear(torch.cat((x, tmp),dim=-1))

        elif self.mode == 'init':
            x = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)

        else:
            tmp = Variable(self.pe[:,:x.shape[1],:x.shape[2]],requires_grad=False)
            tmp = tmp.repeat(x.shape[0], 1, 1, 1)
            x = x + tmp

        return self.dropout(x)

class PositionalEncoding_spatial_custom(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model_feature, d_model_hidden, dropout,mode='cat'):
        super(PositionalEncoding_spatial_custom, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = int(d_model_hidden/2)
        # Compute the positional encodings once in log space.
        
        self.linear = nn.Linear(d_model_feature + d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x, nonego_idx):#x(batch,num,feature)
   
        pe_h = torch.zeros((x.shape[1], self.d_model)).to(x.device)
        pe_w = torch.zeros((x.shape[1], self.d_model)).to(x.device)
        position_h, position_w = nonego_idx[0],nonego_idx[1]
        position_h = position_h.unsqueeze(1)
        position_w = position_w.unsqueeze(1)

        div_term = torch.exp(torch.arange(0, self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(position_h.device)
        pe_h[:, 0::2] = torch.sin(position_h * div_term)
        pe_h[:, 1::2] = torch.cos(position_h * div_term)

        pe_w[:, 0::2] = torch.sin(position_w * div_term)
        pe_w[:, 1::2] = torch.cos(position_w * div_term)
       
        pe = torch.cat((pe_h,pe_w),dim=-1).unsqueeze(0).repeat(x.shape[0],1,1)

        if self.mode == 'cat':
            x = self.linear(torch.cat((x, pe),dim=-1))

        else:
            x = x + Variable(pe,requires_grad=False)

        return self.dropout(x)


class PositionalEncoding_sensor_dist(nn.Module):

    def __init__(self, d_model, dropout):
        super(PositionalEncoding_sensor_dist, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = d_model
        self.linear = nn.Linear(2*d_model,d_model)
        self.mode = 'cat'
         
    def forward(self, x, sensor_dist):#x.shape=[batch,H,W,F],senfor_dist.shape=[batch,H,W]

        x = x.reshape(-1,x.shape[-1])
        sensor_dist = sensor_dist.reshape(-1)

        pe = torch.zeros(x.shape[0], self.d_model).to(x.device)
        position = sensor_dist.unsqueeze(1) 
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(x.device)
         
        pe[:, 0::2] = torch.sin(position * div_term)  
        pe[:, 1::2] = torch.cos(position * div_term)   

        if self.mode == 'cat':
            x = self.linear(torch.cat((x,Variable(pe,require_grad=False)),dim=-1))

        else:
            x = x + Variable(pe,requires_grad=False)
        
        return self.dropout(x)

class PositionalEncoding_irregular_time(nn.Module):

    def __init__(self, d_model_feature, d_model_hidden, dropout):
        super(PositionalEncoding_irregular_time, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.d_model = d_model_feature
        self.linear = nn.Linear(d_model_feature+d_model_hidden, d_model_feature)
        self.mode = 'plus'

    def forward(self, x, time):#x.shape=[batch,height,width,feature],time.shape=[batch]

        pe = torch.zeros(x.shape[0],self.d_model).to(x.device)
        position = time.unsqueeze(1) 
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).to(x.device)
         
        pe[:, 0::2] = torch.sin(position * div_term)  
        pe[:, 1::2] = torch.cos(position * div_term)   

        if self.mode == 'cat':
            if len(x.shape) == 4:
                pe = pe.unsqueeze(1).unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
            else:
                pe = pe.unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], 1)
            x = self.linear(torch.cat((x, tmp),dim=-1))

        else:
            if len(x.shape) == 4:
                pe = pe.unsqueeze(1).unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], x.shape[2], 1)
            else:
                pe = pe.unsqueeze(1)
                tmp = Variable(pe, requires_grad=False)
                tmp = tmp.repeat(1, x.shape[1], 1)
            x = x + tmp

        return self.dropout(x)

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    
    # ipdb.set_trace()
    
    # print('score',scores.shape,'mask',mask.shape)

    if mask is not None:
        # mask = mask.transpose(1,2).to(query.device)
        mask = mask.to(query.device)
        scores = scores.masked_fill(mask == 0, -1e10) 

    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)

    return torch.matmul(p_attn, value), p_attn


class MultiHeadedAttention(nn.Module):
    def __init__(self, h, q_dim, k_dim,v_dim, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h  
        self.h = h 
        self.q_dim = q_dim
        self.k_dim = k_dim
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.linear_q = nn.Linear(q_dim,d_model)
        self.linear_k = nn.Linear(k_dim,d_model)
        self.linear_v = nn.Linear(v_dim,d_model)
        self.linear_out = nn.Linear(d_model,q_dim)
        self.attn = None  
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        # 1) Do all the linear projections in batch from d_model => h x d_k 

        query = self.linear_q(query).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
        key = self.linear_k(key).view(nbatches,-1,self.h,self.d_k).transpose(1,2)
        value = self.linear_v(value).view(nbatches,-1,self.h,self.d_k).transpose(1,2)

        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(query, key, value, mask=mask,
                                 dropout=self.dropout)

        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(1, 2).contiguous() \
            .view(nbatches ,-1, self.h * self.d_k)

        return self.linear_out(x)

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
 
    def forward(self, x):
        "Norm"
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, x, sublayer):
        # print('sub',x.shape,self.dropout(sublayer(self.norm(x))).shape)
        return x + self.dropout(sublayer(self.norm(x))) #残差连接

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model,d_out, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_out)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

class EncoderLayer_temporal(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer_temporal, self).__init__()
        self.self_attn = self_attn #sublayer 1 
        self.feed_forward = feed_forward #sublayer 2 
        self.sublayer = clones(SublayerConnection(size, dropout), 2) #拷贝两个SublayerConnection，一个为了attention，一个是独自作为简单的神经网络
        self.size = size

    def forward(self, q, k, v, mask):
        "Follow Figure 1 (left) for connections."
        # print('x',x.shape,'attn',self.self_attn(x,x,x,mask).shape)
        x = self.sublayer[0](q, lambda q: self.self_attn(q, k, v, mask)) #attention层 对应层1
        return self.sublayer[1](x, self.feed_forward) # 对应 层2

class Spatial_conv(nn.Module):
    def __init__(self, input_channels, output_channels, kernel_size):
        super(Spatial_conv,self).__init__()
        self.conv = nn.Conv2d(in_channels = input_channels,
                              out_channels = output_channels,
                              kernel_size = (kernel_size,kernel_size),
                              stride = (1,1),
                              padding = (int((kernel_size-1)/2),int((kernel_size-1)/2))
        )

        self.relu = nn.ReLU()
        # self.bn = nn.BatchNorm2d(num_features=output_channels)

    def forward(self, x):
        init_x = x.clone()
        x = self.conv(x)
        x = self.relu(x)
        # x = self.bn(x)
        return x+init_x

class SpatialAttention_mtf(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpatialAttention_mtf, self).__init__()

        self.conv1 = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, curr, prev):#####(batch, channel, h, w)
        avg_out = torch.mean(curr, dim=1, keepdim=True)
        max_out, _ = torch.max(curr, dim=1, keepdim=True)
        y = torch.cat([avg_out, max_out], dim=1)
        y = self.conv1(y)
        return self.sigmoid(y) * prev ####(batch-1, channel, h, w)

class ScaledDotProductAttention(nn.Module):
    """
    Scaled Dot-Product Attention proposed in "Attention Is All You Need"
    Compute the dot products of the query with all keys, divide each by sqrt(dim),
    and apply a softmax function to obtain the weights on the values
    Args: dim, mask
        dim (int): dimention of attention
        mask (torch.Tensor): tensor containing indices to be masked
    Inputs: query, key, value, mask
        - **query** (batch, q_len, d_model): tensor containing projection
          vector for decoder.
        - **key** (batch, k_len, d_model): tensor containing projection
          vector for encoder.
        - **value** (batch, v_len, d_model): tensor containing features of the
          encoded input sequence.
        - **mask** (-): tensor containing indices to be masked
    Returns: context, attn
        - **context**: tensor containing the context vector from
          attention mechanism.
        - **attn**: tensor containing the attention (alignment) from the
          encoder outputs.
    """

    def __init__(self, dim):
        super(ScaledDotProductAttention, self).__init__()
        self.sqrt_dim = np.sqrt(dim)

    def forward(self, query, key, value, mask):
        score = torch.bmm(query, key.transpose(1, 2)) / self.sqrt_dim
        if mask != None:
            score = score.masked_fill(mask == 0, -1e10) 
        attn = F.softmax(score, -1)
        context = torch.bmm(attn, value)
        return context

class AttentionFusion(nn.Module):
    def __init__(self, settings, input_dim, hidden_dim, num_heads, num_layers, neighbor_range, height, width, dropout, max_len=300,compensation=True):
        super(AttentionFusion, self).__init__()
        self.sweep_length = settings['sweep_length']
        self.if_max_aggre = settings['if_max_aggre']
        self.if_spatial_encoding = settings['if_spatial_encoding']
        self.if_ego_time_encoding = settings['if_ego_time_encoding']
        self.if_nonego_time_encoding = settings['if_nonego_time_encoding']
        self.if_sensor_encoding = settings['if_sensor_encoding']
        self.if_time_attn_aggre = settings['if_time_attn_aggre']
        self.if_spatial_conv = settings['if_spatial_conv']
        self.if_dotproductattn = settings['if_dotproductattn']
        self.use_mask = settings['use_mask']
        self.confidence_fetch = settings['confidence_fetch']
        self.if_conv_aggre = settings['if_conv_aggre']
        self.sup_individual = settings['sup_individual']

        self.input_dim = input_dim 
        self.hidden_dim = hidden_dim
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.neighbor_range = neighbor_range
        self.height = height
        self.width = width
        self.compensation=compensation

        self.syncnet = SyncLSTM(channel_size=self.input_dim,
                                h=self.height,
                                w=self.width, 
                                k=self.sweep_length,
                                TM_Flag=False)            
        
        print("================ Rain attention Init ================")


    

    def forward(self, features, sensor_dist):
        """
        Parameters:
        -----------
        features: torch.Tensor, shape (B, C, H, W)
            neighbor features in ego view
            B : 按照顺序排列: ego[curr], cav_1[curr], cav_1[past_1], ..., cav_2[past_k], ....
            C : feature 维度，一般为 256
            H : height, 一般为 100
            W : width, 一般为 252

        sensor_dist: torch.Tensor, shape (B, H, W)
            distance between each block and the sensor

        valid_mask: (B, 1, H, W)
            为0/1,表示是否在该位置有效

        Returns:
        --------
        output_feature: torch.Tensor, shape (C, H, W)
            ego cav's fused feature at current frame.
        """
        if self.confidence_fetch == False:
            if self.sup_individual:
                curr_features = features[1::self.sweep_length+1].clone()    # [n, C, H, W]
                latency_features = features[2::self.sweep_length+1].clone()  # [n, C, H, W]
                if self.compensation==False:
                    
                    # final_fuse_feature = torch.max(torch.cat((features[:1],curr_features),dim=0),dim=0)[0]
                    final_fuse_feature_latency = torch.max(torch.cat((features[:1],latency_features),dim=0),dim=0)[0]

                    return final_fuse_feature_latency, 0

                cav_num, C, H, W = features.shape
                individual_recon_loss = 0
                latency_recon_loss = 0                
                if if_save_pt:
                    torch.save(features[1:], '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'features'+'.pt')
                
                num_nonego = int((cav_num-1)/(self.sweep_length+1)) 
                # print('###################',num_nonego)
                nonego_features = features[1:].clone()                  # [nonego_cavs(1+k), C, H, W]
                
                nonego_splits = nonego_features.split(self.sweep_length+1,0) # ([1+k, C, H, W], [], ..., [])
                nonego_splits = torch.stack(nonego_splits,dim=0)        # [nonego_cars, 1+k, C, H, W]
                if if_save_pt:
                    torch.save(nonego_splits, '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'nonego_splits'+'.pt')
                nonego_history = nonego_splits[:,1:]                    # [nonego_cars, k, C, H, W]
                
                nonego_history = torch.flip(nonego_history,dims=[1])
                hidden = self.syncnet(nonego_history,[1])
                non_ego_mask_past_union = torch.ones((num_nonego, C, H, W)).to(hidden)
                final_estimations = hidden.squeeze(1)                   # [non_ego_cars, C, H, W]
                if if_save_pt:
                    torch.save(final_estimations, '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'nonego_estimations'+'.pt')

                individual_recon_loss_pos = torch.torch.nn.functional.smooth_l1_loss(final_estimations,nonego_splits[:,0],reduction='sum')

                individual_recon_loss_pos /= final_estimations.shape[0]
                
                # if non_ego_mask_past_union.sum() != 0:
                #     individual_recon_loss_pos /= (non_ego_mask_past_union.sum())

                # individual_recon_loss = individual_recon_loss_pos + individual_recon_loss_neg
                individual_recon_loss = individual_recon_loss_pos

                latency_recon_loss = torch.torch.nn.functional.smooth_l1_loss(latency_features, nonego_splits[:,0],reduction='sum')/C
                if non_ego_mask_past_union.sum() != 0:
                    latency_recon_loss /= (non_ego_mask_past_union.sum())

                final_fuse_estimation = torch.max(torch.cat((features[:1],final_estimations),dim=0), dim=0)[0]
                final_fuse_feature = torch.max(torch.cat((features[:1],curr_features),dim=0),dim=0)[0]
                final_fuse_feature_latency = torch.max(torch.cat((features[:1],latency_features),dim=0),dim=0)[0]
             
                return final_fuse_estimation, individual_recon_loss # +fuse_recon_loss
                # return final_fuse_feature, 0


class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_syncnet(nn.Module):
    def __init__(self, args):
        super(raindrop_syncnet, self).__init__()
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.agg_mode == 'RAIN':
            d_ff = args['d_ff']
            h = args['h']
            dropout = args['dropout']
            input_dim = args['input_dim']
            num_layers = args['num_layers']
            neighbor_range = args['neighbor_range']
            max_len = args['max_len']
            height = 100
            width = 352

            self.if_syncnet = True

            self.fuse_layer1 = False
            if self.multi_scale:
                layer_nums = args['layer_nums']
                num_filters = args['num_filters']
                self.num_levels = len(layer_nums)
                self.fuse_modules = nn.ModuleList()
                for idx in range(self.num_levels):
                    if idx==0:
                        fuse_network = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len, compensation=True)
                    else:
                        fuse_network = AttentionFusion(args['fusion_setting'], input_dim, d_ff, h, num_layers, neighbor_range, height, width, dropout, max_len, compensation=False)
                    for p in fuse_network.parameters():
                        if p.dim() > 1:
                            nn.init.xavier_uniform_(p)

                    height = int(height/2)
                    width = int(width/2)
                    input_dim = input_dim * 2
                    d_ff = d_ff * 2
                    self.fuse_modules.append(fuse_network)
            else:
                self.fuse_modules = AttentionFusion(input_dim, d_ff, h, num_layers, neighbor_range, 100, 252, dropout, max_len)
                # This was important from their code. 
                # Initialize parameters with Glorot / fan_avg.
                for p in self.fuse_modules.parameters():
                    if p.dim() > 1:
                        nn.init.xavier_uniform_(p)

        else:
            if self.multi_scale:
                layer_nums = args['layer_nums']
                num_filters = args['num_filters']
                self.num_levels = len(layer_nums)
                self.fuse_modules = nn.ModuleList()
                for idx in range(self.num_levels):
                    if self.agg_mode == 'MAX':
                        fuse_network = MaxFusion()
                    else:
                        # self.agg_mode == 'ATTEN':
                        fuse_network = AttenFusion(num_filters[idx])
                    self.fuse_modules.append(fuse_network)
            else: 
                if self.agg_mode == 'MAX': # max fusion, debug use
                    self.fuse_modules = MaxFusion()

        print("##### Raindrop + ** SyncNet ** w.o. single supervision Init! #####")

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k-(k-1), dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        
        return split_x

    def forward(self, x, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]
        # print('num',x.shape[0],pairwise_t_matrix.shape)
        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2
        
        if self.multi_scale:
            
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(x)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))

            ups = []
            all_recon_loss = 0
            batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
            for i in range(self.num_levels):
                x = feats[i] if with_resnet else backbone.blocks[i](x)
                ############ 2. Split the confidence map #######################
                # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                # print('x',x.shape)
                batch_node_features = self.regroup(x, record_len, K)
                
                ############ 3. Fusion ####################################
                x_fuse = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) # [Nx(k+1), 2, 3]
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    # print('node and t mat',node_features.shape,t_matrix.shape,record_len)
                    t_matrix = t_matrix[K-1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))

                    # TODO: 这个位置
                    record_frames = np.ones((N))*K
                    
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        features, recon_loss = self.fuse_modules[i](neighbor_feature, sensor_dist)
                        if if_save_pt:
                            torch.save(batch_time_intervals[b], '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/'+ 'timestamps'+'.pt')
                        x_fuse.append(features)
                        all_recon_loss += recon_loss

                    else: # ATTEN, MAX, Transformer
                        x_fuse.append(self.fuse_modules[i](neighbor_feature))

                x_fuse = torch.stack(x_fuse)

                ############ 4. Deconv ####################################
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[i](x_fuse))
                else:
                    ups.append(x_fuse)
                
            if len(ups) > 1:
                x_fuse = torch.cat(ups, dim=1)
            elif len(ups) == 1:
                x_fuse = ups[0]
            
            if len(backbone.deblocks) > self.num_levels:
                x_fuse = backbone.deblocks[-1](x_fuse)
                
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                # _, (B, 1, H, W), float
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix,
                                                (H, W))
                record_frames = np.ones((N))*K
                # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                if self.agg_mode == "RAIN":
                    # for sensor embedding
                    sensor_dist = -1# (B, H, W)
                    x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)
        
        return x_fuse, all_recon_loss, {}
opencood/models/fuse_modules/raindrop_syncnet.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# Date: 2023/04/13
# for [ where2comm + syncnet ]
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
from opencood.models.sub_modules.convgru import ConvGRU
from icecream import ic
from matplotlib import pyplot as plt
# from opencood.models.sub_modules.SyncLSTM import SyncLSTM
from opencood.models.sub_modules.MotionNet import STPN, MotionPrediction, StateEstimation, FlowUncPrediction
# from opencood.models.sub_modules.dcn_net import DCNNet
from opencood.models.comm_modules.where2comm_multisweep import Communication
from opencood.models.fuse_modules.SyncNet import SyncLSTM

# def get_warped_feature_mask(flow, updated_grid):
#     ''' 
#     get the mask of the warped feature map, where the value is 1 means the location is valid, and 0 means the location is invalid.
#     ----------
#     Args:
#         flow: (2, H, W)
#         updated_grid: (2, H, W)

#     Returns:
#         mask: (H, W)
#     '''
#     def get_large_flow_indices(flow, threshold=0):
#         '''
#         find the indices of the flow map where the flow is larger than the threshold, which means these locations have been moved.
#         ----------
#         Args:
#             flow: (2, H, W)
            
#         '''
#         max_values, max_indices = torch.max(torch.abs(flow[:2]), dim=0)
#         large_indices = torch.nonzero(max_values > threshold, as_tuple=False)
#         return large_indices

#     def remove_duplicate_points(points):
#         unique_points, inverse_indices = torch.unique(points, sorted=True, return_inverse=True, dim=0)
#         return unique_points, inverse_indices
    
#     def get_nonzero_idx(flow, idx):
#         flow_values = flow[:, idx[:, 0], idx[:, 1]]
#         nonzero_idx = torch.nonzero(torch.abs(flow_values).sum(dim=0) == 0, as_tuple=False).squeeze()
#         return idx[nonzero_idx]
    
#     flow_idx = get_large_flow_indices(flow)

#     mask = torch.ones(flow.shape[-2], flow.shape[-1])
#     if flow_idx.shape[0] == 0:
#         return mask

#     # print(flow_idx)
#     updated_grid_points_tmp = updated_grid[:, flow_idx[:,0], flow_idx[:,1]].to(torch.int64).T
#     # change the order of dim 1
#     updated_grid_points = torch.zeros_like(updated_grid_points_tmp)
#     updated_grid_points[:, 0] = updated_grid_points_tmp[:, 1]
#     updated_grid_points[:, 1] = updated_grid_points_tmp[:, 0]
#     # print(updated_grid_points)

#     unique_points_idx, _ = remove_duplicate_points(updated_grid_points)
#     # print(unique_points_idx)

#     nonzero_idx = get_nonzero_idx(flow, unique_points_idx)
#     # print(nonzero_idx)
    
#     # TODO: mask out the outlier idx

#     if len(nonzero_idx.shape) > 1:
#         mask[nonzero_idx[:, 0], nonzero_idx[:, 1]] = 0
#     else: 
#         mask[nonzero_idx[0], nonzero_idx[1]] = 0

#     return mask

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_fuse(nn.Module):
    def __init__(self, args):
        super(raindrop_fuse, self).__init__()
        
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.multi_scale:
            layer_nums = args['layer_nums']
            num_filters = args['num_filters']
            self.num_levels = len(layer_nums)
            self.fuse_modules = nn.ModuleList()
            for idx in range(self.num_levels):
                if self.agg_mode == 'MAX':
                    fuse_network = MaxFusion()
                else:
                    # self.agg_mode == 'ATTEN':
                    fuse_network = AttenFusion(num_filters[idx])
                self.fuse_modules.append(fuse_network)
        else: 
            if self.agg_mode == 'MAX': # max fusion, debug use
                self.fuse_modules = MaxFusion()

        self.syncnet = SyncLSTM(channel_size = 64, h = 200, w = 704, k = 2, TM_Flag = False)
        print('*** channel_size = 64, h = 200, w = 704, k = 2, TM_Flag = False ***')

        # self.stpn = STPN(args['channel_size'])
        # self.motion_pred = MotionPrediction(seq_len=1)
        # self.state_classify = StateEstimation(motion_category_num=1)

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k-(k-1), dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        
        return split_x

    def generate_estimated_feats(self, feats, pairwise_t_matrix, record_len):
        '''
        update the feature by SyncNet
        Note:
            1. ego feature does not need to be updated

        params:
            feats: (sum(B,N,K), C, H, W)
            pairwise_t_matrix: (B, L, K, 2, 3)
            record_len: (B)

        return:
            updated feature: [N, C, H, W]
        '''
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        batch_node_features = self.regroup(feats, record_len, k=K)

        updated_features_list = []
        flow_list = []
        all_recon_loss = 0
        for b in range(B):
            all_feats = batch_node_features[b]
            non_ego_feats = all_feats[1:].clone()
            non_ego_feats = non_ego_feats.reshape(-1, K, C, H, W)
            estimated_non_ego_feats = self.syncnet(non_ego_feats[:, 1:], [1]) # [N-1, 1, C, H, W]
            # estimated_non_ego_feats = non_ego_feats[:, 1:2]
            updated_features_list.append(torch.cat([all_feats[0].unsqueeze(0), estimated_non_ego_feats[:,0]], dim=0))
            recon_loss = F.smooth_l1_loss(estimated_non_ego_feats, non_ego_feats[:,:1])
            all_recon_loss += recon_loss

        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)
        
        return updated_features_all, all_recon_loss

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        # 2. feature compensation with flow
        # 2.1 generate flow, 在同一个坐标系内，计算每个cav的flow
        # 2.2 compensation
        # x: (BxNxK, C, H, W) -> (BxN, C, H, W)
        updated_features, all_recon_loss = self.generate_estimated_feats(x, pairwise_t_matrix, record_len) # (BxN, C, H, W)
        
        # 3. feature fusion
        if self.multi_scale:
            ups = []
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(updated_features)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
            for i in range(self.num_levels):
                x = feats[i] if with_resnet else backbone.blocks[i](x)  # (BxN, C', H, W)

                ############ 1. Communication (Mask the features) #########
                if i==0:
                    if self.communication:
                        batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                        batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                        _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                        communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                        # x = x * communication_masks_tensor
                    else:
                        communication_rates = torch.tensor(0).to(x.device)
                else:
                    if self.communication:
                        communication_masks_tensor = F.max_pool2d(communication_masks_tensor, kernel_size=2)
                        # TODO: scale = 1, 2 不加 mask
                        # x = x * communication_masks_tensor
                
                ############ 2. Split the confidence map #######################
                # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                batch_node_features = self.regroup(x, record_len, 1)

                ############ 3. Fusion ####################################
                x_fuse = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    t_matrix = pairwise_t_matrix[b][:N, 0, :, :] #(N, 2, 3)
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))
                    record_frames = np.ones((N))*K
                    # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                    
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])
                    else: # ATTEN, MAX, Transformer
                        x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])

                x_fuse = torch.stack(x_fuse)

                ############ 4. Deconv ####################################
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[i](x_fuse))
                else:
                    ups.append(x_fuse)
                
            if len(ups) > 1:
                x_fuse = torch.cat(ups, dim=1)
            elif len(ups) == 1:
                x_fuse = ups[0]
            
            if len(backbone.deblocks) > self.num_levels:
                x_fuse = backbone.deblocks[-1](x_fuse)
                
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_confidence_maps = self.regroup(rm, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                # _, (B, 1, H, W), float
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix,
                                                (H, W))
                record_frames = np.ones((N))*K
                # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                if self.agg_mode == "RAIN":
                    # for sensor embedding
                    sensor_dist = -1# (B, H, W)
                    x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)

            # self.fuse_modsules(x_fuse, record_len)
        
        return x_fuse, communication_rates, {}, all_recon_loss

class TemporalFusion(nn.Module):
    def __init__(self, args):
        super(TemporalFusion, self).__init__()
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        # channel_size = args['channel_size']
        # spatial_size = args['spatial_size']
        # TM_Flag = args['TM_Flag']
        # compressed_size = args['compressed_size']
        # self.compensation_net = SyncLSTM(channel_size, spatial_size, TM_Flag, compressed_size)
        self.stpn = STPN(height_feat_size=args['channel_size'])
        self.flow_thre = args['flow_thre']
        self.motion_pred = MotionPrediction(seq_len=1)
        self.state_classify = StateEstimation(motion_category_num=1)

        self.flow_unc_flag = False
        if 'flow_unc_flag' in args:
            self.flow_unc_flag = True
            self.flow_unc_pred = FlowUncPrediction(seq_len=1)

        self.dcn = False
        if 'dcn' in args:
            self.dcn = True
            self.dcn_net = DCNNet(args['dcn'])

        if 'FlowPredictionFix' in args.keys() and args['FlowPredictionFix']:
            self.FlowPredictionFix()

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def FlowPredictionFix(self):
        print('------ Flow Prediction Fix, Only train DCN ------')
        for p in self.stpn.parameters():
            p.requires_grad = False
        for p in self.motion_pred.parameters():
            p.requires_grad = False
        for p in self.state_classify.parameters():
            p.requires_grad = False

    def forward(self, x, record_len, pairwise_t_matrix, data_dict=None):
        """
        Fusion forwarding.

        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        feat_seqs = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # iteratively warp feature to current timestamp
        batch_feat_seqs = []
        for b in range(B):
            # number of valid timestamps
            K = record_len[b]
            t_matrix = pairwise_t_matrix[b][:K, :K, :, :]
            curr_feat_seq = warp_affine_simple(feat_seqs[b],
                                               t_matrix[0, :, :, :],
                                               (H, W))
            batch_feat_seqs.append(curr_feat_seq[None, ...])
        batch_feat_seqs = torch.cat(batch_feat_seqs, dim=0)  # b, K, c, h, w
        batch_hist_feat_seqs = batch_feat_seqs[:, 1:].flip(1)

        # Backbone network
        bevs = self.stpn(batch_hist_feat_seqs)  # b, K, c, h, w

        # Motion Displacement prediction
        flow = self.motion_pred(bevs)
        flow = flow.view(-1, 2, bevs.size(-2), bevs.size(-1))

        flow_unc = None
        if self.flow_unc_flag:
            flow_unc = self.flow_unc_pred(bevs)
            flow_unc = flow_unc.view(-1, 2, bevs.size(-2), bevs.size(-1))

        # flow = data_dict['flow_gt']

        # Motion State Classification head
        state_class_pred = self.state_classify(bevs)

        # Given disp shift feature
        x_coord = torch.arange(bevs.size(-1)).float()
        y_coord = torch.arange(bevs.size(-2)).float()
        y, x = torch.meshgrid(y_coord, x_coord)
        grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(flow.shape[0], -1, -1, -1).to(
            flow.device)
        # updated_grid = grid + flow * (state_class_pred.sigmoid() > self.flow_thre)
        updated_grid = grid - flow
        updated_grid[:, 0, :, :] = updated_grid[:, 0, :, :] / (bevs.size(-1) / 2.0) - 1.0
        updated_grid[:, 1, :, :] = updated_grid[:, 1, :, :] / (bevs.size(-2) / 2.0) - 1.0
        out = F.grid_sample(batch_feat_seqs[:, 1], grid=updated_grid.permute(0, 2, 3, 1), mode='bilinear')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=updated_grid.permute(0,2,3,1), mode='nearest')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=grid.permute(0,2,3,1), mode='bilinear')
        # out = batch_feat_seqs[:,1]
        if self.dcn:
            out = self.dcn_net(out)

        flow_dict = {'flow':flow, 'flow_unc': flow_unc}
        return flow_dict, state_class_pred, out

    def forward_debug(self, x, origin_x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding
        Used for debug and visualization


        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        origin_x: torch.Tensor
            pillars (sum(n_cav), C, H * downsample_rate, W * downsample_rate)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        from matplotlib import pyplot as plt

        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        split_x = self.regroup(x, record_len)
        split_origin_x = self.regroup(origin_x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L, 1, H, W)).to(x)
                roi_mask[b, i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :], (H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        # visualize warped feature map
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0  # ego
            mask = roi_mask[b, i, :N, ...]
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H, W))
            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/feature_{b}_{idx}")
                plt.clf()
                plt.imshow(mask[idx][0].detach().cpu().numpy())
                plt.savefig(
                    f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/mask_feature_{b}_{idx}")
                plt.clf()

        # visualize origin pillar feature
        origin_node_features = split_origin_x

        for b in range(B):
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            i = 0  # ego
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(origin_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H * self.downsample_rate, W * self.downsample_rate))

            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/origin_{b}_{idx}")
                plt.clf()
opencood/models/fuse_modules/raindrop_w_syncnet.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# Date: 2023/04/13
# for [ where2comm + syncnet ]
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
from opencood.models.sub_modules.convgru import ConvGRU
from icecream import ic
from matplotlib import pyplot as plt
# from opencood.models.sub_modules.SyncLSTM import SyncLSTM
from opencood.models.sub_modules.MotionNet import STPN, MotionPrediction, StateEstimation, FlowUncPrediction
# from opencood.models.sub_modules.dcn_net import DCNNet
from opencood.models.comm_modules.where2comm_multisweep import Communication
from opencood.models.fuse_modules.SyncNet import SyncLSTM

# def get_warped_feature_mask(flow, updated_grid):
#     ''' 
#     get the mask of the warped feature map, where the value is 1 means the location is valid, and 0 means the location is invalid.
#     ----------
#     Args:
#         flow: (2, H, W)
#         updated_grid: (2, H, W)

#     Returns:
#         mask: (H, W)
#     '''
#     def get_large_flow_indices(flow, threshold=0):
#         '''
#         find the indices of the flow map where the flow is larger than the threshold, which means these locations have been moved.
#         ----------
#         Args:
#             flow: (2, H, W)
            
#         '''
#         max_values, max_indices = torch.max(torch.abs(flow[:2]), dim=0)
#         large_indices = torch.nonzero(max_values > threshold, as_tuple=False)
#         return large_indices

#     def remove_duplicate_points(points):
#         unique_points, inverse_indices = torch.unique(points, sorted=True, return_inverse=True, dim=0)
#         return unique_points, inverse_indices
    
#     def get_nonzero_idx(flow, idx):
#         flow_values = flow[:, idx[:, 0], idx[:, 1]]
#         nonzero_idx = torch.nonzero(torch.abs(flow_values).sum(dim=0) == 0, as_tuple=False).squeeze()
#         return idx[nonzero_idx]
    
#     flow_idx = get_large_flow_indices(flow)

#     mask = torch.ones(flow.shape[-2], flow.shape[-1])
#     if flow_idx.shape[0] == 0:
#         return mask

#     # print(flow_idx)
#     updated_grid_points_tmp = updated_grid[:, flow_idx[:,0], flow_idx[:,1]].to(torch.int64).T
#     # change the order of dim 1
#     updated_grid_points = torch.zeros_like(updated_grid_points_tmp)
#     updated_grid_points[:, 0] = updated_grid_points_tmp[:, 1]
#     updated_grid_points[:, 1] = updated_grid_points_tmp[:, 0]
#     # print(updated_grid_points)

#     unique_points_idx, _ = remove_duplicate_points(updated_grid_points)
#     # print(unique_points_idx)

#     nonzero_idx = get_nonzero_idx(flow, unique_points_idx)
#     # print(nonzero_idx)
    
#     # TODO: mask out the outlier idx

#     if len(nonzero_idx.shape) > 1:
#         mask[nonzero_idx[:, 0], nonzero_idx[:, 1]] = 0
#     else: 
#         mask[nonzero_idx[0], nonzero_idx[1]] = 0

#     return mask

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_fuse(nn.Module):
    def __init__(self, args):
        super(raindrop_fuse, self).__init__()
        
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.multi_scale:
            layer_nums = args['layer_nums']
            num_filters = args['num_filters']
            self.num_levels = len(layer_nums)
            self.fuse_modules = nn.ModuleList()
            for idx in range(self.num_levels):
                if self.agg_mode == 'MAX':
                    fuse_network = MaxFusion()
                else:
                    # self.agg_mode == 'ATTEN':
                    fuse_network = AttenFusion(num_filters[idx])
                self.fuse_modules.append(fuse_network)
        else: 
            if self.agg_mode == 'MAX': # max fusion, debug use
                self.fuse_modules = MaxFusion()

        # TODO: OPV2V/IRV2V 200x704, DAIR-V2X 200x504
        self.syncnet = SyncLSTM(channel_size = 64, h = 200, w = 704, k = 3, TM_Flag = False)
        print('*** channel_size = 64, h = 200, w = 504, k = 3, TM_Flag = False ***')

        # self.stpn = STPN(args['channel_size'])
        # self.motion_pred = MotionPrediction(seq_len=1)
        # self.state_classify = StateEstimation(motion_category_num=1)

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        
        return split_x

    def generate_estimated_feats(self, feats, pairwise_t_matrix, record_len):
        '''
        update the feature by SyncNet
        Note:
            1. ego feature does not need to be updated

        params:
            feats: (sum(B,N,K), C, H, W)
            pairwise_t_matrix: (B, L, K, 2, 3)
            record_len: (B)

        return:
            updated feature: [N, C, H, W]
        '''
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        batch_node_features = self.regroup(feats, record_len, k=K)

        updated_features_list = []
        flow_list = []
        all_recon_loss = 0
        for b in range(B):
            # TODO: IRV2V
            # all_feats = batch_node_features[b]
            # non_ego_feats = all_feats[1:].clone()
            # non_ego_feats = non_ego_feats.reshape(-1, K, C, H, W)
            # estimated_non_ego_feats = self.syncnet(non_ego_feats[:, 1:], [1]) # [N-1, 1, C, H, W]
            # # estimated_non_ego_feats = non_ego_feats[:, 1:2]
            # updated_features_list.append(torch.cat([all_feats[0].unsqueeze(0), estimated_non_ego_feats[:,0]], dim=0))
            # recon_loss = F.smooth_l1_loss(estimated_non_ego_feats, non_ego_feats[:,:1])
            # all_recon_loss += recon_loss

            # TODO: DAIR-V2X
            all_feats = batch_node_features[b]
            non_ego_feats = all_feats[K:].clone()
            non_ego_feats = non_ego_feats.reshape(-1, K, C, H, W)
            estimated_non_ego_feats_to_bp = self.syncnet(non_ego_feats[:, 1:], [1]) # [N-1, 1, C, H, W]

            # TODO: DAIR-V2X, test 用所有帧，train用前两帧
            # estimated_non_ego_feats = self.syncnet(non_ego_feats[:, :-1], [1]) # [N-1, 1, C, H, W]
            estimated_non_ego_feats = self.syncnet(non_ego_feats, [1]) # [N-1, 1, C, H, W]

            # estimated_non_ego_feats = non_ego_feats[:, 1:2]
            updated_features_list.append(torch.cat([all_feats[0].unsqueeze(0), estimated_non_ego_feats[:,0]], dim=0))
            recon_loss = F.smooth_l1_loss(estimated_non_ego_feats_to_bp, non_ego_feats[:,:1])
            all_recon_loss += recon_loss

        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)
        
        return updated_features_all, all_recon_loss

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        # 2. feature compensation with flow
        # 2.1 generate flow, 在同一个坐标系内，计算每个cav的flow
        # 2.2 compensation
        # x: (BxNxK, C, H, W) -> (BxN, C, H, W)
        updated_features, all_recon_loss = self.generate_estimated_feats(x, pairwise_t_matrix, record_len) # (BxN, C, H, W)
        
        # 3. feature fusion
        if self.multi_scale:
            ups = []
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(updated_features)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
            for i in range(self.num_levels):
                x = feats[i] if with_resnet else backbone.blocks[i](x)  # (BxN, C', H, W)

                ############ 1. Communication (Mask the features) #########
                if i==0:
                    if self.communication:
                        batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                        batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                        _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                        communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                        # x = x * communication_masks_tensor
                    else:
                        communication_rates = torch.tensor(0).to(x.device)
                else:
                    if self.communication:
                        communication_masks_tensor = F.max_pool2d(communication_masks_tensor, kernel_size=2)
                        # TODO: scale = 1, 2 不加 mask
                        # x = x * communication_masks_tensor
                
                ############ 2. Split the confidence map #######################
                # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                batch_node_features = self.regroup(x, record_len, 1)

                ############ 3. Fusion ####################################
                x_fuse = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    t_matrix = pairwise_t_matrix[b][:N, 0, :, :] #(N, 2, 3)
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))
                    record_frames = np.ones((N))*K
                    # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                    
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])
                    else: # ATTEN, MAX, Transformer
                        x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])

                x_fuse = torch.stack(x_fuse)

                ############ 4. Deconv ####################################
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[i](x_fuse))
                else:
                    ups.append(x_fuse)
                
            if len(ups) > 1:
                x_fuse = torch.cat(ups, dim=1)
            elif len(ups) == 1:
                x_fuse = ups[0]
            
            if len(backbone.deblocks) > self.num_levels:
                x_fuse = backbone.deblocks[-1](x_fuse)
                
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_confidence_maps = self.regroup(rm, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                # _, (B, 1, H, W), float
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix,
                                                (H, W))
                record_frames = np.ones((N))*K
                # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                if self.agg_mode == "RAIN":
                    # for sensor embedding
                    sensor_dist = -1# (B, H, W)
                    x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)

            # self.fuse_modsules(x_fuse, record_len)
        
        return x_fuse, communication_rates, {}, all_recon_loss

class TemporalFusion(nn.Module):
    def __init__(self, args):
        super(TemporalFusion, self).__init__()
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        # channel_size = args['channel_size']
        # spatial_size = args['spatial_size']
        # TM_Flag = args['TM_Flag']
        # compressed_size = args['compressed_size']
        # self.compensation_net = SyncLSTM(channel_size, spatial_size, TM_Flag, compressed_size)
        self.stpn = STPN(height_feat_size=args['channel_size'])
        self.flow_thre = args['flow_thre']
        self.motion_pred = MotionPrediction(seq_len=1)
        self.state_classify = StateEstimation(motion_category_num=1)

        self.flow_unc_flag = False
        if 'flow_unc_flag' in args:
            self.flow_unc_flag = True
            self.flow_unc_pred = FlowUncPrediction(seq_len=1)

        self.dcn = False
        if 'dcn' in args:
            self.dcn = True
            self.dcn_net = DCNNet(args['dcn'])

        if 'FlowPredictionFix' in args.keys() and args['FlowPredictionFix']:
            self.FlowPredictionFix()

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def FlowPredictionFix(self):
        print('------ Flow Prediction Fix, Only train DCN ------')
        for p in self.stpn.parameters():
            p.requires_grad = False
        for p in self.motion_pred.parameters():
            p.requires_grad = False
        for p in self.state_classify.parameters():
            p.requires_grad = False

    def forward(self, x, record_len, pairwise_t_matrix, data_dict=None):
        """
        Fusion forwarding.

        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        feat_seqs = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # iteratively warp feature to current timestamp
        batch_feat_seqs = []
        for b in range(B):
            # number of valid timestamps
            K = record_len[b]
            t_matrix = pairwise_t_matrix[b][:K, :K, :, :]
            curr_feat_seq = warp_affine_simple(feat_seqs[b],
                                               t_matrix[0, :, :, :],
                                               (H, W))
            batch_feat_seqs.append(curr_feat_seq[None, ...])
        batch_feat_seqs = torch.cat(batch_feat_seqs, dim=0)  # b, K, c, h, w
        batch_hist_feat_seqs = batch_feat_seqs[:, 1:].flip(1)

        # Backbone network
        bevs = self.stpn(batch_hist_feat_seqs)  # b, K, c, h, w

        # Motion Displacement prediction
        flow = self.motion_pred(bevs)
        flow = flow.view(-1, 2, bevs.size(-2), bevs.size(-1))

        flow_unc = None
        if self.flow_unc_flag:
            flow_unc = self.flow_unc_pred(bevs)
            flow_unc = flow_unc.view(-1, 2, bevs.size(-2), bevs.size(-1))

        # flow = data_dict['flow_gt']

        # Motion State Classification head
        state_class_pred = self.state_classify(bevs)

        # Given disp shift feature
        x_coord = torch.arange(bevs.size(-1)).float()
        y_coord = torch.arange(bevs.size(-2)).float()
        y, x = torch.meshgrid(y_coord, x_coord)
        grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(flow.shape[0], -1, -1, -1).to(
            flow.device)
        # updated_grid = grid + flow * (state_class_pred.sigmoid() > self.flow_thre)
        updated_grid = grid - flow
        updated_grid[:, 0, :, :] = updated_grid[:, 0, :, :] / (bevs.size(-1) / 2.0) - 1.0
        updated_grid[:, 1, :, :] = updated_grid[:, 1, :, :] / (bevs.size(-2) / 2.0) - 1.0
        out = F.grid_sample(batch_feat_seqs[:, 1], grid=updated_grid.permute(0, 2, 3, 1), mode='bilinear')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=updated_grid.permute(0,2,3,1), mode='nearest')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=grid.permute(0,2,3,1), mode='bilinear')
        # out = batch_feat_seqs[:,1]
        if self.dcn:
            out = self.dcn_net(out)

        flow_dict = {'flow':flow, 'flow_unc': flow_unc}
        return flow_dict, state_class_pred, out

    def forward_debug(self, x, origin_x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding
        Used for debug and visualization


        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        origin_x: torch.Tensor
            pillars (sum(n_cav), C, H * downsample_rate, W * downsample_rate)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        from matplotlib import pyplot as plt

        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        split_x = self.regroup(x, record_len)
        split_origin_x = self.regroup(origin_x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L, 1, H, W)).to(x)
                roi_mask[b, i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :], (H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        # visualize warped feature map
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0  # ego
            mask = roi_mask[b, i, :N, ...]
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H, W))
            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/feature_{b}_{idx}")
                plt.clf()
                plt.imshow(mask[idx][0].detach().cpu().numpy())
                plt.savefig(
                    f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/mask_feature_{b}_{idx}")
                plt.clf()

        # visualize origin pillar feature
        origin_node_features = split_origin_x

        for b in range(B):
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            i = 0  # ego
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(origin_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H * self.downsample_rate, W * self.downsample_rate))

            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/origin_{b}_{idx}")
                plt.clf()
opencood/models/fuse_modules/raindrop_w_syncnet_dair.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# Date: 2023/04/13
# for [ where2comm + syncnet ]
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
from opencood.models.sub_modules.convgru import ConvGRU
from icecream import ic
from matplotlib import pyplot as plt
# from opencood.models.sub_modules.SyncLSTM import SyncLSTM
from opencood.models.sub_modules.MotionNet import STPN, MotionPrediction, StateEstimation, FlowUncPrediction
# from opencood.models.sub_modules.dcn_net import DCNNet
from opencood.models.comm_modules.where2comm_multisweep import Communication
from opencood.models.fuse_modules.SyncNet import SyncLSTM

# def get_warped_feature_mask(flow, updated_grid):
#     ''' 
#     get the mask of the warped feature map, where the value is 1 means the location is valid, and 0 means the location is invalid.
#     ----------
#     Args:
#         flow: (2, H, W)
#         updated_grid: (2, H, W)

#     Returns:
#         mask: (H, W)
#     '''
#     def get_large_flow_indices(flow, threshold=0):
#         '''
#         find the indices of the flow map where the flow is larger than the threshold, which means these locations have been moved.
#         ----------
#         Args:
#             flow: (2, H, W)
            
#         '''
#         max_values, max_indices = torch.max(torch.abs(flow[:2]), dim=0)
#         large_indices = torch.nonzero(max_values > threshold, as_tuple=False)
#         return large_indices

#     def remove_duplicate_points(points):
#         unique_points, inverse_indices = torch.unique(points, sorted=True, return_inverse=True, dim=0)
#         return unique_points, inverse_indices
    
#     def get_nonzero_idx(flow, idx):
#         flow_values = flow[:, idx[:, 0], idx[:, 1]]
#         nonzero_idx = torch.nonzero(torch.abs(flow_values).sum(dim=0) == 0, as_tuple=False).squeeze()
#         return idx[nonzero_idx]
    
#     flow_idx = get_large_flow_indices(flow)

#     mask = torch.ones(flow.shape[-2], flow.shape[-1])
#     if flow_idx.shape[0] == 0:
#         return mask

#     # print(flow_idx)
#     updated_grid_points_tmp = updated_grid[:, flow_idx[:,0], flow_idx[:,1]].to(torch.int64).T
#     # change the order of dim 1
#     updated_grid_points = torch.zeros_like(updated_grid_points_tmp)
#     updated_grid_points[:, 0] = updated_grid_points_tmp[:, 1]
#     updated_grid_points[:, 1] = updated_grid_points_tmp[:, 0]
#     # print(updated_grid_points)

#     unique_points_idx, _ = remove_duplicate_points(updated_grid_points)
#     # print(unique_points_idx)

#     nonzero_idx = get_nonzero_idx(flow, unique_points_idx)
#     # print(nonzero_idx)
    
#     # TODO: mask out the outlier idx

#     if len(nonzero_idx.shape) > 1:
#         mask[nonzero_idx[:, 0], nonzero_idx[:, 1]] = 0
#     else: 
#         mask[nonzero_idx[0], nonzero_idx[1]] = 0

#     return mask

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_fuse(nn.Module):
    def __init__(self, args):
        super(raindrop_fuse, self).__init__()
        
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.multi_scale:
            layer_nums = args['layer_nums']
            num_filters = args['num_filters']
            self.num_levels = len(layer_nums)
            self.fuse_modules = nn.ModuleList()
            for idx in range(self.num_levels):
                if self.agg_mode == 'MAX':
                    fuse_network = MaxFusion()
                else:
                    # self.agg_mode == 'ATTEN':
                    fuse_network = AttenFusion(num_filters[idx])
                self.fuse_modules.append(fuse_network)
        else: 
            if self.agg_mode == 'MAX': # max fusion, debug use
                self.fuse_modules = MaxFusion()

        # TODO: OPV2V/IRV2V 200x704, DAIR-V2X 200x504
        self.syncnet = SyncLSTM(channel_size = 64, h = 200, w = 504, k = 2, TM_Flag = False)
        print('*** channel_size = 64, h = 200, w = 504, k = 2, TM_Flag = False ***')

        # self.stpn = STPN(args['channel_size'])
        # self.motion_pred = MotionPrediction(seq_len=1)
        # self.state_classify = StateEstimation(motion_category_num=1)

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        
        return split_x

    def generate_estimated_feats(self, feats, pairwise_t_matrix, record_len):
        '''
        update the feature by SyncNet
        Note:
            1. ego feature does not need to be updated

        params:
            feats: (sum(B,N,K), C, H, W)
            pairwise_t_matrix: (B, L, K, 2, 3)
            record_len: (B)

        return:
            updated feature: [N, C, H, W]
        '''
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        batch_node_features = self.regroup(feats, record_len, k=K)

        updated_features_list = []
        flow_list = []
        all_recon_loss = 0
        for b in range(B):
            # TODO: IRV2V
            # all_feats = batch_node_features[b]
            # non_ego_feats = all_feats[1:].clone()
            # non_ego_feats = non_ego_feats.reshape(-1, K, C, H, W)
            # estimated_non_ego_feats = self.syncnet(non_ego_feats[:, 1:], [1]) # [N-1, 1, C, H, W]
            # # estimated_non_ego_feats = non_ego_feats[:, 1:2]
            # updated_features_list.append(torch.cat([all_feats[0].unsqueeze(0), estimated_non_ego_feats[:,0]], dim=0))
            # recon_loss = F.smooth_l1_loss(estimated_non_ego_feats, non_ego_feats[:,:1])
            # all_recon_loss += recon_loss

            # TODO: DAIR-V2X
            all_feats = batch_node_features[b]
            non_ego_feats = all_feats[K:].clone()
            non_ego_feats = non_ego_feats.reshape(-1, K, C, H, W)
            estimated_non_ego_feats_to_bp = self.syncnet(non_ego_feats[:, 1:], [1]) # [N-1, 1, C, H, W]

            # TODO: DAIR-V2X, test 用所有帧，train用前两帧
            # estimated_non_ego_feats = self.syncnet(non_ego_feats[:, :-1], [1]) # [N-1, 1, C, H, W]
            estimated_non_ego_feats = self.syncnet(non_ego_feats, [1]) # [N-1, 1, C, H, W]

            # estimated_non_ego_feats = non_ego_feats[:, 1:2]
            updated_features_list.append(torch.cat([all_feats[0].unsqueeze(0), estimated_non_ego_feats[:,0]], dim=0))
            recon_loss = F.smooth_l1_loss(estimated_non_ego_feats_to_bp, non_ego_feats[:,:1])
            all_recon_loss += recon_loss

        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)
        
        return updated_features_all, all_recon_loss

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        # 2. feature compensation with flow
        # 2.1 generate flow, 在同一个坐标系内，计算每个cav的flow
        # 2.2 compensation
        # x: (BxNxK, C, H, W) -> (BxN, C, H, W)
        updated_features, all_recon_loss = self.generate_estimated_feats(x, pairwise_t_matrix, record_len) # (BxN, C, H, W)
        
        # 3. feature fusion
        if self.multi_scale:
            ups = []
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(updated_features)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
            for i in range(self.num_levels):
                x = feats[i] if with_resnet else backbone.blocks[i](x)  # (BxN, C', H, W)

                ############ 1. Communication (Mask the features) #########
                if i==0:
                    if self.communication:
                        batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                        batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                        _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                        communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                        # x = x * communication_masks_tensor
                    else:
                        communication_rates = torch.tensor(0).to(x.device)
                else:
                    if self.communication:
                        communication_masks_tensor = F.max_pool2d(communication_masks_tensor, kernel_size=2)
                        # TODO: scale = 1, 2 不加 mask
                        # x = x * communication_masks_tensor
                
                ############ 2. Split the confidence map #######################
                # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                batch_node_features = self.regroup(x, record_len, 1)

                ############ 3. Fusion ####################################
                x_fuse = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    t_matrix = pairwise_t_matrix[b][:N, 0, :, :] #(N, 2, 3)
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))
                    record_frames = np.ones((N))*K
                    # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                    
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])
                    else: # ATTEN, MAX, Transformer
                        x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])

                x_fuse = torch.stack(x_fuse)

                ############ 4. Deconv ####################################
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[i](x_fuse))
                else:
                    ups.append(x_fuse)
                
            if len(ups) > 1:
                x_fuse = torch.cat(ups, dim=1)
            elif len(ups) == 1:
                x_fuse = ups[0]
            
            if len(backbone.deblocks) > self.num_levels:
                x_fuse = backbone.deblocks[-1](x_fuse)
                
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_confidence_maps = self.regroup(rm, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                # _, (B, 1, H, W), float
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix,
                                                (H, W))
                record_frames = np.ones((N))*K
                # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                if self.agg_mode == "RAIN":
                    # for sensor embedding
                    sensor_dist = -1# (B, H, W)
                    x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)

            # self.fuse_modsules(x_fuse, record_len)
        
        return x_fuse, communication_rates, {}, all_recon_loss

class TemporalFusion(nn.Module):
    def __init__(self, args):
        super(TemporalFusion, self).__init__()
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        # channel_size = args['channel_size']
        # spatial_size = args['spatial_size']
        # TM_Flag = args['TM_Flag']
        # compressed_size = args['compressed_size']
        # self.compensation_net = SyncLSTM(channel_size, spatial_size, TM_Flag, compressed_size)
        self.stpn = STPN(height_feat_size=args['channel_size'])
        self.flow_thre = args['flow_thre']
        self.motion_pred = MotionPrediction(seq_len=1)
        self.state_classify = StateEstimation(motion_category_num=1)

        self.flow_unc_flag = False
        if 'flow_unc_flag' in args:
            self.flow_unc_flag = True
            self.flow_unc_pred = FlowUncPrediction(seq_len=1)

        self.dcn = False
        if 'dcn' in args:
            self.dcn = True
            self.dcn_net = DCNNet(args['dcn'])

        if 'FlowPredictionFix' in args.keys() and args['FlowPredictionFix']:
            self.FlowPredictionFix()

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def FlowPredictionFix(self):
        print('------ Flow Prediction Fix, Only train DCN ------')
        for p in self.stpn.parameters():
            p.requires_grad = False
        for p in self.motion_pred.parameters():
            p.requires_grad = False
        for p in self.state_classify.parameters():
            p.requires_grad = False

    def forward(self, x, record_len, pairwise_t_matrix, data_dict=None):
        """
        Fusion forwarding.

        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        feat_seqs = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # iteratively warp feature to current timestamp
        batch_feat_seqs = []
        for b in range(B):
            # number of valid timestamps
            K = record_len[b]
            t_matrix = pairwise_t_matrix[b][:K, :K, :, :]
            curr_feat_seq = warp_affine_simple(feat_seqs[b],
                                               t_matrix[0, :, :, :],
                                               (H, W))
            batch_feat_seqs.append(curr_feat_seq[None, ...])
        batch_feat_seqs = torch.cat(batch_feat_seqs, dim=0)  # b, K, c, h, w
        batch_hist_feat_seqs = batch_feat_seqs[:, 1:].flip(1)

        # Backbone network
        bevs = self.stpn(batch_hist_feat_seqs)  # b, K, c, h, w

        # Motion Displacement prediction
        flow = self.motion_pred(bevs)
        flow = flow.view(-1, 2, bevs.size(-2), bevs.size(-1))

        flow_unc = None
        if self.flow_unc_flag:
            flow_unc = self.flow_unc_pred(bevs)
            flow_unc = flow_unc.view(-1, 2, bevs.size(-2), bevs.size(-1))

        # flow = data_dict['flow_gt']

        # Motion State Classification head
        state_class_pred = self.state_classify(bevs)

        # Given disp shift feature
        x_coord = torch.arange(bevs.size(-1)).float()
        y_coord = torch.arange(bevs.size(-2)).float()
        y, x = torch.meshgrid(y_coord, x_coord)
        grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(flow.shape[0], -1, -1, -1).to(
            flow.device)
        # updated_grid = grid + flow * (state_class_pred.sigmoid() > self.flow_thre)
        updated_grid = grid - flow
        updated_grid[:, 0, :, :] = updated_grid[:, 0, :, :] / (bevs.size(-1) / 2.0) - 1.0
        updated_grid[:, 1, :, :] = updated_grid[:, 1, :, :] / (bevs.size(-2) / 2.0) - 1.0
        out = F.grid_sample(batch_feat_seqs[:, 1], grid=updated_grid.permute(0, 2, 3, 1), mode='bilinear')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=updated_grid.permute(0,2,3,1), mode='nearest')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=grid.permute(0,2,3,1), mode='bilinear')
        # out = batch_feat_seqs[:,1]
        if self.dcn:
            out = self.dcn_net(out)

        flow_dict = {'flow':flow, 'flow_unc': flow_unc}
        return flow_dict, state_class_pred, out

    def forward_debug(self, x, origin_x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding
        Used for debug and visualization


        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        origin_x: torch.Tensor
            pillars (sum(n_cav), C, H * downsample_rate, W * downsample_rate)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        from matplotlib import pyplot as plt

        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        split_x = self.regroup(x, record_len)
        split_origin_x = self.regroup(origin_x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L, 1, H, W)).to(x)
                roi_mask[b, i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :], (H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        # visualize warped feature map
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0  # ego
            mask = roi_mask[b, i, :N, ...]
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H, W))
            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/feature_{b}_{idx}")
                plt.clf()
                plt.imshow(mask[idx][0].detach().cpu().numpy())
                plt.savefig(
                    f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/mask_feature_{b}_{idx}")
                plt.clf()

        # visualize origin pillar feature
        origin_node_features = split_origin_x

        for b in range(B):
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            i = 0  # ego
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(origin_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H * self.downsample_rate, W * self.downsample_rate))

            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/origin_{b}_{idx}")
                plt.clf()
opencood/models/fuse_modules/raindrop_w_syncnet_dairv2x.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# Date: 2023/04/13
# for [ where2comm + syncnet ]
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
from opencood.models.sub_modules.convgru import ConvGRU
from icecream import ic
from matplotlib import pyplot as plt
# from opencood.models.sub_modules.SyncLSTM import SyncLSTM
from opencood.models.sub_modules.MotionNet import STPN, MotionPrediction, StateEstimation, FlowUncPrediction
# from opencood.models.sub_modules.dcn_net import DCNNet
from opencood.models.comm_modules.where2comm_multisweep import Communication
from opencood.models.fuse_modules.SyncNet import SyncLSTM

# def get_warped_feature_mask(flow, updated_grid):
#     ''' 
#     get the mask of the warped feature map, where the value is 1 means the location is valid, and 0 means the location is invalid.
#     ----------
#     Args:
#         flow: (2, H, W)
#         updated_grid: (2, H, W)

#     Returns:
#         mask: (H, W)
#     '''
#     def get_large_flow_indices(flow, threshold=0):
#         '''
#         find the indices of the flow map where the flow is larger than the threshold, which means these locations have been moved.
#         ----------
#         Args:
#             flow: (2, H, W)
            
#         '''
#         max_values, max_indices = torch.max(torch.abs(flow[:2]), dim=0)
#         large_indices = torch.nonzero(max_values > threshold, as_tuple=False)
#         return large_indices

#     def remove_duplicate_points(points):
#         unique_points, inverse_indices = torch.unique(points, sorted=True, return_inverse=True, dim=0)
#         return unique_points, inverse_indices
    
#     def get_nonzero_idx(flow, idx):
#         flow_values = flow[:, idx[:, 0], idx[:, 1]]
#         nonzero_idx = torch.nonzero(torch.abs(flow_values).sum(dim=0) == 0, as_tuple=False).squeeze()
#         return idx[nonzero_idx]
    
#     flow_idx = get_large_flow_indices(flow)

#     mask = torch.ones(flow.shape[-2], flow.shape[-1])
#     if flow_idx.shape[0] == 0:
#         return mask

#     # print(flow_idx)
#     updated_grid_points_tmp = updated_grid[:, flow_idx[:,0], flow_idx[:,1]].to(torch.int64).T
#     # change the order of dim 1
#     updated_grid_points = torch.zeros_like(updated_grid_points_tmp)
#     updated_grid_points[:, 0] = updated_grid_points_tmp[:, 1]
#     updated_grid_points[:, 1] = updated_grid_points_tmp[:, 0]
#     # print(updated_grid_points)

#     unique_points_idx, _ = remove_duplicate_points(updated_grid_points)
#     # print(unique_points_idx)

#     nonzero_idx = get_nonzero_idx(flow, unique_points_idx)
#     # print(nonzero_idx)
    
#     # TODO: mask out the outlier idx

#     if len(nonzero_idx.shape) > 1:
#         mask[nonzero_idx[:, 0], nonzero_idx[:, 1]] = 0
#     else: 
#         mask[nonzero_idx[0], nonzero_idx[1]] = 0

#     return mask

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_fuse(nn.Module):
    def __init__(self, args):
        super(raindrop_fuse, self).__init__()
        
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.multi_scale:
            layer_nums = args['layer_nums']
            num_filters = args['num_filters']
            self.num_levels = len(layer_nums)
            self.fuse_modules = nn.ModuleList()
            for idx in range(self.num_levels):
                if self.agg_mode == 'MAX':
                    fuse_network = MaxFusion()
                else:
                    # self.agg_mode == 'ATTEN':
                    fuse_network = AttenFusion(num_filters[idx])
                self.fuse_modules.append(fuse_network)
        else: 
            if self.agg_mode == 'MAX': # max fusion, debug use
                self.fuse_modules = MaxFusion()

        # TODO: OPV2V/IRV2V 200x704, DAIR-V2X 200x504
        self.syncnet = SyncLSTM(channel_size = 64, h = 200, w = 504, k = 2, TM_Flag = False)
        print('*** channel_size = 64, h = 200, w = 504, k = 2, TM_Flag = False ***')

        # self.stpn = STPN(args['channel_size'])
        # self.motion_pred = MotionPrediction(seq_len=1)
        # self.state_classify = StateEstimation(motion_category_num=1)

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        
        return split_x

    def generate_estimated_feats(self, feats, pairwise_t_matrix, record_len):
        '''
        update the feature by SyncNet
        Note:
            1. ego feature does not need to be updated

        params:
            feats: (sum(B,N,K), C, H, W)
            pairwise_t_matrix: (B, L, K, 2, 3)
            record_len: (B)

        return:
            updated feature: [N, C, H, W]
        '''
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        batch_node_features = self.regroup(feats, record_len, k=K)

        updated_features_list = []
        flow_list = []
        all_recon_loss = 0
        for b in range(B):
            # TODO: IRV2V
            # all_feats = batch_node_features[b]
            # non_ego_feats = all_feats[1:].clone()
            # non_ego_feats = non_ego_feats.reshape(-1, K, C, H, W)
            # estimated_non_ego_feats = self.syncnet(non_ego_feats[:, 1:], [1]) # [N-1, 1, C, H, W]
            # # estimated_non_ego_feats = non_ego_feats[:, 1:2]
            # updated_features_list.append(torch.cat([all_feats[0].unsqueeze(0), estimated_non_ego_feats[:,0]], dim=0))
            # recon_loss = F.smooth_l1_loss(estimated_non_ego_feats, non_ego_feats[:,:1])
            # all_recon_loss += recon_loss

            # TODO: DAIR-V2X
            all_feats = batch_node_features[b]
            non_ego_feats = all_feats[K:].clone()
            non_ego_feats = non_ego_feats.reshape(-1, K, C, H, W)
            estimated_non_ego_feats_to_bp = self.syncnet(non_ego_feats[:, 1:], [1]) # [N-1, 1, C, H, W]

            # TODO: DAIR-V2X, test 用所有帧，train用前两帧
            # estimated_non_ego_feats = self.syncnet(non_ego_feats[:, :-1], [1]) # [N-1, 1, C, H, W]
            estimated_non_ego_feats = self.syncnet(non_ego_feats, [1]) # [N-1, 1, C, H, W]

            # estimated_non_ego_feats = non_ego_feats[:, 1:2]
            updated_features_list.append(torch.cat([all_feats[0].unsqueeze(0), estimated_non_ego_feats[:,0]], dim=0))
            recon_loss = F.smooth_l1_loss(estimated_non_ego_feats_to_bp, non_ego_feats[:,:1])
            all_recon_loss += recon_loss

        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)
        
        return updated_features_all, all_recon_loss

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        # 2. feature compensation with flow
        # 2.1 generate flow, 在同一个坐标系内，计算每个cav的flow
        # 2.2 compensation
        # x: (BxNxK, C, H, W) -> (BxN, C, H, W)
        updated_features, all_recon_loss = self.generate_estimated_feats(x, pairwise_t_matrix, record_len) # (BxN, C, H, W)
        
        # 3. feature fusion
        if self.multi_scale:
            ups = []
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(updated_features)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
            for i in range(self.num_levels):
                x = feats[i] if with_resnet else backbone.blocks[i](x)  # (BxN, C', H, W)

                ############ 1. Communication (Mask the features) #########
                if i==0:
                    if self.communication:
                        batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                        batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                        _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                        communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                        # x = x * communication_masks_tensor
                    else:
                        communication_rates = torch.tensor(0).to(x.device)
                else:
                    if self.communication:
                        communication_masks_tensor = F.max_pool2d(communication_masks_tensor, kernel_size=2)
                        # TODO: scale = 1, 2 不加 mask
                        # x = x * communication_masks_tensor
                
                ############ 2. Split the confidence map #######################
                # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                batch_node_features = self.regroup(x, record_len, 1)

                ############ 3. Fusion ####################################
                x_fuse = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    t_matrix = pairwise_t_matrix[b][:N, 0, :, :] #(N, 2, 3)
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))
                    record_frames = np.ones((N))*K
                    # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                    
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])
                    else: # ATTEN, MAX, Transformer
                        x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])

                x_fuse = torch.stack(x_fuse)

                ############ 4. Deconv ####################################
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[i](x_fuse))
                else:
                    ups.append(x_fuse)
                
            if len(ups) > 1:
                x_fuse = torch.cat(ups, dim=1)
            elif len(ups) == 1:
                x_fuse = ups[0]
            
            if len(backbone.deblocks) > self.num_levels:
                x_fuse = backbone.deblocks[-1](x_fuse)
                
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_confidence_maps = self.regroup(rm, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                # _, (B, 1, H, W), float
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix,
                                                (H, W))
                record_frames = np.ones((N))*K
                # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                if self.agg_mode == "RAIN":
                    # for sensor embedding
                    sensor_dist = -1# (B, H, W)
                    x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)

            # self.fuse_modsules(x_fuse, record_len)
        
        return x_fuse, communication_rates, {}, all_recon_loss

class TemporalFusion(nn.Module):
    def __init__(self, args):
        super(TemporalFusion, self).__init__()
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        # channel_size = args['channel_size']
        # spatial_size = args['spatial_size']
        # TM_Flag = args['TM_Flag']
        # compressed_size = args['compressed_size']
        # self.compensation_net = SyncLSTM(channel_size, spatial_size, TM_Flag, compressed_size)
        self.stpn = STPN(height_feat_size=args['channel_size'])
        self.flow_thre = args['flow_thre']
        self.motion_pred = MotionPrediction(seq_len=1)
        self.state_classify = StateEstimation(motion_category_num=1)

        self.flow_unc_flag = False
        if 'flow_unc_flag' in args:
            self.flow_unc_flag = True
            self.flow_unc_pred = FlowUncPrediction(seq_len=1)

        self.dcn = False
        if 'dcn' in args:
            self.dcn = True
            self.dcn_net = DCNNet(args['dcn'])

        if 'FlowPredictionFix' in args.keys() and args['FlowPredictionFix']:
            self.FlowPredictionFix()

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def FlowPredictionFix(self):
        print('------ Flow Prediction Fix, Only train DCN ------')
        for p in self.stpn.parameters():
            p.requires_grad = False
        for p in self.motion_pred.parameters():
            p.requires_grad = False
        for p in self.state_classify.parameters():
            p.requires_grad = False

    def forward(self, x, record_len, pairwise_t_matrix, data_dict=None):
        """
        Fusion forwarding.

        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        feat_seqs = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # iteratively warp feature to current timestamp
        batch_feat_seqs = []
        for b in range(B):
            # number of valid timestamps
            K = record_len[b]
            t_matrix = pairwise_t_matrix[b][:K, :K, :, :]
            curr_feat_seq = warp_affine_simple(feat_seqs[b],
                                               t_matrix[0, :, :, :],
                                               (H, W))
            batch_feat_seqs.append(curr_feat_seq[None, ...])
        batch_feat_seqs = torch.cat(batch_feat_seqs, dim=0)  # b, K, c, h, w
        batch_hist_feat_seqs = batch_feat_seqs[:, 1:].flip(1)

        # Backbone network
        bevs = self.stpn(batch_hist_feat_seqs)  # b, K, c, h, w

        # Motion Displacement prediction
        flow = self.motion_pred(bevs)
        flow = flow.view(-1, 2, bevs.size(-2), bevs.size(-1))

        flow_unc = None
        if self.flow_unc_flag:
            flow_unc = self.flow_unc_pred(bevs)
            flow_unc = flow_unc.view(-1, 2, bevs.size(-2), bevs.size(-1))

        # flow = data_dict['flow_gt']

        # Motion State Classification head
        state_class_pred = self.state_classify(bevs)

        # Given disp shift feature
        x_coord = torch.arange(bevs.size(-1)).float()
        y_coord = torch.arange(bevs.size(-2)).float()
        y, x = torch.meshgrid(y_coord, x_coord)
        grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(flow.shape[0], -1, -1, -1).to(
            flow.device)
        # updated_grid = grid + flow * (state_class_pred.sigmoid() > self.flow_thre)
        updated_grid = grid - flow
        updated_grid[:, 0, :, :] = updated_grid[:, 0, :, :] / (bevs.size(-1) / 2.0) - 1.0
        updated_grid[:, 1, :, :] = updated_grid[:, 1, :, :] / (bevs.size(-2) / 2.0) - 1.0
        out = F.grid_sample(batch_feat_seqs[:, 1], grid=updated_grid.permute(0, 2, 3, 1), mode='bilinear')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=updated_grid.permute(0,2,3,1), mode='nearest')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=grid.permute(0,2,3,1), mode='bilinear')
        # out = batch_feat_seqs[:,1]
        if self.dcn:
            out = self.dcn_net(out)

        flow_dict = {'flow':flow, 'flow_unc': flow_unc}
        return flow_dict, state_class_pred, out

    def forward_debug(self, x, origin_x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding
        Used for debug and visualization


        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        origin_x: torch.Tensor
            pillars (sum(n_cav), C, H * downsample_rate, W * downsample_rate)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        from matplotlib import pyplot as plt

        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        split_x = self.regroup(x, record_len)
        split_origin_x = self.regroup(origin_x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L, 1, H, W)).to(x)
                roi_mask[b, i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :], (H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        # visualize warped feature map
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0  # ego
            mask = roi_mask[b, i, :N, ...]
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H, W))
            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/feature_{b}_{idx}")
                plt.clf()
                plt.imshow(mask[idx][0].detach().cpu().numpy())
                plt.savefig(
                    f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/mask_feature_{b}_{idx}")
                plt.clf()

        # visualize origin pillar feature
        origin_node_features = split_origin_x

        for b in range(B):
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            i = 0  # ego
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(origin_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H * self.downsample_rate, W * self.downsample_rate))

            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/origin_{b}_{idx}")
                plt.clf()
opencood/models/fuse_modules/raindrop_w_syncnet_dariv2x.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yuxi Wei, Sizhe Wei <sizhewei@sjtu.edu.cn> 
# Date: 2023/04/13
# for [ where2comm + syncnet ]
# License: TDG-Attribution-NonCommercial-NoDistrib

import torch
import torch.nn as nn
import easydict
import copy
import math
from torch.autograd import Variable
import torch.nn.functional as F
import random
# from dgl.nn.pytorch.factory import KNNGraph
# import dgl
import numpy as np
import os

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
from opencood.models.sub_modules.convgru import ConvGRU
from icecream import ic
from matplotlib import pyplot as plt
# from opencood.models.sub_modules.SyncLSTM import SyncLSTM
from opencood.models.sub_modules.MotionNet import STPN, MotionPrediction, StateEstimation, FlowUncPrediction
# from opencood.models.sub_modules.dcn_net import DCNNet
from opencood.models.comm_modules.where2comm_multisweep import Communication
from opencood.models.fuse_modules.SyncNet import SyncLSTM

# def get_warped_feature_mask(flow, updated_grid):
#     ''' 
#     get the mask of the warped feature map, where the value is 1 means the location is valid, and 0 means the location is invalid.
#     ----------
#     Args:
#         flow: (2, H, W)
#         updated_grid: (2, H, W)

#     Returns:
#         mask: (H, W)
#     '''
#     def get_large_flow_indices(flow, threshold=0):
#         '''
#         find the indices of the flow map where the flow is larger than the threshold, which means these locations have been moved.
#         ----------
#         Args:
#             flow: (2, H, W)
            
#         '''
#         max_values, max_indices = torch.max(torch.abs(flow[:2]), dim=0)
#         large_indices = torch.nonzero(max_values > threshold, as_tuple=False)
#         return large_indices

#     def remove_duplicate_points(points):
#         unique_points, inverse_indices = torch.unique(points, sorted=True, return_inverse=True, dim=0)
#         return unique_points, inverse_indices
    
#     def get_nonzero_idx(flow, idx):
#         flow_values = flow[:, idx[:, 0], idx[:, 1]]
#         nonzero_idx = torch.nonzero(torch.abs(flow_values).sum(dim=0) == 0, as_tuple=False).squeeze()
#         return idx[nonzero_idx]
    
#     flow_idx = get_large_flow_indices(flow)

#     mask = torch.ones(flow.shape[-2], flow.shape[-1])
#     if flow_idx.shape[0] == 0:
#         return mask

#     # print(flow_idx)
#     updated_grid_points_tmp = updated_grid[:, flow_idx[:,0], flow_idx[:,1]].to(torch.int64).T
#     # change the order of dim 1
#     updated_grid_points = torch.zeros_like(updated_grid_points_tmp)
#     updated_grid_points[:, 0] = updated_grid_points_tmp[:, 1]
#     updated_grid_points[:, 1] = updated_grid_points_tmp[:, 0]
#     # print(updated_grid_points)

#     unique_points_idx, _ = remove_duplicate_points(updated_grid_points)
#     # print(unique_points_idx)

#     nonzero_idx = get_nonzero_idx(flow, unique_points_idx)
#     # print(nonzero_idx)
    
#     # TODO: mask out the outlier idx

#     if len(nonzero_idx.shape) > 1:
#         mask[nonzero_idx[:, 0], nonzero_idx[:, 1]] = 0
#     else: 
#         mask[nonzero_idx[0], nonzero_idx[1]] = 0

#     return mask

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]
        
class raindrop_fuse(nn.Module):
    def __init__(self, args):
        super(raindrop_fuse, self).__init__()
        
        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']

        #################################################
        if self.multi_scale:
            layer_nums = args['layer_nums']
            num_filters = args['num_filters']
            self.num_levels = len(layer_nums)
            self.fuse_modules = nn.ModuleList()
            for idx in range(self.num_levels):
                if self.agg_mode == 'MAX':
                    fuse_network = MaxFusion()
                else:
                    # self.agg_mode == 'ATTEN':
                    fuse_network = AttenFusion(num_filters[idx])
                self.fuse_modules.append(fuse_network)
        else: 
            if self.agg_mode == 'MAX': # max fusion, debug use
                self.fuse_modules = MaxFusion()

        self.syncnet = SyncLSTM(channel_size = 64, h = 200, w = 704, k = 2, TM_Flag = False)
        print('*** channel_size = 64, h = 200, w = 704, k = 2, TM_Flag = False ***')

        # self.stpn = STPN(args['channel_size'])
        # self.motion_pred = MotionPrediction(seq_len=1)
        # self.state_classify = StateEstimation(motion_category_num=1)

    def regroup(self, x, len, k=0):
        """
        split into different batch and time k

        Parameters
        ----------
        x : torch.Tensor
            input data, (B, ...)

        len : list 
            cav num in differnt batch, eg [3, 3]

        k : torch.Tensor
            num of past frames
        
        Returns:
        --------
        split_x : list
            different cav's lidar feature
            for example: k=4, then return [(3x4, C, H, W), (2x4, C, H, W), ...]
        """
        cum_sum_len = torch.cumsum(len*k-(k-1), dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        
        return split_x

    def generate_estimated_feats(self, feats, pairwise_t_matrix, record_len):
        '''
        update the feature by SyncNet
        Note:
            1. ego feature does not need to be updated

        params:
            feats: (sum(B,N,K), C, H, W)
            pairwise_t_matrix: (B, L, K, 2, 3)
            record_len: (B)

        return:
            updated feature: [N, C, H, W]
        '''
        _, C, H, W = feats.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        batch_node_features = self.regroup(feats, record_len, k=K)

        updated_features_list = []
        flow_list = []
        all_recon_loss = 0
        for b in range(B):
            all_feats = batch_node_features[b]
            non_ego_feats = all_feats[1:].clone()
            non_ego_feats = non_ego_feats.reshape(-1, K, C, H, W)
            estimated_non_ego_feats = self.syncnet(non_ego_feats[:, 1:], [1]) # [N-1, 1, C, H, W]
            # estimated_non_ego_feats = non_ego_feats[:, 1:2]
            updated_features_list.append(torch.cat([all_feats[0].unsqueeze(0), estimated_non_ego_feats[:,0]], dim=0))
            recon_loss = F.smooth_l1_loss(estimated_non_ego_feats, non_ego_feats[:,:1])
            all_recon_loss += recon_loss

        updated_features_all = torch.cat(updated_features_list, dim=0)  # (sum(B,N), C, H, W)
        
        return updated_features_all, all_recon_loss

    def forward(self, x, rm, record_len, pairwise_t_matrix, time_diffs, backbone=None, heads=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
        
        rm : torch.Tensor
            confidence map, (sum(n_cav), 2, H, W)

        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, K, 2, 3) 

        time_diffs : torch.Tensor
            time interval, (sum(n_cav), )
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L, K = pairwise_t_matrix.shape[:3]

        # (B,L,k,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
        # [B, L, L, 2, 3], 只要x,y这两个维度就可以(z只有一层)，所以提取[0,1,3]作为仿射变换矩阵, 大小(2, 3)
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        # 2. feature compensation with flow
        # 2.1 generate flow, 在同一个坐标系内，计算每个cav的flow
        # 2.2 compensation
        # x: (BxNxK, C, H, W) -> (BxN, C, H, W)
        updated_features, all_recon_loss = self.generate_estimated_feats(x, pairwise_t_matrix, record_len) # (BxN, C, H, W)
        
        # 3. feature fusion
        if self.multi_scale:
            ups = []
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(updated_features)  # tuple((B, C, H, W), (B, 2C, H/2, W/2), (B, 4C, H/4, W/4))
            
            for i in range(self.num_levels):
                x = feats[i] if with_resnet else backbone.blocks[i](x)  # (BxN, C', H, W)

                ############ 1. Communication (Mask the features) #########
                if i==0:
                    if self.communication:
                        batch_confidence_maps = self.regroup(rm, record_len, K) # [[2*3, 2, 100/2, 252/2], [3*3, 2, 100/2, 252/2], ...]
                        batch_time_intervals = self.regroup(time_diffs, record_len, K) # [[2*3], [3*3], ...]
                        _, communication_masks_list, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                        communication_masks_tensor = torch.concat(communication_masks_list, dim=0) 
                        # x = x * communication_masks_tensor
                    else:
                        communication_rates = torch.tensor(0).to(x.device)
                else:
                    if self.communication:
                        communication_masks_tensor = F.max_pool2d(communication_masks_tensor, kernel_size=2)
                        # TODO: scale = 1, 2 不加 mask
                        # x = x * communication_masks_tensor
                
                ############ 2. Split the confidence map #######################
                # split x:[(L1, C*2^i, H/2^i, W/2^i), (L2, C*2^i, H/2^i, W/2^i), ...]
                # for example, i=1, past_k=3, b_1 has 2 cav, b_2 has 3 cav ... :
                # [[2*3, 256*2, 100/2, 252/2], [3*3, 256*2, 100/2, 252/2], ...]
                batch_node_features = self.regroup(x, record_len, 1)

                ############ 3. Fusion ####################################
                x_fuse = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # t_matrix[i, j]-> from i to j
                    t_matrix = pairwise_t_matrix[b][:N, 0, :, :] #(N, 2, 3)
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix,
                                                    (H, W))
                    record_frames = np.ones((N))*K
                    # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                    
                    if self.agg_mode == 'RAIN':
                        # for sensor embedding
                        sensor_dist = -1# (B, H, W)
                        x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature, sensor_dist, batch_time_intervals[b], self.regroup(communication_masks_tensor, record_len, K)[b]))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])
                    else: # ATTEN, MAX, Transformer
                        x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # # TODO for scale debug
                        # if i==self.num_levels-1:
                        #     x_fuse.append(self.fuse_modules[i](neighbor_feature))
                        # else:
                        #     x_fuse.append(neighbor_feature[0])

                x_fuse = torch.stack(x_fuse)

                ############ 4. Deconv ####################################
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[i](x_fuse))
                else:
                    ups.append(x_fuse)
                
            if len(ups) > 1:
                x_fuse = torch.cat(ups, dim=1)
            elif len(ups) == 1:
                x_fuse = ups[0]
            
            if len(backbone.deblocks) > self.num_levels:
                x_fuse = backbone.deblocks[-1](x_fuse)
                
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len, K)
            batch_confidence_maps = self.regroup(rm, record_len, K)
            batch_time_intervals = self.regroup(time_diffs, record_len, K)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                # _, (B, 1, H, W), float
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :, :, :].view(-1, 2, 3) #(Nxk, 2, 3)
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix,
                                                (H, W))
                record_frames = np.ones((N))*K
                # def forward(self, feartures, comm_mask, sensor_dist, record_frames, time_diffs):
                if self.agg_mode == "RAIN":
                    # for sensor embedding
                    sensor_dist = -1# (B, H, W)
                    x_fuse.append(self.fuse_modules(neighbor_feature, sensor_dist, batch_time_intervals[b], communication_masks[b]))
                else: # ATTEN, MAX, Transformer
                    x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)

            # self.fuse_modsules(x_fuse, record_len)
        
        return x_fuse, communication_rates, {}, all_recon_loss

class TemporalFusion(nn.Module):
    def __init__(self, args):
        super(TemporalFusion, self).__init__()
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        # channel_size = args['channel_size']
        # spatial_size = args['spatial_size']
        # TM_Flag = args['TM_Flag']
        # compressed_size = args['compressed_size']
        # self.compensation_net = SyncLSTM(channel_size, spatial_size, TM_Flag, compressed_size)
        self.stpn = STPN(height_feat_size=args['channel_size'])
        self.flow_thre = args['flow_thre']
        self.motion_pred = MotionPrediction(seq_len=1)
        self.state_classify = StateEstimation(motion_category_num=1)

        self.flow_unc_flag = False
        if 'flow_unc_flag' in args:
            self.flow_unc_flag = True
            self.flow_unc_pred = FlowUncPrediction(seq_len=1)

        self.dcn = False
        if 'dcn' in args:
            self.dcn = True
            self.dcn_net = DCNNet(args['dcn'])

        if 'FlowPredictionFix' in args.keys() and args['FlowPredictionFix']:
            self.FlowPredictionFix()

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def FlowPredictionFix(self):
        print('------ Flow Prediction Fix, Only train DCN ------')
        for p in self.stpn.parameters():
            p.requires_grad = False
        for p in self.motion_pred.parameters():
            p.requires_grad = False
        for p in self.state_classify.parameters():
            p.requires_grad = False

    def forward(self, x, record_len, pairwise_t_matrix, data_dict=None):
        """
        Fusion forwarding.

        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        feat_seqs = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # iteratively warp feature to current timestamp
        batch_feat_seqs = []
        for b in range(B):
            # number of valid timestamps
            K = record_len[b]
            t_matrix = pairwise_t_matrix[b][:K, :K, :, :]
            curr_feat_seq = warp_affine_simple(feat_seqs[b],
                                               t_matrix[0, :, :, :],
                                               (H, W))
            batch_feat_seqs.append(curr_feat_seq[None, ...])
        batch_feat_seqs = torch.cat(batch_feat_seqs, dim=0)  # b, K, c, h, w
        batch_hist_feat_seqs = batch_feat_seqs[:, 1:].flip(1)

        # Backbone network
        bevs = self.stpn(batch_hist_feat_seqs)  # b, K, c, h, w

        # Motion Displacement prediction
        flow = self.motion_pred(bevs)
        flow = flow.view(-1, 2, bevs.size(-2), bevs.size(-1))

        flow_unc = None
        if self.flow_unc_flag:
            flow_unc = self.flow_unc_pred(bevs)
            flow_unc = flow_unc.view(-1, 2, bevs.size(-2), bevs.size(-1))

        # flow = data_dict['flow_gt']

        # Motion State Classification head
        state_class_pred = self.state_classify(bevs)

        # Given disp shift feature
        x_coord = torch.arange(bevs.size(-1)).float()
        y_coord = torch.arange(bevs.size(-2)).float()
        y, x = torch.meshgrid(y_coord, x_coord)
        grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0).unsqueeze(0).expand(flow.shape[0], -1, -1, -1).to(
            flow.device)
        # updated_grid = grid + flow * (state_class_pred.sigmoid() > self.flow_thre)
        updated_grid = grid - flow
        updated_grid[:, 0, :, :] = updated_grid[:, 0, :, :] / (bevs.size(-1) / 2.0) - 1.0
        updated_grid[:, 1, :, :] = updated_grid[:, 1, :, :] / (bevs.size(-2) / 2.0) - 1.0
        out = F.grid_sample(batch_feat_seqs[:, 1], grid=updated_grid.permute(0, 2, 3, 1), mode='bilinear')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=updated_grid.permute(0,2,3,1), mode='nearest')
        # out = F.grid_sample(batch_feat_seqs[:,1], grid=grid.permute(0,2,3,1), mode='bilinear')
        # out = batch_feat_seqs[:,1]
        if self.dcn:
            out = self.dcn_net(out)

        flow_dict = {'flow':flow, 'flow_unc': flow_unc}
        return flow_dict, state_class_pred, out

    def forward_debug(self, x, origin_x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding
        Used for debug and visualization


        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)

        origin_x: torch.Tensor
            pillars (sum(n_cav), C, H * downsample_rate, W * downsample_rate)

        record_len : list
            shape: (B)

        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego,
            shape: (B, L, L, 4, 4)

        Returns
        -------
        Fused feature.
        """
        from matplotlib import pyplot as plt

        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        split_x = self.regroup(x, record_len)
        split_origin_x = self.regroup(origin_x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:, :, :, [0, 1], :][:, :, :, :, [0, 1, 3]]  # [B, L, L, 2, 3]
        pairwise_t_matrix[..., 0, 1] = pairwise_t_matrix[..., 0, 1] * H / W
        pairwise_t_matrix[..., 1, 0] = pairwise_t_matrix[..., 1, 0] * W / H
        pairwise_t_matrix[..., 0, 2] = pairwise_t_matrix[..., 0, 2] / (
                    self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[..., 1, 2] = pairwise_t_matrix[..., 1, 2] / (
                    self.downsample_rate * self.discrete_ratio * H) * 2

        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L, 1, H, W)).to(x)
                roi_mask[b, i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :], (H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        # visualize warped feature map
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0  # ego
            mask = roi_mask[b, i, :N, ...]
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H, W))
            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/feature_{b}_{idx}")
                plt.clf()
                plt.imshow(mask[idx][0].detach().cpu().numpy())
                plt.savefig(
                    f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/mask_feature_{b}_{idx}")
                plt.clf()

        # visualize origin pillar feature
        origin_node_features = split_origin_x

        for b in range(B):
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            i = 0  # ego
            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(origin_node_features[b],
                                                  t_matrix[i, :, :, :],
                                                  (H * self.downsample_rate, W * self.downsample_rate))

            for idx in range(N):
                plt.imshow(torch.max(neighbor_feature[idx], 0)[0].detach().cpu().numpy())
                plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/origin_{b}_{idx}")
                plt.clf()
opencood/models/fuse_modules/raindrop_w_syncnet_irv2v.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>, Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple

DEBUG=False

class ScaledDotProductAttention(nn.Module):
    """
    Scaled Dot-Product Attention proposed in "Attention Is All You Need"
    Compute the dot products of the query with all keys, divide each by sqrt(dim),
    and apply a softmax function to obtain the weights on the values
    Args: dim, mask
        dim (int): dimention of attention
        mask (torch.Tensor): tensor containing indices to be masked
    Inputs: query, key, value, mask
        - **query** (batch, q_len, d_model): tensor containing projection
          vector for decoder.
        - **key** (batch, k_len, d_model): tensor containing projection
          vector for encoder.
        - **value** (batch, v_len, d_model): tensor containing features of the
          encoded input sequence.
        - **mask** (-): tensor containing indices to be masked
    Returns: context, attn
        - **context**: tensor containing the context vector from
          attention mechanism.
        - **attn**: tensor containing the attention (alignment) from the
          encoder outputs.
    """

    def __init__(self, dim):
        super(ScaledDotProductAttention, self).__init__()
        self.sqrt_dim = np.sqrt(dim)

    def forward(self, query, key, value):
        score = torch.bmm(query, key.transpose(1, 2)) / self.sqrt_dim
        attn = F.softmax(score, -1)
        context = torch.bmm(attn, value)
        return context


class AttFusion(nn.Module):
    def __init__(self, feature_dim):
        super(AttFusion, self).__init__()
        self.att = ScaledDotProductAttention(feature_dim)

    def forward(self, x, record_len, pairwise_t_matrix):
        """
        pairwise_t_matrix : [N,N,2,3]
        """
        split_x = self.regroup(x, record_len)
        batch_size = len(record_len)
        C, H, W = split_x[0].shape[1:]  # C, W, H before
        out = []
        for b, xx in enumerate(split_x):
          N = xx.shape[0]
          t_matrix = pairwise_t_matrix[b][:N, :N, :, :]
          i = 0 # ego
          xx = warp_affine_simple(xx, t_matrix[i, :, :, :], (H, W))

          cav_num = xx.shape[0]
          xx = xx.view(cav_num, C, -1).permute(2, 0, 1) #  (H*W, cav_num, C), perform self attention on each pixel.
          h = self.att(xx, xx, xx)
          h = h.permute(1, 2, 0).view(cav_num, C, H, W)[0, ...].unsqueeze(0)  # C, W, H before
          out.append(h)
        return torch.cat(out, dim=0)

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x



    def forward_debug(self, x, origin_x, record_len, pairwise_t_matrix):
        split_x = self.regroup(x, record_len)
        split_origin_x = self.regroup(origin_x, record_len)
        batch_size = len(record_len)
        C, H, W = split_x[0].shape[1:]  # C, W, H before
        H_origin, W_origin = split_origin_x[0].shape[2:]
        out = []
        from matplotlib import pyplot as plt
        for b, xx in enumerate(split_x):
          N = xx.shape[0]
          t_matrix = pairwise_t_matrix[b][:N, :N, :, :]
          i = 0
          xx = warp_affine_simple(xx, t_matrix[i, :, :, :], (H, W))
          origin_xx = warp_affine_simple(split_origin_x[b], t_matrix[i, :, :, :], (H_origin, W_origin))

          for idx in range(N):
            plt.imshow(torch.max(xx[idx],0)[0].detach().cpu().numpy())
            plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/feature_{b}_{idx}")
            plt.clf()
            plt.imshow(torch.max(origin_xx[idx],0)[0].detach().cpu().numpy())
            plt.savefig(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/debug_warp_feature/origin_feature_{b}_{idx}")
            plt.clf()
        raise
opencood/models/fuse_modules/self_attn.py
MediaBrain-SJTU/CoBEVFlow
"""
Implementation of Simple transformer fusion.


Provided by Yue Hu. 
Have not modified for this codebase.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import kornia


class TransformerMessage(nn.Module):
    def __init__(self, 
                 in_channels=64,
                 trans_layer=[3]):
        super(TransformerMessage, self).__init__()
        self.in_channels = in_channels

        self.trans_layer = trans_layer

        dropout = 0
        nhead = 8
        for c_layer in self.trans_layer:
            d_model = in_channels
            cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
            # Implementation of Feedforward model
            linear1 = nn.Linear(d_model, d_model)
            linear2 = nn.Linear(d_model, d_model)

            norm1 = nn.LayerNorm(d_model)
            norm2 = nn.LayerNorm(d_model)
            dropout0 = nn.Dropout(dropout)
            dropout1 = nn.Dropout(dropout)
            dropout2 = nn.Dropout(dropout)
            self.__setattr__('cross_attn'+str(c_layer), cross_attn)
            self.__setattr__('linear1_'+str(c_layer), linear1)
            self.__setattr__('linear2_'+str(c_layer), linear2)
            self.__setattr__('norm1_'+str(c_layer), norm1)
            self.__setattr__('norm2_'+str(c_layer), norm2)
            self.__setattr__('dropout0_'+str(c_layer), dropout0)
            self.__setattr__('dropout1_'+str(c_layer), dropout1)
            self.__setattr__('dropout2_'+str(c_layer), dropout2)

    def add_pe_map(self, x, normalized=True):
        """ Add positional encoding to feature map.
        Args:
            x: torch.Tensor
                [N, C, H, W]
        
        """
        # scale = 2 * math.pi
        temperature = 10000 
        num_pos_feats = x.shape[-3] // 2  # d

        mask = torch.zeros([x.shape[-2], x.shape[-1]], dtype=torch.bool, device=x.device)
        not_mask = ~mask
        y_embed = not_mask.cumsum(0, dtype=torch.float32)
        x_embed = not_mask.cumsum(1, dtype=torch.float32)

        dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=x.device)
        dim_t = temperature ** (2 * (dim_t // 2) / num_pos_feats)

        pos_x = x_embed[:, :, None] / dim_t
        pos_y = y_embed[:, :, None] / dim_t
        pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2)
        pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2)
        pos = torch.cat((pos_y, pos_x), dim=2).permute(2, 0, 1)

        if len(x.shape) == 5:
            x = x + pos[None,None,:,:,:]
        elif len(x.shape) == 6:
            x = x + pos[None,None,None,:,:,:]
        return x

    def forward(self, x, shift_mats, shift_mats_rev, agent_mask):
        batch, max_agent_num, c, h, w = x.shape

        # ================================
        # First, transform to each coord
        feat_shifted = []
        for agent_i in range(max_agent_num):
            # shift_mat_i = shift_mats[:, agent_i, :, :]
            shift_mat_rev_i = shift_mats_rev[:, agent_i, :, :]
            feat_i = x[:, agent_i, :, :, :]
            feat_shifted_i = []
            for agent_j in range(max_agent_num):
                shift_mat_j = shift_mats[:, agent_j, :, :]
                shift_mat = shift_mat_j.view(batch, 3, 3) @ shift_mat_rev_i.view(batch, 3, 3)
                feat = kornia.warp_perspective(feat_i, shift_mat, dsize=(100 * 2, 100 * 2), align_corners=False)
                feat_shifted_i.append(feat)
            feat_shifted_i = torch.cat([f.unsqueeze(1) for f in feat_shifted_i], dim=1)
            feat_shifted.append(feat_shifted_i)
        feat_shifted = torch.cat([f.unsqueeze(1) for f in feat_shifted], dim=1)


        # ================================
        # x_fuse, _, _ = self.TRANSFORMER_MESSAGE([[],[],[],local_com_mat], [transformed_feature], num_agent_tensor)

        for i, c_layer in enumerate(self.trans_layer):
            batch_updated_features = torch.zeros(batch, max_agent_num, c, h, w).to(shift_mats.device)
            for batch_i in range(batch):
                N = int(torch.sum(agent_mask[batch_i]))
                feat_map = x[batch_i:batch_i+1, :N, :, :, :]
                val_feat = feat_shifted[batch_i:batch_i+1, :N, :N, :, :, :]

                feat_map = self.add_pe_map(feat_map)
                # [b,N,C,H,W] -> [b,N,H,W,C]
                # [b,N,N,C,H,W] -> [N,b,N,H,W,C]
                src = feat_map.permute(0,1,3,4,2).contiguous().view(N*h*w,c).contiguous().unsqueeze(0)
                tgt = val_feat.permute(1,0,2,4,5,3).contiguous().view(N, N*h*w,c).contiguous()
                # print(src.shape)  # torch.Size([1, 120000, 64])
                # print(tgt.shape)  # torch.Size([N, 120000, 64])

                src2, weight_mat = eval('self.cross_attn'+str(c_layer))(src, tgt, value=tgt, attn_mask=None, key_padding_mask=None)
                src = src + eval('self.dropout1_'+str(c_layer))(src2)
                src = eval('self.norm1_'+str(c_layer))(src)
                src2 = eval('self.linear2_'+str(c_layer))(eval('self.dropout0_'+str(c_layer))(F.relu(eval('self.linear1_'+str(c_layer))(src))))
                src = src + eval('self.dropout2_'+str(c_layer))(src2)
                src = eval('self.norm2_'+str(c_layer))(src)

                feat_fuse = src.view(1, N, h, w, c).contiguous().permute(0, 1, 4, 2, 3).contiguous()
                # print(feat_fuse.shape)  # torch.Size([1, N, 64, 200, 200])
                batch_updated_features[batch_i, :N, :, :, :] = feat_fuse.squeeze(0)
        
        return batch_updated_features, None


if __name__=="__main__":
    from icecream import ic
    x = torch.rand((64,6,8))  # [C,H,W]
    temperature = 10000
    num_pos_feats = x.shape[-3] // 2  # [d]

    mask = torch.zeros([x.shape[-2], x.shape[-1]], dtype=torch.bool, device=x.device)  #[H, W]
    not_mask = ~mask
    y_embed = not_mask.cumsum(0, dtype=torch.float32)  # [H, W]
    x_embed = not_mask.cumsum(1, dtype=torch.float32)  # [H, W]

    dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=x.device)  # [0,1,2,...,d]
    dim_t = temperature ** (2 * (dim_t // 2) / num_pos_feats)  # 10000^(2k/d), k is [0,0,1,1,...,d/2,d/2]

    pos_x = x_embed[:, :, None] / dim_t
    pos_y = y_embed[:, :, None] / dim_t

    pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2)
    pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2)
    pos = torch.cat((pos_y, pos_x), dim=2).permute(2, 0, 1)

opencood/models/fuse_modules/transformer.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu


"""
Implementation of transformer encoder fusion.
"""

import torch
import torch.nn as nn

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
import torch.nn.functional as F
from icecream import ic
from matplotlib import pyplot as plt

class MultiheadAttBlock(nn.Module):
    def __init__(self, channels, n_head=8, dropout=0):
        super(MultiheadAttBlock, self).__init__()
        self.attn = nn.MultiheadAttention(channels, n_head, dropout)

    def forward(self, q, k, v):
        """
        order (seq, batch, feature)
        Args:
            q: (1, H*W, C)
            k: (N, H*W, C)
            v: (N, H*W, C)
        Returns:
            outputs: ()
        """
        context, weight = self.attn(q,k,v) # (1, H*W, C)

        return context


class EncodeLayer(nn.Module):
    def __init__(self, channels, n_head=8, dropout=0):
        super(EncodeLayer, self).__init__()
        self.attn = nn.MultiheadAttention(channels, n_head, dropout)
        self.linear1 = nn.Linear(channels, channels)
        self.linear2 = nn.Linear(channels, channels)

        self.norm1 = nn.LayerNorm(channels)
        self.norm2 = nn.LayerNorm(channels)

        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.relu = nn.ReLU()

    def forward(self, q, k, v):
        """
        order (seq, batch, feature)
        Args:
            q: (1, H*W, C)
            k: (N, H*W, C)
            v: (N, H*W, C)
        Returns:
            outputs: ()
        """
        residual = q
        context, weight = self.attn(q,k,v) # (1, H*W, C)
        context = self.dropout1(context)
        output1 = self.norm1(residual + context)

        # feed forward net
        residual = output1 # (1, H*W, C)
        context = self.linear2(self.relu(self.linear1(output1)))
        context = self.dropout2(context)
        output2 = self.norm2(residual + context)

        return output2





class TransformerFusion(nn.Module):
    def __init__(self, args):
        super(TransformerFusion, self).__init__()
        
        self.channels = args['in_channels']
        self.n_head = args['n_head']
        self.dropout = args['dropout_rate']

        self.discrete_ratio = args['voxel_size'][0]
        self.downsample_rate = args['downsample_rate']

        self.encode_layer = EncodeLayer(self.channels, self.n_head, self.dropout)
    
    def add_pe_map(self, x, normalized=True):
        # scale = 2 * math.pi
        temperature = 10000
        num_pos_feats = x.shape[-3] // 2  # positional encoding dimension. C = 2d

        mask = torch.zeros([x.shape[-2], x.shape[-1]], dtype=torch.bool, device=x.device)  #[H, W]
        not_mask = ~mask
        y_embed = not_mask.cumsum(0, dtype=torch.float32)  # [H, W]
        x_embed = not_mask.cumsum(1, dtype=torch.float32)  # [H, W]

        dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=x.device)  # [0,1,2,...,d]
        dim_t = temperature ** (2 * (dim_t // 2) / num_pos_feats)  # 10000^(2k/d), k is [0,0,1,1,...,d/2,d/2]

        pos_x = x_embed[:, :, None] / dim_t
        pos_y = y_embed[:, :, None] / dim_t
        pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2)
        pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2)
        pos = torch.cat((pos_y, pos_x), dim=2).permute(2, 0, 1)  # [C, H, W]

        if len(x.shape) == 4:
            x = x + pos[None,:,:,:]
        elif len(x.shape) == 5:
            x = x + pos[None,None,:,:,:]
        return x

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
            
        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        split_x = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L,1,H,W)).to(x)
                roi_mask[b,i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :],(H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times

        out = []
        # iterate each batch
        for b in range(B):

            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            updated_node_features = []

            # update each node i
            i = 0 # ego
            # (N,1,H,W)
            mask = roi_mask[b, i, :N, ...]

            # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                            t_matrix[i, :, :, :],
                                            (H, W))

            neighbor_feature_flat = neighbor_feature.view(N,C,H*W)  # (N, C, H*W)
            neighbor_feature_flat_pe = self.add_pe_map(neighbor_feature).view(N,C,H*W)  # (N, C, H*W)
            
            query = neighbor_feature_flat_pe[0:1,...].permute(0,2,1)  # (1, H*W, C)
            key = neighbor_feature_flat_pe.permute(0,2,1)  # (N, H*W, C)
            value = neighbor_feature_flat.permute(0,2,1)



            fusion_result = self.encode_layer(query, key, value)  # (1, H*W, C)
            fusion_result = fusion_result.permute(0,2,1).reshape(1, C, H, W)[0]

            out.append(fusion_result)
            
        out = torch.stack(out)
        
        return out









opencood/models/fuse_modules/transformer_fuse.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
Implementation of V2VNet Fusion
"""

from email import message_from_binary_file
import torch
import torch.nn as nn

from opencood.models.sub_modules.torch_transformation_utils import \
    get_discretized_transformation_matrix, get_transformation_matrix, \
    warp_affine_simple, get_rotated_roi
from opencood.models.sub_modules.convgru import ConvGRU


class V2VNetFusion(nn.Module):
    def __init__(self, args):
        super(V2VNetFusion, self).__init__()
        
        in_channels = args['in_channels']
        H, W = args['conv_gru']['H'], args['conv_gru']['W'] # remember to modify for v2xsim dataset
        kernel_size = args['conv_gru']['kernel_size']
        num_gru_layers = args['conv_gru']['num_layers']

        self.discrete_ratio = args['voxel_size'][0]  
        self.downsample_rate = args['downsample_rate']  
        self.num_iteration = args['num_iteration']
        self.gru_flag = args['gru_flag']
        self.agg_operator = args['agg_operator']

        self.msg_cnn = nn.Conv2d(in_channels * 2, in_channels, kernel_size=3,
                                 stride=1, padding=1)
        self.conv_gru = ConvGRU(input_size=(H, W),
                                input_dim=in_channels * 2,
                                hidden_dim=[in_channels] * num_gru_layers,
                                kernel_size=kernel_size,
                                num_layers=num_gru_layers,
                                batch_first=True,
                                bias=True,
                                return_all_layers=False)
        self.mlp = nn.Linear(in_channels, in_channels)

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, x, record_len, pairwise_t_matrix, weight=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
            
        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 
        
        weight: torch.Tensor
            Weight of aggregating coming message
            shape: (B, L, L)
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # split x:[(L1, C, H, W), (L2, C, H, W), ...]
        # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
        split_x = self.regroup(x, record_len)

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        # (B*L,L,1,H,W)
        roi_mask = torch.zeros((B, L, L, 1, H, W)).to(x)
        for b in range(B):
            N = record_len[b]
            for i in range(N):
                one_tensor = torch.ones((L,1,H,W)).to(x)
                roi_mask[b,i] = warp_affine_simple(one_tensor, pairwise_t_matrix[b][i, :, :, :],(H, W))

        batch_node_features = split_x
        # iteratively update the features for num_iteration times
        for l in range(self.num_iteration):

            batch_updated_node_features = []
            # iterate each batch
            for b in range(B):

                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

                updated_node_features = []

                # update each node i
                for i in range(N):
                    # (N,1,H,W)
                    mask = roi_mask[b, i, :N, ...]
                    # (N,C,H,W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
                    # Notice we put i one the first dim of t_matrix. Different from original.
                    # t_matrix[i,j] = Tji
                    neighbor_feature = warp_affine_simple(batch_node_features[b],
                                                   t_matrix[i, :, :, :],
                                                   (H, W))

                    # (N,C,H,W)
                    ego_agent_feature = batch_node_features[b][i].unsqueeze(
                        0).repeat(N, 1, 1, 1)
                    #(N,2C,H,W)
                    neighbor_feature = torch.cat(
                        [neighbor_feature, ego_agent_feature], dim=1)
                    # (N,C,H,W)
                    # message contains all feature map from j to ego i.
                    message = self.msg_cnn(neighbor_feature) * mask

                    # (C,H,W)
                    if self.agg_operator=="avg":
                        agg_feature = torch.mean(message, dim=0)
                    elif self.agg_operator=="max":
                        agg_feature = torch.max(message, dim=0)[0]
                    elif self.agg_operator=='weight':
                        agg_feature = torch.sum(message * weight[b][i,:N].view(-1,1,1,1), dim=0)
                    else:
                        raise ValueError("agg_operator has wrong value")
                    # (2C, H, W)
                    cat_feature = torch.cat(
                        [batch_node_features[b][i, ...], agg_feature], dim=0)
                    # (C,H,W)
                    if self.gru_flag:
                        gru_out = \
                            self.conv_gru(cat_feature.unsqueeze(0).unsqueeze(0))[
                                0][
                                0].squeeze(0).squeeze(0)
                    else:
                        gru_out = batch_node_features[b][i, ...] + agg_feature
                    updated_node_features.append(gru_out.unsqueeze(0))
                # (N,C,H,W)
                batch_updated_node_features.append(
                    torch.cat(updated_node_features, dim=0))
            batch_node_features = batch_updated_node_features
        # (B,C,H,W)
        out = torch.cat(
            [itm[0, ...].unsqueeze(0) for itm in batch_node_features], dim=0)
        # (B,C,H,W) -> (B, H, W, C) -> (B,C,H,W)
        out = self.mlp(out.permute(0, 2, 3, 1)).permute(0, 3, 1, 2)

        return out



opencood/models/fuse_modules/v2v_fuse.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
Implementation of V2VNet Fusion
"""

from turtle import update
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np

from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
from opencood.models.comm_modules.where2comm import Communication


class ScaledDotProductAttention(nn.Module):
    """
    Scaled Dot-Product Attention proposed in "Attention Is All You Need"
    Compute the dot products of the query with all keys, divide each by sqrt(dim),
    and apply a softmax function to obtain the weights on the values
    Args: dim, mask
        dim (int): dimention of attention
        mask (torch.Tensor): tensor containing indices to be masked
    Inputs: query, key, value, mask
        - **query** (batch, q_len, d_model): tensor containing projection
          vector for decoder.
        - **key** (batch, k_len, d_model): tensor containing projection
          vector for encoder.
        - **value** (batch, v_len, d_model): tensor containing features of the
          encoded input sequence.
        - **mask** (-): tensor containing indices to be masked
    Returns: context, attn
        - **context**: tensor containing the context vector from
          attention mechanism.
        - **attn**: tensor containing the attention (alignment) from the
          encoder outputs.
    """

    def __init__(self, dim):
        super(ScaledDotProductAttention, self).__init__()
        self.sqrt_dim = np.sqrt(dim)

    def forward(self, query, key, value):
        score = torch.bmm(query, key.transpose(1, 2)) / self.sqrt_dim
        attn = F.softmax(score, -1)
        context = torch.bmm(attn, value)
        return context

class AttenFusion(nn.Module):
    def __init__(self, feature_dim):
        super(AttenFusion, self).__init__()
        self.att = ScaledDotProductAttention(feature_dim)

    def forward(self, x):
        cav_num, C, H, W = x.shape
        x = x.view(cav_num, C, -1).permute(2, 0, 1) #  (H*W, cav_num, C), perform self attention on each pixel.
        x = self.att(x, x, x)
        x = x.permute(1, 2, 0).view(cav_num, C, H, W)[0]  # C, W, H before
        return x

class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def forward(self, x):
        return torch.max(x, dim=0)[0]


class EncodeLayer(nn.Module):
    def __init__(self, channels, n_head=8, dropout=0):
        super(EncodeLayer, self).__init__()
        self.attn = nn.MultiheadAttention(channels, n_head, dropout)
        self.linear1 = nn.Linear(channels, channels)
        self.linear2 = nn.Linear(channels, channels)

        self.norm1 = nn.LayerNorm(channels)
        self.norm2 = nn.LayerNorm(channels)

        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

        self.relu = nn.ReLU()

    def forward(self, q, k, v, confidence_map=None):
        """
        order (seq, batch, feature)
        Args:
            q: (1, H*W, C)
            k: (N, H*W, C)
            v: (N, H*W, C)
        Returns:
            outputs: ()
        """
        residual = q
        if confidence_map is not None:
            context, weight = self.attn(q,k,v, quality_map=confidence_map) # (1, H*W, C)
        else:
            context, weight = self.attn(q,k,v) # (1, H*W, C)
        context = self.dropout1(context)
        output1 = self.norm1(residual + context)

        # feed forward net
        residual = output1 # (1, H*W, C)
        context = self.linear2(self.relu(self.linear1(output1)))
        context = self.dropout2(context)
        output2 = self.norm2(residual + context)

        return output2

class TransformerFusion(nn.Module):
    def __init__(self, channels=256, n_head=8, with_spe=True, with_scm=True, dropout=0):
        super(TransformerFusion, self).__init__()

        self.encode_layer = EncodeLayer(channels, n_head, dropout)
        self.with_spe = with_spe
        self.with_scm = with_scm
        
    def forward(self, batch_neighbor_feature, batch_neighbor_feature_pe, batch_confidence_map, record_len):
        x_fuse = []
        B = len(record_len)
        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            neighbor_feature = batch_neighbor_feature[b]
            _, C, H, W = neighbor_feature.shape
            neighbor_feature_flat = neighbor_feature.view(N,C,H*W)  # (N, C, H*W)

            if self.with_spe:
                neighbor_feature_pe = batch_neighbor_feature_pe[b]
                neighbor_feature_flat_pe = neighbor_feature_pe.view(N,C,H*W)  # (N, C, H*W)
                query = neighbor_feature_flat_pe[0:1,...].permute(0,2,1)  # (1, H*W, C)
                key = neighbor_feature_flat_pe.permute(0,2,1)  # (N, H*W, C)
            else:
                query = neighbor_feature_flat[0:1,...].permute(0,2,1)  # (1, H*W, C)
                key = neighbor_feature_flat.permute(0,2,1)  # (N, H*W, C)
            
            value = neighbor_feature_flat.permute(0,2,1)

            if self.with_scm:
                confidence_map = batch_confidence_map[b]
                fused_feature = self.encode_layer(query, key, value, confidence_map)  # (1, H*W, C)
            else:
                fused_feature = self.encode_layer(query, key, value)  # (1, H*W, C)
            
            fused_feature = fused_feature.permute(0,2,1).reshape(1, C, H, W)

            x_fuse.append(fused_feature)
        x_fuse = torch.concat(x_fuse, dim=0)
        return x_fuse

def add_pe_map(x):
    # scale = 2 * math.pi
    temperature = 10000
    num_pos_feats = x.shape[-3] // 2  # positional encoding dimension. C = 2d

    mask = torch.zeros([x.shape[-2], x.shape[-1]], dtype=torch.bool, device=x.device)  #[H, W]
    not_mask = ~mask
    y_embed = not_mask.cumsum(0, dtype=torch.float32)  # [H, W]
    x_embed = not_mask.cumsum(1, dtype=torch.float32)  # [H, W]

    dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=x.device)  # [0,1,2,...,d]
    dim_t = temperature ** (2 * (dim_t // 2) / num_pos_feats)  # 10000^(2k/d), k is [0,0,1,1,...,d/2,d/2]

    pos_x = x_embed[:, :, None] / dim_t
    pos_y = y_embed[:, :, None] / dim_t
    pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2)
    pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2)
    pos = torch.cat((pos_y, pos_x), dim=2).permute(2, 0, 1)  # [C, H, W]

    if len(x.shape) == 4:
        x_pe = x + pos[None,:,:,:]
    elif len(x.shape) == 5:
        x_pe = x + pos[None,None,:,:,:]
    return x_pe


class Where2comm(nn.Module):
    def __init__(self, args):
        super(Where2comm, self).__init__()

        self.communication = False
        self.round = 1
        if 'communication' in args:
            self.communication = True
            self.naive_communication = Communication(args['communication'])
            if 'round' in args['communication']:
                self.round = args['communication']['round']
        self.discrete_ratio = args['voxel_size'][0]  # voxel_size[0]=0.4    
        self.downsample_rate = args['downsample_rate']  # 2/4, downsample rate from original feature map [200, 704]
        
        self.agg_mode = args['agg_operator']['mode']
        self.multi_scale = args['multi_scale']
        if self.multi_scale:
            layer_nums = args['layer_nums']
            num_filters = args['num_filters']
            self.num_levels = len(layer_nums)
            self.fuse_modules = nn.ModuleList()
            for idx in range(self.num_levels):
                if self.agg_mode == 'ATTEN':
                    fuse_network = AttenFusion(num_filters[idx])
                elif self.agg_mode == 'MAX':
                    fuse_network = MaxFusion()
                elif self.agg_mode == 'Transformer':
                    fuse_network = TransformerFusion(
                                                channels=num_filters[idx], 
                                                n_head=args['agg_operator']['n_head'], 
                                                with_spe=args['agg_operator']['with_spe'], 
                                                with_scm=args['agg_operator']['with_scm'])
                self.fuse_modules.append(fuse_network)
        else:
            if self.agg_mode == 'ATTEN':
                self.fuse_modules = AttenFusion(args['agg_operator']['feature_dim'])
            elif self.agg_mode == 'MAX':
                self.fuse_modules = MaxFusion()   
            elif self.agg_mode == 'Transformer':
                self.fuse_network = TransformerFusion(
                                            channels=args['agg_operator']['feature_dim'], 
                                            n_head=args['agg_operator']['n_head'], 
                                            with_spe=args['agg_operator']['with_spe'], 
                                            with_scm=args['agg_operator']['with_scm'])     

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, x, rm, record_len, pairwise_t_matrix, backbone=None, heads=None):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
            
        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The transformation matrix from each cav to ego, 
            shape: (B, L, L, 4, 4) 
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

        # rm = torch.ones_like(rm) # 所有地块进行通信，debug用，需注释掉
        if self.multi_scale:
            ups = []
            # backbone.__dict__()
            with_resnet = True if hasattr(backbone, 'resnet') else False
            if with_resnet:
                feats = backbone.resnet(x)
            
            for i in range(self.num_levels):
                x = feats[i] if with_resnet else backbone.blocks[i](x)

                ############ 1. Communication (Mask the features) #########
                if i==0:
                    if self.communication:
                        batch_confidence_maps = self.regroup(rm, record_len)# cls_head 出来的 map 
                        _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
                        x = x * communication_masks
                    else:
                        communication_rates = torch.tensor(0).to(x.device)
                
                ############ 2. Split the confidence map #######################
                # split x:[(L1, C, H, W), (L2, C, H, W), ...]
                # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
                batch_node_features = self.regroup(x, record_len)
                
                ############ 3. Fusion ####################################
                x_fuse = []
                for b in range(B):
                    # number of valid agent
                    N = record_len[b]
                    # (N,N,4,4)
                    # t_matrix[i, j]-> from i to j
                    t_matrix = pairwise_t_matrix[b][:N, :N, :, :]
                    node_features = batch_node_features[b]
                    C, H, W = node_features.shape[1:]
                    neighbor_feature = warp_affine_simple(node_features,
                                                    t_matrix[0, :, :, :],
                                                    (H, W))
                    x_fuse.append(self.fuse_modules[i](neighbor_feature))
                x_fuse = torch.stack(x_fuse)

                ############ 4. Deconv ####################################
                if len(backbone.deblocks) > 0:
                    ups.append(backbone.deblocks[i](x_fuse))
                else:
                    ups.append(x_fuse)
                
            if len(ups) > 1:
                x_fuse = torch.cat(ups, dim=1)
            elif len(ups) == 1:
                x_fuse = ups[0]
            
            if len(backbone.deblocks) > self.num_levels:
                x_fuse = backbone.deblocks[-1](x_fuse)
        else:
            ############ 1. Split the features #######################
            # split x:[(L1, C, H, W), (L2, C, H, W), ...]
            # for example [[2, 256, 50, 176], [1, 256, 50, 176], ...]
            batch_node_features = self.regroup(x, record_len)
            batch_confidence_maps = self.regroup(rm, record_len)

            ############ 2. Communication (Mask the features) #########
            if self.communication:
                _, communication_masks, communication_rates = self.naive_communication(batch_confidence_maps, record_len, pairwise_t_matrix)
            else:
                communication_rates = torch.tensor(0).to(x.device)
            
            ############ 3. Fusion ####################################
            x_fuse = []
            for b in range(B):
                # number of valid agent
                N = record_len[b]
                # (N,N,4,4)
                # t_matrix[i, j]-> from i to j
                t_matrix = pairwise_t_matrix[b][:N, :N, :, :]
                node_features = batch_node_features[b]
                if self.communication:
                    node_features = node_features * communication_masks[b]
                neighbor_feature = warp_affine_simple(node_features,
                                                t_matrix[0, :, :, :],
                                                (H, W))
                x_fuse.append(self.fuse_modules(neighbor_feature))
            x_fuse = torch.stack(x_fuse)
        
        return x_fuse, communication_rates, {}


opencood/models/fuse_modules/where2comm_attn.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import math

import torch.nn as nn

from opencood.models.fuse_modules.self_attn import AttFusion
from opencood.models.pixor import Bottleneck, BackBone, Header


class BackBoneIntermediate(BackBone):

    def __init__(self, block, num_block, geom, use_bn=True):
        super(BackBoneIntermediate, self).__init__(block,
                                                   num_block,
                                                   geom, use_bn)

        self.fusion_net3 = AttFusion(192)
        self.fusion_net4 = AttFusion(256)
        self.fusion_net5 = AttFusion(384)

    def forward(self, x, record_len):
        # Here c3, c4, c5 includes all cav
        c3, c4, c5 = self.encode(x)

        # Here c3, c4, c5 only include ego 
        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


class PIXORIntermediate(nn.Module):
    """
    The Pixor backbone. The input of PIXOR nn module is a tensor of
    [batch_size, height, weight, channel], The output of PIXOR nn module
    is also a tensor of [batch_size, height/4, weight/4, channel].  Note that
     we convert the dimensions to [C, H, W] for PyTorch's nn.Conv2d functions

    Parameters
    ----------
    args : dict
        The arguments of the model.

    Attributes
    ----------
    backbone : opencood.object
        The backbone used to extract features.
    header : opencood.object
        Header used to predict the classification and coordinates.
    """

    def __init__(self, args):
        super(PIXORIntermediate, self).__init__()
        geom = args["geometry_param"]
        use_bn = args["use_bn"]
        self.backbone = BackBoneIntermediate(Bottleneck, [3, 6, 6, 3],
                                             geom,
                                             use_bn)
        self.header = Header(use_bn)

        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

        prior = 0.01
        self.header.clshead.weight.data.fill_(-math.log((1.0 - prior) / prior))
        self.header.clshead.bias.data.fill_(0)
        self.header.reghead.weight.data.fill_(0)
        self.header.reghead.bias.data.fill_(0)

    def forward(self, data_dict):
        bev_input = data_dict['processed_lidar']["bev_input"]
        record_len = data_dict['record_len']

        features = self.backbone(bev_input, record_len)
        # cls -- (N, 1, W/4, L/4)
        # reg -- (N, 6, W/4, L/4)
        cls, reg = self.header(features)

        output_dict = {
            "cls": cls,
            "reg": reg
        }

        return output_dict

opencood/models/pixor_intermediate.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, OpenPCDet
# License: TDG-Attribution-NonCommercial-NoDistrib


import torch
import torch.nn as nn


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
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv

class PointPillar(nn.Module):
    def __init__(self, args):
        super(PointPillar, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        
        # self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        # self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
        #                           kernel_size=1)
        # self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
        #                           kernel_size=1)

        # Base BEV Backbone -> ResNetBEVBackbone
        # self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        if 'resnet' in args['base_bev_backbone'] and args['base_bev_backbone']['resnet']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        ############################################
        
        self.out_channel = sum(args['base_bev_backbone']['num_upsample_filter'])

        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
            self.out_channel = args['shrink_header']['dim'][-1]

        self.cls_head = nn.Conv2d(self.out_channel, args['anchor_number'], # 384
                                kernel_size=1)
        self.reg_head = nn.Conv2d(self.out_channel, 7 * args['anchor_number'], # 384
                                kernel_size=1)

        if 'dir_args' in args.keys():
            self.use_dir = True
            self.dir_head = nn.Conv2d(self.out_channel, args['dir_args']['num_bins'] * args['anchor_number'],
                                kernel_size=1) # BIN_NUM = 2， # 384
        else:
            self.use_dir = False

    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_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points}

        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)

        spatial_features_2d = batch_dict['spatial_features_2d']

        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)

        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)

        output_dict = {'psm': psm,
                       'rm': rm}
                       
        if self.use_dir:
            dm = self.dir_head(spatial_features_2d)
            output_dict.update({'dm': dm})

        return output_dict
opencood/models/point_pillar.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>, Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import torch.nn as nn
import torch

from opencood.models.sub_modules.pillar_vfe import PillarVFE
from opencood.models.sub_modules.point_pillar_scatter import PointPillarScatter
from opencood.models.sub_modules.deformable_transformer_backbone import DeformableTransformerBackbone
from opencood.models.sub_modules.base_bev_backbone import BaseBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
from opencood.data_utils.post_processor import SimpleVoxelPostprocessor
from opencood.models.point_pillar import PointPillar
from opencood.utils.transformation_utils import get_pairwise_transformation_torch
from opencood.utils.model_utils import weight_init


class PointPillarDeformTransformer(nn.Module):
    def __init__(self, args):
        super(PointPillarDeformTransformer, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = DeformableTransformerBackbone(args['deformable_transfomer_backbone'])
        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])

        self.cls_head = nn.Conv2d(128 * 2, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 2, 7 * args['anchor_number'],
                                  kernel_size=1)
        if args['backbone_fix']:
            self.backbone_fix()
        
        self.apply(weight_init)
            

    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']
        record_len = data_dict['record_len']
        lidar_pose = data_dict['lidar_pose']  # [sum(cav), 6]

        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len,
                      'pairwise_t_matrix': pairwise_t_matrix,
                      'lidar_pose': lidar_pose}


        
        # n, 4 -> n, c
        batch_dict = self.pillar_vfe(batch_dict)
        # n, c -> N, C, H, W
        batch_dict = self.scatter(batch_dict)

        batch_dict = self.backbone(batch_dict)
        # N, C, H', W'. [N, 384, 100, 352]
        spatial_features_2d = batch_dict['spatial_features_2d']

        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)

        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)


        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict

opencood/models/point_pillar_deform_transformer.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu, Yiming Li, Dekun Ma

import torch
import torch.nn as nn
import numpy as np
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
import torch.nn.functional as F
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
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
from opencood.data_utils.post_processor import UncertaintyVoxelPostprocessor
from opencood.models.sub_modules.downsample_conv import DownsampleConv

def generate_noise(pos_std, rot_std, pos_mean=0, rot_mean=0):
    """ Add localization error to the 6dof pose
        Noise includes position (x,y) and rotation (yaw).
        We use gaussian distribution to generate noise.

    Args:

        pos_std : float 
            std of gaussian dist, in meter

        rot_std : float
            std of gaussian dist, in degree

        pos_mean : float
            mean of gaussian dist, in meter

        rot_mean : float
            mean of gaussian dist, in degree

    Returns:
        pose_noise: np.ndarray, [6,]
            [x, y, z, roll, yaw, pitch]
    """

    xy = np.random.normal(pos_mean, pos_std, size=(2))
    yaw = np.random.normal(rot_mean, rot_std, size=(1))

    pose_noise = np.array([xy[0], xy[1], 0, 0, yaw[0], 0])

    return pose_noise

def get_pairwise_transformation2ego(past_k_lidar_pose, noise_level, max_cav=5):
    """
    Get transformation matrixes accross different agents to curr ego at all past timestamps.

    Parameters
    ----------
    past_k_lidar_pose: np.array [N, k=1, 6]
    
    ego_pose : list
        ego pose

    max_cav : int
        The maximum number of cav, default 5

    Return
    ------
    pairwise_t_matrix : np.array
        The transformation matrix each cav to curr ego at past k frames.
        shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
        pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
    """
    pos_std = noise_level['pos_std']
    rot_std = noise_level['rot_std']
    pos_mean = 0 
    rot_mean = 0
    
    pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

    ego_pose = past_k_lidar_pose[0, 0]

    t_list = []

    # save all transformation matrix in a list in order first.
    for cav_id in range(past_k_lidar_pose.shape[0]):
        loc_noise = generate_noise(pos_std, rot_std)
        t_list.append(x_to_world(past_k_lidar_pose[cav_id, 0].cpu().numpy()+loc_noise)) # Twx
    
    ego_pose = x_to_world(ego_pose.cpu().numpy())

    for i in range(len(t_list)):
        for j in range(len(t_list)):
            # identity matrix to self
            if i != j:
                # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                pairwise_t_matrix[i, j] = t_matrix

    pairwise_t_matrix = torch.tensor(pairwise_t_matrix).to(past_k_lidar_pose.device)
    return pairwise_t_matrix


def regroup(x, record_len):
    cum_sum_len = torch.cumsum(record_len, dim=0)
    split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
    return split_x

class PointPillarDiscoNet(nn.Module):
    def __init__(self, args):
        super(PointPillarDiscoNet, self).__init__()
        self.discrete_ratio = args['voxel_size'][0]
        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        self.pixel_weight_layer = PixelWeightLayer(128 * 3)
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_num'],
                                  kernel_size=1)

        self.noise_flag = 0
        if 'noise' in args.keys():
            self.noise_flag = True
            self.noise_level = {'pos_std': args['noise']['pos_std'], 'rot_std': args['noise']['rot_std'], 'pos_mean': args['noise']['pos_mean'], 'rot_mean': args['noise']['rot_mean']}
            print(f'=== noise level : {self.noise_level} ===')

    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']

        # teacher_voxel_features = data_dict['teacher_processed_lidar']['voxel_features']
        # teacher_voxel_coords = data_dict['teacher_processed_lidar']['voxel_coords']
        # teacher_voxel_num_points = data_dict['teacher_processed_lidar']['voxel_num_points']

        record_len = data_dict['record_len']
        # lidar_pose = data_dict['lidar_pose']
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len,
                      'pairwise_t_matrix': pairwise_t_matrix}



        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)

        _, _, H0, W0 = batch_dict['spatial_features'].shape

        batch_dict = self.backbone(batch_dict)


        spatial_features_2d = batch_dict['spatial_features_2d']

        noise_pairwise_t_matrix = None
        B = record_len.shape[0]
        if self.noise_flag and B==1:
            # noise_level = {'pos_std': 0.5, 'rot_std': 0, 'pos_mean': 0, 'rot_mean': 0}
            noise_pairwise_t_matrix = get_pairwise_transformation2ego(data_dict['past_lidar_pose'], self.noise_level, max_cav=5).unsqueeze(0)
            
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)
        ########## FUSION SRART ##########
        # we concat ego's feature with other agent
        # first transform feature to ego's coordinate
        split_x = regroup(spatial_features_2d, record_len)

        B = pairwise_t_matrix.shape[0]
        _, C, H, W = spatial_features_2d.shape

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (W0 * self.discrete_ratio) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (H0 * self.discrete_ratio) * 2

        if noise_pairwise_t_matrix is not None:
            noise_pairwise_t_matrix = noise_pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] 
            noise_pairwise_t_matrix[...,0,1] = noise_pairwise_t_matrix[...,0,1] * H / W
            noise_pairwise_t_matrix[...,1,0] = noise_pairwise_t_matrix[...,1,0] * W / H
            noise_pairwise_t_matrix[...,0,2] = noise_pairwise_t_matrix[...,0,2] / (W0 * self.discrete_ratio) * 2
            noise_pairwise_t_matrix[...,1,2] = noise_pairwise_t_matrix[...,1,2] / (H0 * self.discrete_ratio) * 2

        out = []

        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            if noise_pairwise_t_matrix is not None:
                t_matrix = noise_pairwise_t_matrix[b][:N, :N, :, :] #(N, 2, 3)
            else:
                t_matrix = pairwise_t_matrix[b][:N, :N, :, :] #(N, 2, 3)
            # t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0 # ego
            # (N, C, H, W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(split_x[b],
                                            t_matrix[i, :, :, :],
                                            (H, W))
            # (N, C, H, W)
            ego_feature = split_x[b][0].view(1, C, H, W).expand(N, -1, -1, -1)
            # (N, 2C, H, W)
            neighbor_feature_cat = torch.cat((neighbor_feature, ego_feature), dim=1)
            # (N, 1, H, W)
            agent_weight = self.pixel_weight_layer(neighbor_feature_cat) 
            # (N, 1, H, W)
            agent_weight = F.softmax(agent_weight, dim=0)

            agent_weight = agent_weight.expand(-1, C, -1, -1)
            # (N, C, H, W)
            feature_fused = torch.sum(agent_weight * neighbor_feature, dim=0)
            out.append(feature_fused)

        spatial_features_2d = torch.stack(out)


        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)

        output_dict = {'feature': spatial_features_2d,
                       'psm': psm,
                       'rm': rm}

        return output_dict




class PixelWeightLayer(nn.Module):
    def __init__(self, channel):
        super(PixelWeightLayer, self).__init__()

        self.conv1_1 = nn.Conv2d(channel * 2, 128, kernel_size=1, stride=1, padding=0)
        self.bn1_1 = nn.BatchNorm2d(128)

        self.conv1_2 = nn.Conv2d(128, 32, kernel_size=1, stride=1, padding=0)
        self.bn1_2 = nn.BatchNorm2d(32)

        self.conv1_3 = nn.Conv2d(32, 8, kernel_size=1, stride=1, padding=0)
        self.bn1_3 = nn.BatchNorm2d(8)

        self.conv1_4 = nn.Conv2d(8, 1, kernel_size=1, stride=1, padding=0)
        # self.bn1_4 = nn.BatchNorm2d(1)

    def forward(self, x):
        x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
        x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
        x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
        x_1 = F.relu(self.bn1_3(self.conv1_3(x_1)))
        x_1 = F.relu(self.conv1_4(x_1))

        return x_1
opencood/models/point_pillar_disconet.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu, Yiming Li, Dekun Ma
# Modified: Sizhe Wei

import torch
import torch.nn as nn


import torch.nn.functional as F
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
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
from opencood.data_utils.post_processor import UncertaintyVoxelPostprocessor
from opencood.models.sub_modules.downsample_conv import DownsampleConv

def regroup(x, record_len):
    cum_sum_len = torch.cumsum(record_len, dim=0)
    split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
    return split_x

class PointPillarDiscoNetSyncNet(nn.Module):
    def __init__(self, args):
        super(PointPillarDiscoNetSyncNet, self).__init__()
        self.discrete_ratio = args['voxel_size'][0]
        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        self.pixel_weight_layer = PixelWeightLayer(128 * 3)
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_num'],
                                  kernel_size=1)

    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']

        teacher_voxel_features = data_dict['teacher_processed_lidar']['voxel_features']
        teacher_voxel_coords = data_dict['teacher_processed_lidar']['voxel_coords']
        teacher_voxel_num_points = data_dict['teacher_processed_lidar']['voxel_num_points']

        record_len = data_dict['record_len']
        lidar_pose = data_dict['lidar_pose']
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len,
                      'pairwise_t_matrix': pairwise_t_matrix}



        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)

        _, _, H0, W0 = batch_dict['spatial_features'].shape

        batch_dict = self.backbone(batch_dict)


        spatial_features_2d = batch_dict['spatial_features_2d']
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)
        ########## FUSION SRART ##########
        # we concat ego's feature with other agent
        # first transform feature to ego's coordinate
        split_x = regroup(spatial_features_2d, record_len)

        B = pairwise_t_matrix.shape[0]
        _, C, H, W = spatial_features_2d.shape

        # (B,L,L,2,3)
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (W0 * self.discrete_ratio) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (H0 * self.discrete_ratio) * 2

        out = []

        for b in range(B):
            # number of valid agent
            N = record_len[b]
            # (N,N,4,4)
            # t_matrix[i, j]-> from i to j
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0 # ego
            # (N, C, H, W) neighbor_feature is agent i's neighborhood warping to agent i's perspective
            # Notice we put i one the first dim of t_matrix. Different from original.
            # t_matrix[i,j] = Tji
            neighbor_feature = warp_affine_simple(split_x[b],
                                            t_matrix[i, :, :, :],
                                            (H, W))
            # (N, C, H, W)
            ego_feature = split_x[b][0].view(1, C, H, W).expand(N, -1, -1, -1)
            # (N, 2C, H, W)
            neighbor_feature_cat = torch.cat((neighbor_feature, ego_feature), dim=1)
            # (N, 1, H, W)
            agent_weight = self.pixel_weight_layer(neighbor_feature_cat) 
            # (N, 1, H, W)
            agent_weight = F.softmax(agent_weight, dim=0)

            agent_weight = agent_weight.expand(-1, C, -1, -1)
            # (N, C, H, W)
            feature_fused = torch.sum(agent_weight * neighbor_feature, dim=0)
            out.append(feature_fused)

        spatial_features_2d = torch.stack(out)


        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)

        output_dict = {'feature': spatial_features_2d,
                       'psm': psm,
                       'rm': rm}

        return output_dict




class PixelWeightLayer(nn.Module):
    def __init__(self, channel):
        super(PixelWeightLayer, self).__init__()

        self.conv1_1 = nn.Conv2d(channel * 2, 128, kernel_size=1, stride=1, padding=0)
        self.bn1_1 = nn.BatchNorm2d(128)

        self.conv1_2 = nn.Conv2d(128, 32, kernel_size=1, stride=1, padding=0)
        self.bn1_2 = nn.BatchNorm2d(32)

        self.conv1_3 = nn.Conv2d(32, 8, kernel_size=1, stride=1, padding=0)
        self.bn1_3 = nn.BatchNorm2d(8)

        self.conv1_4 = nn.Conv2d(8, 1, kernel_size=1, stride=1, padding=0)
        # self.bn1_4 = nn.BatchNorm2d(1)

    def forward(self, x):
        x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
        x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
        x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
        x_1 = F.relu(self.bn1_3(self.conv1_3(x_1)))
        x_1 = F.relu(self.conv1_4(x_1))

        return x_1
opencood/models/point_pillar_disconet_syncnet.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu

import torch
import torch.nn as nn


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


class PointPillarDiscoNetTeacher(nn.Module):
    def __init__(self, args):
        super(PointPillarDiscoNetTeacher, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        self.cls_head = nn.Conv2d(128 * 3, args['anchor_num'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_num'],
                                  kernel_size=1)

    def forward(self, data_dict):

        voxel_features = data_dict['teacher_processed_lidar']['voxel_features']
        voxel_coords = data_dict['teacher_processed_lidar']['voxel_coords']
        voxel_num_points = data_dict['teacher_processed_lidar']['voxel_num_points']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points}
        

        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)

        spatial_features_2d = batch_dict['spatial_features_2d']

        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)

        output_dict = {'teacher_feature': spatial_features_2d,
                       'teacher_psm': psm,
                       'teacher_rm': rm}

        return output_dict
opencood/models/point_pillar_disconet_teacher.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

# In CoAlign paper, I use point_pillar_max.py for F-Cooper Implementation.


import torch.nn as nn

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
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
from opencood.models.fuse_modules.f_cooper_fuse import SpatialFusion


class PointPillarFCooper(nn.Module):
    """
    F-Cooper implementation with point pillar backbone.
    """
    def __init__(self, args):
        super(PointPillarFCooper, self).__init__()

        self.max_cav = args['max_cav']
        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        self.fusion_net = SpatialFusion()

        self.cls_head = nn.Conv2d(128 * 2, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 2, 7 * args['anchor_number'],
                                  kernel_size=1)

        if args['backbone_fix']:
            self.backbone_fix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay。
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False

    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']
        record_len = data_dict['record_len']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c
        batch_dict = self.pillar_vfe(batch_dict)
        # n, c -> N, C, H, W
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)

        spatial_features_2d = batch_dict['spatial_features_2d']
        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        fused_feature = self.fusion_net(spatial_features_2d, record_len)

        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)

        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict

opencood/models/point_pillar_fcooper.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import torch
import torch.nn as nn


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):
    def __init__(self, args):
        super(PointPillarIntermediate, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = AttBEVBackbone(args['base_bev_backbone'], 64)


        self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_num'],
                                  kernel_size=1)

    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']
        record_len = data_dict['record_len']
        lidar_pose = data_dict['lidar_pose']
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len,
                      'pairwise_t_matrix': pairwise_t_matrix}
            


        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)

        batch_dict = self.backbone(batch_dict)
        spatial_features_2d = batch_dict['spatial_features_2d']

        if self.dcn and not self.before_backbone:
            spatial_features_2d = self.dcn_net(spatial_features_2d)

        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)

        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict
opencood/models/point_pillar_intermediate.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import torch
import torch.nn as nn


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_resnet_bev_backbone import AttResNetBEVBackbone
from opencood.models.sub_modules.dcn_net import DCNNet
from opencood.utils.transformation_utils import get_pairwise_transformation_torch



class PointPillarIntermediateResNet(nn.Module):
    def __init__(self, args):
        super(PointPillarIntermediateResNet, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = AttResNetBEVBackbone(args['base_bev_backbone'], 64)


        self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
                                  kernel_size=1)

    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']
        record_len = data_dict['record_len']
        lidar_pose = data_dict['lidar_pose']  # [sum(cav), 6]
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len,
                      'pairwise_t_matrix': pairwise_t_matrix}
            


        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)

        batch_dict = self.backbone(batch_dict)
        spatial_features_2d = batch_dict['spatial_features_2d']


        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)

        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict
opencood/models/point_pillar_intermediate_resnet.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import torch
import torch.nn as nn


from opencood.models.sub_modules.pillar_vfe import PillarVFE
from opencood.models.sub_modules.point_pillar_scatter import PointPillarScatter
from opencood.models.sub_modules.max_resnet_bev_backbone import MaxResNetBEVBackbone
from opencood.models.sub_modules.dcn_net import DCNNet
from opencood.utils.transformation_utils import get_pairwise_transformation_torch



class PointPillarIntermediateResNetMax(nn.Module):
    def __init__(self, args):
        super(PointPillarIntermediateResNetMax, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = MaxResNetBEVBackbone(args['base_bev_backbone'], 64)

        self.dcn = False
        if 'dcn' in args:
            self.dcn = True
            self.before_backbone = args['dcn']['before_backbone']
            self.dcn_net = DCNNet(args['dcn'])

        self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
                                  kernel_size=1)

    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']
        record_len = data_dict['record_len']
        lidar_pose = data_dict['lidar_pose']  # [sum(cav), 6]
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len,
                      'pairwise_t_matrix': pairwise_t_matrix}
            


        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)

        if self.dcn and self.before_backbone:
            batch_dict['spatial_features'] = self.dcn_net(batch_dict['spatial_features'])

        batch_dict = self.backbone(batch_dict)
        spatial_features_2d = batch_dict['spatial_features_2d']


        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)

        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict
opencood/models/point_pillar_intermediate_resnet_max.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import torch
import torch.nn as nn


from opencood.models.sub_modules.pillar_vfe import PillarVFE
from opencood.models.sub_modules.point_pillar_scatter import PointPillarScatter
from opencood.models.sub_modules.max_resnet_bev_backbone import MaxResNetBEVBackbone
from opencood.utils.transformation_utils import get_pairwise_transformation_torch
from opencood.models.sub_modules.downsample_conv import DownsampleConv


class PointPillarIntermediateResNetMaxylu(nn.Module):
    def __init__(self, args):
        super(PointPillarIntermediateResNetMaxylu, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = MaxResNetBEVBackbone(args['base_bev_backbone'], 64)

        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])

        self.cls_head = nn.Conv2d(128 * 2, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 2, 7 * args['anchor_number'],
                                  kernel_size=1)
        if 'dir_args' in args.keys():
            self.use_dir = True
            self.dir_head = nn.Conv2d(128 * 2, args['dir_args']['num_bins'] * args['anchor_number'],
                                  kernel_size=1) # BIN_NUM = 2， # 384
        else:
            self.use_dir = False


    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']
        record_len = data_dict['record_len']
        lidar_pose = data_dict['lidar_pose']  # [sum(cav), 6]
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len,
                      'pairwise_t_matrix': pairwise_t_matrix}
            


        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)

        batch_dict = self.backbone(batch_dict)
        spatial_features_2d = batch_dict['spatial_features_2d']
        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)

        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)

        output_dict = {'psm': psm,
                       'rm': rm}

        if self.use_dir:
            dm = self.dir_head(spatial_features_2d)
            output_dict.update({'dm': dm})

        return output_dict
opencood/models/point_pillar_intermediate_resnet_max_ylu.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>, Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


from numpy import record
import torch
import torch.nn as nn
import torch.nn.functional as F

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
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
from opencood.models.sub_modules.mash_utils import QueryEncoder, KeyEncoder, SmoothingNetwork

from icecream import ic


class PointPillarMash(nn.Module):
    def __init__(self, args):
        super(PointPillarMash, self).__init__()

        self.max_cav = args['max_cav']
        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        
        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        self.cls_head = nn.Conv2d(128 * 2, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 2, 7 * args['anchor_number'],
                                  kernel_size=1)

        mash_args = args['mash']
        self.query_encoder = QueryEncoder(mash_args["feature_dim"], mash_args['query_dim'])
        self.key_encoder = KeyEncoder(mash_args["feature_dim"], mash_args['key_dim'])
        self.queryKeySim = nn.Conv2d(mash_args['query_dim'],  mash_args['key_dim'], 1, 1)
        self.smoothing_net = SmoothingNetwork(in_ch=mash_args['H'] * mash_args['W'] + 1)
        self.H = mash_args['H']
        self.W = mash_args['W']
        self.downsample_rate = mash_args['downsample_rate']
        self.discrete_ratio = args['voxel_size'][0]

        if args['backbone_fix']:
            self.backbone_fix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay。
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False

    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']
        record_len = data_dict['record_len']

        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}

        # n, 4 -> n, c
        batch_dict = self.pillar_vfe(batch_dict)
        # n, c -> N, C, H, W
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)
        # N, C, H', W'. [N, 256, 50, 176]
        spatial_features_2d = batch_dict['spatial_features_2d']

        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]

        B = len(record_len)

        querys = self.queryKeySim(self.query_encoder(spatial_features_2d))
        keys = self.key_encoder(spatial_features_2d)

        split_query = self.regroup(querys, record_len)
        split_key = self.regroup(keys, record_len)
        split_feature = self.regroup(spatial_features_2d, record_len)
        
        fuse_features = []
        estimate_volumes = []
        for b in range(B):
            # N, C, H, W
            feature = split_feature[b]
            key = split_key[b]
            query = split_query[b]

            ego = 0
            fuse_feature = [feature[ego]]
            N = record_len[b]

            for i in range(1, N):
                corr_volume = self.computeCorrespondenceVolume(query[ego], key[i])
                corr_volume_decoded = self.smoothCorrespondenceVolume(corr_volume) # (Hs*Ws+1, Ht, Wt)
                grid, mask = self.idx2grid(corr_volume_decoded) # (1, H, W, 2)
                weight = torch.max(corr_volume_decoded, dim=0, keepdim=True)[0]
                estimate_volumes.append(corr_volume_decoded)
                
                warp_feature = F.grid_sample(feature[i].unsqueeze(0), grid).squeeze()
                warp_feature *= weight
                warp_feature *= mask
                fuse_feature.append(warp_feature)
            
            # max / sum
            fuse_features.append(torch.max(torch.stack(fuse_feature), dim = 0)[0])
        
        # B,C,H,W
        out_feature = torch.stack(fuse_features)
        if estimate_volumes:
            corr_vol = torch.stack(estimate_volumes)
        else:
            corr_vol = None

        psm = self.cls_head(out_feature)
        rm = self.reg_head(out_feature)

        output_dict = {'psm': psm,
                       'rm': rm,
                       'corr_vol': corr_vol}

        return output_dict
    
    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def computeCorrespondenceVolume(self, featA, featB):
        """compute the similarity volume

        Args:
            featA: [C, H, W], the query vectors of target agent 
            featB: [C, H, W], the key vectors of support agent

        Returns:
            distAB: [Hs*Ws+1, Ht, Wt]
        """
        C, H, W = featA.shape
        
        distAB = torch.zeros((self.H * self.W +1, self.H, self.W), device=featA.device)
        fA = featA.permute(1,2,0).reshape(-1,C) # (H*W, C)
        fB = featB.permute(1,2,0).reshape(-1,C) # (H*W, C)
        
        fA2 = torch.pow(torch.norm(fA,dim=-1),2).view(-1,1).repeat(1,fA.shape[0]) # (H*W, H*W)
        fB2 = torch.pow(torch.norm(fB,dim=-1),2).view(-1,1).repeat(1,fB.shape[0]) # (H*W, H*W)

 
        normA = torch.pow( fA2 + fB2.t() - 2.*torch.matmul(fA,fB.t()), 0.5 ) # (H*W, H*W)


        distAB[:-1,...] = normA.permute(1,0).reshape(-1, H, W)
        distAB[-1,:,:] = torch.norm(featA,p=2,dim=0)
        distAB = -distAB # two pixel is similar, then distAB[pixel1,pixel2] is low. We want it high 

        return distAB

    def smoothCorrespondenceVolume(self, distAB):
        """ smooth the correspondence Volume

        Args:
            distAB: (Hs*Ws+1, Ht, Wt)
        Returns:
            smoothed distAB
        """
        distAB = distAB.unsqueeze(0)
        output = self.smoothing_net(distAB)
        output.squeeze_(0)

        return output

    def idx2grid(self, matches):
        """
        Args:
            matches: (Hs*Ws + 1, Ht, Wt)
        """
        # should rewrite because H!=W
        # matches = matches.unsqueeze(0) # [1, Hs*Ws + 1, Ht, Wt]


        H, W = matches.shape[-2:]
        X = torch.arange(W).view(1,-1).repeat(H,1).type(torch.long).view(-1).to(matches.device) # (Ht * Wt)
        Y = torch.arange(H).view(-1,1).repeat(1,W).type(torch.long).view(-1).to(matches.device) # (Ht * Wt)
        X = torch.cat([X,torch.tensor([0],device=matches.device)],0)
        Y = torch.cat([Y,torch.tensor([0],device=matches.device)],0)

        idx = torch.argmax(matches.detach(),0).view(-1) # (Ht*Wt), the value is the index in supporting map
        
        # idx has no gradient
        # mask select those have no correspondence.
        # that means, ego's feature is used.
        mask = (idx == (matches.shape[0] - 1)).view(H, W).to(matches.device)

        x = torch.index_select(X,0,idx).view(1,H,W) # x_src in affine_grid
        y = torch.index_select(Y,0,idx).view(1,H,W) # y_src in affine_grid
        x = 2*((1.*x/W)-0.5) # (1, H, W)
        y = 2*((1.*y/H)-0.5) # (1, H, W)

        grid = torch.cat([x.unsqueeze(-1),y.unsqueeze(-1)],-1) # (1, 32, 32, 2)

        return grid, mask

opencood/models/point_pillar_mash.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu


from numpy import record
import torch.nn as nn
import torch

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
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
from opencood.models.fuse_modules.max_fuse import MaxFusion
from opencood.utils.transformation_utils import get_pairwise_transformation_torch
from opencood.data_utils.post_processor import SimpleVoxelPostprocessor
from opencood.models.point_pillar import PointPillar

from icecream import ic

DEBUG = False

class PointPillarMax(nn.Module):
    def __init__(self, args):
        super(PointPillarMax, self).__init__()

        self.max_cav = args['max_cav']
        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        
        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        self.fusion_net = MaxFusion(args['v2vfusion'])

        self.cls_head = nn.Conv2d(128 * 2, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 2, 7 * args['anchor_number'],
                                  kernel_size=1)
        if args['backbone_fix']:
            self.backbone_fix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay。
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False

    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']
        record_len = data_dict['record_len']
        lidar_pose = data_dict['lidar_pose']  # [sum(cav), 6]

        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}


        # n, 4 -> n, c
        batch_dict = self.pillar_vfe(batch_dict)
        # n, c -> N, C, H, W
        batch_dict = self.scatter(batch_dict)
        if DEBUG:
            origin_feature = torch.clone(batch_dict['spatial_features'])

        batch_dict = self.backbone(batch_dict)
        # N, C, H', W'. [N, 256, 50, 176]
        spatial_features_2d = batch_dict['spatial_features_2d']

        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]
        if DEBUG:
            self.fusion_net.forward_debug(spatial_features_2d, origin_feature,record_len, pairwise_t_matrix)
            raise


        fused_feature = self.fusion_net(spatial_features_2d,
                                        record_len,
                                        pairwise_t_matrix)


        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)

        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict

opencood/models/point_pillar_max.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, OpenPCDet
# License: TDG-Attribution-NonCommercial-NoDistrib


import torch
import torch.nn as nn


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
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.fuse_modules.raindrop_attn import raindrop_fuse
from opencood.models.sub_modules.downsample_conv import DownsampleConv

from opencood.models.fuse_modules.max_fuse import MaxFusion

class PointPillarMaxfusionFlow(nn.Module):
    def __init__(self, args):
        super(PointPillarMaxfusionFlow, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        # self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        if 'resnet' in args['base_bev_backbone'] and args['base_bev_backbone']['resnet']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        self.out_channel = sum(args['base_bev_backbone']['num_upsample_filter'])

        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
            self.out_channel = args['shrink_header']['dim'][-1]

        self.rain_fusion = raindrop_fuse(args['fusion_args'])

        self.cls_head = nn.Conv2d(self.out_channel, args['anchor_number'], # 384
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(self.out_channel, 7 * args['anchor_number'], # 384
                                  kernel_size=1)

        # self.cls_head_fused = nn.Conv2d(self.out_channel, args['anchor_number'], # 384
        #                           kernel_size=1)
        # self.reg_head_fused = nn.Conv2d(self.out_channel, 7 * args['anchor_number'], # 384
        #                           kernel_size=1)

        self.fusion_net = MaxFusion(args['fusion_args'])

        if 'dir_args' in args.keys():
            self.use_dir = True
            self.dir_head = nn.Conv2d(self.out_channel, args['dir_args']['num_bins'] * args['anchor_number'],
                                  kernel_size=1) # BIN_NUM = 2， # 384
        else:
            self.use_dir = False

    def forward(self, data_dict, stage = 1, pairwise_t_matrix=None):
        if stage == 1:
            return self.single_forward(data_dict)
        elif stage == 2:
            return self.fuse_forward(data_dict, pairwise_t_matrix)
        else:
            print("File point_pillar_maxfusion_flow.py ERROR: \
                stage option must be in 1(for single detection) or 2(for fuesed detection).")
        
    def single_forward(self, data_dict):    
        output_dict = {}
        intermediate_feature_flow = True

        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_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points}

        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)

        spatial_features_2d = batch_dict['spatial_features_2d']  # B, C, H, W

        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)

        if intermediate_feature_flow:
            output_dict.update({
                'spatial_features_2d': spatial_features_2d  # B, C, H, W
            })
        
        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)

        output_dict.update({
            'spatial_features': batch_dict['spatial_features'],
            'psm': psm,
            'rm': rm
        })
                       
        if self.use_dir:
            dm = self.dir_head(spatial_features_2d)
            output_dict.update({'dm': dm})

        return output_dict

    def fuse_forward(self, updated_dict, pairwise_t_matrix):
        """ 参考 where2comm 的写法 如何将多车的feature 融合起来 返回最终的output dict
        Parameters:
        -----------
        updated_dict : 
            'ego' / cav_id : {
                'updated_spatial_features_2d' : torch.tensor, [C, H, W]
            }

        pairwise_t_matrix: 

        Returns:
        ---------
        output_dict : 

        """
        device = updated_dict['ego']['spatial_features_2d'].device
        spatial_features_2d_list = []
        spatial_features_list = []
        for cav_id, cav_content in updated_dict.items():
            spatial_features_2d_list.append(cav_content['updated_spatial_features_2d'])
            spatial_features_list.append(cav_content['updated_spatial_features'])

        record_len = torch.tensor([len(spatial_features_2d_list)]).to(device)
        spatial_features_2d = torch.stack(spatial_features_2d_list, dim=0).to(device)  # (sum(cav), C, H, W)
        spatial_features = torch.stack(spatial_features_list, dim=0).to(device)  # (sum(cav), C', H', W')

        pairwise_t_matrix = pairwise_t_matrix.unsqueeze(0)
        psm_single = self.cls_head(spatial_features_2d)
        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]
        # fused_feature = self.fusion_net(spatial_features_2d,
        #                                 record_len,
        #                                 pairwise_t_matrix)

        '''
        change pariwise_t_matrix from (B, L, L, 4, 4) to (B, L, K, 4, 4)
        '''
        new_pairwise_t_matrix = pairwise_t_matrix[:, 0, :, :, :].unsqueeze(2) # (B, L, 1, 4, 4)
        record_frames = torch.ones((len(spatial_features_2d_list))).to(device)
        fused_feature, communication_rates, result_dict = self.rain_fusion(spatial_features,
                    psm_single,
                    record_len,
                    new_pairwise_t_matrix, 
                    record_frames,
                    self.backbone,
                    [self.shrink_conv, self.cls_head, self.reg_head])
        if self.shrink_flag:
                fused_feature = self.shrink_conv(fused_feature)

        # ###### debug use, viz updated feature of each cav
        # from matplotlib import pyplot as plt
        # viz_save_path = '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/debug_4_feature_flow'
        # viz_content = torch.max(fused_feature[0], dim=0)[0].detach().cpu()
        # plt.imshow(viz_content)
        # plt.savefig(viz_save_path+'/updated_fused_feature.png')
        # ##############

        # psm = self.cls_head_fused(fused_feature)
        # rm = self.reg_head_fused(fused_feature)

        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)

        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict
opencood/models/point_pillar_maxfusion_flow.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>, Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


from numpy import record
import torch.nn as nn

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
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
from opencood.models.fuse_modules.mean_fuse import MeanFusion

from icecream import ic


class PointPillarMean(nn.Module):
    def __init__(self, args):
        super(PointPillarMean, self).__init__()

        self.max_cav = args['max_cav']
        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        
        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        self.fusion_net = MeanFusion(args['v2vfusion'])

        self.cls_head = nn.Conv2d(128 * 2, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 2, 7 * args['anchor_number'],
                                  kernel_size=1)
        if args['backbone_fix']:
            self.backbone_fix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay。
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False

    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']
        record_len = data_dict['record_len']

        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c
        batch_dict = self.pillar_vfe(batch_dict)
        # n, c -> N, C, H, W
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)
        # N, C, H', W'. [N, 384, 100, 352]
        spatial_features_2d = batch_dict['spatial_features_2d']

        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]


        fused_feature = self.fusion_net(spatial_features_2d,
                                        record_len,
                                        pairwise_t_matrix)


        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)

        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict

opencood/models/point_pillar_mean.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu

# This is not V2X-ViT!! Please use point_pillar_v2xvit.py 
# I have forgotten its performance. Not sure about the transformer_fuse

from numpy import record
import torch.nn as nn

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
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
from opencood.models.sub_modules.dcn_net import DCNNet
from opencood.models.fuse_modules.transformer_fuse import TransformerFusion



class PointPillarTransformer(nn.Module):
    def __init__(self, args):
        super(PointPillarTransformer, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        
        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        self.dcn = False
        if 'dcn' in args:
            self.dcn = True
            self.dcn_net = DCNNet(args['dcn'])


        self.fusion_net = TransformerFusion(args['fusion_args'])

        self.cls_head = nn.Conv2d(128 * 2, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 2, 7 * args['anchor_number'],
                                  kernel_size=1)
        if args['backbone_fix']:
            self.backbone_fix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay。
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False

    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']
        record_len = data_dict['record_len']

        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c
        batch_dict = self.pillar_vfe(batch_dict)
        # n, c -> N, C, H, W
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)
        # N, C, H', W'. [N, 384, 100, 352]
        spatial_features_2d = batch_dict['spatial_features_2d']

        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)
        # dcn
        if self.dcn:
            spatial_features_2d = self.dcn_net(spatial_features_2d)
        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]


        fused_feature = self.fusion_net(spatial_features_2d,
                                        record_len,
                                        pairwise_t_matrix)

        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)


        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict

opencood/models/point_pillar_transformer.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu

import torch
import torch.nn as nn


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
from opencood.utils.model_utils import weight_init

class PointPillarUncertainty(nn.Module):
    def __init__(self, args):
        super(PointPillarUncertainty, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        self.uncertainty_dim = args['uncertainty_dim'] # dim=3 means x, y, yaw, dim=2 means x, y

        self.cls_head = nn.Conv2d(128 * 3, args['anchor_num'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_num'],
                                  kernel_size=1)

        self.s_head = nn.Conv2d(128 * 3, self.uncertainty_dim * args['anchor_num'],
                                  kernel_size=1) 

        if 'dir_args' in args.keys():
            self.use_dir = True
            self.dir_head = nn.Conv2d(128 * 3, args['dir_args']['num_bins'] * args['anchor_num'],
                                  kernel_size=1) # BIN_NUM = 2
        else:
            self.use_dir = False

        
        self.apply(weight_init)
    

    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_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points}

        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)

        spatial_features_2d = batch_dict['spatial_features_2d']

        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)
        sm = self.s_head(spatial_features_2d) # s is log(b) or log(sigma^2)

        output_dict = {'psm': psm,
                       'rm': rm,
                       'sm': sm}
        
        if self.use_dir:
            dm = self.dir_head(spatial_features_2d)
            output_dict.update({'dm': dm})

        return output_dict
opencood/models/point_pillar_uncertainty.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>, Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


from numpy import record
import torch.nn as nn

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
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
from opencood.models.fuse_modules.v2v_fuse import V2VNetFusion

from icecream import ic
import torch
import numpy as np
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world

def generate_noise(pos_std, rot_std, pos_mean=0, rot_mean=0):
    """ Add localization error to the 6dof pose
        Noise includes position (x,y) and rotation (yaw).
        We use gaussian distribution to generate noise.

    Args:

        pos_std : float 
            std of gaussian dist, in meter

        rot_std : float
            std of gaussian dist, in degree

        pos_mean : float
            mean of gaussian dist, in meter

        rot_mean : float
            mean of gaussian dist, in degree

    Returns:
        pose_noise: np.ndarray, [6,]
            [x, y, z, roll, yaw, pitch]
    """

    xy = np.random.normal(pos_mean, pos_std, size=(2))
    yaw = np.random.normal(rot_mean, rot_std, size=(1))

    pose_noise = np.array([xy[0], xy[1], 0, 0, yaw[0], 0])

    return pose_noise

def get_pairwise_transformation2ego(past_k_lidar_pose, noise_level, max_cav=5):
    """
    Get transformation matrixes accross different agents to curr ego at all past timestamps.

    Parameters
    ----------
    past_k_lidar_pose: np.array [N, k=1, 6]
    
    ego_pose : list
        ego pose

    max_cav : int
        The maximum number of cav, default 5

    Return
    ------
    pairwise_t_matrix : np.array
        The transformation matrix each cav to curr ego at past k frames.
        shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
        pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
    """
    pos_std = noise_level['pos_std']
    rot_std = noise_level['rot_std']
    pos_mean = 0 
    rot_mean = 0
    
    pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

    ego_pose = past_k_lidar_pose[0, 0]

    t_list = []

    # save all transformation matrix in a list in order first.
    for cav_id in range(past_k_lidar_pose.shape[0]):
        loc_noise = generate_noise(pos_std, rot_std)
        t_list.append(x_to_world(past_k_lidar_pose[cav_id, 0].cpu().numpy()+loc_noise)) # Twx
    
    ego_pose = x_to_world(ego_pose.cpu().numpy())

    for i in range(len(t_list)):
        for j in range(len(t_list)):
            # identity matrix to self
            if i != j:
                # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                pairwise_t_matrix[i, j] = t_matrix

    pairwise_t_matrix = torch.tensor(pairwise_t_matrix).to(past_k_lidar_pose.device)
    return pairwise_t_matrix

class PointPillarV2VNet(nn.Module):
    def __init__(self, args):
        super(PointPillarV2VNet, self).__init__()

        self.max_cav = args['max_cav']
        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        
        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        self.noise_flag = 0
        if 'noise' in args.keys():
            self.noise_flag = True
            self.noise_level = {'pos_std': args['noise']['pos_std'], 'rot_std': args['noise']['rot_std'], 'pos_mean': args['noise']['pos_mean'], 'rot_mean': args['noise']['rot_mean']}
            print(f'=== noise level : {self.noise_level} ===')

        self.fusion_net = V2VNetFusion(args['v2vfusion'])

        self.cls_head = nn.Conv2d(128 * 2, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 2, 7 * args['anchor_number'],
                                  kernel_size=1)
        if args['backbone_fix']:
            self.backbone_fix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay。
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False

    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']
        record_len = data_dict['record_len']

        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c
        batch_dict = self.pillar_vfe(batch_dict)
        # n, c -> N, C, H, W
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)
        # N, C, H', W'. [N, 256, 50, 176]
        spatial_features_2d = batch_dict['spatial_features_2d']

        noise_pairwise_t_matrix = None
        B = record_len.shape[0]
        if self.noise_flag and B==1:
            # noise_level = {'pos_std': 0.5, 'rot_std': 0, 'pos_mean': 0, 'rot_mean': 0}
            noise_pairwise_t_matrix = get_pairwise_transformation2ego(data_dict['past_lidar_pose'], self.noise_level, max_cav=5).unsqueeze(0)
            # diff = pairwise_t_matrix - noise_pairwise_t_matrix
            # print(f'=== diff.max() = {diff.max()} === diff.min() = {diff.min()} ===')

        if noise_pairwise_t_matrix is not None:
            pairwise_t_matrix = noise_pairwise_t_matrix

        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]

        fused_feature = self.fusion_net(spatial_features_2d,
                                        record_len,
                                        pairwise_t_matrix)


        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)

        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict


        # import pickle
        # with open("/GPFS/rhome/yifanlu/workspace/OpenCOOD/vis_result/v2vfuse/v2v_proj_first.pkl", 'wb') as f:
        #     pickle.dump({"spatial_features":batch_dict['spatial_features'].detach().cpu().numpy(), 
        #             "spatial_features_2d":spatial_features_2d.detach().cpu().numpy(), 
        #             "record_len":record_len.detach().cpu().numpy(), 
        #             "pairwise_t_matrix":pairwise_t_matrix.detach().cpu().numpy()}, f)
        # raise
opencood/models/point_pillar_v2vnet.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Yifan Lu

from numpy import record
import torch
import torch.nn as nn

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
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
from opencood.models.sub_modules.v2v_robust_module import AttentionWrapper, PoseRegressionWraper, WeightedEM, get_intersection, regroup
from opencood.utils.pose_utils import generate_noise_torch
from opencood.utils.transformation_utils import get_pairwise_transformation_torch
from opencood.models.fuse_modules.v2v_fuse import V2VNetFusion

from icecream import ic
from opencood.utils.model_utils import weight_init

class PointPillarV2VNetRobust(nn.Module):
    def __init__(self, args):
        super(PointPillarV2VNetRobust, self).__init__()

        self.max_cav = args['max_cav']
        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        
        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        self.fusion_net = V2VNetFusion(args['v2vfusion'])

        self.cls_head = nn.Conv2d(128 * 2, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 2, 7 * args['anchor_number'],
                                  kernel_size=1)


        self.downsample_rate = args['robust']['downsample_rate']
        self.discrete_ratio = args['robust']['discrete_ratio']
        self.H = args['robust']['H']
        self.W = args['robust']['W']

        self.affine_parameter = {"H":self.H, "W": self.W, "downsample_rate": self.downsample_rate, "discrete_ratio": self.discrete_ratio}
        learnable_alpha = True if 'learnable_alpha' not in args['robust'] else args['robust']['learnable_alpha']

        self.pose_reg_net = PoseRegressionWraper(args['robust']['feature_dim']*2, 
                                                args['robust']['hidden_dim'],
                                                self.affine_parameter
                                                )   

        self.attention_net = AttentionWrapper(args['robust']['feature_dim']*2, 
                                              args['robust']['hidden_dim'],
                                              self.affine_parameter,
                                              learnable_alpha,
                                              )

        self.stage = args['stage'] # [0, 1, 2]

        self.apply(weight_init)

        if self.stage == 1:
            self.backbone_fix()
        if self.stage == 2:
            self.backbone_unfix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone for stage 1 
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.fusion_net.parameters():
            p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False

        for p in self.attention_net.parameters():
            p.requires_grad = False

    def backbone_unfix(self):
        """
        unfix for stage 2
        """

        for p in self.pillar_vfe.parameters():
            p.requires_grad = True

        for p in self.scatter.parameters():
            p.requires_grad = True

        for p in self.backbone.parameters():
            p.requires_grad = True

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = True
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = True

        for p in self.fusion_net.parameters():
            p.requires_grad = True

        for p in self.cls_head.parameters():
            p.requires_grad = True
        for p in self.reg_head.parameters():
            p.requires_grad = True

        for p in self.attention_net.parameters():
            p.requires_grad = True

    def pose_correction(self, features, record_len, pairwise_t_matrix):
        """ use pose regression module to correct relative pose
        Args:

        Returns:
            pairwise_t_matrix_new: 
                [B, L, L, 4, 4], the relative pose after correction.
        """
        return self.pose_reg_net(features, record_len, pairwise_t_matrix)

    def global_correction(self, lidar_pose, pairwise_t_matrix, record_len):
        """
        Args:
            lidar_pose: [N, 3]
                input noisy lidar pose
            pairwise_t_matrix: [B, L, L, 4, 4]
                relative pose after pose regression module
            record_len: list,
                shape [B]
        
        Returns:
            lidar_pose_new: [N, 3]
                refined lidar pose
        """

        B = len(record_len)
        lidar_pose_new = []

        # [[N1,3], [N2, 3], ...]
        lidar_pose_split = regroup(lidar_pose, record_len)

        for b in range(B):
            if record_len[b] == 1:
                lidar_pose_new.append(lidar_pose_split[b])
                continue
            lidar_pose = lidar_pose_split[b]
            intersection_matrix = get_intersection(pairwise_t_matrix[b], self.affine_parameter)
            lidar_pose_corrected = WeightedEM(lidar_pose, pairwise_t_matrix[b],intersection_matrix)

            lidar_pose_new.append(lidar_pose_corrected)

        lidar_pose_new = torch.cat(lidar_pose_new, dim=0)

        return lidar_pose_new
        

    def noise_generator(self, lidar_pose, all_strong=False):
        noise_s = generate_noise_torch(lidar_pose, pos_std=0.4, rot_std=4)  # (N, 6)
        noise_w = generate_noise_torch(lidar_pose, pos_std=0.01, rot_std=0.1)  # (N, 6)
        N = lidar_pose.shape[0]
        
        if all_strong:
            choice = torch.zeros((N, 1), device=lidar_pose.device) # (N, 1) 0 choose strong, 1 choose weak
            noise = noise_s
        else:
            choice = torch.randint(0, 2, (N, 1), device=lidar_pose.device) # (N, 1) 0 choose strong, 1 choose weak
            noise = choice * noise_w + (1-choice) * noise_s
        
        return noise, choice


    def train_forward(self, spatial_features_2d, record_len, lidar_pose, pairwise_t_matrix, stage):
        """
        stage = 0, only training attentive_aggregation and v2vnet, strong noise and weak noise are used.
        stage = 1, only training pose correction module. all strong noise
        stage = 2, all component are used. all strong noise

        Args:
            spatial_features_2d: (N, C, H, W)
            record_len: list
            lidar_pose: (N, 6), it will turn to [N, 3] quickly
        """
        if stage == 0:
            noise, choice = self.noise_generator(lidar_pose, all_strong=False)

        if stage == 1 or stage == 2:
            noise, choice = self.noise_generator(lidar_pose, all_strong=True)

        lidar_pose += noise
        lidar_pose = lidar_pose[:,[0,1,4]] # [N, 3]
        pairwise_t_matrix = get_pairwise_transformation_torch(lidar_pose, self.max_cav, record_len, dof=3)
        
        # when training pairwise_t_matrix, pairwise_t_matrix carries given noise.
        if self.stage == 0:
            scores, weight = self.attention_net(spatial_features_2d, record_len, pairwise_t_matrix)
            fused_feature = self.fusion_net(spatial_features_2d, record_len, pairwise_t_matrix, weight)
            psm = self.cls_head(fused_feature)
            rm = self.reg_head(fused_feature)
            print("scores:", scores)
            print("weight:", weight)
            print("alpha:", self.attention_net.alpha)

            output_dict = {'stage': stage,
                            'scores': scores,
                            'choice': choice,
                            'psm': psm,
                            'rm': rm}

        if self.stage == 1:
            pairwise_corr, _ = self.pose_correction(spatial_features_2d, record_len, pairwise_t_matrix)
            output_dict = {'stage': stage,
                            'pairwise_corr' : pairwise_corr,
                            'pairwise_t_matrix': pairwise_t_matrix}
        
        if self.stage == 2:
            pairwise_corr, pairwise_t_matrix_new = self.pose_correction(spatial_features_2d, record_len, pairwise_t_matrix)
            lidar_pose_corrected = self.global_correction(lidar_pose, pairwise_t_matrix_new, record_len) # [N, 3]

            pairwise_t_matrix_corrected = get_pairwise_transformation_torch(lidar_pose_corrected, self.max_cav, record_len, dof=3)
            scores, weight = self.attention_net(spatial_features_2d, record_len, pairwise_t_matrix_corrected)
            fused_feature = self.fusion_net(spatial_features_2d, record_len, pairwise_t_matrix_corrected, weight)
            psm = self.cls_head(fused_feature)
            rm = self.reg_head(fused_feature)

            output_dict = {'stage': stage,
                            'scores': scores,
                            'psm': psm,
                            'rm': rm,
                            'pairwise_corr' : pairwise_corr,
                            'pairwise_t_matrix': pairwise_t_matrix}

        return output_dict



    def eval_forward(self, spatial_features_2d, record_len, lidar_pose, pairwise_t_matrix, stage):
        """
            same as stage=2 in training, but no noise added.
        """
        lidar_pose = lidar_pose[:,[0,1,4]] # [N, 3]
        pairwise_t_matrix = get_pairwise_transformation_torch(lidar_pose, self.max_cav, record_len, dof=3)

        pairwise_corr, pairwise_t_matrix = self.pose_correction(spatial_features_2d, record_len, pairwise_t_matrix)

        lidar_pose_corrected = self.global_correction(lidar_pose, pairwise_t_matrix, record_len) # [N, 3]
        pairwise_t_matrix_corrected = get_pairwise_transformation_torch(lidar_pose_corrected, self.max_cav, record_len, dof=3)


        scores, weight = self.attention_net(spatial_features_2d, record_len, pairwise_t_matrix_corrected)
        fused_feature = self.fusion_net(spatial_features_2d, record_len, pairwise_t_matrix_corrected, weight)
        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)

        output_dict = {'stage': stage,
                        'scores': scores,
                        'psm': psm,
                        'rm': rm,
                        'pairwise_t_matrix': pairwise_t_matrix}

        return output_dict


    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']
        record_len = data_dict['record_len']

        pairwise_t_matrix = data_dict['pairwise_t_matrix']
        lidar_pose = data_dict['lidar_pose'] # [sum(cav), 6]

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c
        batch_dict = self.pillar_vfe(batch_dict)
        # n, c -> N, C, H, W
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)
        # N, C, H', W'. [N, 256, 50, 176]
        spatial_features_2d = batch_dict['spatial_features_2d']

        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]

        # lidar_pose -> pairwise_t_matrix , same content

        return self.train_forward(spatial_features_2d, record_len, lidar_pose, pairwise_t_matrix, self.stage)
opencood/models/point_pillar_v2vnet_robust.py
MediaBrain-SJTU/CoBEVFlow
"""
This is from
https://github.com/DerrickXuNu/v2x-vit/blob/main/v2xvit/models/point_pillar_transformer.py


"""

import torch
import torch.nn as nn

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
from opencood.models.fuse_modules.fuse_utils import regroup
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
from opencood.models.sub_modules.v2xvit_basic import V2XTransformer
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple

import numpy as np
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world

def generate_noise(pos_std, rot_std, pos_mean=0, rot_mean=0):
    """ Add localization error to the 6dof pose
        Noise includes position (x,y) and rotation (yaw).
        We use gaussian distribution to generate noise.

    Args:

        pos_std : float 
            std of gaussian dist, in meter

        rot_std : float
            std of gaussian dist, in degree

        pos_mean : float
            mean of gaussian dist, in meter

        rot_mean : float
            mean of gaussian dist, in degree

    Returns:
        pose_noise: np.ndarray, [6,]
            [x, y, z, roll, yaw, pitch]
    """

    xy = np.random.normal(pos_mean, pos_std, size=(2))
    yaw = np.random.normal(rot_mean, rot_std, size=(1))

    pose_noise = np.array([xy[0], xy[1], 0, 0, yaw[0], 0])

    return pose_noise

def get_pairwise_transformation2ego(past_k_lidar_pose, noise_level, max_cav=5):
    """
    Get transformation matrixes accross different agents to curr ego at all past timestamps.

    Parameters
    ----------
    past_k_lidar_pose: np.array [N, k=1, 6]
    
    ego_pose : list
        ego pose

    max_cav : int
        The maximum number of cav, default 5

    Return
    ------
    pairwise_t_matrix : np.array
        The transformation matrix each cav to curr ego at past k frames.
        shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
        pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
    """
    pos_std = noise_level['pos_std']
    rot_std = noise_level['rot_std']
    pos_mean = 0 
    rot_mean = 0
    
    pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

    ego_pose = past_k_lidar_pose[0, 0]

    t_list = []

    # save all transformation matrix in a list in order first.
    for cav_id in range(past_k_lidar_pose.shape[0]):
        loc_noise = generate_noise(pos_std, rot_std)
        t_list.append(x_to_world(past_k_lidar_pose[cav_id, 0].cpu().numpy()+loc_noise)) # Twx
    
    ego_pose = x_to_world(ego_pose.cpu().numpy())

    for i in range(len(t_list)):
        for j in range(len(t_list)):
            # identity matrix to self
            if i != j:
                # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                pairwise_t_matrix[i, j] = t_matrix

    pairwise_t_matrix = torch.tensor(pairwise_t_matrix).to(past_k_lidar_pose.device)
    return pairwise_t_matrix


def get_past_k_pairwise_transformation2ego(past_k_lidar_pose, noise_level, k=3, max_cav=5):
    """
    Get transformation matrixes accross different agents to curr ego at all past timestamps.

    Parameters
    ----------
    base_data_dict : dict
        Key : cav id, item: transformation matrix to ego, lidar points.
    
    ego_pose : list
        ego pose

    max_cav : int
        The maximum number of cav, default 5

    Return
    ------
    pairwise_t_matrix : np.array
        The transformation matrix each cav to curr ego at past k frames.
        shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
        pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
    """
    pos_std = noise_level['pos_std']
    rot_std = noise_level['rot_std']
    pos_mean = 0 
    rot_mean = 0
    
    pairwise_t_matrix = np.tile(np.eye(4), (max_cav, k, 1, 1)) # (L, k, 4, 4)

    ego_pose = past_k_lidar_pose[0, 0]

    t_list = []

    # save all transformation matrix in a list in order first.
    for cav_id in range(past_k_lidar_pose.shape[0]):
        past_k_poses = []
        for time_id in range(k):
            loc_noise = generate_noise(pos_std, rot_std)
            past_k_poses.append(x_to_world(past_k_lidar_pose[cav_id, time_id].cpu().numpy()+loc_noise))
        t_list.append(past_k_poses) # Twx
    
    ego_pose = x_to_world(ego_pose.cpu().numpy())
    for i in range(len(t_list)): # different cav
        if i!=0 :
            for j in range(len(t_list[i])): # different time
                # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                pairwise_t_matrix[i, j] = t_matrix
    pairwise_t_matrix = torch.tensor(pairwise_t_matrix).to(past_k_lidar_pose.device)
    return pairwise_t_matrix


class PointPillarV2XVit(nn.Module):
    def __init__(self, args):
        super(PointPillarV2XVit, self).__init__()

        self.max_cav = args['max_cav']
        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        self.noise_flag = 0
        if 'noise' in args.keys():
            self.noise_flag = True
            self.noise_level = {'pos_std': args['noise']['pos_std'], 'rot_std': args['noise']['rot_std'], 'pos_mean': args['noise']['pos_mean'], 'rot_mean': args['noise']['rot_mean']}
            print(f'=== noise level : {self.noise_level} ===')

        self.fusion_net = V2XTransformer(args['transformer'])

        self.cls_head = nn.Conv2d(128 * 2, args['anchor_number'],
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(128 * 2, 7 * args['anchor_number'],
                                  kernel_size=1)
        self.discrete_ratio = args['voxel_size'][0]

        if args['backbone_fix']:
            self.backbone_fix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay。
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False

    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']
        record_len = data_dict['record_len']
        pairwise_t_matrix = data_dict['pairwise_t_matrix']
        """
        We don't consider the time delay
        So spatial_correction_matrix is just identity matrix
        """
        # B, L, 4, 4
        spatial_correction_matrix = torch.eye(4).expand(len(record_len), self.max_cav, 4, 4).to(record_len.device)

        # B, max_cav, 3(dt dv infra), 1, 1
        # prior_encoding =\
        #     data_dict['prior_encoding'].unsqueeze(-1).unsqueeze(-1)
        prior_encoding = \
            torch.zeros(len(record_len), self.max_cav, 3, 1, 1).to(record_len.device)
              

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c
        batch_dict = self.pillar_vfe(batch_dict)
        # n, c -> N, C, H, W
        batch_dict = self.scatter(batch_dict)

        _, _, H, W = batch_dict['spatial_features'].shape
        self.downsample_rate = 1

        batch_dict = self.backbone(batch_dict)

        spatial_features_2d = batch_dict['spatial_features_2d']

        noise_pairwise_t_matrix = None
        B = record_len.shape[0]
        if self.noise_flag and B==1:
            # noise_level = {'pos_std': 0.5, 'rot_std': 0, 'pos_mean': 0, 'rot_mean': 0}
            # noise_pairwise_t_matrix = get_past_k_pairwise_transformation2ego(data_dict['past_lidar_pose'], self.noise_level, k=1, max_cav=5).unsqueeze(0)
            noise_pairwise_t_matrix = get_pairwise_transformation2ego(data_dict['past_lidar_pose'], self.noise_level, max_cav=5).unsqueeze(0)
            
            # diff = pairwise_t_matrix - noise_pairwise_t_matrix
            # print(f'=== diff.max() = {diff.max()} === diff.min() = {diff.min()} ===')

        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)
        # N, C, H, W -> B, L, C, H, W
        regroup_feature, mask = regroup(spatial_features_2d,
                                        record_len,
                                        self.max_cav)

        
        # prior encoding added
        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)

        ######## Added by Yifan Lu 2022.9.5
        regroup_feature_new = []
        B = regroup_feature.shape[0]
        if noise_pairwise_t_matrix is not None:
            pairwise_t_matrix = noise_pairwise_t_matrix
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]
        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2
        H_, W_ = spatial_features_2d.shape[-2:]
        for b in range(B):
            # (B,L,L,2,3)
            ego = 0
            # Original OPV2V
            regroup_feature_new.append(warp_affine_simple(regroup_feature[b], pairwise_t_matrix[b, ego], (H_, W_)))

            # OPV2V Irregular version
            # regroup_feature_new.append(warp_affine_simple(regroup_feature[b], pairwise_t_matrix[b, :, 0], (H_, W_)))
        regroup_feature = torch.stack(regroup_feature_new)
        ###################################

        # b l c h w -> b l h w c
        regroup_feature = regroup_feature.permute(0, 1, 3, 4, 2)
        # transformer fusion
        fused_feature = self.fusion_net(regroup_feature, mask, spatial_correction_matrix)
        # b h w c -> b c h w
        fused_feature = fused_feature.permute(0, 3, 1, 2)

        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)

        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict
opencood/models/point_pillar_v2xvit.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, OpenPCDet
# License: TDG-Attribution-NonCommercial-NoDistrib


import torch
import torch.nn as nn


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
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv

class PointPillarWFlow(nn.Module):
    def __init__(self, args):
        super(PointPillarWFlow, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        # self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        if 'resnet' in args['base_bev_backbone'] and args['base_bev_backbone']['resnet']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)
        
        self.out_channel = sum(args['base_bev_backbone']['num_upsample_filter'])

        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
            self.out_channel = args['shrink_header']['dim'][-1]

        self.cls_head = nn.Conv2d(self.out_channel, args['anchor_number'], # 384
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(self.out_channel, 7 * args['anchor_number'], # 384
                                  kernel_size=1)

        if 'dir_args' in args.keys():
            self.use_dir = True
            self.dir_head = nn.Conv2d(self.out_channel, args['dir_args']['num_bins'] * args['anchor_number'],
                                  kernel_size=1) # BIN_NUM = 2， # 384
        else:
            self.use_dir = False

    def forward(self, data_dict):
        output_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_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points}

        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)

        spatial_features_2d = batch_dict['spatial_features_2d']  # B, C, H, W

        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)

        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)

        output_dict.update({
            'psm': psm,
            'rm': rm
        })
                       
        if self.use_dir:
            dm = self.dir_head(spatial_features_2d)
            output_dict.update({'dm': dm})

        return output_dict
opencood/models/point_pillar_w_flow.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>, Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


from numpy import record
import torch.nn as nn

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
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
# from opencood.models.sub_modules.dcn_net import DCNNet
# from opencood.models.fuse_modules.where2comm import Where2comm
from opencood.models.fuse_modules.where2comm_attn import Where2comm
import torch
import numpy as np
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world

def generate_noise(pos_std, rot_std, pos_mean=0, rot_mean=0):
    """ Add localization error to the 6dof pose
        Noise includes position (x,y) and rotation (yaw).
        We use gaussian distribution to generate noise.

    Args:

        pos_std : float 
            std of gaussian dist, in meter

        rot_std : float
            std of gaussian dist, in degree

        pos_mean : float
            mean of gaussian dist, in meter

        rot_mean : float
            mean of gaussian dist, in degree

    Returns:
        pose_noise: np.ndarray, [6,]
            [x, y, z, roll, yaw, pitch]
    """

    xy = np.random.normal(pos_mean, pos_std, size=(2))
    yaw = np.random.normal(rot_mean, rot_std, size=(1))

    pose_noise = np.array([xy[0], xy[1], 0, 0, yaw[0], 0])

    return pose_noise

def get_pairwise_transformation2ego(lidar_pose, noise_level, max_cav=5):
    """
    Get transformation matrixes accross different agents to curr ego at all past timestamps.

    Parameters
    ----------
    lidar_pose: np.array [N, k=1, 6]
    
    ego_pose : list
        ego pose

    max_cav : int
        The maximum number of cav, default 5

    Return
    ------
    pairwise_t_matrix : np.array
        The transformation matrix each cav to curr ego at past k frames.
        shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
        pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
    """
    pos_std = noise_level['pos_std']
    rot_std = noise_level['rot_std']
    pos_mean = 0 
    rot_mean = 0
    
    pairwise_t_matrix = np.tile(np.eye(4), (max_cav, max_cav, 1, 1)) # (L, L, 4, 4)

    ego_pose = lidar_pose[0]

    t_list = []

    # save all transformation matrix in a list in order first.
    for cav_id in range(lidar_pose.shape[0]):
        loc_noise = generate_noise(pos_std, rot_std)
        t_list.append(x_to_world(lidar_pose[cav_id].cpu().numpy()+loc_noise)) # Twx
    
    ego_pose = x_to_world(ego_pose.cpu().numpy())

    for i in range(len(t_list)):
        for j in range(len(t_list)):
            # identity matrix to self
            if i != j:
                # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                t_matrix = np.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                pairwise_t_matrix[i, j] = t_matrix

    pairwise_t_matrix = torch.tensor(pairwise_t_matrix).to(lidar_pose.device)
    return pairwise_t_matrix


class PointPillarWhere2comm(nn.Module):
    def __init__(self, args):
        super(PointPillarWhere2comm, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        if 'resnet' in args['base_bev_backbone']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        self.noise_flag = 0
        if 'noise' in args.keys():
            self.noise_flag = True
            self.noise_level = {'pos_std': args['noise']['pos_std'], 'rot_std': args['noise']['rot_std'], 'pos_mean': args['noise']['pos_mean'], 'rot_mean': args['noise']['rot_mean']}
            print(f'=== noise level : {self.noise_level} ===')

        self.dcn = False
        # if 'dcn' in args:
        #     self.dcn = True
        #     self.dcn_net = DCNNet(args['dcn'])

        # self.fusion_net = TransformerFusion(args['fusion_args'])
        # TODO: 
        # self.fusion_net = Where2comm(args['fusion_args'])
        self.rain_fusion = Where2comm(args['fusion_args'])
        self.multi_scale = args['fusion_args']['multi_scale']

        if self.shrink_flag:
            dim = args['shrink_header']['dim'][0]
            self.cls_head = nn.Conv2d(int(dim), args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
                                    kernel_size=1)
            # self.fused_cls_head = nn.Conv2d(int(dim), args['anchor_number'],
            #                         kernel_size=1)
            # self.fused_reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
            #                         kernel_size=1)
        else:
            self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
                                    kernel_size=1)
        if args['backbone_fix']:
            self.backbone_fix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False
    
    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    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']
        record_len = data_dict['record_len']

        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c
        batch_dict = self.pillar_vfe(batch_dict)
        # (n, c) -> (batch_cav_size, C, H, W) put pillars into spatial feature map
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)
        # N, C, H', W'. [N, 384, 100, 252]
        spatial_features_2d = batch_dict['spatial_features_2d']

        B = record_len.shape[0]
        if self.noise_flag and B==1:
            # noise_level = {'pos_std': 0.5, 'rot_std': 0, 'pos_mean': 0, 'rot_mean': 0}
            noise_pairwise_t_matrix = get_pairwise_transformation2ego(data_dict['lidar_pose'], self.noise_level, max_cav=5).unsqueeze(0)
            # diff = pairwise_t_matrix - noise_pairwise_t_matrix
            # print(f'=== diff.max() = {diff.max()} === diff.min() = {diff.min()} ===')
            pairwise_t_matrix = noise_pairwise_t_matrix
        
        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)  # (B, 256, H', W')
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)
        # dcn
        # if self.dcn:
        #     spatial_features_2d = self.dcn_net(spatial_features_2d)
        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]
        psm_single = self.cls_head(spatial_features_2d)
        rm_single = self.reg_head(spatial_features_2d)

        # TODO: do not collabrate, only use ego info
        # [Bx2, C, H, W] dair-v2x has 2 cavs
        # print('spatial_features: ', batch_dict['spatial_features'].shape)
        # batch_dict['spatial_features'][1, :, :, :] = 0.

        # print('spatial_features_2d: ', spatial_features_2d.shape)
        if self.multi_scale:
            fused_feature, communication_rates, result_dict = self.rain_fusion(batch_dict['spatial_features'],
                                            psm_single,
                                            record_len,
                                            pairwise_t_matrix, 
                                            self.backbone,
                                            [self.shrink_conv, self.cls_head, self.reg_head])
            # downsample feature to reduce memory
            if self.shrink_flag:
                fused_feature = self.shrink_conv(fused_feature)
        else:
            fused_feature, communication_rates, result_dict = self.fusion_net(spatial_features_2d,
                                            psm_single,
                                            record_len,
                                            pairwise_t_matrix)
            
            
        # print('fused_feature: ', fused_feature.shape)
        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)
        # psm = self.fused_cls_head(fused_feature)
        # rm = self.fused_reg_head(fused_feature)

        output_dict = {'psm': psm,
                       'rm': rm
                       }

        output_dict.update({'psm_single': psm_single,
                       'rm_single': rm_single,
                       'comm_rate': communication_rates
                       })

        output_dict.update(result_dict)
        '''
        split_psm_single = self.regroup(psm_single, record_len)
        split_rm_single = self.regroup(rm_single, record_len)
        psm_single_v = []
        psm_single_i = []
        rm_single_v = []
        rm_single_i = []
        for b in range(len(split_psm_single)):
            psm_single_v.append(split_psm_single[b][0:1])
            psm_single_i.append(split_psm_single[b][1:2])
            rm_single_v.append(split_rm_single[b][0:1])
            rm_single_i.append(split_rm_single[b][1:2])
        psm_single_v = torch.cat(psm_single_v, dim=0)
        psm_single_i = torch.cat(psm_single_i, dim=0)
        rm_single_v = torch.cat(rm_single_v, dim=0)
        rm_single_i = torch.cat(rm_single_i, dim=0)
        output_dict.update({'psm_single_v': psm_single_v,
                       'psm_single_i': psm_single_i,
                       'rm_single_v': rm_single_v,
                       'rm_single_i': rm_single_i,
                       'comm_rate': communication_rates
                       })
        '''
        return output_dict

opencood/models/point_pillar_where2comm.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Sizhe Wei <sizhewei@sjtu.edu.cn>
# License: TDG-Attribution-NonCommercial-NoDistrib


from numpy import record
import torch.nn as nn

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
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
# from opencood.models.sub_modules.dcn_net import DCNNet
# from opencood.models.fuse_modules.where2comm import Where2comm
from opencood.models.fuse_modules.where2comm_attn import Where2comm
from opencood.models.fuse_modules.raindrop_attn import raindrop_fuse
from opencood.models.fuse_modules.raindrop_swin import raindrop_swin
from opencood.models.fuse_modules.raindrop_swin_w_single import raindrop_swin_w_single
import torch

import numpy as np
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world

def generate_noise(pos_std, rot_std, pos_mean=0, rot_mean=0):
    """ Add localization error to the 6dof pose
        Noise includes position (x,y) and rotation (yaw).
        We use gaussian distribution to generate noise.

    Args:

        pos_std : float 
            std of gaussian dist, in meter

        rot_std : float
            std of gaussian dist, in degree

        pos_mean : float
            mean of gaussian dist, in meter

        rot_mean : float
            mean of gaussian dist, in degree

    Returns:
        pose_noise: np.ndarray, [6,]
            [x, y, z, roll, yaw, pitch]
    """

    xy = np.random.normal(pos_mean, pos_std, size=(2))
    yaw = np.random.normal(rot_mean, rot_std, size=(1))

    pose_noise = np.array([xy[0], xy[1], 0, 0, yaw[0], 0])

    return pose_noise

def get_past_k_pairwise_transformation2ego(past_k_lidar_pose, noise_level, k=3, max_cav=5):
    """
    Get transformation matrixes accross different agents to curr ego at all past timestamps.

    Parameters
    ----------
    base_data_dict : dict
        Key : cav id, item: transformation matrix to ego, lidar points.
    
    ego_pose : list
        ego pose

    max_cav : int
        The maximum number of cav, default 5

    Return
    ------
    pairwise_t_matrix : np.array
        The transformation matrix each cav to curr ego at past k frames.
        shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
        pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
    """
    pos_std = noise_level['pos_std']
    rot_std = noise_level['rot_std']
    pos_mean = 0 
    rot_mean = 0
    
    pairwise_t_matrix = np.tile(np.eye(4), (max_cav, k, 1, 1)) # (L, k, 4, 4)

    ego_pose = past_k_lidar_pose[0, 0]

    t_list = []

    # save all transformation matrix in a list in order first.
    for cav_id in range(past_k_lidar_pose.shape[0]):
        past_k_poses = []
        for time_id in range(k):
            loc_noise = generate_noise(pos_std, rot_std)
            past_k_poses.append(x_to_world(past_k_lidar_pose[cav_id, time_id].cpu().numpy()+loc_noise))
        t_list.append(past_k_poses) # Twx
    
    ego_pose = x_to_world(ego_pose.cpu().numpy())
    for i in range(len(t_list)): # different cav
        if i!=0 :
            for j in range(len(t_list[i])): # different time
                # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                pairwise_t_matrix[i, j] = t_matrix
    pairwise_t_matrix = torch.tensor(pairwise_t_matrix).to(past_k_lidar_pose.device)
    return pairwise_t_matrix


class PointPillarWhere2commAttn(nn.Module):
    def __init__(self, args):
        super(PointPillarWhere2commAttn, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        if 'resnet' in args['base_bev_backbone'] and args['base_bev_backbone']['resnet']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        if 'num_sweep_frames' in args:    # number of frames we use in LSTM
            self.k = args['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in args:          # number of time delay
            self.tau = args['time_delay'] 
        else:
            self.tau = 0
        
        self.noise_flag = 0
        if 'noise' in args.keys():
            self.noise_flag = True
            self.noise_level = {'pos_std': args['noise']['pos_std'], 'rot_std': args['noise']['rot_std'], 'pos_mean': args['noise']['pos_mean'], 'rot_mean': args['noise']['rot_mean']}
            print(f'=== noise level : {self.noise_level} ===')

        self.dcn = False
        # if 'dcn' in args:
        #     self.dcn = True
        #     self.dcn_net = DCNNet(args['dcn'])

        # 用于 single 部分的监督
        self.single_supervise = False
        if 'with_compensation' in args and args['with_compensation']:
            self.compensation = True
            if 'with_single_supervise' in args and args['with_single_supervise']:
                self.rain_fusion = raindrop_swin_w_single(args['rain_model'])
                self.single_supervise = True
            else:
                self.rain_fusion = raindrop_swin(args['rain_model'])
        else: 
            self.compensation = False
            self.rain_fusion = raindrop_fuse(args['rain_model'])

        self.multi_scale = args['rain_model']['multi_scale']
        
        if self.shrink_flag:
            dim = args['shrink_header']['dim'][0]
            self.cls_head = nn.Conv2d(int(dim), args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
                                    kernel_size=1)
            self.fused_cls_head = nn.Conv2d(int(dim), args['anchor_number'],
                                    kernel_size=1)
            self.fused_reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
                                    kernel_size=1)
        else:
            self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
                                    kernel_size=1)

        if 'dir_args' in args.keys():
            self.use_dir = True
            self.dir_head = nn.Conv2d(128 * 2, args['dir_args']['num_bins'] * args['anchor_number'],
                                  kernel_size=1) # BIN_NUM = 2， # 384
            self.fused_dir_head = nn.Conv2d(128 * 2, args['dir_args']['num_bins'] * args['anchor_number'],
                                    kernel_size=1)
        else:
            self.use_dir = False

        if args['backbone_fix']:
            self.backbone_fix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False

        # TODO: for only tune dir header use.
        for p in self.fused_cls_head.parameters():
            p.requires_grad = False
        for p in self.fused_reg_head.parameters():
            p.requires_grad = False

        for p in self.rain_fusion.parameters():
            p.requires_grad = False
    
    def regroup(self, x, record_len, k=1):
        cum_sum_len = torch.cumsum(record_len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, data_dict):
        voxel_features = data_dict['processed_lidar']['voxel_features']         #(M, 32, 4)
        voxel_coords = data_dict['processed_lidar']['voxel_coords']             #(M, 4)
        voxel_num_points = data_dict['processed_lidar']['voxel_num_points']     #(M, )
        record_len = data_dict['record_len']
        record_frames = data_dict['past_k_time_interval']                       #(B, )
        pairwise_t_matrix = data_dict['pairwise_t_matrix']                      #(B, L, k, 4, 4)
        
        # debug = 0
        B, _, k, _, _ = pairwise_t_matrix.shape
        # for i in range(B):
        #     debug += record_len[i]*k

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c  ('pillar_features')
        batch_dict = self.pillar_vfe(batch_dict)
        # (n, c) -> (batch_cav_size, C, H, W) put pillars into spatial feature map ('spatial_features')
        # import ipdb; ipdb.set_trace()
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict) # 'spatial_features_2d': (batch_cav_size, 128*3, H/2, W/2)
        # N, C, H', W'. [N, 384, 100, 352]
        spatial_features_2d = batch_dict['spatial_features_2d']

        noise_pairwise_t_matrix = None
        if self.noise_flag and B==1:
            # noise_level = {'pos_std': 0.5, 'rot_std': 0, 'pos_mean': 0, 'rot_mean': 0}
            noise_pairwise_t_matrix = get_past_k_pairwise_transformation2ego(data_dict['past_lidar_pose'], self.noise_level, k=self.k, max_cav=5)[:record_len[0]].unsqueeze(0)
        
        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)  # (B, 256, H', W')
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        ####### debug use, viz feature of each cav
        # from matplotlib import pyplot as plt
        # viz_save_path = '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/debug_4_feature_flow/where2comm'
        # for i in range(spatial_features_2d.shape[0]):
        #     viz_content = torch.max(spatial_features_2d[i], dim=0)[0].detach().cpu()
        #     plt.imshow(viz_content)
        #     plt.savefig(viz_save_path+f'/updated_feature_{i}.png')
        ##############
        
        # dcn
        # if self.dcn:
        #     spatial_features_2d = self.dcn_net(spatial_features_2d)
        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]
        psm_single = self.cls_head(spatial_features_2d)
        rm_single = self.reg_head(spatial_features_2d)
        if self.use_dir:
            dm_single = self.dir_head(spatial_features_2d)

        # rain attention:
        if self.multi_scale:
            if self.compensation:
                if self.single_supervise:
                    fused_feature, single_feature, communication_rates, all_recon_loss, result_dict = self.rain_fusion(batch_dict['spatial_features'],
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])
                else:
                    fused_feature,fused_feature_curr,fused_feature_latency, communication_rates, all_recon_loss, all_latency_recon_loss, result_dict = self.rain_fusion(batch_dict['spatial_features'],
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])
            else: 
                fused_feature, communication_rates, result_dict = self.rain_fusion(batch_dict['spatial_features'],
                    psm_single,
                    record_len,
                    pairwise_t_matrix, 
                    record_frames,
                    self.backbone,
                    [self.shrink_conv, self.cls_head, self.reg_head],
                    noise_pairwise_t_matrix=noise_pairwise_t_matrix)
            # downsample feature to reduce memory
            if self.shrink_flag:
                fused_feature = self.shrink_conv(fused_feature)
                if self.single_supervise:
                    single_feature = self.shrink_conv(single_feature)
                # if self.compensation:
                    
                #         # single_feature = self.shrink_conv(single_feature)
                
                #         fused_feature_curr = self.shrink_conv(fused_feature_curr)
                #         fused_feature_latency = self.shrink_conv(fused_feature_latency)
        else:
            fused_feature, communication_rates, result_dict = self.rain_fusion(spatial_features_2d,
                                            psm_single,
                                            record_len,
                                            pairwise_t_matrix,
                                            record_frames)
            if self.compensation:
                if self.single_supervise:
                    fused_feature, single_feature, communication_rates, all_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])
                else:
                    fused_feature,fused_feature_curr,fused_feature_latency, communication_rates, all_recon_loss, all_latency_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])            

        ####### debug use, viz fused feature of where2comm
        # viz_content = torch.max(fused_feature[0], dim=0)[0].detach().cpu()
        # plt.imshow(viz_content)
        # plt.savefig(viz_save_path+'/fused_feature.png')  
        ##############
        
        # print('fused_feature: ', fused_feature.shape)
        psm = self.fused_cls_head(fused_feature)
        rm = self.fused_reg_head(fused_feature)
        
        # fuse 之后的 feature (ego)
        output_dict = {'psm': psm,
                       'rm': rm}
        
        if self.use_dir:
            dm = self.fused_dir_head(fused_feature)
            output_dict.update({'dm': dm})

        if self.compensation:
            if self.single_supervise:
                psm_nonego_single = self.cls_head(single_feature)
                rm_nonego_single = self.reg_head(single_feature)
                output_dict.update({
                    'psm_nonego_single': psm_nonego_single,
                    'rm_nonego_single': rm_nonego_single
                })
            
            output_dict.update({
                'recon_loss': all_recon_loss, 
                'record_len': record_len
            })

        output_dict.update({'psm_single': psm_single,
                       'rm_single': rm_single,
                       'comm_rate': communication_rates
                       })
        if self.use_dir:
            output_dict.update({'dm_single': dm_single})
        
        output_dict.update(result_dict) 
        
        return output_dict
opencood/models/point_pillar_where2comm_attn.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Sizhe Wei <sizhewei@sjtu.edu.cn>
# License: TDG-Attribution-NonCommercial-NoDistrib
# Date: 2023-04-10

''' 
where2comm as backbone
use MotionNet to generate flow
use GT flow as supervision
''' 
from numpy import record
import torch.nn as nn

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
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
# from opencood.models.sub_modules.dcn_net import DCNNet
# from opencood.models.fuse_modules.where2comm import Where2comm
from opencood.models.fuse_modules.where2comm_attn import Where2comm
from opencood.models.fuse_modules.raindrop_attn import raindrop_fuse
from opencood.models.fuse_modules.raindrop_swin import raindrop_swin
from opencood.models.fuse_modules.raindrop_swin_w_single import raindrop_swin_w_single
from opencood.models.fuse_modules.raindrop_flow import raindrop_fuse
import torch

class PointPillarWhere2commFlow(nn.Module):
    def __init__(self, args):
        super(PointPillarWhere2commFlow, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        if 'resnet' in args['base_bev_backbone'] and args['base_bev_backbone']['resnet']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        if 'num_sweep_frames' in args:    # number of frames we use in LSTM
            self.k = args['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in args:          # number of time delay
            self.tau = args['time_delay'] 
        else:
            self.tau = 0

        self.dcn = False
        # if 'dcn' in args:
        #     self.dcn = True
        #     self.dcn_net = DCNNet(args['dcn'])

        # 用于 single 部分的监督
        self.single_supervise = False
        if 'with_compensation' in args and args['with_compensation']:
            self.compensation = True
            if 'with_single_supervise' in args and args['with_single_supervise']:
                self.rain_fusion = raindrop_swin_w_single(args['rain_model'])
                self.single_supervise = True
            else:
                self.rain_fusion = raindrop_swin(args['rain_model'])
        else: 
            self.compensation = False
            self.rain_fusion = raindrop_fuse(args['rain_model'])

        self.multi_scale = args['rain_model']['multi_scale']
        
        if self.shrink_flag:
            dim = args['shrink_header']['dim'][0]
            self.cls_head = nn.Conv2d(int(dim), args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
                                    kernel_size=1)
        else:
            self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
                                    kernel_size=1)

        if 'backbone_fix' in args.keys() and args['backbone_fix']:
            self.backbone_fix()
            print('=== backbone fixed ===')

    # def where2comm_fix(self):
    #     """
    #     Fix the parameters of backbone during finetune on timedelay
    #     """
    #     for p in self.pillar_vfe.parameters():
    #         p.requires_grad = False

    #     for p in self.scatter.parameters():
    #         p.requires_grad = False

    #     for p in self.backbone.parameters():
    #         p.requires_grad = False

    #     if self.compression:
    #         for p in self.naive_compressor.parameters():
    #             p.requires_grad = False
    #     if self.shrink_flag:
    #         for p in self.shrink_conv.parameters():
    #             p.requires_grad = False

    #     for p in self.cls_head.parameters():
    #         p.requires_grad = False
    #     for p in self.reg_head.parameters():
    #         p.requires_grad = False
    
    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False
    
    def regroup(self, x, record_len, k=1):
        cum_sum_len = torch.cumsum(record_len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, data_dict):
        voxel_features = data_dict['processed_lidar']['voxel_features']         #(M, 32, 4)
        voxel_coords = data_dict['processed_lidar']['voxel_coords']             #(M, 4)
        voxel_num_points = data_dict['processed_lidar']['voxel_num_points']     #(M, )
        record_len = data_dict['record_len']
        record_frames = data_dict['past_k_time_interval']                       #(B, )
        pairwise_t_matrix = data_dict['pairwise_t_matrix']                      #(B, L, k, 4, 4)
        
        # debug = 0
        B, _, k, _, _ = pairwise_t_matrix.shape
        # for i in range(B):
        #     debug += record_len[i]*k

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c  ('pillar_features')
        batch_dict = self.pillar_vfe(batch_dict)
        # (n, c) -> (batch_cav_size, C, H, W) put pillars into spatial feature map ('spatial_features')
        # import ipdb; ipdb.set_trace()
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict) # 'spatial_features_2d': (batch_cav_size, 128*3, H/2, W/2)
        # N, C, H', W'. [N, 384, 100, 352]
        spatial_features_2d = batch_dict['spatial_features_2d']
        
        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)  # (B, 256, H', W')
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        ####### debug use, viz feature of each cav
        # from matplotlib import pyplot as plt
        # viz_save_path = '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/debug_4_feature_flow/where2comm'
        # for i in range(spatial_features_2d.shape[0]):
        #     viz_content = torch.max(spatial_features_2d[i], dim=0)[0].detach().cpu()
        #     plt.imshow(viz_content)
        #     plt.savefig(viz_save_path+f'/updated_feature_{i}.png')
        ##############
        
        # dcn
        # if self.dcn:
        #     spatial_features_2d = self.dcn_net(spatial_features_2d)
        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]
        psm_single = self.cls_head(spatial_features_2d)
        rm_single = self.reg_head(spatial_features_2d)

        # debug use
        flow_gt = data_dict['label_dict']['flow_gt']

        # rain attention:
        if self.multi_scale:
            fused_feature, communication_rates, result_dict, flow, state_preds = self.rain_fusion(batch_dict['spatial_features'],
                psm_single,
                record_len,
                pairwise_t_matrix, 
                record_frames,
                self.backbone,
                [self.shrink_conv, self.cls_head, self.reg_head],
                flow_gt=flow_gt)
            # downsample feature to reduce memory
            if self.shrink_flag:
                fused_feature = self.shrink_conv(fused_feature)
                if self.single_supervise:
                    single_feature = self.shrink_conv(single_feature)
                # if self.compensation:
                    
                #         # single_feature = self.shrink_conv(single_feature)
                
                #         fused_feature_curr = self.shrink_conv(fused_feature_curr)
                #         fused_feature_latency = self.shrink_conv(fused_feature_latency)
        else:
            fused_feature, communication_rates, result_dict = self.rain_fusion(spatial_features_2d,
                                            psm_single,
                                            record_len,
                                            pairwise_t_matrix,
                                            record_frames)
            if self.compensation:
                if self.single_supervise:
                    fused_feature, single_feature, communication_rates, all_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])
                else:
                    fused_feature,fused_feature_curr,fused_feature_latency, communication_rates, all_recon_loss, all_latency_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])            

        ####### debug use, viz fused feature of where2comm
        # viz_content = torch.max(fused_feature[0], dim=0)[0].detach().cpu()
        # plt.imshow(viz_content)
        # plt.savefig(viz_save_path+'/fused_feature.png')  
        ##############
        
        # print('fused_feature: ', fused_feature.shape)
        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)
        
        # fuse 之后的 feature (ego)
        output_dict = {'psm': psm,
                       'rm': rm}

        output_dict.update({'flow_preds': flow,
                            'state_preds': state_preds})

        if self.compensation:
            if self.single_supervise:
                psm_nonego_single = self.cls_head(single_feature)
                rm_nonego_single = self.reg_head(single_feature)
                output_dict.update({
                    'psm_nonego_single': psm_nonego_single,
                    'rm_nonego_single': rm_nonego_single
                })
            
            output_dict.update({
                'recon_loss': all_recon_loss, 
                'record_len': record_len
            })

        output_dict.update({'psm_single': psm_single,
                       'rm_single': rm_single,
                       'comm_rate': communication_rates
                       })
        
        output_dict.update(result_dict) 
        
        return output_dict

opencood/models/point_pillar_where2comm_flow.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Sizhe Wei <sizhewei@sjtu.edu.cn>
# License: TDG-Attribution-NonCommercial-NoDistrib
# Date: 2023-04-10

'''
Use where2comm as backbone
We just need to learn the flow generator for each cav
Training part is MotionNet
'''
from numpy import record
import torch.nn as nn

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
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
from opencood.models.fuse_modules.raindrop_flow_regular import raindrop_fuse
import torch

class PointPillarWhere2commFlowRegular(nn.Module):
    def __init__(self, args):
        super(PointPillarWhere2commFlowRegular, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        if 'resnet' in args['base_bev_backbone'] and args['base_bev_backbone']['resnet']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        if 'num_sweep_frames' in args:    # number of frames we use in LSTM
            self.k = args['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in args:          # number of time delay
            self.tau = args['time_delay'] 
        else:
            self.tau = 0

        self.dcn = False
        # if 'dcn' in args:
        #     self.dcn = True
        #     self.dcn_net = DCNNet(args['dcn'])

        self.compensation = False
        self.rain_fusion = raindrop_fuse(args['rain_model'])

        self.multi_scale = args['rain_model']['multi_scale']

        self.single_supervise = False
        
        if self.shrink_flag:
            dim = args['shrink_header']['dim'][0]
            self.cls_head = nn.Conv2d(int(dim), args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
                                    kernel_size=1)
        else:
            self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
                                    kernel_size=1)


        self.backbone_fix_flag = False
        if 'backbone_fix' in args.keys() and args['backbone_fix']:
            self.backbone_fix_flag = True
            self.backbone_fix()
            print('=== backbone fixed ===')

        self.only_tune_header_flag = False
        if 'only_tune_header' in args.keys() and args['only_tune_header']:
            self.only_tune_header_flag = True
            self.only_tune_header()
            print('=== only tune header ===')

        assert self.backbone_fix_flag == False or self.only_tune_header_flag == False, 'backbone_fix and only_tune_header cannot be True at the same time'
    
    def only_tune_header(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        for p in self.rain_fusion.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False
    
    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False
    
    def regroup(self, x, record_len, k=1):
        cum_sum_len = torch.cumsum(record_len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, data_dict):
        voxel_features = data_dict['processed_lidar']['voxel_features']         #(M, 32, 4)
        voxel_coords = data_dict['processed_lidar']['voxel_coords']             #(M, 4)
        voxel_num_points = data_dict['processed_lidar']['voxel_num_points']     #(M, )
        record_len = data_dict['record_len']
        record_frames = data_dict['past_k_time_interval']                       #(B, )
        pairwise_t_matrix = data_dict['pairwise_t_matrix']                      #(B, L, k, 4, 4)
        pastk_2_past0_tr_mats = data_dict['pastk_2_past0_tr_mats']              #(sum(N_b), k, 4, 4)
        
        # debug = 0
        B, _, k, _, _ = pairwise_t_matrix.shape
        # for i in range(B):
        #     debug += record_len[i]*k

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c  ('pillar_features')
        batch_dict = self.pillar_vfe(batch_dict)
        # (n, c) -> (batch_cav_size, C, H, W) put pillars into spatial feature map ('spatial_features')
        # import ipdb; ipdb.set_trace()
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict) # 'spatial_features_2d': (batch_cav_size, 128*3, H/2, W/2)
        # N, C, H', W'. [N, 384, 100, 352]
        spatial_features_2d = batch_dict['spatial_features_2d']
        
        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)  # (B, 256, H', W')
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        psm_single = self.cls_head(spatial_features_2d)
        rm_single = self.reg_head(spatial_features_2d)

        # debug use
        flow_gt = data_dict['label_dict']['flow_gt']
        warp_mask = data_dict['label_dict']['warp_mask']

        # rain attention:
        if self.multi_scale:
            fused_feature, communication_rates, result_dict, flow, state_preds = self.rain_fusion(batch_dict['spatial_features'],
                psm_single,
                record_len,
                pairwise_t_matrix, 
                record_frames,
                self.backbone,
                [self.shrink_conv, self.cls_head, self.reg_head],
                pastk_2_past0_tr_mats=pastk_2_past0_tr_mats,
                flow_gt=flow_gt, 
                warp_mask=warp_mask)
            # downsample feature to reduce memory
            if self.shrink_flag:
                fused_feature = self.shrink_conv(fused_feature)
                if self.single_supervise:
                    single_feature = self.shrink_conv(single_feature)
        else:
            fused_feature, communication_rates, result_dict = self.rain_fusion(spatial_features_2d,
                                            psm_single,
                                            record_len,
                                            pairwise_t_matrix,
                                            record_frames)
            if self.compensation:
                if self.single_supervise:
                    fused_feature, single_feature, communication_rates, all_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])
                else:
                    fused_feature,fused_feature_curr,fused_feature_latency, communication_rates, all_recon_loss, all_latency_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])            
        
        # print('fused_feature: ', fused_feature.shape)
        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)
        
        # fuse 之后的 feature (ego)
        output_dict = {'psm': psm,
                       'rm': rm}

        output_dict.update({'flow_preds': flow,
                            'state_preds': state_preds})

        if self.compensation:
            if self.single_supervise:
                psm_nonego_single = self.cls_head(single_feature)
                rm_nonego_single = self.reg_head(single_feature)
                output_dict.update({
                    'psm_nonego_single': psm_nonego_single,
                    'rm_nonego_single': rm_nonego_single
                })
            
            output_dict.update({
                'recon_loss': all_recon_loss, 
                'record_len': record_len
            })

        output_dict.update({'psm_single': psm_single,
                       'rm_single': rm_single,
                       'comm_rate': communication_rates
                       })
        
        output_dict.update(result_dict) 
        
        return output_dict

opencood/models/point_pillar_where2comm_flow_regular.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Sizhe Wei <sizhewei@sjtu.edu.cn>
# License: TDG-Attribution-NonCommercial-NoDistrib
'''
Use where2comm as backbone
First generate the box
then use the box to generate the flow
利用box生成flow
将该flow通过某种方法warp feature 
将补偿好的feature融合 得到最终结果
'''

from numpy import record
import torch.nn as nn

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
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
# from opencood.models.sub_modules.dcn_net import DCNNet
# from opencood.models.fuse_modules.where2comm import Where2comm
from opencood.models.fuse_modules.where2comm_attn import Where2comm
from opencood.models.fuse_modules.raindrop_swin import raindrop_swin
from opencood.models.fuse_modules.raindrop_swin_w_single import raindrop_swin_w_single
from opencood.models.fuse_modules.raindrop_flow import raindrop_fuse
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.tools.matcher import Matcher
from collections import OrderedDict
import torch
import numpy as np

def generate_noise(pos_std, rot_std, pos_mean=0, rot_mean=0):
    """ Add localization error to the 6dof pose
        Noise includes position (x,y) and rotation (yaw).
        We use gaussian distribution to generate noise.

    Args:

        pos_std : float 
            std of gaussian dist, in meter

        rot_std : float
            std of gaussian dist, in degree

        pos_mean : float
            mean of gaussian dist, in meter

        rot_mean : float
            mean of gaussian dist, in degree

    Returns:
        pose_noise: np.ndarray, [6,]
            [x, y, z, roll, yaw, pitch]
    """

    xy = np.random.normal(pos_mean, pos_std, size=(2))
    yaw = np.random.normal(rot_mean, rot_std, size=(1))

    pose_noise = np.array([xy[0], xy[1], 0, 0, yaw[0], 0])

    return pose_noise

def get_past_k_pairwise_transformation2ego(past_k_lidar_pose, noise_level, k=3, max_cav=5):
    """
    Get transformation matrixes accross different agents to curr ego at all past timestamps.

    Parameters
    ----------
    base_data_dict : dict
        Key : cav id, item: transformation matrix to ego, lidar points.
    
    ego_pose : list
        ego pose

    max_cav : int
        The maximum number of cav, default 5

    Return
    ------
    pairwise_t_matrix : np.array
        The transformation matrix each cav to curr ego at past k frames.
        shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
        pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
    """
    pos_std = noise_level['pos_std']
    rot_std = noise_level['rot_std']
    pos_mean = 0 
    rot_mean = 0
    
    pairwise_t_matrix = np.tile(np.eye(4), (max_cav, k, 1, 1)) # (L, k, 4, 4)

    ego_pose = past_k_lidar_pose[0, 0]

    t_list = []

    # save all transformation matrix in a list in order first.
    for cav_id in range(past_k_lidar_pose.shape[0]):
        past_k_poses = []
        for time_id in range(k):
            loc_noise = generate_noise(pos_std, rot_std)
            past_k_poses.append(x_to_world(past_k_lidar_pose[cav_id, time_id].cpu().numpy()+loc_noise))
        t_list.append(past_k_poses) # Twx
    
    ego_pose = x_to_world(ego_pose.cpu().numpy())
    for i in range(len(t_list)): # different cav
        if i!=0 :
            for j in range(len(t_list[i])): # different time
                # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                pairwise_t_matrix[i, j] = t_matrix
    pairwise_t_matrix = torch.tensor(pairwise_t_matrix).to(past_k_lidar_pose.device)
    return pairwise_t_matrix



class PointPillarWhere2commFlowTwoStage(nn.Module):
    def __init__(self, args):
        super(PointPillarWhere2commFlowTwoStage, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        if 'resnet' in args['base_bev_backbone'] and args['base_bev_backbone']['resnet']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        if 'num_sweep_frames' in args:    # number of frames we use in LSTM
            self.k = args['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in args:          # number of time delay
            self.tau = args['time_delay'] 
        else:
            self.tau = 0

        self.num_roi_thres = -1
        if 'num_roi_thres' in args.keys():
            self.num_roi_thres = args['num_roi_thres']
            print(f'=== num_roi_thres : {self.num_roi_thres} ===')

        self.dcn = False
        # if 'dcn' in args:
        #     self.dcn = True
        #     self.dcn_net = DCNNet(args['dcn'])

        self.design_mode = 0
        if 'design_mode' in args.keys():
            self.design_mode = args['design_mode']
            print(f'=== design mode : {self.design_mode} ===')

        self.noise_flag = 0
        if 'noise' in args.keys():
            self.noise_flag = True
            self.noise_level = {'pos_std': args['noise']['pos_std'], 'rot_std': args['noise']['rot_std'], 'pos_mean': args['noise']['pos_mean'], 'rot_mean': args['noise']['rot_mean']}
            print(f'=== noise level : {self.noise_level} ===')

        # 用于 single 部分的监督
        self.single_supervise = False
        if 'with_compensation' in args and args['with_compensation']:
            self.compensation = True
            if 'with_single_supervise' in args and args['with_single_supervise']:
                self.rain_fusion = raindrop_swin_w_single(args['rain_model'])
                self.single_supervise = True
            else:
                self.rain_fusion = raindrop_swin(args['rain_model'])
        else: 
            self.compensation = False
            self.rain_fusion = raindrop_fuse(args['rain_model'], self.design_mode)

        self.multi_scale = args['rain_model']['multi_scale']
        
        if self.shrink_flag:
            dim = args['shrink_header']['dim'][0]
            self.cls_head = nn.Conv2d(int(dim), args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
                                    kernel_size=1)
            self.fused_cls_head = nn.Conv2d(int(dim), args['anchor_number'],
                                    kernel_size=1)
            self.fused_reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
                                    kernel_size=1)
        else:
            self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
                                    kernel_size=1)

        if self.design_mode == 0:
            self.matcher = Matcher('flow')
        else:
            self.matcher = Matcher('linear')

        self.backbone_fix_flag = False
        if 'backbone_fix' in args.keys() and args['backbone_fix']:
            self.backbone_fix_flag = True
            self.backbone_fix()
            print('=== backbone fixed ===')

        self.only_tune_header_flag = False
        if 'only_tune_header' in args.keys() and args['only_tune_header']:
            self.only_tune_header_flag = True
            self.only_tune_header()
            print('=== only tune header ===')

        self.viz_bbx_flag = False
        if 'viz_bbx_flag' in args.keys() and args['viz_bbx_flag']:
            self.viz_bbx_flag = True
        
        assert self.backbone_fix_flag == False or self.only_tune_header_flag == False, 'backbone_fix and only_tune_header cannot be True at the same time'
    
    def only_tune_header(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        for p in self.rain_fusion.parameters():
            p.requires_grad = False

        for p in self.matcher.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False
    
    def regroup(self, x, record_len, k=1):
        cum_sum_len = torch.cumsum(record_len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def generate_box_flow(self, data_dict, pred_dict, dataset, device): 
        """
        data_dict : 

        pred_dict : 
            {
                psm_single_list: len = B, each element's shape is (N_b x k, 2, H, W) 
                rm_single_list: len = B, each element's shape is (N_b x k, 14, H, W) 
            }
        """
        # for b in range(B):
        # 1. get box results
        lidar_pose_batch = self.regroup(data_dict['past_lidar_pose'], data_dict['record_len'], k=1)
        past_k_time_diff = self.regroup(data_dict['past_k_time_interval'], data_dict['record_len'], k=3)
        anchor_box = data_dict['anchor_box'] # (H, W, 2, 7)
        psm_single_list = pred_dict['psm_single_list']
        rm_single_list = pred_dict['rm_single_list']

        # H, W = psm_single_list[0].shape[-2:]
        shape_list = torch.tensor([64, 200, 704]).to(device)
        
        trans_mat_pastk_2_past0_batch = []
        B = len(lidar_pose_batch)
        box_flow_map_list = []
        reserved_mask_list = []

        if self.viz_bbx_flag:
            single_box_results = None
        
        # for all batches
        for b in range(B):
            box_results = OrderedDict()
            psm_single = psm_single_list[b].reshape(-1, self.k, 2, psm_single_list[b].shape[-2], psm_single_list[b].shape[-1]) # (N_b, k, 2, H, W)
            rm_single = rm_single_list[b].reshape(-1, self.k, 14, rm_single_list[b].shape[-2], rm_single_list[b].shape[-1]) # (N_b, k, 14, H, W)
            cav_past_k_time_diff = past_k_time_diff[b]
            cav_trans_mat_pastk_2_past0 = []
            # for all cavs
            '''
            box_result : dict for each cav at each time
            {
                cav_idx : {
                    'past_k_time_diff' : 
                    [0] : {
                        pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. (n, 8, 3)
                        pred_box_center_tensor : (n, 7)
                        scores: (n, )
                    },
                    ...
                    [k-1] : { ... }
                }
            }
            '''
            for cav_idx in range(data_dict['record_len'][b]):
                # generate one cav's trans_mat_pastk_2_past0
                pastk_trans_mat_pastk_2_past0 = []
                for i in range(self.k):
                    unit_mat = x1_to_x2(lidar_pose_batch[b][cav_idx, i, :].cpu().numpy(), lidar_pose_batch[b][cav_idx, 0].cpu()) # (4, 4)
                    pastk_trans_mat_pastk_2_past0.append(unit_mat)
                pastk_trans_mat_pastk_2_past0 = torch.from_numpy(np.stack(pastk_trans_mat_pastk_2_past0, axis=0)).to(device) # (k, 4, 4)
                
                # 1. generate one cav's box results
                box_results[cav_idx] = dataset.generate_pred_bbx_frames(psm_single[cav_idx], rm_single[cav_idx], pastk_trans_mat_pastk_2_past0, cav_past_k_time_diff[cav_idx*self.k:cav_idx*self.k+self.k], anchor_box)

            cav_trans_mat_pastk_2_past0.append(pastk_trans_mat_pastk_2_past0)
            
            # 2. generate box flow in one batch
            if self.viz_bbx_flag:
                box_flow_map, mask, matched_idx_list, compensated_results_list = self.matcher(box_results, shape_list=shape_list, viz_flag=self.viz_bbx_flag)
            else:
                box_flow_map, mask = self.matcher(box_results, shape_list=shape_list, viz_flag=self.viz_bbx_flag)
            box_flow_map_list.append(box_flow_map)
            reserved_mask_list.append(mask)

            if self.viz_bbx_flag:
                single_box_results = box_results
        
        final_flow_map = torch.concat(box_flow_map_list, dim=0)
        final_reserved_mask = torch.concat(reserved_mask_list, dim=0)

        if self.viz_bbx_flag:
            return final_flow_map, final_reserved_mask, single_box_results, matched_idx_list, compensated_results_list

        return final_flow_map, final_reserved_mask

    def forward(self, data_dict, dataset=None):
        voxel_features = data_dict['processed_lidar']['voxel_features']         #(M, 32, 4)
        voxel_coords = data_dict['processed_lidar']['voxel_coords']             #(M, 4)
        voxel_num_points = data_dict['processed_lidar']['voxel_num_points']     #(M, )
        record_len = data_dict['record_len']
        record_frames = data_dict['past_k_time_interval']                       #(B, )
        pairwise_t_matrix = data_dict['pairwise_t_matrix']                      #(B, L, k, 4, 4)
        
        # debug = 0
        B, _, k, _, _ = pairwise_t_matrix.shape
        # for i in range(B):
        #     debug += record_len[i]*k

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c  ('pillar_features')
        batch_dict = self.pillar_vfe(batch_dict)
        # (n, c) -> (batch_cav_size, C, H, W) put pillars into spatial feature map ('spatial_features')
        # import ipdb; ipdb.set_trace()
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict) # 'spatial_features_2d': (batch_cav_size, 128*3, H/2, W/2)
        # N, C, H', W'. [N, 384, 100, 352]
        spatial_features_2d = batch_dict['spatial_features_2d']

        noise_pairwise_t_matrix = None
        if self.noise_flag and B==1:
            # noise_level = {'pos_std': 0.5, 'rot_std': 0, 'pos_mean': 0, 'rot_mean': 0}
            noise_pairwise_t_matrix = get_past_k_pairwise_transformation2ego(data_dict['past_lidar_pose'], self.noise_level, k=self.k, max_cav=5)[:record_len[0]].unsqueeze(0)
        
        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)  # (B, 256, H', W')
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        ####### debug use, viz feature of each cav
        # from matplotlib import pyplot as plt
        # viz_save_path = '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/debug_4_feature_flow/where2comm'
        # for i in range(spatial_features_2d.shape[0]):
        #     viz_content = torch.max(spatial_features_2d[i], dim=0)[0].detach().cpu()
        #     plt.imshow(viz_content)
        #     plt.savefig(viz_save_path+f'/updated_feature_{i}.png')
        ##############
        
        # dcn
        # if self.dcn:
        #     spatial_features_2d = self.dcn_net(spatial_features_2d)
        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]
        psm_single = self.cls_head(spatial_features_2d).detach()
        rm_single = self.reg_head(spatial_features_2d).detach()
        single_detection_bbx = None

        # if self.design_mode != 4 and not self.only_tune_header_flag:
        if self.design_mode != 4:
            # generate box flow
            # [B, 256, 50, 176]
            single_output = {}
            single_output.update({'psm_single_list': self.regroup(psm_single, record_len, k), 
            'rm_single_list': self.regroup(rm_single, record_len, k)})
            if self.viz_bbx_flag:
                box_flow_map, reserved_mask, single_detection_bbx, matched_idx_list, compensated_results_list = self.generate_box_flow(data_dict, single_output, dataset, psm_single.device)
            else:
                box_flow_map, reserved_mask = self.generate_box_flow(data_dict, single_output, dataset, psm_single.device)

        if 'flow_gt' in data_dict['label_dict']:
            flow_gt = data_dict['label_dict']['flow_gt']
            mask_gt = data_dict['label_dict']['warp_mask']
        else:
            flow_gt = None
        
        # if self.only_tune_header_flag:
        #     box_flow_map = flow_gt
        #     reserved_mask = mask_gt

        # rain attention:
        if self.multi_scale:
            if self.design_mode == 0 or self.design_mode==5:
                fused_feature, communication_rates, result_dict = self.rain_fusion(batch_dict['spatial_features'],
                    psm_single,
                    record_len,
                    pairwise_t_matrix, 
                    record_frames,
                    self.backbone,
                    [self.shrink_conv, self.cls_head, self.reg_head],
                    box_flow=box_flow_map, reserved_mask=reserved_mask,
                    flow_gt=flow_gt)
            elif self.design_mode == 4:
                fused_feature, communication_rates, result_dict = self.rain_fusion(batch_dict['spatial_features'],
                    psm_single,
                    record_len,
                    pairwise_t_matrix, 
                    record_frames,
                    self.backbone,
                    [self.shrink_conv, self.cls_head, self.reg_head],
                    None, num_roi_thres= self.num_roi_thres, noise_pairwise_t_matrix=noise_pairwise_t_matrix)
            else: 
                fused_feature, communication_rates, result_dict, flow_recon_loss = self.rain_fusion(batch_dict['spatial_features'],
                    psm_single,
                    record_len,
                    pairwise_t_matrix, 
                    record_frames,
                    self.backbone,
                    [self.shrink_conv, self.cls_head, self.reg_head],
                    box_flow=box_flow_map, reserved_mask=reserved_mask,
                    flow_gt=flow_gt)
            # downsample feature to reduce memory
            if self.shrink_flag:
                fused_feature = self.shrink_conv(fused_feature)
                if self.single_supervise:
                    single_feature = self.shrink_conv(single_feature)
                # if self.compensation:
                    
                #         # single_feature = self.shrink_conv(single_feature)
                
                #         fused_feature_curr = self.shrink_conv(fused_feature_curr)
                #         fused_feature_latency = self.shrink_conv(fused_feature_latency)
        else:
            fused_feature, communication_rates, result_dict = self.rain_fusion(spatial_features_2d,
                                            psm_single,
                                            record_len,
                                            pairwise_t_matrix,
                                            record_frames)
            if self.compensation:
                if self.single_supervise:
                    fused_feature, single_feature, communication_rates, all_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])
                else:
                    fused_feature,fused_feature_curr,fused_feature_latency, communication_rates, all_recon_loss, all_latency_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])            

        ####### debug use, viz fused feature of where2comm
        # viz_content = torch.max(fused_feature[0], dim=0)[0].detach().cpu()
        # plt.imshow(viz_content)
        # plt.savefig(viz_save_path+'/fused_feature.png')  
        ##############
        
        # print('fused_feature: ', fused_feature.shape)
        if self.only_tune_header_flag:
            psm = self.fused_cls_head(fused_feature)
            rm = self.fused_reg_head(fused_feature)
        else: 
            psm = self.cls_head(fused_feature)
            rm = self.reg_head(fused_feature)
        
        # fuse 之后的 feature (ego)
        output_dict = {'psm': psm,
                       'rm': rm}

        if self.compensation:
            if self.single_supervise:
                psm_nonego_single = self.cls_head(single_feature)
                rm_nonego_single = self.reg_head(single_feature)
                output_dict.update({
                    'psm_nonego_single': psm_nonego_single,
                    'rm_nonego_single': rm_nonego_single
                })
            
            output_dict.update({
                'recon_loss': all_recon_loss, 
                'record_len': record_len
            })

        output_dict.update({'psm_single': psm_single,
                       'rm_single': rm_single,
                       'comm_rate': communication_rates
                       })

        if self.viz_bbx_flag:
            output_dict.update({
                'single_detection_bbx': single_detection_bbx, 
                'matched_idx_list': matched_idx_list, 
                'compensated_results_list': compensated_results_list
            })

        if self.design_mode == 1:
            output_dict.update({'flow_recon_loss': flow_recon_loss})
        
        output_dict.update(result_dict) 
        
        return output_dict

opencood/models/point_pillar_where2comm_flow_two_stage.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Sizhe Wei <sizhewei@sjtu.edu.cn>
# License: TDG-Attribution-NonCommercial-NoDistrib
'''
Use where2comm as backbone
First generate the box
then use the box to generate the flow
利用box生成flow
将该flow通过某种方法warp feature 
将补偿好的feature融合 得到最终结果
'''

from numpy import record
import torch.nn as nn

from opencood.models.sub_modules.pillar_vfe import PillarVFE
from opencood.models.sub_modules.point_pillar_scatter import PointPillarScatter
# from opencood.models.sub_modules.max_resnet_bev_backbone import MaxResNetBEVBackbone
from opencood.models.sub_modules.base_bev_backbone import BaseBEVBackbone
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
# from opencood.models.sub_modules.dcn_net import DCNNet
# from opencood.models.fuse_modules.where2comm import Where2comm
from opencood.models.fuse_modules.where2comm_attn import Where2comm
from opencood.models.fuse_modules.raindrop_swin import raindrop_swin
from opencood.models.fuse_modules.raindrop_swin_w_single import raindrop_swin_w_single
from opencood.models.fuse_modules.raindrop_flow import raindrop_fuse
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world
from opencood.tools.matcher import Matcher
from collections import OrderedDict
import torch
import numpy as np

def generate_noise(pos_std, rot_std, pos_mean=0, rot_mean=0):
    """ Add localization error to the 6dof pose
        Noise includes position (x,y) and rotation (yaw).
        We use gaussian distribution to generate noise.

    Args:

        pos_std : float 
            std of gaussian dist, in meter

        rot_std : float
            std of gaussian dist, in degree

        pos_mean : float
            mean of gaussian dist, in meter

        rot_mean : float
            mean of gaussian dist, in degree

    Returns:
        pose_noise: np.ndarray, [6,]
            [x, y, z, roll, yaw, pitch]
    """

    xy = np.random.normal(pos_mean, pos_std, size=(2))
    yaw = np.random.normal(rot_mean, rot_std, size=(1))

    pose_noise = np.array([xy[0], xy[1], 0, 0, yaw[0], 0])

    return pose_noise

def get_past_k_pairwise_transformation2ego(past_k_lidar_pose, noise_level, k=3, max_cav=5):
    """
    Get transformation matrixes accross different agents to curr ego at all past timestamps.

    Parameters
    ----------
    base_data_dict : dict
        Key : cav id, item: transformation matrix to ego, lidar points.
    
    ego_pose : list
        ego pose

    max_cav : int
        The maximum number of cav, default 5

    Return
    ------
    pairwise_t_matrix : np.array
        The transformation matrix each cav to curr ego at past k frames.
        shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
        pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
    """
    pos_std = noise_level['pos_std']
    rot_std = noise_level['rot_std']
    pos_mean = 0 
    rot_mean = 0
    
    pairwise_t_matrix = np.tile(np.eye(4), (max_cav, k, 1, 1)) # (L, k, 4, 4)

    ego_pose = past_k_lidar_pose[0, 0]

    t_list = []

    # save all transformation matrix in a list in order first.
    for cav_id in range(past_k_lidar_pose.shape[0]):
        past_k_poses = []
        for time_id in range(k):
            loc_noise = generate_noise(pos_std, rot_std)
            past_k_poses.append(x_to_world(past_k_lidar_pose[cav_id, time_id].cpu().numpy()+loc_noise))
        t_list.append(past_k_poses) # Twx
    
    ego_pose = x_to_world(ego_pose.cpu().numpy())
    for i in range(len(t_list)): # different cav
        if i!=0 :
            for j in range(len(t_list[i])): # different time
                # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                pairwise_t_matrix[i, j] = t_matrix
    pairwise_t_matrix = torch.tensor(pairwise_t_matrix).to(past_k_lidar_pose.device)
    return pairwise_t_matrix


class PointPillarWhere2commFlowTwoStageWDir(nn.Module):
    def __init__(self, args):
        super(PointPillarWhere2commFlowTwoStageWDir, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        # self.backbone = MaxResNetBEVBackbone(args['base_bev_backbone'], 64)
        if 'resnet' in args['base_bev_backbone'] and args['base_bev_backbone']['resnet']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        # if args['compression'] > 0:
        #     self.compression = True
        #     self.naive_compressor = NaiveCompressor(256, args['compression'])

        if 'num_sweep_frames' in args:    # number of frames we use in LSTM
            self.k = args['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in args:          # number of time delay
            self.tau = args['time_delay'] 
        else:
            self.tau = 0

        self.dcn = False
        # if 'dcn' in args:
        #     self.dcn = True
        #     self.dcn_net = DCNNet(args['dcn'])

        self.design_mode = 0
        if 'design_mode' in args.keys():
            self.design_mode = args['design_mode']
            print(f'=== design mode : {self.design_mode} ===')

        self.noise_flag = 0
        if 'noise' in args.keys():
            self.noise_flag = True
            self.noise_level = {'pos_std': args['noise']['pos_std'], 'rot_std': args['noise']['rot_std'], 'pos_mean': args['noise']['pos_mean'], 'rot_mean': args['noise']['rot_mean']}
            print(f'=== noise level : {self.noise_level} ===')

        self.num_roi_thres = -1
        if 'num_roi_thres' in args.keys():
            self.num_roi_thres = args['num_roi_thres']
            print(f'=== num_roi_thres : {self.num_roi_thres} ===')

        # 用于 single 部分的监督
        self.single_supervise = False
        if 'with_compensation' in args and args['with_compensation']:
            self.compensation = True
            if 'with_single_supervise' in args and args['with_single_supervise']:
                self.rain_fusion = raindrop_swin_w_single(args['rain_model'])
                self.single_supervise = True
            else:
                self.rain_fusion = raindrop_swin(args['rain_model'])
        else: 
            self.compensation = False
            self.rain_fusion = raindrop_fuse(args['rain_model'], self.design_mode)

        self.multi_scale = args['rain_model']['multi_scale']
        
        if self.shrink_flag:
            dim = args['shrink_header']['dim'][0]
            self.cls_head = nn.Conv2d(int(dim), args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
                                    kernel_size=1)
            self.fused_cls_head = nn.Conv2d(int(dim), args['anchor_number'],
                                    kernel_size=1)
            self.fused_reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
                                    kernel_size=1)
        else:
            self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
                                    kernel_size=1)

        if 'dir_args' in args.keys():
            self.use_dir = True
            self.dir_head = nn.Conv2d(128 * 2, args['dir_args']['num_bins'] * args['anchor_number'],
                                    kernel_size=1) # BIN_NUM = 2， # 384
            self.fused_dir_head = nn.Conv2d(128 * 2, args['dir_args']['num_bins'] * args['anchor_number'],
                                    kernel_size=1)
        else:
            self.use_dir = False

        if self.design_mode == 0:
            self.matcher = Matcher('flow')
        else:
            self.matcher = Matcher('linear')

        self.backbone_fix_flag = False
        if 'backbone_fix' in args.keys() and args['backbone_fix']:
            self.backbone_fix_flag = True
            self.backbone_fix()
            print('=== backbone fixed ===')

        self.only_tune_header_flag = False
        if 'only_tune_header' in args.keys() and args['only_tune_header']:
            self.only_tune_header_flag = True
            self.only_tune_header()
            print('=== only tune header ===')

        self.viz_bbx_flag = False
        if 'viz_bbx_flag' in args.keys() and args['viz_bbx_flag']:
            self.viz_bbx_flag = True
        
        assert self.backbone_fix_flag == False or self.only_tune_header_flag == False, 'backbone_fix and only_tune_header cannot be True at the same time'
    
    def only_tune_header(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        for p in self.rain_fusion.parameters():
            p.requires_grad = False

        for p in self.matcher.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False
    
    def regroup(self, x, record_len, k=1):
        cum_sum_len = torch.cumsum(record_len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def bandwidth_filter(self, input, num_box):
        topk_idx = torch.argsort(input['scores'], descending=True)[:num_box]

        output = {}
        output['scores'] = input['scores'][topk_idx]
        output['pred_box_3dcorner_tensor']  = input['pred_box_3dcorner_tensor'][topk_idx]
        output['pred_box_center_tensor'] = input['pred_box_center_tensor'][topk_idx]

        return output

    def generate_box_flow(self, data_dict, pred_dict, dataset, shape_list, device): 
        """
        data_dict : 

        pred_dict : 
            {
                psm_single_list: len = B, each element's shape is (N_b x k, 2, H, W) 
                rm_single_list: len = B, each element's shape is (N_b x k, 14, H, W) 
            }
        """
        # for b in range(B):
        # 1. get box results
        lidar_pose_batch = self.regroup(data_dict['past_lidar_pose'], data_dict['record_len'], k=1)
        past_k_time_diff = self.regroup(data_dict['past_k_time_interval'], data_dict['record_len'], k=3)
        anchor_box = data_dict['anchor_box'] # (H, W, 2, 7)
        psm_single_list = pred_dict['psm_single_list']
        rm_single_list = pred_dict['rm_single_list']
        dm_single_list = pred_dict['dm_single_list']

        # H, W = psm_single_list[0].shape[-2:]
        # shape_list = torch.tensor([64, H, W]).to(device)
        
        trans_mat_pastk_2_past0_batch = []
        B = len(lidar_pose_batch)
        box_flow_map_list = []
        reserved_mask_list = []

        if self.viz_bbx_flag:
            ori_reserved_mask_list = []
            single_box_results = None
        
        # for all batches
        for b in range(B):
            box_results = OrderedDict()
            psm_single = psm_single_list[b].reshape(-1, self.k, 2, psm_single_list[b].shape[-2], psm_single_list[b].shape[-1]) # (N_b, k, 2, H, W)
            rm_single = rm_single_list[b].reshape(-1, self.k, 14, rm_single_list[b].shape[-2], rm_single_list[b].shape[-1]) # (N_b, k, 14, H, W)
            if self.use_dir:
                dm_single = dm_single_list[b].reshape(-1, self.k, 4, dm_single_list[b].shape[-2], dm_single_list[b].shape[-1]) # (N_b, k, 4, H, W)

            cav_past_k_time_diff = past_k_time_diff[b]
            cav_trans_mat_pastk_2_past0 = []
            # for all cavs
            '''
            box_result : dict for each cav at each time
            {
                cav_idx : {
                    'past_k_time_diff' : 
                    [0] : {
                        pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. (n, 8, 3)
                        pred_box_center_tensor : (n, 7)
                        scores: (n, )
                    },
                    ...
                    [k-1] : { ... }
                }
            }
            '''
            for cav_idx in range(data_dict['record_len'][b]):
                # generate one cav's trans_mat_pastk_2_past0
                pastk_trans_mat_pastk_2_past0 = []
                for i in range(self.k):
                    unit_mat = x1_to_x2(lidar_pose_batch[b][cav_idx, i, :].cpu().numpy(), lidar_pose_batch[b][cav_idx, 0].cpu()) # (4, 4)
                    pastk_trans_mat_pastk_2_past0.append(unit_mat)
                pastk_trans_mat_pastk_2_past0 = torch.from_numpy(np.stack(pastk_trans_mat_pastk_2_past0, axis=0)).to(device) # (k, 4, 4)
                
                m_single = {}
                m_single['psm_single'] = psm_single[cav_idx]
                m_single['rm_single'] = rm_single[cav_idx]
                if self.use_dir:
                    m_single['dm_single'] = dm_single[cav_idx]
                # 1. generate one cav's box results
                box_results[cav_idx] = dataset.generate_pred_bbx_frames(m_single, pastk_trans_mat_pastk_2_past0, cav_past_k_time_diff[cav_idx*self.k:cav_idx*self.k+self.k], anchor_box)

            cav_trans_mat_pastk_2_past0.append(pastk_trans_mat_pastk_2_past0)
            
            # 2. generate box flow in one batch
            if self.viz_bbx_flag:
                box_flow_map, mask, ori_mask, matched_idx_list, compensated_results_list = self.matcher(box_results, shape_list=shape_list, viz_flag=self.viz_bbx_flag)
                ori_reserved_mask_list.append(ori_mask)
            else:
                box_flow_map, mask = self.matcher(box_results, shape_list=shape_list, viz_flag=self.viz_bbx_flag)
            box_flow_map_list.append(box_flow_map)
            reserved_mask_list.append(mask)

            if self.viz_bbx_flag:
                single_box_results = box_results
        
        final_flow_map = torch.concat(box_flow_map_list, dim=0)
        final_reserved_mask = torch.concat(reserved_mask_list, dim=0)

        if self.viz_bbx_flag:
            ori_reserved_mask = torch.concat(ori_reserved_mask_list, dim=0)
            return final_flow_map, final_reserved_mask, ori_reserved_mask, single_box_results, matched_idx_list, compensated_results_list

        return final_flow_map, final_reserved_mask

    def forward(self, data_dict, dataset=None):
        voxel_features = data_dict['processed_lidar']['voxel_features']         #(M, 32, 4)
        voxel_coords = data_dict['processed_lidar']['voxel_coords']             #(M, 4)
        voxel_num_points = data_dict['processed_lidar']['voxel_num_points']     #(M, )
        record_len = data_dict['record_len']
        record_frames = data_dict['past_k_time_interval']                       #(B, )
        pairwise_t_matrix = data_dict['pairwise_t_matrix']                      #(B, L, k, 4, 4)
        
        # debug = 0
        B, _, k, _, _ = pairwise_t_matrix.shape
        # for i in range(B):
        #     debug += record_len[i]*k

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c  ('pillar_features')
        batch_dict = self.pillar_vfe(batch_dict)
        # (n, c) -> (batch_cav_size, C, H, W) put pillars into spatial feature map ('spatial_features')
        # import ipdb; ipdb.set_trace()
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict) # 'spatial_features_2d': (batch_cav_size, 128*3, H/2, W/2)
        # N, C, H', W'. [N, 384, 100, 352]
        spatial_features_2d = batch_dict['spatial_features_2d']
        
        shape_list = torch.tensor(batch_dict['spatial_features'].shape[-3:]).to(pairwise_t_matrix.device)

        noise_pairwise_t_matrix = None
        if self.noise_flag and B==1:
            # noise_level = {'pos_std': 0.5, 'rot_std': 0, 'pos_mean': 0, 'rot_mean': 0}
            noise_pairwise_t_matrix = get_past_k_pairwise_transformation2ego(data_dict['past_lidar_pose'], self.noise_level, k=self.k, max_cav=5)[:record_len[0]].unsqueeze(0)

        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)  # (B, 256, H', W')
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        ####### debug use, viz feature of each cav
        # from matplotlib import pyplot as plt
        # viz_save_path = '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/debug_4_feature_flow/where2comm'
        # for i in range(spatial_features_2d.shape[0]):
        #     viz_content = torch.max(spatial_features_2d[i], dim=0)[0].detach().cpu()
        #     plt.imshow(viz_content)
        #     plt.savefig(viz_save_path+f'/updated_feature_{i}.png')
        ##############
        
        # dcn
        # if self.dcn:
        #     spatial_features_2d = self.dcn_net(spatial_features_2d)
        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]
        psm_single = self.cls_head(spatial_features_2d).detach()
        rm_single = self.reg_head(spatial_features_2d).detach()

        if self.use_dir:
            dm_single = self.dir_head(spatial_features_2d)

        single_detection_bbx = None

        # if self.design_mode != 4 and not self.only_tune_header_flag:
        if self.design_mode != 4:
            # generate box flow
            # [B, 256, 50, 176]
            single_output = {}
            single_output.update({'psm_single_list': self.regroup(psm_single, record_len, k), 
            'rm_single_list': self.regroup(rm_single, record_len, k)})
            if self.use_dir:
                single_output.update({
                    'dm_single_list': self.regroup(dm_single, record_len, k)
                })
            if self.viz_bbx_flag:
                box_flow_map, reserved_mask, ori_reserved_mask, single_detection_bbx, matched_idx_list, compensated_results_list = self.generate_box_flow(data_dict, single_output, dataset, shape_list, psm_single.device)
            else:
                box_flow_map, reserved_mask = self.generate_box_flow(data_dict, single_output, dataset, shape_list, psm_single.device)

        if 'flow_gt' in data_dict['label_dict']:
            flow_gt = data_dict['label_dict']['flow_gt']
            mask_gt = data_dict['label_dict']['warp_mask']
        else:
            flow_gt = None
        
        # if self.only_tune_header_flag:
        #     box_flow_map = flow_gt
        #     reserved_mask = mask_gt

        # rain attention:
        if self.multi_scale:
            if self.design_mode == 0 or self.design_mode==5:
                if self.viz_bbx_flag:
                    fused_feature, communication_rates, result_dict, single_updated_feature = self.rain_fusion(batch_dict['spatial_features'],
                    psm_single,
                    record_len,
                    pairwise_t_matrix, 
                    record_frames,
                    self.backbone,
                    [self.shrink_conv, self.cls_head, self.reg_head],
                    box_flow=box_flow_map, reserved_mask=reserved_mask,
                    flow_gt=flow_gt, viz_bbx_flag=self.viz_bbx_flag)
                else:
                    fused_feature, communication_rates, result_dict = self.rain_fusion(batch_dict['spatial_features'],
                        psm_single,
                        record_len,
                        pairwise_t_matrix, 
                        record_frames,
                        self.backbone,
                        [self.shrink_conv, self.cls_head, self.reg_head],
                        box_flow=box_flow_map, reserved_mask=reserved_mask,
                        flow_gt=flow_gt, viz_bbx_flag=self.viz_bbx_flag, noise_pairwise_t_matrix=noise_pairwise_t_matrix)
            elif self.design_mode == 4:
                fused_feature, communication_rates, result_dict = self.rain_fusion(batch_dict['spatial_features'],
                    psm_single,
                    record_len,
                    pairwise_t_matrix, 
                    record_frames,
                    self.backbone,
                    [self.shrink_conv, self.cls_head, self.reg_head])
            else: 
                fused_feature, communication_rates, result_dict, flow_recon_loss = self.rain_fusion(batch_dict['spatial_features'],
                    psm_single,
                    record_len,
                    pairwise_t_matrix, 
                    record_frames,
                    self.backbone,
                    [self.shrink_conv, self.cls_head, self.reg_head],
                    box_flow=box_flow_map, reserved_mask=reserved_mask,
                    flow_gt=flow_gt)
            # downsample feature to reduce memory
            if self.shrink_flag:
                fused_feature = self.shrink_conv(fused_feature)
                if self.single_supervise:
                    single_feature = self.shrink_conv(single_feature)
                # if self.compensation:
                    
                #         # single_feature = self.shrink_conv(single_feature)
                
                #         fused_feature_curr = self.shrink_conv(fused_feature_curr)
                #         fused_feature_latency = self.shrink_conv(fused_feature_latency)
        else:
            fused_feature, communication_rates, result_dict = self.rain_fusion(spatial_features_2d,
                                            psm_single,
                                            record_len,
                                            pairwise_t_matrix,
                                            record_frames)
            if self.compensation:
                if self.single_supervise:
                    fused_feature, single_feature, communication_rates, all_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])
                else:
                    fused_feature,fused_feature_curr,fused_feature_latency, communication_rates, all_recon_loss, all_latency_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])            

        ####### debug use, viz fused feature of where2comm
        # viz_content = torch.max(fused_feature[0], dim=0)[0].detach().cpu()
        # plt.imshow(viz_content)
        # plt.savefig(viz_save_path+'/fused_feature.png')  
        ##############
        
        # print('fused_feature: ', fused_feature.shape)
        if self.only_tune_header_flag:
            psm = self.fused_cls_head(fused_feature)
            rm = self.fused_reg_head(fused_feature)
        else: 
            psm = self.cls_head(fused_feature)
            rm = self.reg_head(fused_feature)
            # psm = self.fused_cls_head(fused_feature)
            # rm = self.fused_reg_head(fused_feature)
        
        # fuse 之后的 feature (ego)
        output_dict = {'psm': psm,
                       'rm': rm}

        if self.use_dir:
            # dm = self.dir_head(fused_feature)
            dm = self.fused_dir_head(fused_feature)
            output_dict.update({'dm': dm})

        if self.compensation:
            if self.single_supervise:
                psm_nonego_single = self.cls_head(single_feature)
                rm_nonego_single = self.reg_head(single_feature)
                output_dict.update({
                    'psm_nonego_single': psm_nonego_single,
                    'rm_nonego_single': rm_nonego_single
                })
            
            output_dict.update({
                'recon_loss': all_recon_loss, 
                'record_len': record_len
            })

        output_dict.update({'psm_single': psm_single,
                       'rm_single': rm_single,
                       'comm_rate': communication_rates
                       })

        if self.viz_bbx_flag:
            output_dict.update({
                'single_detection_bbx': single_detection_bbx, 
                'matched_idx_list': matched_idx_list, 
                'compensated_results_list': compensated_results_list
            })
            _, C, H, W = batch_dict['spatial_features'].shape
            output_dict.update({
                'single_updated_feature': single_updated_feature,
                'single_original_feature': batch_dict['spatial_features'].reshape(-1, self.k, C, H, W)[:, 0, :, :, :], 
                'single_flow_map': box_flow_map, 
                'single_reserved_mask': reserved_mask, 
                'single_original_reserved_mask': ori_reserved_mask
            })

        if self.design_mode == 1:
            output_dict.update({'flow_recon_loss': flow_recon_loss})
        
        output_dict.update(result_dict) 
        
        return output_dict

opencood/models/point_pillar_where2comm_flow_two_stage_w_dir.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Hao Xiang <haxiang@g.ucla.edu>, Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib
'''
Use pre-trained late-fusioin model, 
'''

from numpy import record
import torch.nn as nn

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
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
# from opencood.models.sub_modules.dcn_net import DCNNet
# from opencood.models.fuse_modules.where2comm import Where2comm
from opencood.models.fuse_modules.where2comm_attn import Where2comm
from opencood.models.fuse_modules.raindrop_attn_single_2_fused import raindrop_fuse
import torch

class PointPillarWhere2commSingle2Fused(nn.Module):
    def __init__(self, args):
        super(PointPillarWhere2commSingle2Fused, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        if 'resnet' in args['base_bev_backbone'] and args['base_bev_backbone']['resnet']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])
        if 'num_sweep_frames' in args:    # number of frames we use in LSTM
            self.k = args['num_sweep_frames']
        else:
            self.k = 0
        self.dcn = False

        # self.fusion_net = TransformerFusion(args['fusion_args'])
        self.fusion_net = raindrop_fuse(args['fusion_args'])
        self.multi_scale = args['rain_model']['multi_scale']

        if self.shrink_flag:
            self.cls_head = nn.Conv2d(int(dim), args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
                                    kernel_size=1)
        else:
            self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
                                    kernel_size=1)

        self.cls_head_fused = nn.Conv2d(int(dim), args['anchor_number'], kernel_size=1)
        self.reg_head_fused = nn.Conv2d(int(dim), 7 * args['anchor_number'], kernel_size=1)
        
        if args['backbone_fix']:
            self.backbone_fix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False
    
    def regroup(self, x, record_len, k=1):
        cum_sum_len = torch.cumsum(record_len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    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']
        record_len = data_dict['record_len']
        record_frames = data_dict['past_k_time_interval']                       #(B, )
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        batch_dict = {'voxel_features': voxel_features,
                        'voxel_coords': voxel_coords,
                        'voxel_num_points': voxel_num_points,
                        'record_len': record_len}
        # n, 4 -> n, c
        batch_dict = self.pillar_vfe(batch_dict)
        # (n, c) -> (batch_cav_size, C, H, W) put pillars into spatial feature map
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)
        # N, C, H', W'. [N, 384, 100, 252]
        spatial_features_2d = batch_dict['spatial_features_2d']
        
        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)  # (B, 256, H', W')
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)
        # dcn
        # if self.dcn:
        #     spatial_features_2d = self.dcn_net(spatial_features_2d)
        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]
        psm_single = self.cls_head(spatial_features_2d)
        rm_single = self.reg_head(spatial_features_2d)

        # print('spatial_features_2d: ', spatial_features_2d.shape)
        if self.multi_scale:
            fused_feature, communication_rates, result_dict = self.fusion_net(batch_dict['spatial_features'],
                                            psm_single,
                                            record_len,
                                            pairwise_t_matrix, 
                                            self.backbone,
                                            [self.shrink_conv, self.cls_head, self.reg_head])
            # downsample feature to reduce memory
            if self.shrink_flag:
                fused_feature = self.shrink_conv(fused_feature)
        else:
            fused_feature, communication_rates, result_dict = self.fusion_net(spatial_features_2d,
                                            psm_single,
                                            record_len,
                                            pairwise_t_matrix,
                                            record_frames)
        
        # print('fused_feature: ', fused_feature.shape)
        # psm = self.cls_head(fused_feature)
        # rm = self.reg_head(fused_feature)
        psm = self.cls_head_fused(fused_feature)
        rm = self.reg_head_fused(fused_feature)

        output_dict = {'psm': psm,
                        'rm': rm
                        }
        output_dict.update(result_dict)
        output_dict['comm_rate'] = communication_rates

        return output_dict

opencood/models/point_pillar_where2comm_single_2_fused.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Sizhe Wei <sizhewei@sjtu.edu.cn>
# License: TDG-Attribution-NonCommercial-NoDistrib


from numpy import record
import torch.nn as nn

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
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
# from opencood.models.sub_modules.dcn_net import DCNNet
# from opencood.models.fuse_modules.where2comm import Where2comm
from opencood.models.fuse_modules.where2comm_attn import Where2comm
from opencood.models.fuse_modules.raindrop_syncnet import raindrop_syncnet
import torch

class PointPillarWhere2commSyncnet(nn.Module):
    def __init__(self, args):
        super(PointPillarWhere2commSyncnet, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        if 'resnet' in args['base_bev_backbone'] and args['base_bev_backbone']['resnet']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        if 'num_sweep_frames' in args:    # number of frames we use in LSTM
            self.k = args['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in args:          # number of time delay
            self.tau = args['time_delay'] 
        else:
            self.tau = 0

        self.compensation = True

        self.dcn = False

        self.rain_fusion = raindrop_syncnet(args['rain_model'])

        self.multi_scale = args['rain_model']['multi_scale']
        
        if self.shrink_flag:
            self.cls_head = nn.Conv2d(128 * 2, args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(128 * 2, 7 * args['anchor_number'],
                                    kernel_size=1)
        else:
            self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
                                    kernel_size=1)

        if args['backbone_fix']:
            self.backbone_fix()

    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False
    
    def regroup(self, x, record_len, k=1):
        cum_sum_len = torch.cumsum(record_len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, data_dict):
        voxel_features = data_dict['processed_lidar']['voxel_features']         #(M, 32, 4)
        voxel_coords = data_dict['processed_lidar']['voxel_coords']             #(M, 4)
        voxel_num_points = data_dict['processed_lidar']['voxel_num_points']     #(M, )
        record_len = data_dict['record_len']
        record_frames = data_dict['past_k_time_interval']                       #(B, )
        pairwise_t_matrix = data_dict['pairwise_t_matrix']                      #(B, L, k, 4, 4)
        
        # debug = 0
        B, _, k, _, _ = pairwise_t_matrix.shape
        # for i in range(B):
        #     debug += record_len[i]*k

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c  ('pillar_features')
        batch_dict = self.pillar_vfe(batch_dict)
        # (n, c) -> (batch_cav_size, C, H, W) put pillars into spatial feature map ('spatial_features')
        # import ipdb; ipdb.set_trace()
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict) # 'spatial_features_2d': (batch_cav_size, 128*3, H/2, W/2)
        # N, C, H', W'. [N, 384, 100, 352]
        spatial_features_2d = batch_dict['spatial_features_2d']
        
        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)  # (B, 256, H', W')
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        # output only contains ego
        # [B, 256, 50, 176]
        psm_single = self.cls_head(spatial_features_2d)
        rm_single = self.reg_head(spatial_features_2d)

        # rain attention:
        if self.multi_scale:
            fused_feature, all_recon_loss, result_dict = self.rain_fusion(batch_dict['spatial_features'],
                                                    record_len,
                                                    pairwise_t_matrix, 
                                                    record_frames,
                                                    self.backbone,
                                                    [self.shrink_conv, self.cls_head, self.reg_head])
            # downsample feature to reduce memory
            if self.shrink_flag:
                fused_feature = self.shrink_conv(fused_feature)
        else:
            fused_feature, result_dict = self.fusion_net(spatial_features_2d,
                                            psm_single,
                                            record_len,
                                            pairwise_t_matrix,
                                            record_frames)
            
            
        # print('fused_feature: ', fused_feature.shape)
        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)
        
        # fuse 之后的 feature (ego)
        output_dict = {'psm': psm,
                       'rm': rm}

        if self.compensation:
            output_dict.update({
                'recon_loss': all_recon_loss, 
                'record_len': record_len
            })

        output_dict.update(result_dict)  # TODO: 这个现在是空字典  
        
        return output_dict

opencood/models/point_pillar_where2comm_syncnet.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Sizhe Wei <sizhewei@sjtu.edu.cn>
# License: TDG-Attribution-NonCommercial-NoDistrib


from numpy import record
import torch.nn as nn

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
from opencood.models.sub_modules.base_bev_backbone_resnet import ResNetBEVBackbone
from opencood.models.sub_modules.downsample_conv import DownsampleConv
from opencood.models.sub_modules.naive_compress import NaiveCompressor
# from opencood.models.sub_modules.dcn_net import DCNNet
# from opencood.models.fuse_modules.where2comm import Where2comm
from opencood.models.fuse_modules.where2comm_attn import Where2comm
from opencood.models.fuse_modules.raindrop_swin import raindrop_swin
from opencood.models.fuse_modules.raindrop_swin_w_single import raindrop_swin_w_single
from opencood.models.fuse_modules.raindrop_w_syncnet import raindrop_fuse
import torch

import numpy as np
from opencood.utils.transformation_utils import tfm_to_pose, x1_to_x2, x_to_world

def generate_noise(pos_std, rot_std, pos_mean=0, rot_mean=0):
    """ Add localization error to the 6dof pose
        Noise includes position (x,y) and rotation (yaw).
        We use gaussian distribution to generate noise.

    Args:

        pos_std : float 
            std of gaussian dist, in meter

        rot_std : float
            std of gaussian dist, in degree

        pos_mean : float
            mean of gaussian dist, in meter

        rot_mean : float
            mean of gaussian dist, in degree

    Returns:
        pose_noise: np.ndarray, [6,]
            [x, y, z, roll, yaw, pitch]
    """

    xy = np.random.normal(pos_mean, pos_std, size=(2))
    yaw = np.random.normal(rot_mean, rot_std, size=(1))

    pose_noise = np.array([xy[0], xy[1], 0, 0, yaw[0], 0])

    return pose_noise

def get_past_k_pairwise_transformation2ego(past_k_lidar_pose, noise_level, k=3, max_cav=5):
    """
    Get transformation matrixes accross different agents to curr ego at all past timestamps.

    Parameters
    ----------
    base_data_dict : dict
        Key : cav id, item: transformation matrix to ego, lidar points.
    
    ego_pose : list
        ego pose

    max_cav : int
        The maximum number of cav, default 5

    Return
    ------
    pairwise_t_matrix : np.array
        The transformation matrix each cav to curr ego at past k frames.
        shape: (L, k, 4, 4), L is the max cav number in a scene, k is the num of past frames
        pairwise_t_matrix[i, j] is T i_to_ego at past_j frame
    """
    pos_std = noise_level['pos_std']
    rot_std = noise_level['rot_std']
    pos_mean = 0 
    rot_mean = 0
    
    pairwise_t_matrix = np.tile(np.eye(4), (max_cav, k, 1, 1)) # (L, k, 4, 4)

    ego_pose = past_k_lidar_pose[0, 0]

    t_list = []

    # save all transformation matrix in a list in order first.
    for cav_id in range(past_k_lidar_pose.shape[0]):
        past_k_poses = []
        for time_id in range(k):
            loc_noise = generate_noise(pos_std, rot_std)
            past_k_poses.append(x_to_world(past_k_lidar_pose[cav_id, time_id].cpu().numpy()+loc_noise))
        t_list.append(past_k_poses) # Twx
    
    ego_pose = x_to_world(ego_pose.cpu().numpy())
    for i in range(len(t_list)): # different cav
        if i!=0 :
            for j in range(len(t_list[i])): # different time
                # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                t_matrix = np.linalg.solve(t_list[i][j], ego_pose)  # Tjw*Twi = Tji
                pairwise_t_matrix[i, j] = t_matrix
    pairwise_t_matrix = torch.tensor(pairwise_t_matrix).to(past_k_lidar_pose.device)
    return pairwise_t_matrix


class PointPillarWhere2commSyncnetNew(nn.Module):
    def __init__(self, args):
        super(PointPillarWhere2commSyncnetNew, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        if 'resnet' in args['base_bev_backbone'] and args['base_bev_backbone']['resnet']:
            self.backbone = ResNetBEVBackbone(args['base_bev_backbone'], 64)
        else:
            self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        # used to downsample the feature map for efficient computation
        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
        self.compression = False

        if args['compression'] > 0:
            self.compression = True
            self.naive_compressor = NaiveCompressor(256, args['compression'])

        if 'num_sweep_frames' in args:    # number of frames we use in LSTM
            self.k = args['num_sweep_frames']
        else:
            self.k = 0

        if 'time_delay' in args:          # number of time delay
            self.tau = args['time_delay'] 
        else:
            self.tau = 0

        self.noise_flag = 0
        if 'noise' in args.keys():
            self.noise_flag = True
            self.noise_level = {'pos_std': args['noise']['pos_std'], 'rot_std': args['noise']['rot_std'], 'pos_mean': args['noise']['pos_mean'], 'rot_mean': args['noise']['rot_mean']}
            print(f'=== noise level : {self.noise_level} ===')

        self.dcn = False
        # if 'dcn' in args:
        #     self.dcn = True
        #     self.dcn_net = DCNNet(args['dcn'])

        # 用于 single 部分的监督
        self.single_supervise = False
        if 'with_compensation' in args and args['with_compensation']:
            self.compensation = True
            if 'with_single_supervise' in args and args['with_single_supervise']:
                self.rain_fusion = raindrop_swin_w_single(args['rain_model'])
                self.single_supervise = True
            else:
                self.rain_fusion = raindrop_swin(args['rain_model'])
        else: 
            self.compensation = False
            self.rain_fusion = raindrop_fuse(args['rain_model'])

        self.multi_scale = args['rain_model']['multi_scale']
        
        if self.shrink_flag:
            dim = args['shrink_header']['dim'][0]
            self.cls_head = nn.Conv2d(int(dim), args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(int(dim), 7 * args['anchor_number'],
                                    kernel_size=1)
        else:
            self.cls_head = nn.Conv2d(128 * 3, args['anchor_number'],
                                    kernel_size=1)
            self.reg_head = nn.Conv2d(128 * 3, 7 * args['anchor_number'],
                                    kernel_size=1)

        if 'backbone_fix' in args.keys() and args['backbone_fix']:
            self.backbone_fix()
            print('=== backbone fixed ===')

    # def where2comm_fix(self):
    #     """
    #     Fix the parameters of backbone during finetune on timedelay
    #     """
    #     for p in self.pillar_vfe.parameters():
    #         p.requires_grad = False

    #     for p in self.scatter.parameters():
    #         p.requires_grad = False

    #     for p in self.backbone.parameters():
    #         p.requires_grad = False

    #     if self.compression:
    #         for p in self.naive_compressor.parameters():
    #             p.requires_grad = False
    #     if self.shrink_flag:
    #         for p in self.shrink_conv.parameters():
    #             p.requires_grad = False

    #     for p in self.cls_head.parameters():
    #         p.requires_grad = False
    #     for p in self.reg_head.parameters():
    #         p.requires_grad = False
    
    def backbone_fix(self):
        """
        Fix the parameters of backbone during finetune on timedelay
        """
        for p in self.pillar_vfe.parameters():
            p.requires_grad = False

        for p in self.scatter.parameters():
            p.requires_grad = False

        for p in self.backbone.parameters():
            p.requires_grad = False

        if self.compression:
            for p in self.naive_compressor.parameters():
                p.requires_grad = False
        if self.shrink_flag:
            for p in self.shrink_conv.parameters():
                p.requires_grad = False

        for p in self.cls_head.parameters():
            p.requires_grad = False
        for p in self.reg_head.parameters():
            p.requires_grad = False
    
    def regroup(self, x, record_len, k=1):
        cum_sum_len = torch.cumsum(record_len*k, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, data_dict):
        voxel_features = data_dict['processed_lidar']['voxel_features']         #(M, 32, 4)
        voxel_coords = data_dict['processed_lidar']['voxel_coords']             #(M, 4)
        voxel_num_points = data_dict['processed_lidar']['voxel_num_points']     #(M, )
        record_len = data_dict['record_len']
        record_frames = data_dict['past_k_time_interval']                       #(B, )
        pairwise_t_matrix = data_dict['pairwise_t_matrix']                      #(B, L, k, 4, 4)
        
        # debug = 0
        B, _, k, _, _ = pairwise_t_matrix.shape
        # for i in range(B):
        #     debug += record_len[i]*k

        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'record_len': record_len}
        # n, 4 -> n, c  ('pillar_features')
        batch_dict = self.pillar_vfe(batch_dict)
        # (n, c) -> (batch_cav_size, C, H, W) put pillars into spatial feature map ('spatial_features')
        # import ipdb; ipdb.set_trace()
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict) # 'spatial_features_2d': (batch_cav_size, 128*3, H/2, W/2)
        # N, C, H', W'. [N, 384, 100, 352]
        spatial_features_2d = batch_dict['spatial_features_2d']

        noise_pairwise_t_matrix = None
        if self.noise_flag and B==1:
            # noise_level = {'pos_std': 0.5, 'rot_std': 0, 'pos_mean': 0, 'rot_mean': 0}
            noise_pairwise_t_matrix = get_past_k_pairwise_transformation2ego(data_dict['past_lidar_pose'], self.noise_level, k=self.k, max_cav=5)[:record_len[0]].unsqueeze(0)
            pairwise_t_matrix = noise_pairwise_t_matrix
        
        # downsample feature to reduce memory
        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)  # (B, 256, H', W')
        # compressor
        if self.compression:
            spatial_features_2d = self.naive_compressor(spatial_features_2d)

        ####### debug use, viz feature of each cav
        # from matplotlib import pyplot as plt
        # viz_save_path = '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/debug_4_feature_flow/where2comm'
        # for i in range(spatial_features_2d.shape[0]):
        #     viz_content = torch.max(spatial_features_2d[i], dim=0)[0].detach().cpu()
        #     plt.imshow(viz_content)
        #     plt.savefig(viz_save_path+f'/updated_feature_{i}.png')
        ##############
        
        # dcn
        # if self.dcn:
        #     spatial_features_2d = self.dcn_net(spatial_features_2d)
        # spatial_features_2d is [sum(cav_num), 256, 50, 176]
        # output only contains ego
        # [B, 256, 50, 176]
        psm_single = self.cls_head(spatial_features_2d)
        rm_single = self.reg_head(spatial_features_2d)

        # for debug use
        # flow_gt = data_dict['label_dict']['flow_gt']

        # rain attention:
        if self.multi_scale:
            fused_feature, communication_rates, result_dict, all_recon_loss = self.rain_fusion(batch_dict['spatial_features'],
                psm_single,
                record_len,
                pairwise_t_matrix, 
                record_frames,
                self.backbone,
                [self.shrink_conv, self.cls_head, self.reg_head])
            # downsample feature to reduce memory
            if self.shrink_flag:
                fused_feature = self.shrink_conv(fused_feature)
                if self.single_supervise:
                    single_feature = self.shrink_conv(single_feature)
                # if self.compensation:
                    
                #         # single_feature = self.shrink_conv(single_feature)
                
                #         fused_feature_curr = self.shrink_conv(fused_feature_curr)
                #         fused_feature_latency = self.shrink_conv(fused_feature_latency)
        else:
            fused_feature, communication_rates, result_dict = self.rain_fusion(spatial_features_2d,
                                            psm_single,
                                            record_len,
                                            pairwise_t_matrix,
                                            record_frames)
            if self.compensation:
                if self.single_supervise:
                    fused_feature, single_feature, communication_rates, all_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])
                else:
                    fused_feature,fused_feature_curr,fused_feature_latency, communication_rates, all_recon_loss, all_latency_recon_loss, result_dict = self.rain_fusion(spatial_features_2d,
                                                psm_single,
                                                record_len,
                                                pairwise_t_matrix, 
                                                record_frames,
                                                self.backbone,
                                                [self.shrink_conv, self.cls_head, self.reg_head])            

        ####### debug use, viz fused feature of where2comm
        # viz_content = torch.max(fused_feature[0], dim=0)[0].detach().cpu()
        # plt.imshow(viz_content)
        # plt.savefig(viz_save_path+'/fused_feature.png')  
        ##############
        
        # print('fused_feature: ', fused_feature.shape)
        psm = self.cls_head(fused_feature)
        rm = self.reg_head(fused_feature)
        
        # fuse 之后的 feature (ego)
        output_dict = {'psm': psm,
                       'rm': rm}

        output_dict.update({'recon_loss': all_recon_loss})

        # output_dict.update({'flow_preds': flow,
        #                     'state_preds': state_preds})

        if self.compensation:
            if self.single_supervise:
                psm_nonego_single = self.cls_head(single_feature)
                rm_nonego_single = self.reg_head(single_feature)
                output_dict.update({
                    'psm_nonego_single': psm_nonego_single,
                    'rm_nonego_single': rm_nonego_single
                })
            
            output_dict.update({
                'recon_loss': all_recon_loss, 
                'record_len': record_len
            })

        output_dict.update({'psm_single': psm_single,
                       'rm_single': rm_single,
                       'comm_rate': communication_rates
                       })
        
        output_dict.update(result_dict) 
        
        return output_dict

opencood/models/point_pillar_where2comm_syncnet_new.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, OpenPCDet
# License: TDG-Attribution-NonCommercial-NoDistrib


import torch
import torch.nn as nn


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
from opencood.models.sub_modules.downsample_conv import DownsampleConv

class PointPillar(nn.Module):
    def __init__(self, args):
        super(PointPillar, self).__init__()

        # PIllar VFE
        self.pillar_vfe = PillarVFE(args['pillar_vfe'],
                                    num_point_features=4,
                                    voxel_size=args['voxel_size'],
                                    point_cloud_range=args['lidar_range'])
        self.scatter = PointPillarScatter(args['point_pillar_scatter'])
        self.backbone = BaseBEVBackbone(args['base_bev_backbone'], 64)

        self.out_channel = sum(args['base_bev_backbone']['num_upsample_filter'])

        self.shrink_flag = False
        if 'shrink_header' in args:
            self.shrink_flag = True
            self.shrink_conv = DownsampleConv(args['shrink_header'])
            self.out_channel = args['shrink_header']['dim'][-1]

        self.cls_head = nn.Conv2d(self.out_channel, args['anchor_number'], # 384
                                  kernel_size=1)
        self.reg_head = nn.Conv2d(self.out_channel, 7 * args['anchor_number'], # 384
                                  kernel_size=1)
        
        if 'dir_args' in args.keys():
            self.use_dir = True
            self.dir_head = nn.Conv2d(self.out_channel, args['dir_args']['num_bins'] * args['anchor_number'],
                                  kernel_size=1) # BIN_NUM = 2， # 384
        else:
            self.use_dir = False

    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_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points}

        batch_dict = self.pillar_vfe(batch_dict)
        batch_dict = self.scatter(batch_dict)
        batch_dict = self.backbone(batch_dict)

        spatial_features_2d = batch_dict['spatial_features_2d']

        if self.shrink_flag:
            spatial_features_2d = self.shrink_conv(spatial_features_2d)

        psm = self.cls_head(spatial_features_2d)
        rm = self.reg_head(spatial_features_2d)

        output_dict = {'psm': psm,
                       'rm': rm}
                       
        if self.use_dir:
            dm = self.dir_head(spatial_features_2d)
            output_dict.update({'dm': dm})

        return output_dict
opencood/models/point_pillar_ylu.py
MediaBrain-SJTU/CoBEVFlow
# Copyright (c) 2020 Mitsubishi Electric Research Laboratories (MERL). All rights reserved. The software, documentation and/or data in this file is provided on an "as is" basis, and MERL has no obligations to provide maintenance, support, updates, enhancements or modifications. MERL specifically disclaims any warranties, including, but not limited to, the implied warranties of merchantability and fitness for any particular purpose. In no event shall MERL be liable to any party for direct, indirect, special, incidental, or consequential damages, including lost profits, arising out of the use of this software and its documentation, even if MERL has been advised of the possibility of such damages. As more fully described in the license agreement that was required in order to download this software, documentation and/or data, permission to use, copy and modify this software without fee is granted, but only for educational, research and non-commercial purposes.

import torch.nn.functional as F
import torch.nn as nn
import torch


class CellClassification(nn.Module):
    def __init__(self, category_num=5):
        super(CellClassification, self).__init__()
        self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, category_num, kernel_size=1, stride=1, padding=0)

        self.bn1 = nn.BatchNorm2d(32)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.conv2(x)

        return x


class StateEstimation(nn.Module):
    def __init__(self, motion_category_num=2):
        super(StateEstimation, self).__init__()
        self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, motion_category_num, kernel_size=1, stride=1, padding=0)

        self.bn1 = nn.BatchNorm2d(32)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.conv2(x)

        return x


class MotionPrediction(nn.Module):
    def __init__(self, seq_len):
        super(MotionPrediction, self).__init__()
        self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 2 * seq_len, kernel_size=1, stride=1, padding=0)

        self.bn1 = nn.BatchNorm2d(32)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.conv2(x)

        return x


class FlowUncPrediction(nn.Module):
    def __init__(self, seq_len):
        super(FlowUncPrediction, self).__init__()
        self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(32, 2 * seq_len, kernel_size=1, stride=1, padding=0)

        self.bn1 = nn.BatchNorm2d(32)

    def forward(self, x):
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.conv2(x)

        return x


class Conv3D(nn.Module):
    def __init__(self, in_channel, out_channel, kernel_size, stride, padding):
        super(Conv3D, self).__init__()
        self.conv3d = nn.Conv3d(in_channel, out_channel, kernel_size=kernel_size, stride=stride, padding=padding)
        self.bn3d = nn.BatchNorm3d(out_channel)

    def forward(self, x):
        # input x: (batch, seq, c, h, w)
        x = x.permute(0, 2, 1, 3, 4).contiguous()  # (batch, c, seq_len, h, w)
        x = F.relu(self.bn3d(self.conv3d(x)))
        x = x.permute(0, 2, 1, 3, 4).contiguous()  # (batch, seq_len, c, h, w)

        return x


class STPN(nn.Module):
    def __init__(self, height_feat_size=13):
        super(STPN, self).__init__()
        self.conv_pre_1 = nn.Conv2d(height_feat_size, 32, kernel_size=3, stride=1, padding=1)
        self.conv_pre_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
        self.bn_pre_1 = nn.BatchNorm2d(32)
        self.bn_pre_2 = nn.BatchNorm2d(32)

        self.conv3d_1 = Conv3D(64, 64, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
        self.conv3d_2 = Conv3D(128, 128, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))

        self.conv1_1 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
        self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)

        self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
        self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)

        self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
        self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)

        self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
        self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)

        self.conv5_1 = nn.Conv2d(512 + 256, 256, kernel_size=3, stride=1, padding=1)
        self.conv5_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)

        self.conv6_1 = nn.Conv2d(256 + 128, 128, kernel_size=3, stride=1, padding=1)
        self.conv6_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)

        self.conv7_1 = nn.Conv2d(128 + 64, 64, kernel_size=3, stride=1, padding=1)
        self.conv7_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)

        self.conv8_1 = nn.Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1)
        self.conv8_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)

        self.bn1_1 = nn.BatchNorm2d(64)
        self.bn1_2 = nn.BatchNorm2d(64)

        self.bn2_1 = nn.BatchNorm2d(128)
        self.bn2_2 = nn.BatchNorm2d(128)

        self.bn3_1 = nn.BatchNorm2d(256)
        self.bn3_2 = nn.BatchNorm2d(256)

        self.bn4_1 = nn.BatchNorm2d(512)
        self.bn4_2 = nn.BatchNorm2d(512)

        self.bn5_1 = nn.BatchNorm2d(256)
        self.bn5_2 = nn.BatchNorm2d(256)

        self.bn6_1 = nn.BatchNorm2d(128)
        self.bn6_2 = nn.BatchNorm2d(128)

        self.bn7_1 = nn.BatchNorm2d(64)
        self.bn7_2 = nn.BatchNorm2d(64)

        self.bn8_1 = nn.BatchNorm2d(32)
        self.bn8_2 = nn.BatchNorm2d(32)

    def forward(self, x):
        batch, seq, z, h, w = x.size()

        x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
        x = F.relu(self.bn_pre_1(self.conv_pre_1(x)))
        x = F.relu(self.bn_pre_2(self.conv_pre_2(x)))

        # -------------------------------- Encoder Path --------------------------------
        # -- STC block 1
        x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
        x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))

        x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3)).contiguous()  # (batch, seq, c, h, w)
        x_1 = self.conv3d_1(x_1)
        x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()  # (batch * seq, c, h, w)

        # -- STC block 2
        x_2 = F.relu(self.bn2_1(self.conv2_1(x_1)))
        x_2 = F.relu(self.bn2_2(self.conv2_2(x_2)))

        x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3)).contiguous()  # (batch, seq, c, h, w)
        x_2 = self.conv3d_2(x_2)
        x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()  # (batch * seq, c, h, w), seq = 1

        # -- STC block 3
        x_3 = F.relu(self.bn3_1(self.conv3_1(x_2)))
        x_3 = F.relu(self.bn3_2(self.conv3_2(x_3)))

        # -- STC block 4
        x_4 = F.relu(self.bn4_1(self.conv4_1(x_3)))
        x_4 = F.relu(self.bn4_2(self.conv4_2(x_4)))

        # -------------------------------- Decoder Path --------------------------------
        x_5 = F.relu(self.bn5_1(self.conv5_1(torch.cat((F.interpolate(x_4, scale_factor=(2, 2)), x_3), dim=1))))
        x_5 = F.relu(self.bn5_2(self.conv5_2(x_5)))

        x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3))
        x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
        x_2 = F.adaptive_max_pool3d(x_2, (1, None, None))
        x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
        x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()

        x_6 = F.relu(self.bn6_1(self.conv6_1(torch.cat((F.interpolate(x_5, scale_factor=(2, 2)), x_2), dim=1))))
        x_6 = F.relu(self.bn6_2(self.conv6_2(x_6)))

        x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3))
        x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
        x_1 = F.adaptive_max_pool3d(x_1, (1, None, None))
        x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
        x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()

        x_7 = F.relu(self.bn7_1(self.conv7_1(torch.cat((F.interpolate(x_6, scale_factor=(2, 2)), x_1), dim=1))))
        x_7 = F.relu(self.bn7_2(self.conv7_2(x_7)))

        x = x.view(batch, -1, x.size(1), x.size(2), x.size(3))
        x = x.permute(0, 2, 1, 3, 4).contiguous()
        x = F.adaptive_max_pool3d(x, (1, None, None))
        x = x.permute(0, 2, 1, 3, 4).contiguous()
        x = x.view(-1, x.size(2), x.size(3), x.size(4)).contiguous()

        x_8 = F.relu(self.bn8_1(self.conv8_1(torch.cat((F.interpolate(x_7, scale_factor=(2, 2)), x), dim=1))))
        res_x = F.relu(self.bn8_2(self.conv8_2(x_8)))

        return res_x


class MotionNet(nn.Module):
    def __init__(self, out_seq_len=20, motion_category_num=2, height_feat_size=13):
        super(MotionNet, self).__init__()
        self.out_seq_len = out_seq_len

        self.cell_classify = CellClassification()
        self.motion_pred = MotionPrediction(seq_len=self.out_seq_len)
        self.state_classify = StateEstimation(motion_category_num=motion_category_num)
        self.stpn = STPN(height_feat_size=height_feat_size)


    def forward(self, bevs):
        bevs = bevs.permute(0, 1, 4, 2, 3)  # (Batch, seq, z, h, w)

        # Backbone network
        x = self.stpn(bevs)

        # Cell Classification head
        cell_class_pred = self.cell_classify(x)

        # Motion State Classification head
        state_class_pred = self.state_classify(x)

        # Motion Displacement prediction
        disp = self.motion_pred(x)
        disp = disp.view(-1, 2, x.size(-2), x.size(-1))

        return disp, cell_class_pred, state_class_pred


# For MGDA loss computation
class FeatEncoder(nn.Module):
    def __init__(self, height_feat_size=13):
        super(FeatEncoder, self).__init__()
        self.stpn = STPN(height_feat_size=height_feat_size)

    def forward(self, bevs):
        bevs = bevs.permute(0, 1, 4, 2, 3)  # (Batch, seq, z, h, w)
        x = self.stpn(bevs)

        return x


class MotionNetMGDA(nn.Module):
    def __init__(self, out_seq_len=20, motion_category_num=2):
        super(MotionNetMGDA, self).__init__()
        self.out_seq_len = out_seq_len

        self.cell_classify = CellClassification()
        self.motion_pred = MotionPrediction(seq_len=self.out_seq_len)
        self.state_classify = StateEstimation(motion_category_num=motion_category_num)

    def forward(self, stpn_out):
        # Cell Classification head
        cell_class_pred = self.cell_classify(stpn_out)

        # Motion State Classification head
        state_class_pred = self.state_classify(stpn_out)

        # Motion Displacement prediction
        disp = self.motion_pred(stpn_out)
        disp = disp.view(-1, 2, stpn_out.size(-2), stpn_out.size(-1))

        return disp, cell_class_pred, state_class_pred

# # Copyright (c) 2020 Mitsubishi Electric Research Laboratories (MERL). All rights reserved. The software, documentation and/or data in this file is provided on an "as is" basis, and MERL has no obligations to provide maintenance, support, updates, enhancements or modifications. MERL specifically disclaims any warranties, including, but not limited to, the implied warranties of merchantability and fitness for any particular purpose. In no event shall MERL be liable to any party for direct, indirect, special, incidental, or consequential damages, including lost profits, arising out of the use of this software and its documentation, even if MERL has been advised of the possibility of such damages. As more fully described in the license agreement that was required in order to download this software, documentation and/or data, permission to use, copy and modify this software without fee is granted, but only for educational, research and non-commercial purposes.

# import torch.nn.functional as F
# import torch.nn as nn
# import torch


# class CellClassification(nn.Module):
#     def __init__(self, category_num=5):
#         super(CellClassification, self).__init__()
#         self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
#         self.conv2 = nn.Conv2d(32, category_num, kernel_size=1, stride=1, padding=0)

#         self.bn1 = nn.BatchNorm2d(32)

#     def forward(self, x):
#         x = F.relu(self.bn1(self.conv1(x)))
#         x = self.conv2(x)

#         return x


# class StateEstimation(nn.Module):
#     def __init__(self, motion_category_num=2):
#         super(StateEstimation, self).__init__()
#         self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
#         self.conv2 = nn.Conv2d(32, motion_category_num, kernel_size=1, stride=1, padding=0)

#         self.bn1 = nn.BatchNorm2d(32)

#     def forward(self, x):
#         x = F.relu(self.bn1(self.conv1(x)))
#         x = self.conv2(x)

#         return x


# class MotionPrediction(nn.Module):
#     def __init__(self, seq_len):
#         super(MotionPrediction, self).__init__()
#         self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
#         self.conv2 = nn.Conv2d(32, 2 * seq_len, kernel_size=1, stride=1, padding=0)

#         self.bn1 = nn.BatchNorm2d(32)

#     def forward(self, x):
#         x = F.relu(self.bn1(self.conv1(x)))
#         x = self.conv2(x)

#         return x


# class FlowUncPrediction(nn.Module):
#     def __init__(self, seq_len):
#         super(FlowUncPrediction, self).__init__()
#         self.conv1 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
#         self.conv2 = nn.Conv2d(32, 2 * seq_len, kernel_size=1, stride=1, padding=0)

#         self.bn1 = nn.BatchNorm2d(32)

#     def forward(self, x):
#         x = F.relu(self.bn1(self.conv1(x)))
#         x = self.conv2(x)

#         return x


# class Conv3D(nn.Module):
#     def __init__(self, in_channel, out_channel, kernel_size, stride, padding):
#         super(Conv3D, self).__init__()
#         self.conv3d = nn.Conv3d(in_channel, out_channel, kernel_size=kernel_size, stride=stride, padding=padding)
#         self.bn3d = nn.BatchNorm3d(out_channel)

#     def forward(self, x):
#         # input x: (batch, seq, c, h, w)
#         x = x.permute(0, 2, 1, 3, 4).contiguous()  # (batch, c, seq_len, h, w)
#         x = F.relu(self.bn3d(self.conv3d(x)))
#         x = x.permute(0, 2, 1, 3, 4).contiguous()  # (batch, seq_len, c, h, w)

#         return x


# class STPN(nn.Module):
#     def __init__(self, height_feat_size=13):
#         super(STPN, self).__init__()
#         self.conv_pre_1 = nn.Conv2d(height_feat_size, 32, kernel_size=3, stride=1, padding=1)
#         self.conv_pre_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)
#         self.bn_pre_1 = nn.BatchNorm2d(32)
#         self.bn_pre_2 = nn.BatchNorm2d(32)

#         self.conv3d_1 = Conv3D(64, 64, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
#         self.conv3d_2 = Conv3D(128, 128, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))

#         self.conv1_1 = nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1)
#         self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)

#         self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1)
#         self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)

#         self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1)
#         self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)

#         self.conv4_1 = nn.Conv2d(256 + 128, 128, kernel_size=3, stride=1, padding=1)
#         self.conv4_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)

#         self.conv5_1 = nn.Conv2d(128 + 64, 64, kernel_size=3, stride=1, padding=1)
#         self.conv5_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)

#         self.conv6_1 = nn.Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1)
#         self.conv6_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)

#         # self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1)
#         # self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1)

#         # self.conv5_1 = nn.Conv2d(512 + 256, 256, kernel_size=3, stride=1, padding=1)
#         # self.conv5_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1)

#         # self.conv6_1 = nn.Conv2d(256 + 128, 128, kernel_size=3, stride=1, padding=1)
#         # self.conv6_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1)

#         # self.conv7_1 = nn.Conv2d(128 + 64, 64, kernel_size=3, stride=1, padding=1)
#         # self.conv7_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1)

#         # self.conv8_1 = nn.Conv2d(64 + 32, 32, kernel_size=3, stride=1, padding=1)
#         # self.conv8_2 = nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1)

#         self.bn1_1 = nn.BatchNorm2d(64)
#         self.bn1_2 = nn.BatchNorm2d(64)

#         self.bn2_1 = nn.BatchNorm2d(128)
#         self.bn2_2 = nn.BatchNorm2d(128)

#         self.bn3_1 = nn.BatchNorm2d(256)
#         self.bn3_2 = nn.BatchNorm2d(256)

#         self.bn4_1 = nn.BatchNorm2d(128)
#         self.bn4_2 = nn.BatchNorm2d(128)

#         self.bn5_1 = nn.BatchNorm2d(64)
#         self.bn5_2 = nn.BatchNorm2d(64)

#         self.bn6_1 = nn.BatchNorm2d(32)
#         self.bn6_2 = nn.BatchNorm2d(32)
        
#         # self.bn4_1 = nn.BatchNorm2d(512)
#         # self.bn4_2 = nn.BatchNorm2d(512)

#         # self.bn5_1 = nn.BatchNorm2d(256)
#         # self.bn5_2 = nn.BatchNorm2d(256)

#         # self.bn6_1 = nn.BatchNorm2d(128)
#         # self.bn6_2 = nn.BatchNorm2d(128)

#         # self.bn7_1 = nn.BatchNorm2d(64)
#         # self.bn7_2 = nn.BatchNorm2d(64)

#         # self.bn8_1 = nn.BatchNorm2d(32)
#         # self.bn8_2 = nn.BatchNorm2d(32)

#         self.conv_suf_1 = nn.Conv2d(32, height_feat_size, kernel_size=3, stride=1, padding=1)
#         self.conv_suf_2 = nn.Conv2d(height_feat_size, height_feat_size, kernel_size=3, stride=1, padding=1)
#         self.bn_suf_1 = nn.BatchNorm2d(height_feat_size)
#         self.bn_suf_2 = nn.BatchNorm2d(height_feat_size)

#     def forward(self, x):
#         batch, seq, z, h, w = x.size()

#         x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
#         x = F.relu(self.bn_pre_1(self.conv_pre_1(x)))
#         x = F.relu(self.bn_pre_2(self.conv_pre_2(x)))

#         # -------------------------------- Encoder Path --------------------------------
#         # -- STC block 1
#         x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
#         x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))

#         x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3)).contiguous()  # (batch, seq, c, h, w)
#         x_1 = self.conv3d_1(x_1)
#         x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()  # (batch * seq, c, h, w)

#         # -- STC block 2
#         x_2 = F.relu(self.bn2_1(self.conv2_1(x_1)))
#         x_2 = F.relu(self.bn2_2(self.conv2_2(x_2)))

#         # TODO: comment this block when seq<=3
#         x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3)).contiguous()  # (batch, seq, c, h, w)
#         x_2 = self.conv3d_2(x_2)
#         x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()  # (batch * seq, c, h, w), seq = 1
#         ##############################

#         # -- STC block 3
#         x_3 = F.relu(self.bn3_1(self.conv3_1(x_2)))
#         x_3 = F.relu(self.bn3_2(self.conv3_2(x_3)))

#         # -------------------------------- Decoder Path --------------------------------
#         x_4 = F.relu(self.bn4_1(self.conv4_1(torch.cat((F.interpolate(x_3, scale_factor=(2, 2)), x_2), dim=1))))
#         x_4 = F.relu(self.bn4_2(self.conv4_2(x_4)))

#         x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3))
#         x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
#         x_1 = F.adaptive_max_pool3d(x_1, (1, None, None))
#         x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
#         x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()

#         x_5 = F.relu(self.bn5_1(self.conv5_1(torch.cat((F.interpolate(x_4, scale_factor=(2, 2)), x_1), dim=1))))
#         x_5 = F.relu(self.bn5_2(self.conv5_2(x_5)))

#         x = x.view(batch, -1, x.size(1), x.size(2), x.size(3))
#         x = x.permute(0, 2, 1, 3, 4).contiguous()
#         x = F.adaptive_max_pool3d(x, (1, None, None))
#         x = x.permute(0, 2, 1, 3, 4).contiguous()
#         x = x.view(-1, x.size(2), x.size(3), x.size(4)).contiguous()

#         x_6 = F.relu(self.bn6_1(self.conv6_1(torch.cat((F.interpolate(x_5, scale_factor=(2, 2)), x), dim=1))))
#         x = F.relu(self.bn6_2(self.conv6_2(x_6)))

#         # x = F.relu(self.bn_suf_1(self.conv_suf_1(x)))
#         # x = F.relu(self.bn_suf_2(self.conv_suf_2(x)))

#         '''
#         # # -- STC block 4
#         # x_4 = F.relu(self.bn4_1(self.conv4_1(x_3)))
#         # x_4 = F.relu(self.bn4_2(self.conv4_2(x_4)))

#         # -------------------------------- Decoder Path --------------------------------
#         x_5 = F.relu(self.bn5_1(self.conv5_1(torch.cat((F.interpolate(x_4, scale_factor=(2, 2)), x_3), dim=1))))
#         x_5 = F.relu(self.bn5_2(self.conv5_2(x_5)))

#         x_2 = x_2.view(batch, -1, x_2.size(1), x_2.size(2), x_2.size(3))
#         x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
#         x_2 = F.adaptive_max_pool3d(x_2, (1, None, None))
#         x_2 = x_2.permute(0, 2, 1, 3, 4).contiguous()
#         x_2 = x_2.view(-1, x_2.size(2), x_2.size(3), x_2.size(4)).contiguous()

#         x_6 = F.relu(self.bn6_1(self.conv6_1(torch.cat((F.interpolate(x_5, scale_factor=(2, 2)), x_2), dim=1))))
#         x_6 = F.relu(self.bn6_2(self.conv6_2(x_6)))

#         x_1 = x_1.view(batch, -1, x_1.size(1), x_1.size(2), x_1.size(3))
#         x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
#         x_1 = F.adaptive_max_pool3d(x_1, (1, None, None))
#         x_1 = x_1.permute(0, 2, 1, 3, 4).contiguous()
#         x_1 = x_1.view(-1, x_1.size(2), x_1.size(3), x_1.size(4)).contiguous()

#         x_7 = F.relu(self.bn7_1(self.conv7_1(torch.cat((F.interpolate(x_6, scale_factor=(2, 2)), x_1), dim=1))))
#         x_7 = F.relu(self.bn7_2(self.conv7_2(x_7)))

#         x = x.view(batch, -1, x.size(1), x.size(2), x.size(3))
#         x = x.permute(0, 2, 1, 3, 4).contiguous()
#         x = F.adaptive_max_pool3d(x, (1, None, None))
#         x = x.permute(0, 2, 1, 3, 4).contiguous()
#         x = x.view(-1, x.size(2), x.size(3), x.size(4)).contiguous()

#         x_8 = F.relu(self.bn8_1(self.conv8_1(torch.cat((F.interpolate(x_7, scale_factor=(2, 2)), x), dim=1))))
#         res_x = F.relu(self.bn8_2(self.conv8_2(x_8)))
#         '''

#         return x


# class MotionNet(nn.Module):
#     def __init__(self, out_seq_len=20, motion_category_num=2, height_feat_size=13):
#         super(MotionNet, self).__init__()
#         self.out_seq_len = out_seq_len

#         self.cell_classify = CellClassification()
#         self.motion_pred = MotionPrediction(seq_len=self.out_seq_len)
#         self.state_classify = StateEstimation(motion_category_num=motion_category_num)
#         self.stpn = STPN(height_feat_size=height_feat_size)


#     def forward(self, bevs):
#         bevs = bevs.permute(0, 1, 4, 2, 3)  # (Batch, seq, z, h, w)

#         # Backbone network
#         x = self.stpn(bevs)

#         # Cell Classification head
#         cell_class_pred = self.cell_classify(x)

#         # Motion State Classification head
#         state_class_pred = self.state_classify(x)

#         # Motion Displacement prediction
#         disp = self.motion_pred(x)
#         disp = disp.view(-1, 2, x.size(-2), x.size(-1))

#         return disp, cell_class_pred, state_class_pred


# # For MGDA loss computation
# class FeatEncoder(nn.Module):
#     def __init__(self, height_feat_size=13):
#         super(FeatEncoder, self).__init__()
#         self.stpn = STPN(height_feat_size=height_feat_size)

#     def forward(self, bevs):
#         bevs = bevs.permute(0, 1, 4, 2, 3)  # (Batch, seq, z, h, w)
#         x = self.stpn(bevs)

#         return x


# class MotionNetMGDA(nn.Module):
#     def __init__(self, out_seq_len=20, motion_category_num=2):
#         super(MotionNetMGDA, self).__init__()
#         self.out_seq_len = out_seq_len

#         self.cell_classify = CellClassification()
#         self.motion_pred = MotionPrediction(seq_len=self.out_seq_len)
#         self.state_classify = StateEstimation(motion_category_num=motion_category_num)

#     def forward(self, stpn_out):
#         # Cell Classification head
#         cell_class_pred = self.cell_classify(stpn_out)

#         # Motion State Classification head
#         state_class_pred = self.state_classify(stpn_out)

#         # Motion Displacement prediction
#         disp = self.motion_pred(stpn_out)
#         disp = disp.view(-1, 2, stpn_out.size(-2), stpn_out.size(-1))

#         return disp, cell_class_pred, state_class_pred
opencood/models/sub_modules/MotionNet.py
MediaBrain-SJTU/CoBEVFlow
import numpy as np
import torch
import torch.nn as nn

from opencood.models.fuse_modules.self_attn import AttFusion
from opencood.models.sub_modules.auto_encoder import AutoEncoder

DEBUG = False

class AttBEVBackbone(nn.Module):
    def __init__(self, model_cfg, input_channels):
        super().__init__()
        self.model_cfg = model_cfg
        self.compress = False

        self.discrete_ratio = model_cfg['voxel_size'][0]
        self.downsample_rate = 1

            

        if 'compression' in model_cfg and model_cfg['compression'] > 0:
            self.compress = True
            self.compress_layer = model_cfg['compression']

        if 'layer_nums' in self.model_cfg:

            assert len(self.model_cfg['layer_nums']) == \
                   len(self.model_cfg['layer_strides']) == \
                   len(self.model_cfg['num_filters'])

            layer_nums = self.model_cfg['layer_nums']
            layer_strides = self.model_cfg['layer_strides']
            num_filters = self.model_cfg['num_filters']
        else:
            layer_nums = layer_strides = num_filters = []

        if 'upsample_strides' in self.model_cfg:
            assert len(self.model_cfg['upsample_strides']) \
                   == len(self.model_cfg['num_upsample_filter'])

            num_upsample_filters = self.model_cfg['num_upsample_filter']
            upsample_strides = self.model_cfg['upsample_strides']

        else:
            upsample_strides = num_upsample_filters = []

        num_levels = len(layer_nums)
        c_in_list = [input_channels, *num_filters[:-1]]

        self.blocks = nn.ModuleList()
        self.fuse_modules = nn.ModuleList()
        self.deblocks = nn.ModuleList()

        if self.compress:
            self.compression_modules = nn.ModuleList()

        for idx in range(num_levels):
            cur_layers = [
                nn.ZeroPad2d(1),
                nn.Conv2d(
                    c_in_list[idx], num_filters[idx], kernel_size=3,
                    stride=layer_strides[idx], padding=0, bias=False
                ),
                nn.BatchNorm2d(num_filters[idx], eps=1e-3, momentum=0.01),
                nn.ReLU()
            ]

            fuse_network = AttFusion(num_filters[idx])
            self.fuse_modules.append(fuse_network)
            if self.compress and self.compress_layer - idx > 0:
                self.compression_modules.append(AutoEncoder(num_filters[idx],
                                                            self.compress_layer-idx))

            for k in range(layer_nums[idx]):
                cur_layers.extend([
                    nn.Conv2d(num_filters[idx], num_filters[idx],
                              kernel_size=3, padding=1, bias=False),
                    nn.BatchNorm2d(num_filters[idx], eps=1e-3, momentum=0.01),
                    nn.ReLU()
                ])

            self.blocks.append(nn.Sequential(*cur_layers))
            if len(upsample_strides) > 0:
                stride = upsample_strides[idx]
                if stride >= 1:
                    self.deblocks.append(nn.Sequential(
                        nn.ConvTranspose2d(
                            num_filters[idx], num_upsample_filters[idx],
                            upsample_strides[idx],
                            stride=upsample_strides[idx], bias=False
                        ),
                        nn.BatchNorm2d(num_upsample_filters[idx],
                                       eps=1e-3, momentum=0.01),
                        nn.ReLU()
                    ))
                else:
                    stride = np.round(1 / stride).astype(np.int)
                    self.deblocks.append(nn.Sequential(
                        nn.Conv2d(
                            num_filters[idx], num_upsample_filters[idx],
                            stride,
                            stride=stride, bias=False
                        ),
                        nn.BatchNorm2d(num_upsample_filters[idx], eps=1e-3,
                                       momentum=0.01),
                        nn.ReLU()
                    ))

        c_in = sum(num_upsample_filters)
        if len(upsample_strides) > num_levels:
            self.deblocks.append(nn.Sequential(
                nn.ConvTranspose2d(c_in, c_in, upsample_strides[-1],
                                   stride=upsample_strides[-1], bias=False),
                nn.BatchNorm2d(c_in, eps=1e-3, momentum=0.01),
                nn.ReLU(),
            ))

        self.num_bev_features = c_in

    def forward(self, data_dict):
        spatial_features = data_dict['spatial_features']
        if DEBUG:
            origin_features = torch.clone(spatial_features)
        record_len = data_dict['record_len']
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        ups = []
        ret_dict = {}
        x = spatial_features

        H, W = x.shape[2:]   #  200, 704
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]

        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        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)
            if DEBUG:
                self.fuse_modules[i].forward_debug(x, origin_features, record_len, pairwise_t_matrix)
            else:
                x_fuse = self.fuse_modules[i](x, record_len, pairwise_t_matrix)

            stride = int(spatial_features.shape[2] / x.shape[2])
            ret_dict['spatial_features_%dx' % stride] = x

            if len(self.deblocks) > 0:
                ups.append(self.deblocks[i](x_fuse))
            else:
                ups.append(x_fuse)

        if len(ups) > 1:
            x = torch.cat(ups, dim=1)
        elif len(ups) == 1:
            x = ups[0]

        if len(self.deblocks) > len(self.blocks):
            x = self.deblocks[-1](x)

        data_dict['spatial_features_2d'] = x
        return data_dict

opencood/models/sub_modules/att_bev_backbone.py
MediaBrain-SJTU/CoBEVFlow
import numpy as np
import torch
import torch.nn as nn

from opencood.models.sub_modules.resblock import ResNetModified, BasicBlock, Bottleneck
from opencood.models.fuse_modules.self_attn import AttFusion
from opencood.models.sub_modules.auto_encoder import AutoEncoder

DEBUG = False

class AttResNetBEVBackbone(nn.Module):
    def __init__(self, model_cfg, input_channels):
        super().__init__()
        self.model_cfg = model_cfg
        self.compress = False

        
        self.discrete_ratio = model_cfg['voxel_size'][0]
        self.downsample_rate = 1


            
        layer_nums = model_cfg['layer_nums']
        num_filters = model_cfg['num_filters']
        layer_strides = model_cfg['layer_strides']
        upsample_strides = model_cfg['upsample_strides']
        num_upsample_filters = model_cfg['num_upsample_filter']
        self.level_num = len(layer_nums)

        self.resnet = ResNetModified(BasicBlock, 
                                    layer_nums,
                                    layer_strides,
                                    num_filters)



        num_levels = len(layer_nums)
        c_in_list = [input_channels, *num_filters[:-1]]

        self.fuse_modules = nn.ModuleList()
        self.deblocks = nn.ModuleList()

        if self.compress:
            self.compression_modules = nn.ModuleList()

        for idx in range(num_levels):
            
            fuse_network = AttFusion(num_filters[idx])
            self.fuse_modules.append(fuse_network)
            
            if len(upsample_strides) > 0:
                stride = upsample_strides[idx]
                self.deblocks.append(nn.Sequential(
                    nn.ConvTranspose2d(
                        num_filters[idx], num_upsample_filters[idx],
                        upsample_strides[idx],
                        stride=upsample_strides[idx], bias=False
                    ),
                    nn.BatchNorm2d(num_upsample_filters[idx],
                                    eps=1e-3, momentum=0.01),
                    nn.ReLU()
                ))


        c_in = sum(num_upsample_filters)
        if len(upsample_strides) > num_levels:
            self.deblocks.append(nn.Sequential(
                nn.ConvTranspose2d(c_in, c_in, upsample_strides[-1],
                                   stride=upsample_strides[-1], bias=False),
                nn.BatchNorm2d(c_in, eps=1e-3, momentum=0.01),
                nn.ReLU(),
            ))

        self.num_bev_features = c_in

    def forward(self, data_dict):
        spatial_features = data_dict['spatial_features']
        if DEBUG:
            origin_features = torch.clone(spatial_features)

        record_len = data_dict['record_len']
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        ups = []
        ret_dict = {}
        x = spatial_features

        H, W = x.shape[2:]   #  200, 704
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]

        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        features = self.resnet(x)  # tuple of features
        ups = []

        for i in range(self.level_num):
            if DEBUG:
                self.fuse_modules[i].forward_debug(features[i], origin_features, record_len, pairwise_t_matrix)
            else:
                x_fuse = self.fuse_modules[i](features[i], record_len, pairwise_t_matrix)
            if len(self.deblocks) > 0:
                ups.append(self.deblocks[i](x_fuse))
            else:
                ups.append(x_fuse)

        if len(ups) > 1:
            x = torch.cat(ups, dim=1)
        elif len(ups) == 1:
            x = ups[0]

        if len(self.deblocks) > self.level_num:
            x = self.deblocks[-1](x)

        data_dict['spatial_features_2d'] = x
        return data_dict

opencood/models/sub_modules/att_resnet_bev_backbone.py
MediaBrain-SJTU/CoBEVFlow
import numpy as np
import torch
import torch.nn as nn

from opencood.models.sub_modules.resblock import ResNetModified, BasicBlock, Bottleneck
from opencood.models.fuse_modules.self_attn import AttFusion
from opencood.models.sub_modules.auto_encoder import AutoEncoder
from opencood.models.sub_modules.dcn_net import DCNNet

DEBUG = False

class AttResNetBEVBackboneDCN(nn.Module):
    def __init__(self, model_cfg, input_channels):
        super().__init__()
        self.model_cfg = model_cfg
        self.compress = False

        
        self.discrete_ratio = model_cfg['voxel_size'][0]
        self.downsample_rate = 1


        self.dcn_before_upsample = model_cfg['dcn_before_upsample']
        layer_nums = model_cfg['layer_nums']
        num_filters = model_cfg['num_filters']
        layer_strides = model_cfg['layer_strides']
        upsample_strides = model_cfg['upsample_strides']
        num_upsample_filters = model_cfg['num_upsample_filter']
        self.level_num = len(layer_nums)

        self.resnet = ResNetModified(BasicBlock, 
                                    layer_nums,
                                    layer_strides,
                                    num_filters)



        num_levels = len(layer_nums)
        c_in_list = [input_channels, *num_filters[:-1]]

        self.fuse_modules = nn.ModuleList()
        self.deblocks = nn.ModuleList()
        self.dcn_blocks = nn.ModuleList()

        if self.compress:
            self.compression_modules = nn.ModuleList()

        for idx in range(num_levels):
            
            fuse_network = AttFusion(num_filters[idx])
            self.fuse_modules.append(fuse_network)
            
            if len(upsample_strides) > 0:
                stride = upsample_strides[idx]
                self.deblocks.append(nn.Sequential(
                    nn.ConvTranspose2d(
                        num_filters[idx], num_upsample_filters[idx],
                        upsample_strides[idx],
                        stride=upsample_strides[idx], bias=False
                    ),
                    nn.BatchNorm2d(num_upsample_filters[idx],
                                    eps=1e-3, momentum=0.01),
                    nn.ReLU()
                ))
        for idx in range(num_level):
            if self.dcn_before_upsample:
                dcn_args = {}
                dcn_args['n_blocks'] = 1
                dcn_args['in_channels'] = [num_filters[idx]]
                dcn_args['out_channels'] = [num_filters[idx]]
                dcn_args['stride'] = [1]
                dcn_args['kernel_size'] = [3]
                dcn_args['padding'] = [1]
                dcn_network = DCNNet(dcn_args)
            else:
                args = {}
                args['n_blocks'] = 1
                args['in_channels'] = [num_upsample_filters[idx]]
                args['out_channels'] = [num_upsample_filters[idx]]
                args['stride'] = [1]
                args['kernel_size'] = [3]
                args['padding'] = [1]
                dcn_network = DCNNet(dcn_args)
            self.dcn_blocks.append(dcn_network)

        c_in = sum(num_upsample_filters)
        if len(upsample_strides) > num_levels:
            self.deblocks.append(nn.Sequential(
                nn.ConvTranspose2d(c_in, c_in, upsample_strides[-1],
                                   stride=upsample_strides[-1], bias=False),
                nn.BatchNorm2d(c_in, eps=1e-3, momentum=0.01),
                nn.ReLU(),
            ))

        self.num_bev_features = c_in

    def forward(self, data_dict):
        spatial_features = data_dict['spatial_features']
        if DEBUG:
            origin_features = torch.clone(spatial_features)

        record_len = data_dict['record_len']
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        ups = []
        ret_dict = {}
        x = spatial_features

        H, W = x.shape[2:]   #  200, 704
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]

        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        features = self.resnet(x)  # tuple of features
        ups = []

        for i in range(self.level_num):
            x_fuse = self.fuse_modules[i](features[i], record_len, pairwise_t_matrix)
            if len(self.deblocks) > 0:
                if self.dcn_before_upsample:
                    ups.append(self.deblocks[i](self.dcn_blocks[i](x_fuse)))
                else:
                    ups.append(self.dcn_blocks[i](self.deblocks[i](x_fuse)))
            else:
                ups.append(x_fuse)

        if len(ups) > 1:
            x = torch.cat(ups, dim=1)
        elif len(ups) == 1:
            x = ups[0]

        if len(self.deblocks) > self.level_num:
            x = self.deblocks[-1](x)

        data_dict['spatial_features_2d'] = x
        return data_dict

opencood/models/sub_modules/att_resnet_bev_backbone_dcn.py
MediaBrain-SJTU/CoBEVFlow
import numpy as np
import torch
import torch.nn as nn

from opencood.models.sub_modules.resblock import ResNetModified, BasicBlock, Bottleneck
from opencood.models.fuse_modules.self_attn import AttFusion
from opencood.models.sub_modules.auto_encoder import AutoEncoder

DEBUG = False

class AttResNetBEVBackboneSS(nn.Module):
    # single scale
    def __init__(self, model_cfg, input_channels):
        super().__init__()
        self.model_cfg = model_cfg
        self.compress = False

        
        self.discrete_ratio = model_cfg['voxel_size'][0]
        self.downsample_rate = 1


            
        layer_nums = model_cfg['layer_nums']
        num_filters = model_cfg['num_filters']
        layer_strides = model_cfg['layer_strides']
        upsample_strides = model_cfg['upsample_strides']
        num_upsample_filters = model_cfg['num_upsample_filter']
        self.level_num = len(layer_nums)

        self.resnet = ResNetModified(BasicBlock, 
                                    layer_nums,
                                    layer_strides,
                                    num_filters)



        num_levels = len(layer_nums)
        c_in_list = [input_channels, *num_filters[:-1]]

        self.fuse_modules = nn.ModuleList()
        self.deblocks = nn.ModuleList()

        if self.compress:
            self.compression_modules = nn.ModuleList()

        for idx in range(num_levels):
            
            fuse_network = AttFusion(num_filters[idx])
            self.fuse_modules.append(fuse_network)
            
            if len(upsample_strides) > 0:
                stride = upsample_strides[idx]
                self.deblocks.append(nn.Sequential(
                    nn.ConvTranspose2d(
                        num_filters[idx], num_upsample_filters[idx],
                        upsample_strides[idx],
                        stride=upsample_strides[idx], bias=False
                    ),
                    nn.BatchNorm2d(num_upsample_filters[idx],
                                    eps=1e-3, momentum=0.01),
                    nn.ReLU()
                ))


        c_in = sum(num_upsample_filters)
        if len(upsample_strides) > num_levels:
            self.deblocks.append(nn.Sequential(
                nn.ConvTranspose2d(c_in, c_in, upsample_strides[-1],
                                   stride=upsample_strides[-1], bias=False),
                nn.BatchNorm2d(c_in, eps=1e-3, momentum=0.01),
                nn.ReLU(),
            ))

        self.num_bev_features = c_in

    def forward(self, data_dict):
        spatial_features = data_dict['spatial_features']
        if DEBUG:
            origin_features = torch.clone(spatial_features)

        record_len = data_dict['record_len']
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        ups = []
        ret_dict = {}
        x = spatial_features

        H, W = x.shape[2:]   #  200, 704
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]

        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        features = self.resnet(x)  # tuple of features
        ups = []

 
        for i in range(1): # only one scale
            x_fuse = self.fuse_modules[i](features[i], record_len, pairwise_t_matrix)
            if len(self.deblocks) > 0:
                ups.append(self.deblocks[i](x_fuse))
            else:
                ups.append(x_fuse)

        x = ups[0]

        data_dict['spatial_features_2d'] = x
        return data_dict

opencood/models/sub_modules/att_resnet_bev_backbone_ss.py
MediaBrain-SJTU/CoBEVFlow
import numpy as np
import torch
import torch.nn as nn


class BaseBEVBackbone(nn.Module):
    def __init__(self, model_cfg, input_channels):
        super().__init__()
        self.model_cfg = model_cfg

        if 'layer_nums' in self.model_cfg:

            assert len(self.model_cfg['layer_nums']) == \
                   len(self.model_cfg['layer_strides']) == \
                   len(self.model_cfg['num_filters'])

            layer_nums = self.model_cfg['layer_nums']
            layer_strides = self.model_cfg['layer_strides']
            num_filters = self.model_cfg['num_filters']
        else:
            layer_nums = layer_strides = num_filters = []

        if 'upsample_strides' in self.model_cfg:
            assert len(self.model_cfg['upsample_strides']) \
                   == len(self.model_cfg['num_upsample_filter'])

            num_upsample_filters = self.model_cfg['num_upsample_filter']
            upsample_strides = self.model_cfg['upsample_strides']

        else:
            upsample_strides = num_upsample_filters = []

        num_levels = len(layer_nums)
        c_in_list = [input_channels, *num_filters[:-1]]

        self.blocks = nn.ModuleList()
        self.deblocks = nn.ModuleList()

        for idx in range(num_levels):
            cur_layers = [
                nn.ZeroPad2d(1),
                nn.Conv2d(
                    c_in_list[idx], num_filters[idx], kernel_size=3,
                    stride=layer_strides[idx], padding=0, bias=False
                ),
                nn.BatchNorm2d(num_filters[idx], eps=1e-3, momentum=0.01),
                nn.ReLU()
            ]
            for k in range(layer_nums[idx]):
                cur_layers.extend([
                    nn.Conv2d(num_filters[idx], num_filters[idx],
                              kernel_size=3, padding=1, bias=False),
                    nn.BatchNorm2d(num_filters[idx], eps=1e-3, momentum=0.01),
                    nn.ReLU()
                ])

            self.blocks.append(nn.Sequential(*cur_layers))
            if len(upsample_strides) > 0:
                stride = upsample_strides[idx]
                if stride >= 1:
                    self.deblocks.append(nn.Sequential(
                        nn.ConvTranspose2d(
                            num_filters[idx], num_upsample_filters[idx],
                            upsample_strides[idx],
                            stride=upsample_strides[idx], bias=False
                        ),
                        nn.BatchNorm2d(num_upsample_filters[idx],
                                       eps=1e-3, momentum=0.01),
                        nn.ReLU()
                    ))
                else:
                    stride = np.round(1 / stride).astype(np.int)
                    self.deblocks.append(nn.Sequential(
                        nn.Conv2d(
                            num_filters[idx], num_upsample_filters[idx],
                            stride,
                            stride=stride, bias=False
                        ),
                        nn.BatchNorm2d(num_upsample_filters[idx], eps=1e-3,
                                       momentum=0.01),
                        nn.ReLU()
                    ))

        c_in = sum(num_upsample_filters)
        if len(upsample_strides) > num_levels:
            self.deblocks.append(nn.Sequential(
                nn.ConvTranspose2d(c_in, c_in, upsample_strides[-1],
                                   stride=upsample_strides[-1], bias=False),
                nn.BatchNorm2d(c_in, eps=1e-3, momentum=0.01),
                nn.ReLU(),
            ))

        self.num_bev_features = c_in

    def forward(self, data_dict):
        spatial_features = data_dict['spatial_features']

        ups = []
        ret_dict = {}
        x = spatial_features

        for i in range(len(self.blocks)):
            x = self.blocks[i](x)

            stride = int(spatial_features.shape[2] / x.shape[2])
            ret_dict['spatial_features_%dx' % stride] = x

            if len(self.deblocks) > 0:
                ups.append(self.deblocks[i](x))
            else:
                ups.append(x)

        if len(ups) > 1:
            x = torch.cat(ups, dim=1)
        elif len(ups) == 1:
            x = ups[0]

        if len(self.deblocks) > len(self.blocks):
            x = self.deblocks[-1](x)

        data_dict['spatial_features_2d'] = x # [N,C,100,352]

        return data_dict

opencood/models/sub_modules/base_bev_backbone.py
MediaBrain-SJTU/CoBEVFlow
import numpy as np
import torch
import torch.nn as nn

from opencood.models.sub_modules.resblock import ResNetModified, BasicBlock
# from opencood.models.sub_modules.resblock2 import ResNetModified, BasicBlock
DEBUG = False

class ResNetBEVBackbone(nn.Module):
    def __init__(self, model_cfg, input_channels):
        super().__init__()
        self.model_cfg = model_cfg

        if 'layer_nums' in self.model_cfg:

            assert len(self.model_cfg['layer_nums']) == \
                   len(self.model_cfg['layer_strides']) == \
                   len(self.model_cfg['num_filters'])

            layer_nums = self.model_cfg['layer_nums']
            layer_strides = self.model_cfg['layer_strides']
            num_filters = self.model_cfg['num_filters']
        else:
            layer_nums = layer_strides = num_filters = []

        if 'upsample_strides' in self.model_cfg:
            assert len(self.model_cfg['upsample_strides']) \
                   == len(self.model_cfg['num_upsample_filter'])

            num_upsample_filters = self.model_cfg['num_upsample_filter']
            upsample_strides = self.model_cfg['upsample_strides']

        else:
            upsample_strides = num_upsample_filters = []

        self.resnet = ResNetModified(BasicBlock, 
                                        layer_nums,
                                        layer_strides,
                                        num_filters)

        num_levels = len(layer_nums)
        self.num_levels = len(layer_nums)
        self.deblocks = nn.ModuleList()

        for idx in range(num_levels):
            if len(upsample_strides) > 0:
                stride = upsample_strides[idx]
                if stride >= 1:
                    self.deblocks.append(nn.Sequential(
                        nn.ConvTranspose2d(
                            num_filters[idx], num_upsample_filters[idx],
                            upsample_strides[idx],
                            stride=upsample_strides[idx], bias=False
                        ),
                        nn.BatchNorm2d(num_upsample_filters[idx],
                                       eps=1e-3, momentum=0.01),
                        nn.ReLU()
                    ))
                else:
                    stride = np.round(1 / stride).astype(np.int)
                    self.deblocks.append(nn.Sequential(
                        nn.Conv2d(
                            num_filters[idx], num_upsample_filters[idx],
                            stride,
                            stride=stride, bias=False
                        ),
                        nn.BatchNorm2d(num_upsample_filters[idx], eps=1e-3,
                                       momentum=0.01),
                        nn.ReLU()
                    ))

        c_in = sum(num_upsample_filters)
        if len(upsample_strides) > num_levels:
            self.deblocks.append(nn.Sequential(
                nn.ConvTranspose2d(c_in, c_in, upsample_strides[-1],
                                   stride=upsample_strides[-1], bias=False),
                nn.BatchNorm2d(c_in, eps=1e-3, momentum=0.01),
                nn.ReLU(),
            ))

        self.num_bev_features = c_in

    def forward(self, data_dict):
        spatial_features = data_dict['spatial_features'] # (batch_size, feature, H, W)

        x = self.resnet(spatial_features)  # tuple of features
        ups = []

        for i in range(self.num_levels):
            if len(self.deblocks) > 0:
                ups.append(self.deblocks[i](x[i]))
            else:
                ups.append(x[i])

        if len(ups) > 1:
            x = torch.cat(ups, dim=1) # feature dim 
        elif len(ups) == 1:
            x = ups[0]

        if len(self.deblocks) > self.num_levels:
            x = self.deblocks[-1](x)

        data_dict['spatial_features_2d'] = x
        return data_dict

opencood/models/sub_modules/base_bev_backbone_resnet.py
MediaBrain-SJTU/CoBEVFlow
import torch
from torch import nn

from einops import rearrange


class PreNorm(nn.Module):
    def __init__(self, dim, fn):
        super().__init__()
        self.norm = nn.LayerNorm(dim)
        self.fn = fn

    def forward(self, x, **kwargs):
        return self.fn(self.norm(x), **kwargs)


class FeedForward(nn.Module):
    def __init__(self, dim, hidden_dim, dropout=0.):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(dim, hidden_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x):
        return self.net(x)


class CavAttention(nn.Module):
    """
    Vanilla CAV attention.
    """
    def __init__(self, dim, heads, dim_head=64, dropout=0.1):
        super().__init__()
        inner_dim = heads * dim_head

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.attend = nn.Softmax(dim=-1)
        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias=False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, dim),
            nn.Dropout(dropout)
        )

    def forward(self, x, mask, prior_encoding):
        # x: (B, L, H, W, C) -> (B, H, W, L, C)
        # mask: (B, L)
        x = x.permute(0, 2, 3, 1, 4)
        # mask: (B, 1, H, W, L, 1)
        mask = mask.unsqueeze(1)

        # qkv: [(B, H, W, L, C_inner) *3]
        qkv = self.to_qkv(x).chunk(3, dim=-1)
        # q: (B, M, H, W, L, C)
        q, k, v = map(lambda t: rearrange(t, 'b h w l (m c) -> b m h w l c',
                                          m=self.heads), qkv)

        # attention, (B, M, H, W, L, L)
        att_map = torch.einsum('b m h w i c, b m h w j c -> b m h w i j',
                               q, k) * self.scale
        # add mask
        att_map = att_map.masked_fill(mask == 0, -float('inf'))
        # softmax
        att_map = self.attend(att_map)

        # out:(B, M, H, W, L, C_head)
        out = torch.einsum('b m h w i j, b m h w j c -> b m h w i c', att_map,
                           v)
        out = rearrange(out, 'b m h w l c -> b h w l (m c)',
                        m=self.heads)
        out = self.to_out(out)
        # (B L H W C)
        out = out.permute(0, 3, 1, 2, 4)
        return out


class BaseEncoder(nn.Module):
    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout=0.):
        super().__init__()
        self.layers = nn.ModuleList([])
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                PreNorm(dim, CavAttention(dim,
                                          heads=heads,
                                          dim_head=dim_head,
                                          dropout=dropout)),
                PreNorm(dim, FeedForward(dim, mlp_dim, dropout=dropout))
            ]))

    def forward(self, x, mask):
        for attn, ff in self.layers:
            x = attn(x, mask=mask) + x
            x = ff(x) + x
        return x


class BaseTransformer(nn.Module):
    def __init__(self, args):
        super().__init__()

        dim = args['dim']
        depth = args['depth']
        heads = args['heads']
        dim_head = args['dim_head']
        mlp_dim = args['mlp_dim']
        dropout = args['dropout']
        max_cav = args['max_cav']

        self.encoder = BaseEncoder(dim, depth, heads, dim_head, mlp_dim,
                                   dropout)

    def forward(self, x, mask):
        # B, L, H, W, C
        output = self.encoder(x, mask)
        # B, H, W, C
        output = output[:, 0]

        return 
opencood/models/sub_modules/base_transformer.py
MediaBrain-SJTU/CoBEVFlow
"""
This module is designed for box alignment
It should be used for 1-round communication, maybe adapt to 2-round communication latter
i.e, collaborative agent send the full feature map and noisy pose once together

We will use g2o for pose graph optimization.
"""


from cv2 import threshold
from opencood.models.sub_modules.pose_graph_optim import PoseGraphOptimization2D
from opencood.utils.transformation_utils import pose_to_tfm
from opencood.utils import box_utils
from collections import OrderedDict
import numpy as np
import torch
import torch.nn.functional as F
import g2o
from icecream import ic
import os

DEBUG = False

def all_pair_l2(A, B):
    """ All pair L2 distance for A and B
    Args:
        A : torch.Tensor
            shape [N_A, D]
        B : torch.Tensor
            shape [N_B, D]
    Returns:
        C : torch.Tensor
            shape [N_A, N_B]
    """
    TwoAB = 2*A@B.T
    C = torch.sqrt(torch.sum(A * A, 1, keepdim=True).expand_as(TwoAB) \
        + torch.sum(B * B, 1, keepdim=True).T.expand_as(TwoAB) \
        - TwoAB)
    return C

def box_alignment_relative_sample(
            pred_corners_list,
            noisy_lidar_pose, 
            clean_lidar_pose=None, 
            uncertainty_list=None, 
            order='hwl', 
            landmark_SE2=True,
            adaptive_landmark=False):
    """ Perform box alignment for one sample. 
    Correcting the relative pose.

    Args:
        pred_corners_list: in each ego coordinate
            [[N_1, 8, 3], ..., [N_cav1, 8, 3]]

        clean_lidar_poses:
            [N_cav1, 6], in degree
        
        noisy_lidar_poses:
            [N_cav1, 6], in degree

        uncertainty_list:
            [[N_1, 3], [N_2, 3], ..., [N_cav1, 3]]

        landmark_SE2:
            if True, the landmark is SE(2), otherwise R^2
        
        adaptive_landmark: (when landmark_SE2 = True)
            if True, landmark will turn to R^2 if yaw angles differ a lot

    Returns: 
        refined_lidar_poses: np.ndarray
            [N_cav1, 3], 
    """

    ## first transform point from ego coordinate to world coordinate, using lidar_pose.
    N = noisy_lidar_pose.shape[0]
    device = pred_corners_list[0].device
    lidar_pose_noisy_tfm = pose_to_tfm(noisy_lidar_pose, dof=6)

    pred_corners_world_list = \
        [box_utils.project_box3d(pred_corners_list[i], lidar_pose_noisy_tfm[i]) for i in range(N)]  # [[N1, 8, 3], [N2, 8, 3],...]
    pred_box3d_list = \
        [box_utils.corner_to_center_torch(corner, order).to(device) for corner in pred_corners_list]   # [[N1, 7], [N2, 7], ...], angle in radius
    pred_box3d_world_list = \
        [box_utils.corner_to_center_torch(corner, order).to(device) for corner in pred_corners_world_list]   # [[N1, 7], [N2, 7], ...], angle in radius

    pred_center_list = \
        [torch.mean(corner_tensor, dim=[1]) for corner_tensor in pred_corners_list] # [[N1,3], [N2,3], ...]

    pred_center_world_list = \
        [pred_box3d_world[:,:3] for pred_box3d_world in pred_box3d_world_list]

    pred_yaw_world_list = \
        [pred_box3d[:, 6] for pred_box3d in pred_box3d_world_list]

    pred_len = \
        [pred_center.shape[0] for pred_center in pred_center_list] 

    


    box_idx_to_agent = []
    for i in range(N):
        box_idx_to_agent += [i] * pred_len[i] 
    

    pred_center_cat = torch.cat(pred_center_list, dim=0)   # [sum(pred_box), 3]
    pred_center_world_cat = torch.cat(pred_center_world_list, dim=0)  # [sum(pred_box), 3]
    pred_box3d_cat = torch.cat(pred_box3d_list, dim=0)  # [sum(pred_box), 7]
    pred_yaw_world_cat = torch.cat(pred_yaw_world_list)  # [sum(pred_box)]


    w_a = 1.6 # width of anchor
    l_a = 3.9 # length of anchor
    d_a_square = w_a ** 2 + l_a ** 2 # anchor's diag


    if uncertainty_list is not None:
        pred_log_sigma2_cat = torch.cat(uncertainty_list)
        pred_certainty_cat = torch.exp(-pred_log_sigma2_cat)
        pred_certainty_cat[:,:2] /= d_a_square # sigma_delta_x -> sigma_x. 


    pred_center_world_cat_cpu = pred_center_world_cat.cpu() # if use gpu, it will get nan.
    pred_center_allpair_dist = all_pair_l2(pred_center_world_cat_cpu, pred_center_world_cat_cpu) # [sum(pred_box), sum(pred_box)]


    # let pair from one vehicle be max distance
    MAX_DIST = 10000
    cum = 0
    for i in range(N):
        pred_center_allpair_dist[cum: cum + pred_len[i], cum: cum +pred_len[i]] = MAX_DIST
        cum += pred_len[i]


    cluster_id = N # let the vertex id of object start from N
    cluster_dict = OrderedDict()
    remain_box = set(range(cum))
    thres = 0.75  # l2 distance within the threshold, can be considered as one object.
    for box_idx in range(cum): 

        if box_idx not in remain_box:  # already assigned
            continue
        within_thres_idx_tensor = (pred_center_allpair_dist[box_idx] < thres).nonzero().flatten()
        within_thres_idx_list = within_thres_idx_tensor.cpu().numpy().tolist()

        if len(within_thres_idx_list) == 0:  # if it's a single box
            continue

        # start from within_thres_idx_list, find new box added to the cluster
        explored = [box_idx]
        unexplored = [idx for idx in within_thres_idx_list if idx in remain_box]

        while unexplored:
            idx = unexplored[0]
            within_thres_idx_tensor = (pred_center_allpair_dist[idx] < thres).nonzero().flatten()
            within_thres_idx_list = within_thres_idx_tensor.cpu().numpy().tolist()
            for newidx in within_thres_idx_list:
                if (newidx not in explored) and (newidx not in unexplored) and (newidx in remain_box):
                    unexplored.append(newidx)
            unexplored.remove(idx)
            explored.append(idx)
        
        if len(explored) == 1: # it's a single box, neighbors have been assigned
            remain_box.remove(box_idx)
            continue
        
        cluster_box_idxs = explored

        cluster_dict[cluster_id] = OrderedDict()
        cluster_dict[cluster_id]['box_idx'] = [idx for idx in cluster_box_idxs]
        cluster_dict[cluster_id]['box_dist'] = [pred_center_cat[idx].norm() for idx in cluster_box_idxs]  # distance to observer 
        cluster_dict[cluster_id]['box_center_world'] = [pred_center_world_cat[idx] for idx in cluster_box_idxs]  # coordinate in world, [3,]
        cluster_dict[cluster_id]['box_yaw'] = [pred_yaw_world_cat[idx] for idx in cluster_box_idxs]

        yaw_var = torch.var(torch.as_tensor(cluster_dict[cluster_id]['box_yaw']), unbiased=False)
        
        if landmark_SE2:
            if adaptive_landmark and yaw_var > 0.2:
                landmark = pred_center_world_cat[box_idx].clone()[:2]
            else:
                landmark = pred_center_world_cat[box_idx].clone()
                landmark[2] = pred_yaw_world_cat[box_idx]
        else:
            landmark = pred_center_world_cat[box_idx].clone()[:2]


        cluster_dict[cluster_id]['landmark'] = landmark.cpu().numpy()  # [x, y, yaw] or [x, y]
        cluster_dict[cluster_id]['landmark_SE2'] = True if landmark.shape[0] == 3 else False

        DEBUG = False
        if DEBUG:
            from icecream import ic
            ic(cluster_dict[cluster_id]['box_idx'])
            ic(cluster_dict[cluster_id]['box_center_world'])
            ic(cluster_dict[cluster_id]['box_yaw'])
            ic(cluster_dict[cluster_id]['landmark'])
        

        cluster_id += 1
        for idx in cluster_box_idxs:
            remain_box.remove(idx)

    vertex_num = cluster_id
    agent_num = N
    landmark_num = cluster_id - N
    # ic(agent_num)
    # ic(landmark_num)

    """
        Now we have clusters for objects. we can create pose graph.
        First we consider center as landmark.
        Maybe set corner as landmarks in the future.
    """
    pgo = PoseGraphOptimization2D(verbose=False)

    # Add agent to vertexs
    for agent_id in range(N):
        v_id = agent_id
        # notice lidar_pose use degree format, translate it to radians.
        pose_np = noisy_lidar_pose[agent_id, [0,1,4]].cpu().numpy()
        pose_np[2] = np.deg2rad(pose_np[2])  # radians
        v_pose = g2o.SE2(pose_np)
        
        if agent_id == 0:
            pgo.add_vertex(id=v_id, pose=v_pose, fixed=True)
        else:
            pgo.add_vertex(id=v_id, pose=v_pose, fixed=False)

    # Add landmark to vertexs
    for landmark_id in range(N, cluster_id):
        v_id = landmark_id
        landmark = cluster_dict[landmark_id]['landmark'] # (3,) or (2,)
        landmark_SE2 = cluster_dict[landmark_id]['landmark_SE2']

        if landmark_SE2:
            v_pose = g2o.SE2(landmark)
        else:
            v_pose = landmark

        pgo.add_vertex(id=v_id, pose=v_pose, fixed=False, SE2=landmark_SE2)

    # Add agent-landmark edge to edge
    for landmark_id in range(N, cluster_id):
        landmark_SE2 = cluster_dict[landmark_id]['landmark_SE2']

        for box_idx in cluster_dict[landmark_id]['box_idx']:
            agent_id = box_idx_to_agent[box_idx]
            if landmark_SE2:
                e_pose = g2o.SE2(pred_box3d_cat[box_idx][[0,1,6]].cpu().numpy().astype(np.float64))
                info = np.identity(3, dtype=np.float64)
                if uncertainty_list is not None:
                    info[[0,1,2],[0,1,2]] = pred_certainty_cat[box_idx].cpu().numpy()
            else:
                e_pose = pred_box3d_cat[box_idx][[0,1]].cpu().numpy().astype(np.float64)
                info = np.identity(2, dtype=np.float64)
                if uncertainty_list is not None:
                    info[[0,1],[0,1]] = pred_certainty_cat[box_idx][:2].cpu().numpy()


            pgo.add_edge(vertices=[agent_id, landmark_id], measurement=e_pose, information=info, SE2=landmark_SE2)
    
    pgo.optimize()

    pose_new_list = []
    for agent_id in range(N):
        # print(pgo.get_pose(agent_id).vector())
        pose_new_list.append(pgo.get_pose(agent_id).vector())

    refined_pose = np.array(pose_new_list)
    refined_pose[:,2] = np.rad2deg(refined_pose[:,2])  # rad -> degree, same as source

    return refined_pose

def box_alignment_sample(pred_corners_list, lidar_poses_for_tfm, noisy_lidar_poses, uncertainty_list=None, order='hwl'):
    """ Perform box alignment for one sample.
    Args:
        pred_corners_list: in each ego coordinate
            [[N_1, 8, 3], ..., [N_cav1, 8, 3]]

        lidar_poses:
            [N_cav1, 6] , in degree

        scores_list:
            [[N_1, 3], [N_2, 3], ..., [N_cav1, 3]]

    Returns: 
        refined_lidar_poses: np.ndarray
            [N_cav1, 3], 
    """
    
    ## first transform point from ego coordinate to world coordinate, using lidar_pose.
    lidar_poses = lidar_poses_for_tfm
    N = lidar_poses.shape[0]
    device = pred_corners_list[0].device
    lidar_pose_tfm = pose_to_tfm(lidar_poses, dof=6)  # Tw_c



    pred_corners_world_list = \
        [box_utils.project_box3d(pred_corners_list[i], lidar_pose_tfm[i]) for i in range(N)]  # [[N1, 8, 3], [N2, 8, 3],...]
    pred_box3d_list = \
        [box_utils.corner_to_center_torch(corner, order).to(device) for corner in pred_corners_list]   # [[N1, 7], [N2, 7], ...], angle in radius
    pred_box3d_world_list = \
        [box_utils.corner_to_center_torch(corner, order).to(device) for corner in pred_corners_world_list]   # [[N1, 7], [N2, 7], ...], angle in radius

    pred_center_list = \
        [torch.mean(corner_tensor, dim=[1]) for corner_tensor in pred_corners_list] # [[N1,3], [N2,3], ...]

    pred_center_world_list = \
        [pred_box3d_world[:,:3] for pred_box3d_world in pred_box3d_world_list]

    pred_yaw_world_list = \
        [pred_box3d[:, 6] for pred_box3d in pred_box3d_world_list]

    pred_len = \
        [pred_center.shape[0] for pred_center in pred_center_list] 


    box_idx_to_agent = []
    for i in range(N):
        box_idx_to_agent += [i] * pred_len[i] 

    if DEBUG:
        vis_corners_list(pred_corners_world_list,filename="/GPFS/rhome/yifanlu/OpenCOOD/box_align_items/gt_box_noisy_pose.png")
    

    pred_center_cat = torch.cat(pred_center_list, dim=0)   # [sum(pred_box), 3]
    pred_center_world_cat = torch.cat(pred_center_world_list, dim=0)  # [sum(pred_box), 3]
    pred_box3d_cat = torch.cat(pred_box3d_list, dim=0)  # [sum(pred_box), 7]
    pred_yaw_world_cat = torch.cat(pred_yaw_world_list)  # [sum(pred_box)]

    pred_center_world_cat_cpu = pred_center_world_cat.cpu() # if use gpu, it will get nan.
    pred_center_allpair_dist = all_pair_l2(pred_center_world_cat_cpu, pred_center_world_cat_cpu) # [sum(pred_box), sum(pred_box)]


    # let pair from one vehicle be max distance
    MAX_DIST = 10000
    cum = 0
    for i in range(N):
        pred_center_allpair_dist[cum: cum + pred_len[i], cum: cum +pred_len[i]] = MAX_DIST
        cum += pred_len[i]


    cluster_id = N # let the vertex id of object start from N
    cluster_dict = OrderedDict()
    remain_box = set(range(cum))
    thres = 1  # l2 distance within the threshold, can be considered as one object.
    for box_idx in range(cum):
        if box_idx not in remain_box:  # already assigned
            continue
        within_thres_idx_tensor = (pred_center_allpair_dist[box_idx] < thres).nonzero().flatten()
        within_thres_idx_list = within_thres_idx_tensor.cpu().numpy().tolist()

        if len(within_thres_idx_list) == 0:  # if it's a single box
            continue

        # start from within_thres_idx_list, find new box added to the cluster
        explored = [box_idx]
        unexplored = [idx for idx in within_thres_idx_list if idx in remain_box]

        while unexplored:
            idx = unexplored[0]
            within_thres_idx_tensor = (pred_center_allpair_dist[idx] < thres).nonzero().flatten()
            within_thres_idx_list = within_thres_idx_tensor.cpu().numpy().tolist()
            for newidx in within_thres_idx_list:
                if (newidx not in explored) and (newidx not in unexplored) and (newidx in remain_box):
                    unexplored.append(newidx)
            unexplored.remove(idx)
            explored.append(idx)
        
        if len(explored) == 1: # it's a single box, neighbors have been assigned
            remain_box.remove(box_idx)
            continue
        
        cluster_box_idxs = explored

        cluster_dict[cluster_id] = OrderedDict()
        cluster_dict[cluster_id]['box_idx'] = [idx for idx in cluster_box_idxs]
        cluster_dict[cluster_id]['box_dist'] = [pred_center_cat[idx].norm() for idx in cluster_box_idxs]  # distance to observer 
        cluster_dict[cluster_id]['box_center_world'] = [pred_center_world_cat[idx] for idx in cluster_box_idxs]  # coordinate in world, [3,]
        cluster_dict[cluster_id]['box_yaw'] = [pred_yaw_world_cat[idx] for idx in cluster_box_idxs]


        box_dist = torch.as_tensor(cluster_dict[cluster_id]['box_dist']).to(device)
        box_weight = F.normalize(1/box_dist, p=1, dim=0) # [n]
        centers = torch.stack(cluster_dict[cluster_id]['box_center_world'], dim=0) # [n, 3]
        yaws = torch.stack(cluster_dict[cluster_id]['box_yaw'])  # [n]

        weighted_center = torch.sum(box_weight.unsqueeze(-1) * centers, dim=0) # [3,]
        weighted_yaw = torch.sum(box_weight * yaws) # [1,]

        weighted_center[2] = weighted_yaw  # just replace z to yaw

        cluster_dict[cluster_id]['se2'] = weighted_center  # [x, y, yaw]

        # DEBUG = True
        if DEBUG:
            from icecream import ic
            ic(cluster_dict[cluster_id]['box_idx'])
            ic(centers)
            ic(yaws)
            ic(box_weight)
            ic(cluster_dict[cluster_id]['se2'])

        cluster_dict[cluster_id].pop('box_dist')
        cluster_dict[cluster_id].pop('box_center_world')
        cluster_dict[cluster_id].pop('box_yaw')

        cluster_id += 1
        for idx in cluster_box_idxs:
            remain_box.remove(idx)

    vertex_num = cluster_id
    agent_num = N
    landmark_num = cluster_id - N
    # ic(agent_num)
    # ic(landmark_num)

    """
        Now we have clusters for objects. we can create pose graph.
        First we consider center as landmark.
        Maybe set corner as landmarks in the future.
    """
    pgo = PoseGraphOptimization2D(verbose=False)
    if DEBUG:
        pgo = PoseGraphOptimization2D(verbose=True)
    # Add agent to vertexs
    for agent_id in range(N):
        v_id = agent_id
        # notice lidar_pose use degree format, translate it to radius.
        # pose_np = lidar_poses[agent_id, [0,1,4]].cpu().numpy()
        pose_np = noisy_lidar_poses[agent_id, [0,1,4]].cpu().numpy()
        pose_np[2] = np.deg2rad(pose_np[2])  # radius
        v_pose = g2o.SE2(pose_np)
        # if agent_id == 0 and DEBUG:
        #     pgo.add_vertex(id=v_id, pose=v_pose, fixed=True)
        # else:
        pgo.add_vertex(id=v_id, pose=v_pose, fixed=False)

    # Add landmark to vertexs
    for landmark_id in range(N, cluster_id):
        v_id = landmark_id
        v_pose = g2o.SE2(cluster_dict[landmark_id]['se2'].cpu().numpy())
        pgo.add_vertex(id=v_id, pose=v_pose, fixed=False)

    # Add agent-landmark edge to edge
    for landmark_id in range(N, cluster_id):
        for box_idx in cluster_dict[landmark_id]['box_idx']:
            agent_id = box_idx_to_agent[box_idx]
            e_pose = g2o.SE2(pred_box3d_cat[box_idx][[0,1,6]].cpu().numpy())
            pgo.add_edge(vertices=[agent_id, landmark_id], measurement=e_pose, information=np.identity(3))
    
    pgo.optimize()

    pose_new_list = []
    for agent_id in range(N):
        # print(pgo.get_pose(agent_id).vector())
        pose_new_list.append(pgo.get_pose(agent_id).vector())

    refined_pose = np.array(pose_new_list)
    refined_pose[:,2] = np.rad2deg(refined_pose[:,2])  # rad -> degree, same as source

    return refined_pose

def box_alignment(pred_corner3d_list, uncertainty_list, lidar_poses, record_len, proj_first=False):
    """
    Args:
        pred_corner3d_list: list of tensors, with shape [[N1_object, 8, 3], [N2_object, 8, 3], ...,[N_sumcav_object, 8, 3]]
            box in each agent's coordinate. (proj_first=False)
        
        pred_box3d_list: not necessary
            list of tensors, with shape [[N1_object, 7], [N2_object, 7], ...,[N_sumcav_object, 7]]

        scores_list: list of tensor, [[N1_object,], [N2_object,], ...,[N_sumcav_object,]]
            box confidence score.

        lidar_poses: torch.Tensor [sum(cav), 6]

        record_len: torch.Tensor
    Returns:
        refined_lidar_pose: torch.Tensor [sum(cav), 6]
    """
    refined_lidar_pose = []
    start_idx = 0
    for b in record_len:
        refined_lidar_pose.append(
            torch.from_numpy(
                box_alignment_relative_sample(
                    pred_corner3d_list[start_idx: start_idx + b],
                    lidar_poses[start_idx: start_idx + b],
                    clean_lidar_pose=None,
                    uncertainty_list= None if uncertainty_list is None else uncertainty_list[start_idx: start_idx + b]
                )
            )
        )
        start_idx += b

    return torch.cat(refined_lidar_pose, dim=0)

def vis_corners_list(corner3d_list, filename="/GPFS/rhome/yifanlu/OpenCOOD/opencood/corners.png"):
    """
    Args:
        corner3d: list of  torch.Tensor, shape [N, 8, 3]

    """
    COLOR = ['red','springgreen','dodgerblue', 'darkviolet']
    box_idx = 0

    for idx in range(len(corner3d_list)):
        corner3d = corner3d_list[idx]
        if torch.is_tensor(corner3d):
            corner3d = corner3d.cpu().numpy()

        corner2d = corner3d[:,:4,:2]
        import matplotlib.pyplot as plt
        for i in range(corner2d.shape[0]):
            plt.scatter(corner2d[i,[0,1],0], corner2d[i,[0,1], 1], s=2, c=COLOR[idx])
            plt.plot(corner2d[i,[0,1,2,3,0],0], corner2d[i,[0,1,2,3,0], 1], linewidth=1, c=COLOR[idx])
            plt.text(corner2d[i,0,0], corner2d[i,0,1], s=str(box_idx), fontsize="xx-small")
            box_idx += 1
    plt.gca().invert_yaxis()
    plt.axis('equal')
    plt.savefig(filename, dpi=400)
    plt.clf()

def vis_corners(corner3d, filename="/GPFS/rhome/yifanlu/OpenCOOD/opencood/corners.png"):
    """
    Args:
        corner3d: torch.Tensor, shape [N, 8, 3]

        box3d: torch.Tensor shape [N, 7]
    """
    if torch.is_tensor(corner3d):
        corner3d = corner3d.cpu().numpy()


    corner2d = corner3d[:,:4,:2]
    import matplotlib.pyplot as plt
    for i in range(corner2d.shape[0]):
        plt.scatter(corner2d[i,[0,1],0], corner2d[i,[0,1], 1], s=2)
        plt.plot(corner2d[i,[0,1,2,3,0],0], corner2d[i,[0,1,2,3,0], 1])
        # plt.text(corner2d[i,0,0], corner2d[i,0,1], s=f"{box3d[i,0]:.2f},{box3d[i,1]:.2f},{box3d[i,6]:.2f}", fontsize='xx-small')
    plt.axis('equal')
    plt.savefig(filename, dpi=300)
    plt.clf()

def vis_pose(lidar_poses):
    """
    Args:
        lidar_poses: torch.Tensor shape [N_, 6], x,y,z, roll, yaw, pitch
    """
    h = 1.56
    l = 3.9
    w = 1.6
    if torch.is_tensor(lidar_poses):
        lidar_poses = lidar_poses.cpu().numpy()

    box3d = np.zeros((lidar_poses.shape[0], 7))
    box3d[:,0] = lidar_poses[:,0]
    box3d[:,1] = lidar_poses[:,1]
    box3d[:,3] = h # hwl order
    box3d[:,4] = w
    box3d[:,5] = l
    box3d[:,6] = np.deg2rad(lidar_poses[:,4])  # degree -> radius

    corner3d = box_utils.boxes_to_corners_3d(box3d, order='hwl')
    vis_corners(corner3d, box3d, "/GPFS/rhome/yifanlu/OpenCOOD/opencood/pose_corners.png")

def test_pred_gt_box():
    gt_corners_list = torch.load("/GPFS/rhome/yifanlu/workspace/OpenCOOD/box_align_items/previous_items/gt_box_list.pt")
    data = torch.load("/GPFS/rhome/yifanlu/OpenCOOD/box_align_items/stage1_output_02/0.pt")
    pred_corner3d_list, pred_box3d_list, scores_list, record_len, lidar_pose, lidar_pose_clean = data

    lidar_pose_tfm = pose_to_tfm(lidar_pose, dof=6)
    lidar_pose_clean_tfm = pose_to_tfm(lidar_pose_clean, dof=6)  # Tw_c
    N = lidar_pose.shape[0]

    pred_corners_world_list = \
        [box_utils.project_box3d(pred_corner3d_list[i], lidar_pose_tfm[i]) for i in range(N)]  # [[N1, 8, 3], [N2, 8, 3],...]

    gt_corners_world_list = \
        [box_utils.project_box3d(gt_corners_list[i], lidar_pose_clean_tfm[i]) for i in range(N)]

    vis_corners_list([torch.cat(pred_corners_world_list, dim=0), torch.cat(gt_corners_world_list, dim=0)], filename="/GPFS/rhome/yifanlu/OpenCOOD/box_align_items/gt_box_pred_box.png")

    


def test_gt_boxes_world():
    data = torch.load("/GPFS/rhome/yifanlu/OpenCOOD/box_align_items/stage1_output_02/0.pt")
    pred_corner3d_list, pred_box3d_list, scores_list, record_len, lidar_pose, lidar_pose_clean = data
    
    gt_poses_tensor = lidar_pose_clean
    noisy_poses_tensor = lidar_pose

    gt_corners_list = torch.load("/GPFS/rhome/yifanlu/workspace/OpenCOOD/box_align_items/previous_items/gt_box_list.pt")
    
    # refined_poses = box_alignment_sample(gt_corners_list, noisy_poses_tensor, noisy_poses_tensor)
    refined_poses = box_alignment_relative_sample(pred_corner3d_list, noisy_poses_tensor, gt_poses_tensor)
    print("before:\n", noisy_poses_tensor.cpu().numpy()[:,[0,1,4]])
    
    print("after:\n", refined_poses)

    print("gt:\n", gt_poses_tensor.cpu().numpy()[:,[0,1,4]])

    # gt_corners_world_list = \
    #     [box_utils.project_box3d(gt_corners_list[i], lidar_pose_tfm[i]) for i in range(3)]  # [[N1, 8, 3], [N2, 8, 3],...]

    # vis_corners_list(gt_corners_world_list, filename="/GPFS/rhome/yifanlu/workspace/OpenCOOD/box_align_items/gt_corners.png")


def test_box_align_tmp():
    """
        This func input different noise_std pose (load from stored files).
        And run pose graph optimization, compare the localization error w/wo uncertainty/landmark SE2, etc.
    """
    noise_stds = ['02','04','06']
    items = ["16"]
    torch.set_printoptions(precision=3, sci_mode=False)
    np.set_printoptions(precision=3, suppress=True)
    for item in items:
        for noise_std in noise_stds:
            file_dir = f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/box_align_items/stage1_output_{noise_std}_w_uncertainty/{item}.pt"
            data = torch.load(file_dir)
            pred_corner3d_list, pred_box3d_list, uncertainty_list, record_len, lidar_pose, lidar_pose_clean = data
            lidar_pose[0] = lidar_pose_clean[0]
            refined_pose_SE2 = box_alignment_relative_sample(pred_corner3d_list, lidar_pose_clean, lidar_pose, uncertainty_list=uncertainty_list, landmark_SE2=True)
            refined_pose = box_alignment_relative_sample(pred_corner3d_list, lidar_pose_clean, lidar_pose, uncertainty_list=uncertainty_list, landmark_SE2=False)
            # refined_pose = box_alignment_sample(pred_corner3d_list, lidar_pose, lidar_pose)
            lidar_pose_clean = lidar_pose_clean[:,[0,1,4]].cpu().numpy()
            print(f"noise std: {noise_std}: SE2")
            print(np.abs(refined_pose_SE2 - lidar_pose_clean))
            # print(f"PointXY")
            # print(np.abs(refined_pose - lidar_pose_clean))
            print(f"original error:")
            lidar_pose = lidar_pose[:,[0,1,4]].cpu().numpy()
            print(np.abs(lidar_pose - lidar_pose_clean))
            # print(refined_pose_w_u)
            # print(lidar_pose_clean) 


def test_box_align(noise_std="04", relative=True, use_uncertainty=False):
    from glob import glob
    data_dir = f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/box_align_items/stage1_output_{noise_std}_w_uncertainty/*"
    trans_error_before_list = []
    rotate_error_before_list = []
    trans_error_after_list = []
    rotate_error_after_list = []

    full_files = glob(data_dir)
    for filename in full_files:
        data = torch.load(filename)
        
        if relative is False:
            pred_corner3d_list, pred_box3d_list, scores_list, record_len, lidar_pose, lidar_pose_clean = data
            refined_pose = box_alignment_sample(pred_corner3d_list, None, lidar_pose)
        elif relative is True:
            pred_corner3d_list, pred_box3d_list, uncertainty_list, record_len, lidar_pose, lidar_pose_clean = data
            lidar_pose[0] = lidar_pose_clean[0]
            # if not use_uncertainty:
            #     uncertainty_list = None
            refined_pose = box_alignment_relative_sample(pred_corner3d_list, lidar_pose_clean, lidar_pose, uncertainty_list=uncertainty_list)
            uncertainty_list = None
            refined_pose_wo_uncertainty = box_alignment_relative_sample(pred_corner3d_list, lidar_pose_clean, lidar_pose, uncertainty_list=uncertainty_list)

        lidar_pose = lidar_pose.cpu().numpy()[:,[0,1,4]]
        lidar_pose_clean = lidar_pose_clean.cpu().numpy()[:,[0,1,4]]
        np.set_printoptions(suppress=True, precision=4)
        print(lidar_pose[1:])
        print(refined_pose_wo_uncertainty[1:])
        print(refined_pose[1:])
        print(lidar_pose_clean[1:])
        print()

        error_before = np.abs(lidar_pose - lidar_pose_clean)
        error_after = np.abs(refined_pose - lidar_pose_clean)

        trans_error_before_list.append(np.mean(error_before[:,[0,1]]))
        rotate_error_before_list.append(np.mean(error_before[:,2]))

        trans_error_after_list.append(np.mean(error_after[:,[0,1]]))
        rotate_error_after_list.append(np.mean(error_after[:,2]))

    raise


    out_quantile_dict = {0.8:None, 0.5:None, 0.3:None}
    for q in out_quantile_dict.keys():
        out_quantile_dict[q] = (np.quantile(trans_error_before_list, q), 
                                    np.quantile(trans_error_after_list, q), 
                                    np.quantile(rotate_error_before_list, q),
                                    np.quantile(rotate_error_after_list, q))

    return out_quantile_dict
    # return np.mean(trans_error_before_list), np.mean(rotate_error_before_list), np.mean(trans_error_after_list), np.mean(rotate_error_after_list)

def main1():
    """
    This function test the box alignment performance on the subset of training set.
    """
    for noise in ['02', '04', '06']:
        out = test_box_align(noise, relative=True, use_uncertainty=True)
        for k,v in out.items():
            with open(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/box_align_items/rel_quantile{k*100}_{noise}_w_u.txt", "w") as f:
                f.write(f"trans error before: \t {v[0]}\n")
                f.write(f"trans error after:  \t {v[1]}\n\n")

                f.write(f"rotate error before: \t {v[2]}\n")
                f.write(f"rotate error after: \t {v[3]}\n")


        out = test_box_align(noise, relative=True, use_uncertainty=False)
        for k,v in out.items():
            with open(f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/box_align_items/rel_quantile{k*100}_{noise}_wo_u.txt", "w") as f:
                f.write(f"trans error before: \t {v[0]}\n")
                f.write(f"trans error after:  \t {v[1]}\n\n")

                f.write(f"rotate error before: \t {v[2]}\n")
                f.write(f"rotate error after: \t {v[3]}\n")


def vis_pose_graph(
            poses,
            pred_corner3d, 
            save_dir_path="/GPFS/rhome/yifanlu/OpenCOOD/box_align_items/pose_graph_vis",
            ):
    """
    Args:
        poses: list of np.ndarray
            each item is a pose . [pose_before, ..., pose_refined]

        pred_corner3d: list
            predicted box for each agent.

    """
    COLOR = ['red','springgreen','dodgerblue', 'darkviolet', 'orange']
    from opencood.utils.transformation_utils import get_relative_transformation

    if not os.path.exists(save_dir_path):
        os.mkdir(save_dir_path)

    for iter, pose in enumerate(poses):
        box_idx = 0
        # we first transform other agents' box to ego agent's coordinate
        relative_t_matrix = get_relative_transformation(pose)
        N = pose.shape[0]

        pred_corners3d_in_ego = [box_utils.project_box3d(pred_corner3d[i].cpu().numpy(), relative_t_matrix[i]) for i in range(N)]

        for agent_id in range(len(pred_corners3d_in_ego)):
            corner3d = pred_corners3d_in_ego[agent_id]
            agent_pos = relative_t_matrix[agent_id][:2,3] # agent's position in ego's coordinate
            if torch.is_tensor(corner3d):
                corner3d = corner3d.cpu().numpy()

            corner2d = corner3d[:,:4,:2]
            center2d = np.mean(corner2d, axis=1)
            import matplotlib.pyplot as plt
            for i in range(corner2d.shape[0]):
                plt.scatter(corner2d[i,[0,1],0], corner2d[i,[0,1], 1], s=2, c=COLOR[agent_id])
                plt.plot(corner2d[i,[0,1,2,3,0],0], corner2d[i,[0,1,2,3,0], 1], linewidth=1, c=COLOR[agent_id])
                plt.text(corner2d[i,0,0], corner2d[i,0,1], s=str(box_idx), fontsize="xx-small")
                # add a line connecting box center and agent.
                box_center = center2d[i] # [2,]
                connection_x = [agent_pos[0], box_center[0]]
                connection_y = [agent_pos[1], box_center[1]]
                # print(connection_x)
                # print(connection_y)
                # print()
                plt.plot(connection_x, connection_y,'--', linewidth=0.5, c=COLOR[agent_id], alpha=0.3)
                box_idx += 1
        
        filename = os.path.join(save_dir_path, f"{iter}.png")
        plt.gca().invert_yaxis()
        plt.axis('equal')
        plt.savefig(filename, dpi=400)
        plt.clf()

def vis_pose_graphs():
    noise_stds = ['02','04','06']
    items = ["53", "63", "73", "83"]
    torch.set_printoptions(precision=3, sci_mode=False)
    np.set_printoptions(precision=3, suppress=True)
    for item in items:
        for noise_std in noise_stds:
            file_dir = f"/GPFS/rhome/yifanlu/workspace/OpenCOOD/box_align_items/stage1_output_{noise_std}_w_uncertainty/{item}.pt"
            data = torch.load(file_dir)
            pred_corner3d_list, pred_box3d_list, uncertainty_list, record_len, lidar_pose, lidar_pose_clean = data
            lidar_pose[0] = lidar_pose_clean[0]
            refined_pose_SE2 = box_alignment_relative_sample(pred_corner3d_list, lidar_pose_clean, lidar_pose, uncertainty_list=uncertainty_list, landmark_SE2=True)
            ## visualize pred_corner3d with refined_pose. We can set different iteration to animate
            save_dir_path = f"/GPFS/rhome/yifanlu/OpenCOOD/box_align_items/pose_graph_vis/{item}_{noise_std}"
            poses = [lidar_pose.cpu().numpy(), refined_pose_SE2]
            vis_pose_graph(poses, pred_corner3d_list, save_dir_path)




def main2():
    pass

if __name__ == "__main__":
    # vis_pose_graphs()
    test_box_align_tmp()
    # main1()
    # test_gt_boxes_world()
opencood/models/sub_modules/box_align.py
MediaBrain-SJTU/CoBEVFlow
"""
This module is designed for box alignment

We will use g2o for pose graph optimization.
"""


from opencood.models.sub_modules.pose_graph_optim import PoseGraphOptimization2D
from opencood.utils.transformation_utils import pose_to_tfm
from opencood.utils.common_utils import check_torch_to_numpy
from opencood.utils import box_utils
from collections import OrderedDict
import numpy as np
import torch
import torch.nn.functional as F
import g2o
from icecream import ic
import copy
import os

DEBUG = False

def vis_pose_graph(poses, pred_corner3d, save_dir_path):
    """
    Args:
        poses: list of np.ndarray
            each item is a pose . [pose_before, ..., pose_refined]

        pred_corner3d: list
            predicted box for each agent.

    """
    COLOR = ['red','springgreen','dodgerblue', 'darkviolet', 'orange']
    from opencood.utils.transformation_utils import get_relative_transformation

    if not os.path.exists(save_dir_path):
        os.makedirs(save_dir_path)

    for iter, pose in enumerate(poses):
        box_idx = 0
        # we first transform other agents' box to ego agent's coordinate
        relative_t_matrix = get_relative_transformation(pose)
        N = pose.shape[0]
        nonempty_indices = [idx for (idx, corners) in enumerate(pred_corner3d) if len(corners)!=0]
        pred_corners3d_in_ego = [box_utils.project_box3d(pred_corner3d[i], relative_t_matrix[i]) for i in nonempty_indices]

        for agent_id in range(len(pred_corners3d_in_ego)):
            if agent_id not in nonempty_indices:
                continue
            corner3d = pred_corners3d_in_ego[agent_id]
            agent_pos = relative_t_matrix[agent_id][:2,3] # agent's position in ego's coordinate

            corner2d = corner3d[:,:4,:2]
            center2d = np.mean(corner2d, axis=1)
            import matplotlib.pyplot as plt
            for i in range(corner2d.shape[0]):
                plt.scatter(corner2d[i,[0,1],0], corner2d[i,[0,1], 1], s=2, c=COLOR[agent_id])
                plt.plot(corner2d[i,[0,1,2,3,0],0], corner2d[i,[0,1,2,3,0], 1], linewidth=1, c=COLOR[agent_id])
                plt.text(corner2d[i,0,0], corner2d[i,0,1], s=str(box_idx), fontsize="xx-small")
                # add a line connecting box center and agent.
                box_center = center2d[i] # [2,]
                connection_x = [agent_pos[0], box_center[0]]
                connection_y = [agent_pos[1], box_center[1]]

                plt.plot(connection_x, connection_y,'--', linewidth=0.5, c=COLOR[agent_id], alpha=0.3)
                box_idx += 1
        
        filename = os.path.join(save_dir_path, f"{iter}.png")
        plt.gca().invert_yaxis()
        plt.axis('equal')
        plt.savefig(filename, dpi=400)
        plt.clf()


def all_pair_l2(A, B):
    """ All pair L2 distance for A and B
    Args:
        A : np.ndarray
            shape [N_A, D]
        B : np.ndarray
            shape [N_B, D]
    Returns:
        C : np.ndarray
            shape [N_A, N_B]
    """
    TwoAB = 2*A@B.T  # [N_A, N_B]
    C = np.sqrt(
              np.sum(A * A, 1, keepdims=True).repeat(TwoAB.shape[1], axis=1) \
            + np.sum(B * B, 1, keepdims=True).T.repeat(TwoAB.shape[0], axis=0) \
            - TwoAB
        )
    return C




def box_alignment_relative_sample_np(
            pred_corners_list,
            noisy_lidar_pose, 
            uncertainty_list=None, 
            landmark_SE2=True,
            adaptive_landmark=False,
            normalize_uncertainty=False,
            abandon_hard_cases = False,
            drop_hard_boxes = False,
            drop_unsure_edge = False,
            use_uncertainty = True,
            thres = 1.,
            yaw_var_thres = 0.2):
    """ Perform box alignment for one sample. 
    Correcting the relative pose.

    Args:
        pred_corners_list: in each ego coordinate
            [[N_1, 8, 3], ..., [N_cav1, 8, 3]]

        clean_lidar_poses:
            [N_cav1, 6], in degree
        
        noisy_lidar_poses:
            [N_cav1, 6], in degree

        uncertainty_list:
            [[N_1, 3], [N_2, 3], ..., [N_cav1, 3]]

        landmark_SE2:
            if True, the landmark is SE(2), otherwise R^2
        
        adaptive_landmark: (when landmark_SE2 = True)
            if True, landmark will turn to R^2 if yaw angles differ a lot

        normalize_uncertainty: bool
            if True, normalize the uncertainty
        
        abandon_hard_cases: bool
            if True, algorithm will just return original poses for hard cases

        drop_unsure_edge: bool

    Returns: 
        refined_lidar_poses: np.ndarray
            [N_cav1, 3], 
    """
    if not use_uncertainty:
        uncertainty_list = None
    ## first transform point from ego coordinate to world coordinate, using lidar_pose.
    order = 'lwh'  # hwl
    N = noisy_lidar_pose.shape[0]
    lidar_pose_noisy_tfm = pose_to_tfm(noisy_lidar_pose)

    nonempty_indices = [idx for (idx, corners) in enumerate(pred_corners_list) if len(corners)!=0] # if one agent detects no boxes, its corners is just [].
    
    pred_corners_world_list = \
        [box_utils.project_box3d(pred_corners_list[i], lidar_pose_noisy_tfm[i]) for i in nonempty_indices]  # [[N1, 8, 3], [N2, 8, 3],...]
    pred_box3d_list = \
        [box_utils.corner_to_center(corner, order) for corner in pred_corners_list if len(corner)!=0]   # [[N1, 7], [N2, 7], ...], angle in radian
    pred_box3d_world_list = \
        [box_utils.corner_to_center(corner, order) for corner in pred_corners_world_list]   # [[N1, 7], [N2, 7], ...], angle in radian
    pred_center_list = \
        [np.mean(corners, axis=1) for corners in pred_corners_list if len(corners)!=0] # [[N1,3], [N2,3], ...]

    pred_center_world_list = \
        [pred_box3d_world[:,:3] for pred_box3d_world in pred_box3d_world_list]
    pred_yaw_world_list = \
        [pred_box3d[:, 6] for pred_box3d in pred_box3d_world_list]
    pred_len = \
        [len(corners) for corners in pred_corners_list] 


    box_idx_to_agent = []
    for i in range(N):
        box_idx_to_agent += [i] * pred_len[i] 
    
    pred_center_cat = np.concatenate(pred_center_list, axis=0)   # [sum(pred_box), 3]
    pred_center_world_cat =  np.concatenate(pred_center_world_list, axis=0)  # [sum(pred_box), 3]
    pred_box3d_cat =  np.concatenate(pred_box3d_list, axis=0)  # [sum(pred_box), 7]
    pred_yaw_world_cat = np.concatenate(pred_yaw_world_list, axis=0)  # [sum(pred_box)]

    # hard-coded currently
    w_a = 1.6 # width of anchor
    l_a = 3.9 # length of anchor
    d_a_square = w_a ** 2 + l_a ** 2 # anchor's diag


    if uncertainty_list is not None:
        pred_log_sigma2_cat = np.concatenate([i for i in uncertainty_list if len(i)!=0], axis=0)
        # Since the regression target is x_t = (x_g - x_a)/d_a, 
        # var(x) = d_a^2 * var(x_t)
        # so we 1/var(x) = 1/var(x_t) / d_a^2  
        # sigma_{delta_x}^2 -> sigma_x^2. 
        pred_certainty_cat = np.exp(-pred_log_sigma2_cat)
        pred_certainty_cat[:,:2] /= d_a_square 


        if normalize_uncertainty:
            pred_certainty_cat = np.sqrt(pred_certainty_cat)


    pred_center_allpair_dist = all_pair_l2(pred_center_world_cat, pred_center_world_cat) # [sum(pred_box), sum(pred_box)]

    # let pair from one vehicle be max distance
    MAX_DIST = 10000
    cum = 0
    for i in range(N):
        pred_center_allpair_dist[cum: cum + pred_len[i], cum: cum +pred_len[i]] = MAX_DIST   # do not include itself
        cum += pred_len[i]


    cluster_id = N # let the vertex id of object start from N
    cluster_dict = OrderedDict()
    remain_box = set(range(cum))

    for box_idx in range(cum): 

        if box_idx not in remain_box:  # already assigned
            continue
        
        within_thres_idx_tensor = (pred_center_allpair_dist[box_idx] < thres).nonzero()[0]
        within_thres_idx_list = within_thres_idx_tensor.tolist()

        if len(within_thres_idx_list) == 0:  # if it's a single box
            continue

        # start from within_thres_idx_list, find new box added to the cluster
        explored = [box_idx]
        unexplored = [idx for idx in within_thres_idx_list if idx in remain_box]

        while unexplored:
            idx = unexplored[0]
            within_thres_idx_tensor = (pred_center_allpair_dist[box_idx] < thres).nonzero()[0]
            within_thres_idx_list = within_thres_idx_tensor.tolist()
            for newidx in within_thres_idx_list:
                if (newidx not in explored) and (newidx not in unexplored) and (newidx in remain_box):
                    unexplored.append(newidx)
            unexplored.remove(idx)
            explored.append(idx)
        
        if len(explored) == 1: # it's a single box, neighbors have been assigned
            remain_box.remove(box_idx)
            continue
        
        cluster_box_idxs = explored

        cluster_dict[cluster_id] = OrderedDict()
        cluster_dict[cluster_id]['box_idx'] = [idx for idx in cluster_box_idxs]
        cluster_dict[cluster_id]['box_center_world'] = [pred_center_world_cat[idx] for idx in cluster_box_idxs]  # coordinate in world, [3,]
        cluster_dict[cluster_id]['box_yaw'] = [pred_yaw_world_cat[idx] for idx in cluster_box_idxs]

        yaw_var = np.var(cluster_dict[cluster_id]['box_yaw'])
        cluster_dict[cluster_id]['box_yaw_varies'] = yaw_var > yaw_var_thres
        cluster_dict[cluster_id]['active'] = True


        ########### adaptive_landmark ##################
        if landmark_SE2:
            if adaptive_landmark and yaw_var > yaw_var_thres:
                landmark = pred_center_world_cat[box_idx][:2]
                for _box_idx in cluster_box_idxs:
                    pred_certainty_cat[_box_idx] *= 2
            else:
                landmark = copy.deepcopy(pred_center_world_cat[box_idx])
                landmark[2] = pred_yaw_world_cat[box_idx]
        else:
            landmark = pred_center_world_cat[box_idx][:2]
        ##################################################


        cluster_dict[cluster_id]['landmark'] = landmark  # [x, y, yaw] or [x, y]
        cluster_dict[cluster_id]['landmark_SE2'] = True if landmark.shape[0] == 3 else False

        DEBUG = False
        if DEBUG:
            from icecream import ic
            ic(cluster_dict[cluster_id]['box_idx'])
            ic(cluster_dict[cluster_id]['box_center_world'])
            ic(cluster_dict[cluster_id]['box_yaw'])
            ic(cluster_dict[cluster_id]['landmark'])
        

        cluster_id += 1
        for idx in cluster_box_idxs:
            remain_box.remove(idx)

    
    vertex_num = cluster_id
    agent_num = N
    landmark_num = cluster_id - N


    ########### abandon_hard_cases ##########
    """
        We should think what is hard cases for agent-object pose graph optimization
            1. Overlapping boxes are rare (landmark_num <= 3)
            2. Yaw angles differ a lot
    """

    if abandon_hard_cases:
        # case1: object num is smaller than 3
        if landmark_num <= 3:
            return noisy_lidar_pose[:,[0,1,4]]
        
        # case2: more than half of the landmarks yaw varies 
        yaw_varies_cnt = sum([cluster_dict[i]["box_yaw_varies"] for i in range(agent_num, vertex_num)])
        if yaw_varies_cnt >= 0.5 * landmark_num:
            return noisy_lidar_pose[:,[0,1,4]]

    ########### drop hard boxes ############

    if drop_hard_boxes:
        for landmark_id in range(agent_num, vertex_num):
            if cluster_dict[landmark_id]['box_yaw_varies']:
                cluster_dict[landmark_id]['active'] = False




    """
        Now we have clusters for objects. we can create pose graph.
        First we consider center as landmark.
        Maybe set corner as landmarks in the future.
    """
    pgo = PoseGraphOptimization2D(verbose=False)

    # Add agent to vertexs
    for agent_id in range(agent_num):
        v_id = agent_id
        # notice lidar_pose use degree format, translate it to radians.
        pose_np = noisy_lidar_pose[agent_id, [0,1,4]]
        pose_np[2] = np.deg2rad(pose_np[2])  # radians
        v_pose = g2o.SE2(pose_np)
        
        if agent_id == 0:
            pgo.add_vertex(id=v_id, pose=v_pose, fixed=True)
        else:
            pgo.add_vertex(id=v_id, pose=v_pose, fixed=False)

    # Add object to vertexs
    for landmark_id in range(agent_num, vertex_num):
        v_id = landmark_id
        landmark = cluster_dict[landmark_id]['landmark'] # (3,) or (2,)
        landmark_SE2 = cluster_dict[landmark_id]['landmark_SE2']

        if landmark_SE2:
            v_pose = g2o.SE2(landmark)
        else:
            v_pose = landmark

        pgo.add_vertex(id=v_id, pose=v_pose, fixed=False, SE2=landmark_SE2)

    # Add agent-object edge to edge set
    for landmark_id in range(agent_num, vertex_num):
        landmark_SE2 = cluster_dict[landmark_id]['landmark_SE2']

        if not cluster_dict[landmark_id]['active']:
            continue

        for box_idx in cluster_dict[landmark_id]['box_idx']:
            agent_id = box_idx_to_agent[box_idx]
            if landmark_SE2:
                e_pose = g2o.SE2(pred_box3d_cat[box_idx][[0,1,6]].astype(np.float64))
                info = np.identity(3, dtype=np.float64)
                if uncertainty_list is not None:
                    info[[0,1,2],[0,1,2]] = pred_certainty_cat[box_idx]

                    ############ drop_unsure_edge ###########
                    if drop_unsure_edge and sum(pred_certainty_cat[box_idx]) < 100:
                        continue

            else:
                e_pose = pred_box3d_cat[box_idx][[0,1]].astype(np.float64)
                info = np.identity(2, dtype=np.float64)
                if uncertainty_list is not None:
                    info[[0,1],[0,1]] = pred_certainty_cat[box_idx][:2]

                    ############ drop_unsure_edge ############
                    if drop_unsure_edge and sum(pred_certainty_cat[box_idx]) < 100:
                        continue

            pgo.add_edge(vertices=[agent_id, landmark_id], measurement=e_pose, information=info, SE2=landmark_SE2)
    
    pgo.optimize()

    pose_new_list = []
    for agent_id in range(agent_num):
        # print(pgo.get_pose(agent_id).vector())
        pose_new_list.append(pgo.get_pose(agent_id).vector())

    refined_pose = np.array(pose_new_list)
    refined_pose[:,2] = np.rad2deg(refined_pose[:,2])  # rad -> degree, same as source

    return refined_pose

def box_alignment_relative_np(pred_corner3d_list, 
                              uncertainty_list, 
                              lidar_poses, 
                              record_len, 
                              **kwargs):
    """
    Args:
        pred_corner3d_list: list of tensors, with shape [[N1_object, 8, 3], [N2_object, 8, 3], ...,[N_sumcav_object, 8, 3]]
            box in each agent's coordinate. (proj_first=False)
        
        pred_box3d_list: not necessary
            list of tensors, with shape [[N1_object, 7], [N2_object, 7], ...,[N_sumcav_object, 7]]

        scores_list: list of tensor, [[N1_object,], [N2_object,], ...,[N_sumcav_object,]]
            box confidence score.

        lidar_poses: torch.Tensor [sum(cav), 6]

        record_len: torch.Tensor
    Returns:
        refined_lidar_pose: torch.Tensor [sum(cav), 6]
    """
    refined_lidar_pose = []
    start_idx = 0
    for b in record_len:
        refined_lidar_pose.append(
            box_alignment_relative_sample_np(
                pred_corner3d_list[start_idx: start_idx + b],
                lidar_poses[start_idx: start_idx + b],
                uncertainty_list= None if uncertainty_list is None else uncertainty_list[start_idx: start_idx + b],
                **kwargs
            )
        )
        start_idx += b

    return np.cat(refined_lidar_pose, axis=0)



opencood/models/sub_modules/box_align_v2.py
MediaBrain-SJTU/CoBEVFlow
import torch
from torch import nn



class SSFA(nn.Module):
    def __init__(self, args):
        super(SSFA, self).__init__()
        self._num_input_features = args['feature_num']  # 128

        seq = [nn.ZeroPad2d(1)]
        seq += get_conv_layers('Conv2d', 128, 128, n_layers=3, kernel_size=[3, 3, 3],
                                                  stride=[1, 1, 1], padding=[0, 1, 1], sequential=False)
        self.bottom_up_block_0 = nn.Sequential(*seq)
        self.bottom_up_block_1 = get_conv_layers('Conv2d', 128, 256, n_layers=3, kernel_size=[3, 3, 3],
                                                  stride=[2, 1, 1], padding=[1, 1, 1])

        self.trans_0 = get_conv_layers('Conv2d', 128, 128, n_layers=1, kernel_size=[1], stride=[1], padding=[0])
        self.trans_1 = get_conv_layers('Conv2d', 256, 256, n_layers=1, kernel_size=[1], stride=[1], padding=[0])

        self.deconv_block_0 = get_conv_layers('ConvTranspose2d', 256, 128, n_layers=1, kernel_size=[3], stride=[2],
                                              padding=[1], output_padding=[1])
        self.deconv_block_1 = get_conv_layers('ConvTranspose2d', 256, 128, n_layers=1, kernel_size=[3], stride=[2],
                                              padding=[1], output_padding=[1])

        self.conv_0 = get_conv_layers('Conv2d', 128, 128, n_layers=1, kernel_size=[3], stride=[1], padding=[1])
        self.conv_1 = get_conv_layers('Conv2d', 128, 128, n_layers=1, kernel_size=[3], stride=[1], padding=[1])

        self.w_0 = get_conv_layers('Conv2d', 128, 1, n_layers=1, kernel_size=[1], stride=[1], padding=[0], relu_last=False)
        self.w_1 = get_conv_layers('Conv2d', 128, 1, n_layers=1, kernel_size=[1], stride=[1], padding=[0], relu_last=False)

    # default init_weights for conv(msra) and norm in ConvModule
    def init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.xavier_normal_(m.weight, gain=1)
                if hasattr(m, "bias") and m.bias is not None:
                    nn.init.constant_(m.bias, 0)

    def forward(self, x):
        x_0 = self.bottom_up_block_0(x)
        x_1 = self.bottom_up_block_1(x_0)
        x_trans_0 = self.trans_0(x_0)
        x_trans_1 = self.trans_1(x_1)
        x_middle_0 = self.deconv_block_0(x_trans_1) + x_trans_0
        x_middle_1 = self.deconv_block_1(x_trans_1)
        x_output_0 = self.conv_0(x_middle_0)
        x_output_1 = self.conv_1(x_middle_1)

        x_weight_0 = self.w_0(x_output_0)
        x_weight_1 = self.w_1(x_output_1)
        x_weight = torch.softmax(torch.cat([x_weight_0, x_weight_1], dim=1), dim=1)
        x_output = x_output_0 * x_weight[:, 0:1, :, :] + x_output_1 * x_weight[:, 1:, :, :]

        return x_output.contiguous()


def get_conv_layers(conv_name, in_channels, out_channels, n_layers, kernel_size, stride,
                    padding, relu_last=True, sequential=True, **kwargs):
    """
    Build convolutional layers. kernel_size, stride and padding should be a list with the lengths that match n_layers
    """
    seq = []
    for i in range(n_layers):
        seq.extend([getattr(nn, conv_name)(in_channels, out_channels, kernel_size[i], stride=stride[i],
                                           padding=padding[i], bias=False, **{k: v[i] for k, v in kwargs.items()}),
                    nn.BatchNorm2d(out_channels, eps=1e-3, momentum=0.01)])
        if i < n_layers - 1 or relu_last:
            seq.append(nn.ReLU())
        in_channels = out_channels
    if sequential:
        return nn.Sequential(*seq)
    else:
        return seq


class Head(nn.Module):
    def __init__(self, num_input, num_pred, num_cls, num_iou=2, use_dir=False, num_dir=1):
        super(Head, self).__init__()
        self.use_dir = use_dir

        self.conv_box = nn.Conv2d(num_input, num_pred, 1)  # 128 -> 14
        self.conv_cls = nn.Conv2d(num_input, num_cls, 1)   # 128 -> 2
        self.conv_iou = nn.Conv2d(num_input, num_iou, 1, bias=False)

        if self.use_dir:
            self.conv_dir = nn.Conv2d(num_input, num_dir, 1)  # 128 -> 4

    def forward(self, x):
        box_preds = self.conv_box(x)
        cls_preds = self.conv_cls(x)
        ret_dict = {"box_preds": box_preds, "cls_preds": cls_preds}
        if self.use_dir:
            dir_preds = self.conv_dir(x)  # dir_preds.shape=[8, w, h, 4]
            ret_dict["dir_cls_preds"] = dir_preds
        else:
            ret_dict["dir_cls_preds"] = torch.zeros((len(box_preds), 1, 2))

        ret_dict["iou_preds"] = self.conv_iou(x)

        return ret_dict
opencood/models/sub_modules/cia_ssd_utils.py
MediaBrain-SJTU/CoBEVFlow
# from mmcv.ops import DeformConv2dPack as DCN
# import torch
# import torch.nn as nn

# class DCNNet(nn.Module):
#     def __init__(self, args):
#         super(DCNNet,self).__init__()

#         module_list =[]
#         in_channels = args['in_channels']
#         out_channels = args['out_channels']
#         stride = args['stride']
#         kernel_size = args['kernel_size']
#         padding = args['padding']

#         for i in range(args['n_blocks']):
#             module_list.append(DCN(in_channels[i],out_channels[i],kernel_size[i],stride=stride[i],padding=padding[i]))
#         self.model = nn.Sequential(*module_list)

#     def forward(self, x):
#         return self.model(x)
opencood/models/sub_modules/dcn_net.py
MediaBrain-SJTU/CoBEVFlow
import numpy as np
import torch
import math
import torch.nn as nn
from opencood.models.sub_modules.resblock import ResNetModified, BasicBlock, Bottleneck
from opencood.models.sub_modules.detr_module import PositionEmbeddingSine, \
                DeformableTransformerEncoderLayer, DeformableTransformerEncoder
from opencood.models.fuse_modules.self_attn import AttFusion
from opencood.models.fuse_modules.deform_fuse import DeformFusion
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple


DEBUG = True


"""
    Different from MaxFusion in max_fuse.py
    This is a simplified version.
    pairwise_t_matrix is already scaled.
"""
def regroup(x, record_len):
    cum_sum_len = torch.cumsum(record_len, dim=0)
    split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
    return split_x

class MaxFusion(nn.Module):
    def __init__(self):
        super().__init__()
    def forward(self, x, record_len, pairwise_t_matrix):
        """
        pairwise_t_matrix is already normalized [B, L, L, 2, 3]
        """
        split_x = regroup(x, record_len)
        batch_size = len(record_len)
        C, H, W = split_x[0].shape[1:]  # C, W, H before
        out = []
        for b, xx in enumerate(split_x):
            N = xx.shape[0]
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]
            i = 0
            xx = warp_affine_simple(xx, t_matrix[i, :, :, :], (H, W))

            h = torch.max(xx, dim=0)[0]  # C, W, H before
            out.append(h)
        return torch.stack(out, dim=0)




class DeformableTransformerBackbone(nn.Module):
    def __init__(self, model_cfg):
        super().__init__()
        self.model_cfg = model_cfg
        self.compress = False

        
        self.proj_first = True
        if ('proj_first' in model_cfg) and (model_cfg['proj_first'] is False):
            self.proj_first = False
            self.discrete_ratio = model_cfg['voxel_size'][0]
            self.downsample_rate = 1

        self.level_num = len(model_cfg['layer_nums']) # exactly 3 now

        layer_nums = model_cfg['layer_nums']
        num_filters = model_cfg['num_filters']
        layer_strides = model_cfg['layer_strides']
        hidden_dim = model_cfg['hidden_dim']
        upsample_strides = model_cfg['upsample_strides']
        num_upsample_filters = model_cfg['num_upsample_filter']

        self.resnet = ResNetModified(BasicBlock, 
                                        layer_nums,
                                        layer_strides,
                                        num_filters)

        self.position_embedding = PositionEmbeddingSine(hidden_dim//2)

        self.hidden_dim = hidden_dim

        if model_cfg['fusion'] == 'max':
            self.fuse_net = [MaxFusion() for _ in range(self.level_num)]
        elif model_cfg['fusion'] == 'self_att':
            self.fuse_net = [AttFusion(n_filter) for n_filter in num_filters]
        elif model_cfg['fusion'] == 'deform':
            self.fuse_net = DeformFusion(num_filters[0], model_cfg['deform_method'])
        elif model_cfg['fusion'] == 'deform_w_cycle':
            assert self.proj_first is False
            assert model_cfg['deform_method'] == 'rigid'
            self.fuse_net = DeformFusion(num_filters[0], model_cfg['deform_method'], cycle_consist_loss=True)
        else:
            raise

        input_proj_list = []
        for i in range(self.level_num):
            proj_in_channels = num_filters[i]
            input_proj_list.append(nn.Sequential(
                nn.Conv2d(proj_in_channels, self.hidden_dim, kernel_size=1),
                nn.GroupNorm(32, self.hidden_dim),
            ))

        self.input_proj = nn.ModuleList(input_proj_list)
        self.level_embed = nn.Parameter(torch.Tensor(self.level_num, self.hidden_dim))
        self.upsample_strides = model_cfg['upsample_strides']

        encoder_layer = DeformableTransformerEncoderLayer(self.hidden_dim, model_cfg['dim_feedforward'],
                                                          model_cfg['dropout'], model_cfg['activation'],
                                                          self.level_num, model_cfg['n_head'], model_cfg['enc_n_points'])
        self.encoder = DeformableTransformerEncoder(encoder_layer, model_cfg['num_encoder_layers'])

        self.deblocks = nn.ModuleList()
        for idx in range(self.level_num):
            self.deblocks.append(nn.Sequential(
                nn.ConvTranspose2d(
                    self.hidden_dim, num_upsample_filters[idx],
                    upsample_strides[idx],
                    stride=upsample_strides[idx], bias=False
                ),
                nn.BatchNorm2d(num_upsample_filters[idx],
                                eps=1e-3, momentum=0.01),
                nn.ReLU()
                ))

    def forward(self, data_dict):
        spatial_features = data_dict['spatial_features']
        if DEBUG:
            origin_features = torch.clone(spatial_features)

        record_len = data_dict['record_len']
        pairwise_t_matrix = data_dict['pairwise_t_matrix']
        lidar_pose = data_dict['lidar_pose'] # (sum(cav),6 )

        ups = []
        ret_dict = {}
        x = spatial_features

        B = len(record_len)
        H, W = x.shape[2:]  ## this is original feature map [200, 704], not downsampled
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]

        if not self.proj_first:
            pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
            pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
            pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
            pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        features = self.resnet(x)  # feature[i] is (sum(cav), C, H, W), different i, different C, H, W

        if self.model_cfg['fusion'].startswith('deform'):
            x_fuseds = self.fuse_net(features, record_len, pairwise_t_matrix, lidar_pose)
        else:
            x_fuseds = [self.fuse_net[i](features[i], record_len, pairwise_t_matrix) for i in range(len(features))]

        pos_embeds = [self.position_embedding(x_fused) for x_fused in x_fuseds]
        srcs = [self.input_proj[i](x_fuseds[i]) for i in range(len(x_fuseds))]


        # srcs = []
        # pos_embeds = []
        # for i, feat in enumerate(features):
        #     x_fused = self.fuse_net[i](feat, record_len, pairwise_t_matrix)
        #     x_pos = self.position_embedding(x_fused)
        #     x_fused = self.input_proj[i](x_fused)
        #     srcs.append(x_fused)  # (B, hidden_dim, H1, W1)
        #     pos_embeds.append(x_pos)


        src_flatten = []
        mask_flatten = []
        lvl_pos_embed_flatten = []
        spatial_shapes = []
        for lvl, (src, pos_embed) in enumerate(zip(srcs, pos_embeds)):
            bs, c, h, w = src.shape
            spatial_shape = (h, w)
            spatial_shapes.append(spatial_shape)
            src = src.flatten(2).transpose(1, 2)
            pos_embed = pos_embed.flatten(2).transpose(1, 2)
            
            lvl_pos_embed = pos_embed + self.level_embed[lvl].view(1, 1, -1)
            lvl_pos_embed_flatten.append(lvl_pos_embed)
            src_flatten.append(src)
        src_flatten = torch.cat(src_flatten, 1)
        mask_flatten = torch.zeros(src_flatten.shape[:2], device=src_flatten.device, dtype=torch.bool)
        lvl_pos_embed_flatten = torch.cat(lvl_pos_embed_flatten, 1)
        spatial_shapes = torch.as_tensor(spatial_shapes, dtype=torch.long, device=src_flatten.device)
        level_start_index = torch.cat((spatial_shapes.new_zeros((1, )), spatial_shapes.prod(1).cumsum(0)[:-1]))
        valid_ratios = torch.stack([self.get_valid_ratio(m) for m in srcs], 1)


        memory = self.encoder(src_flatten, spatial_shapes, level_start_index, valid_ratios, lvl_pos_embed_flatten, mask_flatten)
        flatten_length = [h*w for (h,w) in spatial_shapes]
        output_split = torch.split(memory, flatten_length, dim=1)
        output_features = [output.reshape(bs,spatial_shapes[i][0], spatial_shapes[i][1],self.hidden_dim).permute(0,3,1,2) for i, output in enumerate(output_split)]
        
        ups = []
        for i, feat in enumerate(output_features):
            feat = self.deblocks[i](feat)
            ups.append(feat)

        ups = torch.cat(ups, dim=1)

        x = ups

        data_dict['spatial_features_2d'] = x
        return data_dict


    def get_valid_ratio(self, x):
        N, _, H, W = x.shape
        mask = torch.zeros((N,H,W),dtype=torch.bool,device=x.device)
        valid_H = torch.sum(~mask[:, :, 0], 1)
        valid_W = torch.sum(~mask[:, 0, :], 1)
        valid_ratio_h = valid_H.float() / H
        valid_ratio_w = valid_W.float() / W
        valid_ratio = torch.stack([valid_ratio_w, valid_ratio_h], -1)
        return valid_ratio
opencood/models/sub_modules/deformable_transformer_backbone.py
MediaBrain-SJTU/CoBEVFlow
import torch
import torch.nn as nn
import math
import copy
from opencood.models.sub_modules.ms_deform_attn import MSDeformAttn
import torch.nn.functional as F

class PositionEmbeddingSine(nn.Module):
    """
    This is a more standard version of the position embedding, very similar to the one
    used by the Attention is all you need paper, generalized to work on images.
    """
    def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None):
        super().__init__()
        self.num_pos_feats = num_pos_feats
        self.temperature = temperature
        self.normalize = normalize
        if scale is not None and normalize is False:
            raise ValueError("normalize should be True if scale is passed")
        if scale is None:
            scale = 2 * math.pi
        self.scale = scale

    def forward(self, x):
        """
        Args:
            x: torch.Tensor
                [N, C, H, W]
        """
        mask = torch.zeros((x.shape[0], x.shape[-2],x.shape[-1]), dtype=torch.bool, device=x.device)
        assert mask is not None
        not_mask = ~mask
        y_embed = not_mask.cumsum(1, dtype=torch.float32)
        x_embed = not_mask.cumsum(2, dtype=torch.float32)
        if self.normalize:
            eps = 1e-6
            y_embed = (y_embed - 0.5) / (y_embed[:, -1:, :] + eps) * self.scale
            x_embed = (x_embed - 0.5) / (x_embed[:, :, -1:] + eps) * self.scale

        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)

        pos_x = x_embed[:, :, :, None] / dim_t
        pos_y = y_embed[:, :, :, None] / dim_t
        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
        return pos


class DeformableTransformerEncoderLayer(nn.Module):
    def __init__(self,
                 d_model=256, d_ffn=1024,
                 dropout=0.1, activation="relu",
                 n_levels=4, n_heads=8, n_points=4):
        super().__init__()

        # self attention
        self.self_attn = MSDeformAttn(d_model, n_levels, n_heads, n_points)
        self.dropout1 = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(d_model)

        # ffn
        self.linear1 = nn.Linear(d_model, d_ffn)
        self.activation = _get_activation_fn(activation)
        self.dropout2 = nn.Dropout(dropout)
        self.linear2 = nn.Linear(d_ffn, d_model)
        self.dropout3 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)

    @staticmethod
    def with_pos_embed(tensor, pos):
        return tensor if pos is None else tensor + pos

    def forward_ffn(self, src):
        src2 = self.linear2(self.dropout2(self.activation(self.linear1(src))))
        src = src + self.dropout3(src2)
        src = self.norm2(src)
        return src

    def forward(self, src, pos, reference_points, spatial_shapes, level_start_index, padding_mask=None):
        # self attention
        src2 = self.self_attn(self.with_pos_embed(src, pos), reference_points, src, spatial_shapes, level_start_index, padding_mask)
        src = src + self.dropout1(src2)
        src = self.norm1(src)

        # ffn
        src = self.forward_ffn(src)

        return src


class DeformableTransformerEncoder(nn.Module):
    def __init__(self, encoder_layer, num_layers):
        super().__init__()
        self.layers = _get_clones(encoder_layer, num_layers)
        self.num_layers = num_layers

    @staticmethod
    def get_reference_points(spatial_shapes, valid_ratios, device):
        reference_points_list = []
        for lvl, (H_, W_) in enumerate(spatial_shapes):

            ref_y, ref_x = torch.meshgrid(torch.linspace(0.5, H_ - 0.5, H_, dtype=torch.float32, device=device),
                                          torch.linspace(0.5, W_ - 0.5, W_, dtype=torch.float32, device=device))
            ref_y = ref_y.reshape(-1)[None] / (valid_ratios[:, None, lvl, 1] * H_)
            ref_x = ref_x.reshape(-1)[None] / (valid_ratios[:, None, lvl, 0] * W_)
            ref = torch.stack((ref_x, ref_y), -1)

            reference_points_list.append(ref)
        reference_points = torch.cat(reference_points_list, 1)
        reference_points = reference_points[:, :, None] * valid_ratios[:, None]
        return reference_points

    def forward(self, src, spatial_shapes, level_start_index, valid_ratios, pos=None, padding_mask=None):
        output = src
        reference_points = self.get_reference_points(spatial_shapes, valid_ratios, device=src.device)
        for _, layer in enumerate(self.layers):
            output = layer(output, pos, reference_points, spatial_shapes, level_start_index, padding_mask)

        return output




def _get_clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])

def _get_activation_fn(activation):
    """Return an activation function given a string"""
    if activation == "relu":
        return F.relu
    if activation == "gelu":
        return F.gelu
    if activation == "glu":
        return F.glu
    raise RuntimeError(F"activation should be relu/gelu, not {activation}.")
opencood/models/sub_modules/detr_module.py
MediaBrain-SJTU/CoBEVFlow
import torch
import torch.nn as nn


class conv2DBatchNormRelu(nn.Module):
    def __init__(
        self,
        in_channels,
        n_filters,
        k_size,
        stride,
        padding,
        bias=True,
        dilation=1,
        is_batchnorm=True,
        shouldUseReLU=True,
    ):
        super(conv2DBatchNormRelu, self).__init__()

        conv_mod = nn.Conv2d(
            int(in_channels),
            int(n_filters),
            kernel_size=k_size,
            padding=padding,
            stride=stride,
            bias=bias,
            dilation=dilation,
        )

        if shouldUseReLU:
            if is_batchnorm:
                self.cbr_unit = nn.Sequential(conv_mod, nn.BatchNorm2d(int(n_filters)), nn.ReLU(inplace=False))
            else:
                self.cbr_unit = nn.Sequential(conv_mod, nn.ReLU(inplace=False))
        else:
            if is_batchnorm:
                self.cbr_unit = nn.Sequential(conv_mod, nn.BatchNorm2d(int(n_filters)))
            else:
                self.cbr_unit = nn.Sequential(conv_mod)


    def forward(self, inputs):
        outputs = self.cbr_unit(inputs)
        return outputs

class segnetDown3(nn.Module):
    def __init__(self, in_size, out_size, indices=False):
        super(segnetDown3, self).__init__()
        self.conv1 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1)
        self.conv2 = conv2DBatchNormRelu(out_size, out_size, 3, 1, 1)
        self.conv3 = conv2DBatchNormRelu(out_size, out_size, 3, 1, 1)

        self.indices = indices
        if indices:
            self.maxpool_with_argmax = nn.MaxPool2d(2, 2, return_indices=True)
        else:
            self.maxpool_without_argmax = nn.MaxPool2d(2, 2, return_indices=False)

    def forward(self, inputs):
        outputs = self.conv1(inputs)
        outputs = self.conv2(outputs)
        outputs = self.conv3(outputs)
        unpooled_shape = outputs.size()

        if self.indices:
            outputs, indices = self.maxpool_with_argmax(outputs)
            return outputs, indices, unpooled_shape
        else:
            outputs = self.maxpool_without_argmax(outputs)
            return outputs

class segnetUp3(nn.Module):
    def __init__(self, in_size, out_size, shouldUseReLU=True):
        super(segnetUp3, self).__init__()
        self.unpool = nn.MaxUnpool2d(2, 2)
        self.up = torch.nn.Upsample(scale_factor=2,mode='bilinear',align_corners=True)
        self.conv1 = conv2DBatchNormRelu(in_size, in_size, 3, 1, 1)
        self.conv2 = conv2DBatchNormRelu(in_size, in_size, 3, 1, 1)
        self.conv3 = conv2DBatchNormRelu(in_size, out_size, 3, 1, 1, shouldUseReLU=shouldUseReLU)

    def forward(self, inputs, indices=None, output_shape=None):
        if indices is not None:
            outputs = self.unpool(input=inputs, indices=indices, output_size=output_shape)
        else:
            outputs = self.up(inputs)

        outputs = self.conv1(outputs)
        outputs = self.conv2(outputs)
        outputs = self.conv3(outputs)
        return outputs

class QueryEncoder(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(QueryEncoder, self).__init__()

        self.model = nn.Sequential(
            conv2DBatchNormRelu(in_ch, 512, 1, 1, 0),
            conv2DBatchNormRelu(512, 512, 1, 1, 0),
            conv2DBatchNormRelu(512, out_ch, 1, 1, 0),
        )
    def forward(self, x):
        return self.model(x)

class KeyEncoder(nn.Module):
    def __init__(self, in_ch, out_ch):
        super(KeyEncoder,self).__init__()
        self.model = nn.Sequential(
            conv2DBatchNormRelu(in_ch, 512, 1, 1, 0),
            # conv2DBatchNormRelu(512, 512, 1, 1, 0),
            conv2DBatchNormRelu(512, out_ch, 1, 1, 0),
        )
    def forward(self, x):
        return self.model(x)

class SmoothingNetwork(nn.Module):
    def __init__(self, in_ch=32*32+1):
        super(SmoothingNetwork, self).__init__()
        out_ch = in_ch
        self.d32to16 = segnetDown3(in_ch,256,indices=True)
        self.d16to08 = segnetDown3(256,128,indices=True)
        self.d08to16 = segnetUp3(128,256)
        self.d16to32 = segnetUp3(256,out_ch)

    def forward(self, distAB):
        Da16d,Da_i16,Da_s16 = self.d32to16(torch.nn.Softmax(1)(distAB))
        Da08d,Da_i08,Da_s08 = self.d16to08(Da16d)
        Da08 = Da08d
        Da16 = self.d08to16(Da08,Da_i08,Da_s08)
        Da32 = self.d16to32(Da16,Da_i16,Da_s16)
        return Da32


opencood/models/sub_modules/mash_utils.py
MediaBrain-SJTU/CoBEVFlow
import torch
from torch import nn

from opencood.pcdet_utils.iou3d_nms.iou3d_nms_utils import boxes_iou3d_gpu
from opencood.utils.box_utils import corner_to_center_torch, boxes_to_corners_3d, project_box3d, project_points_by_matrix_torch
from opencood.utils.transformation_utils import x1_to_x2
from icecream import ic
import copy
pi = 3.141592653


def limit_period(val, offset=0.5, period=2 * pi):
    return val - torch.floor(val / period + offset) * period


class Matcher(nn.Module):
    """Correct localization error and use Algorithm 1:
     BBox matching with scores to fuse the proposal BBoxes"""

    def __init__(self, cfg, pc_range):
        super(Matcher, self).__init__()
        self.pc_range = pc_range

    @torch.no_grad()
    def forward(self, data_dict):
        clusters, scores = self.clustering(data_dict)
        data_dict['boxes_fused'], data_dict[
            'scores_fused'] = self.cluster_fusion(clusters, scores)
        self.merge_keypoints(data_dict)
        return data_dict

    def clustering(self, data_dict):
        """
        Assign predicted boxes to clusters according to their ious with each other
        """
        clusters_batch = []
        scores_batch = []
        record_len = [int(l) for l in data_dict['record_len']]
        lidar_poses = data_dict['lidar_pose'].cpu().numpy()
        for i, l in enumerate(record_len):
            cur_boxes_list = data_dict['det_boxes'][sum(record_len[:i]):sum(record_len[:i])+l]
            
            # Added by Yifan Lu 
            if data_dict['proj_first'] is False:
                cur_boxes_list_ego = []
                # project bounding box to ego coordinate. [x,y,z,l,w,h,yaw]
                cur_boxes_list_ego.append(cur_boxes_list[0])
                for agent_id in range(1, l):
                    tfm = x1_to_x2(lidar_poses[sum(record_len[:i])+agent_id], 
                                   lidar_poses[sum(record_len[:i])])
                    tfm = torch.from_numpy(tfm).to(cur_boxes_list[0].device).float()
                    cur_boxes = cur_boxes_list[agent_id]
                    cur_corners = boxes_to_corners_3d(cur_boxes, order='lwh')
                    cur_corners_ego = project_box3d(cur_corners, tfm)
                    cur_boxes_ego = corner_to_center_torch(cur_corners_ego, order='lwh')
                    cur_boxes_list_ego.append(cur_boxes_ego)
                cur_boxes_list = cur_boxes_list_ego


            cur_scores_list = data_dict['det_scores'][sum(record_len[:i]):sum(record_len[:i])+l]
            cur_boxes_list = [b for b in cur_boxes_list if len(b) > 0]
            cur_scores_list = [s for s in cur_scores_list if len(s) > 0]
            if len(cur_scores_list) == 0:
                clusters_batch.append([torch.Tensor([0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.57]).
                                      to(torch.device('cuda:0')).view(1, 7)])
                scores_batch.append([torch.Tensor([0.01]).to(torch.device('cuda:0')).view(-1)])
                continue

            pred_boxes_cat = torch.cat(cur_boxes_list, dim=0)
            pred_boxes_cat[:, -1] = limit_period(pred_boxes_cat[:, -1])
            pred_scores_cat = torch.cat(cur_scores_list, dim=0)

            ious = boxes_iou3d_gpu(pred_boxes_cat, pred_boxes_cat)
            cluster_indices = torch.zeros(len(ious)).int() # gt assignments of preds
            cur_cluster_id = 1
            while torch.any(cluster_indices == 0):
                cur_idx = torch.where(cluster_indices == 0)[0][0] # find the idx of the first pred which is not assigned yet
                cluster_indices[torch.where(ious[cur_idx] > 0.1)[0]] = cur_cluster_id
                cur_cluster_id += 1
            clusters = []
            scores = []
            for j in range(1, cur_cluster_id):
                clusters.append(pred_boxes_cat[cluster_indices==j])
                scores.append(pred_scores_cat[cluster_indices==j])
            clusters_batch.append(clusters)
            scores_batch.append(scores)

        return clusters_batch, scores_batch

    def cluster_fusion(self, clusters, scores):
        """
        Merge boxes in each cluster with scores as weights for merging
        """
        boxes_fused = []
        scores_fused = []
        for cl, sl in zip(clusters, scores):
            for c, s in zip(cl, sl):
                # reverse direction for non-dominant direction of boxes
                dirs = c[:, -1]
                max_score_idx = torch.argmax(s)
                dirs_diff = torch.abs(dirs - dirs[max_score_idx].item())
                lt_pi = (dirs_diff > pi).int()
                dirs_diff = dirs_diff * (1 - lt_pi) + (
                            2 * pi - dirs_diff) * lt_pi
                score_lt_half_pi = s[dirs_diff > pi / 2].sum()  # larger than
                score_set_half_pi = s[
                    dirs_diff <= pi / 2].sum()  # small equal than
                # select larger scored direction as final direction
                if score_lt_half_pi <= score_set_half_pi:
                    dirs[dirs_diff > pi / 2] += pi
                else:
                    dirs[dirs_diff <= pi / 2] += pi
                dirs = limit_period(dirs)
                s_normalized = s / s.sum()
                sint = torch.sin(dirs) * s_normalized
                cost = torch.cos(dirs) * s_normalized
                theta = torch.atan2(sint.sum(), cost.sum()).view(1, )
                center_dim = c[:, :-1] * s_normalized[:, None]
                
                boxes_fused.append(torch.cat([center_dim.sum(dim=0), theta]))
                s_sorted = torch.sort(s, descending=True).values
                s_fused = 0
                for i, ss in enumerate(s_sorted):
                    s_fused += ss ** (i + 1)
                s_fused = torch.tensor([min(s_fused, 1.0)], device=s.device)
                scores_fused.append(s_fused)

        assert len(boxes_fused) > 0
        boxes_fused = torch.stack(boxes_fused, dim=0)
        len_records = [len(c) for c in clusters]
        boxes_fused = [
            boxes_fused[sum(len_records[:i]):sum(len_records[:i]) + l] for i, l
            in enumerate(len_records)]
        scores_fused = torch.stack(scores_fused, dim=0)
        scores_fused = [
            scores_fused[sum(len_records[:i]):sum(len_records[:i]) + l] for
            i, l in enumerate(len_records)]

        return boxes_fused, scores_fused

    def merge_keypoints(self, data_dict):
        # merge keypoints
        kpts_feat_out = []
        kpts_coor_out = []
        kpts_coor_out_ego = []
        keypoints_features = data_dict['point_features'] # sum(record_len)
        keypoints_coords = data_dict['point_coords'] # [[N,3],...]
        idx = 0
        record_len = data_dict['record_len']
        lidar_poses = data_dict['lidar_pose'].cpu().numpy()
        for l in data_dict['record_len']:
            # Added by Yifan Lu
            # if not project first, first transform the keypoints coords
            if data_dict['proj_first'] is False:
                kpts_coor_cur = []
                for agent_id in range(0, l):
                    tfm = x1_to_x2(lidar_poses[idx+agent_id], lidar_poses[idx])
                    tfm = torch.from_numpy(tfm).to(keypoints_coords[0].device).float()
                    keypoints_coords[idx+agent_id][:, :3] = project_points_by_matrix_torch(keypoints_coords[idx+agent_id][:,:3], tfm)

                kpts_coor_out_ego.append(
                    torch.cat(keypoints_coords[idx:l + idx], dim=0)
                )
                
            kpts_coor_out.append(
                torch.cat(keypoints_coords[idx:l + idx], dim=0))
            kpts_feat_out.append(
                torch.cat(keypoints_features[idx:l + idx], dim=0))
            idx += l
        data_dict['point_features'] = kpts_feat_out
        data_dict['point_coords'] = kpts_coor_out

        if data_dict['proj_first'] is False:
            data_dict['point_coords'] = kpts_coor_out_ego

opencood/models/sub_modules/matcher.py
MediaBrain-SJTU/CoBEVFlow
import numpy as np
import torch
import torch.nn as nn

# from opencood.models.sub_modules.resblock import ResNetModified, BasicBlock, Bottleneck
from opencood.models.sub_modules.resblock2 import ResNetModified, BasicBlock, Bottleneck
from opencood.models.sub_modules.auto_encoder import AutoEncoder
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
DEBUG = False

class MaxResNetBEVBackbone(nn.Module):
    def __init__(self, model_cfg, input_channels):
        super().__init__()
        self.model_cfg = model_cfg
        self.compress = False

        
        self.discrete_ratio = model_cfg['voxel_size'][0]
        self.downsample_rate = 1


            
        layer_nums = model_cfg['layer_nums']
        num_filters = model_cfg['num_filters']
        layer_strides = model_cfg['layer_strides']
        upsample_strides = model_cfg['upsample_strides']
        num_upsample_filters = model_cfg['num_upsample_filter']
        self.level_num = len(layer_nums)

        self.resnet = ResNetModified(BasicBlock, 
                                    layer_nums,
                                    layer_strides,
                                    num_filters)



        num_levels = len(layer_nums)
        c_in_list = [input_channels, *num_filters[:-1]]

        self.fuse_modules = nn.ModuleList()
        self.deblocks = nn.ModuleList()

        if self.compress:
            self.compression_modules = nn.ModuleList()

        for idx in range(num_levels):
            
            fuse_network = MaxFusion()
            self.fuse_modules.append(fuse_network)
            
            if len(upsample_strides) > 0:
                stride = upsample_strides[idx]
                self.deblocks.append(nn.Sequential(
                    nn.ConvTranspose2d(
                        num_filters[idx], num_upsample_filters[idx],
                        upsample_strides[idx],
                        stride=upsample_strides[idx], bias=False
                    ),
                    nn.BatchNorm2d(num_upsample_filters[idx],
                                    eps=1e-3, momentum=0.01),
                    nn.ReLU()
                ))


        c_in = sum(num_upsample_filters)
        if len(upsample_strides) > num_levels:
            self.deblocks.append(nn.Sequential(
                nn.ConvTranspose2d(c_in, c_in, upsample_strides[-1],
                                   stride=upsample_strides[-1], bias=False),
                nn.BatchNorm2d(c_in, eps=1e-3, momentum=0.01),
                nn.ReLU(),
            ))

        self.num_bev_features = c_in

    def forward(self, data_dict):
        spatial_features = data_dict['spatial_features']
        if DEBUG:
            origin_features = torch.clone(spatial_features)

        record_len = data_dict['record_len']
        pairwise_t_matrix = data_dict['pairwise_t_matrix']

        ups = []
        ret_dict = {}
        x = spatial_features

        H, W = x.shape[2:]   #  200, 704
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]

        pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * H / W
        pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * W / H
        pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
        pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2


        features = self.resnet(x)  # tuple of features
        ups = []

        for i in range(self.level_num):
            if DEBUG:
                self.fuse_modules[i].forward_debug(features[i], origin_features, record_len, pairwise_t_matrix)
            else:
                x_fuse = self.fuse_modules[i](features[i], record_len, pairwise_t_matrix)
            if len(self.deblocks) > 0:
                ups.append(self.deblocks[i](x_fuse))
            else:
                ups.append(x_fuse)

        if len(ups) > 1:
            x = torch.cat(ups, dim=1)
        elif len(ups) == 1:
            x = ups[0]

        if len(self.deblocks) > self.level_num:
            x = self.deblocks[-1](x)

        data_dict['spatial_features_2d'] = x
        return data_dict




class MaxFusion(nn.Module):
    def __init__(self):
        super(MaxFusion, self).__init__()

    def regroup(self, x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    def forward(self, x, record_len, pairwise_t_matrix):
        """
        Fusion forwarding.
        
        Parameters
        ----------
        x : torch.Tensor
            input data, (sum(n_cav), C, H, W)
            
        record_len : list
            shape: (B)
            
        pairwise_t_matrix : torch.Tensor
            The affine transformation matrix from each cav to ego, already normalized
            shape: (B, L, L, 2, 3) 
            
        Returns
        -------
        Fused feature.
        """
        _, C, H, W = x.shape
        B, L = pairwise_t_matrix.shape[:2]

        split_x = self.regroup(x, record_len)

        batch_node_features = split_x

        out = []
        # iterate each batch
        for b in range(B):

            N = record_len[b]
            t_matrix = pairwise_t_matrix[b][:N, :N, :, :]

            # update each node i
            i = 0 # ego
            neighbor_feature = warp_affine_simple(batch_node_features[b],
                                            t_matrix[i, :, :, :],
                                            (H, W))
            out.append(torch.max(neighbor_feature, dim=0)[0])
        out = torch.stack(out)
        
        return out

opencood/models/sub_modules/max_resnet_bev_backbone.py
MediaBrain-SJTU/CoBEVFlow
"""
Pillar VFE, credits to OpenPCDet.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F


class PFNLayer(nn.Module):
    def __init__(self,
                 in_channels,
                 out_channels,
                 use_norm=True,
                 last_layer=False):
        super().__init__()

        self.last_vfe = last_layer
        self.use_norm = use_norm
        if not self.last_vfe:
            out_channels = out_channels // 2

        if self.use_norm:
            self.linear = nn.Linear(in_channels, out_channels, bias=False)
            self.norm = nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.01)
        else:
            self.linear = nn.Linear(in_channels, out_channels, bias=True)

        self.part = 50000

    def forward(self, inputs):
        if inputs.shape[0] > self.part:
            # nn.Linear performs randomly when batch size is too large
            num_parts = inputs.shape[0] // self.part
            part_linear_out = [self.linear(
                inputs[num_part * self.part:(num_part + 1) * self.part])
                for num_part in range(num_parts + 1)]
            x = torch.cat(part_linear_out, dim=0)
        else:
            x = self.linear(inputs)
        torch.backends.cudnn.enabled = False
        x = self.norm(x.permute(0, 2, 1)).permute(0, 2,
                                                  1) if self.use_norm else x
        torch.backends.cudnn.enabled = True
        x = F.relu(x)
        x_max = torch.max(x, dim=1, keepdim=True)[0]

        if self.last_vfe:
            return x_max
        else:
            x_repeat = x_max.repeat(1, inputs.shape[1], 1)
            x_concatenated = torch.cat([x, x_repeat], dim=2)
            return x_concatenated


class PillarVFE(nn.Module):
    def __init__(self, model_cfg, num_point_features, voxel_size,
                 point_cloud_range):
        super().__init__()
        self.model_cfg = model_cfg

        self.use_norm = self.model_cfg['use_norm']
        self.with_distance = self.model_cfg['with_distance']

        self.use_absolute_xyz = self.model_cfg['use_absolute_xyz']
        num_point_features += 6 if self.use_absolute_xyz else 3
        if self.with_distance:
            num_point_features += 1

        self.num_filters = self.model_cfg['num_filters']
        assert len(self.num_filters) > 0
        num_filters = [num_point_features] + list(self.num_filters)

        pfn_layers = []
        for i in range(len(num_filters) - 1):
            in_filters = num_filters[i]
            out_filters = num_filters[i + 1]
            pfn_layers.append(
                PFNLayer(in_filters, out_filters, self.use_norm,
                         last_layer=(i >= len(num_filters) - 2))
            )
        self.pfn_layers = nn.ModuleList(pfn_layers)

        self.voxel_x = voxel_size[0] # 每个 voxel 的长/宽/高
        self.voxel_y = voxel_size[1]
        self.voxel_z = voxel_size[2]
        self.x_offset = self.voxel_x / 2 + point_cloud_range[0] # 初始voxel的几何中心的 x/y/z 坐标
        self.y_offset = self.voxel_y / 2 + point_cloud_range[1]
        self.z_offset = self.voxel_z / 2 + point_cloud_range[2]

    def get_output_feature_dim(self):
        return self.num_filters[-1]

    @staticmethod
    def get_paddings_indicator(actual_num, max_num, axis=0):
        actual_num = torch.unsqueeze(actual_num, axis + 1)
        max_num_shape = [1] * len(actual_num.shape)
        max_num_shape[axis + 1] = -1
        max_num = torch.arange(max_num,
                               dtype=torch.int,
                               device=actual_num.device).view(max_num_shape)
        paddings_indicator = actual_num.int() > max_num  # array: (M, max_num_points), 
        return paddings_indicator

    def forward(self, batch_dict):
        """encoding voxel feature using point-pillar method
        Args:
            voxel_features: [M, 32, 4]
            voxel_num_points: [M,]
            voxel_coords: [M, 4]
        Returns:
            features: [M,64], after PFN
        """
        voxel_features, voxel_num_points, coords = \
            batch_dict['voxel_features'], batch_dict['voxel_num_points'], \
            batch_dict['voxel_coords'] # voxel_features: x, y, c, r

        points_mean = \
            voxel_features[:, :, :3].sum(dim=1, keepdim=True) / \
            voxel_num_points.type_as(voxel_features).view(-1, 1, 1)
        f_cluster = voxel_features[:, :, :3] - points_mean  # (M, max_num_points, 3) 每个点与 voxel 中 [重心] 的坐标偏移量 x_c, y_c, z_c
        # coords: 存放每个 voxel 的顺序 -> zyx
        f_center = torch.zeros_like(voxel_features[:, :, :3]) # (M, max_num_points, 3) 每个点与 voxel [center] 的坐标偏移量 x_c, y_c, z_c
        f_center[:, :, 0] = voxel_features[:, :, 0] - (
                coords[:, 3].to(voxel_features.dtype).unsqueeze(
                    1) * self.voxel_x + self.x_offset)      
        f_center[:, :, 1] = voxel_features[:, :, 1] - (
                coords[:, 2].to(voxel_features.dtype).unsqueeze(
                    1) * self.voxel_y + self.y_offset)
        f_center[:, :, 2] = voxel_features[:, :, 2] - (
                coords[:, 1].to(voxel_features.dtype).unsqueeze(
                    1) * self.voxel_z + self.z_offset)

        if self.use_absolute_xyz:
            features = [voxel_features, f_cluster, f_center]
        else:
            features = [voxel_features[..., 3:], f_cluster, f_center]

        if self.with_distance:
            points_dist = torch.norm(voxel_features[:, :, :3], 2, 2,
                                     keepdim=True)
            features.append(points_dist)
        features = torch.cat(features, dim=-1) # (M, max_num_points, 10)

        voxel_count = features.shape[1]
        mask = self.get_paddings_indicator(voxel_num_points, voxel_count,
                                           axis=0)  # (M, max_num_points)
        mask = torch.unsqueeze(mask, -1).type_as(voxel_features)  # (M, max_num_points, 1) float
        features *= mask # 每个voxel块中的补出来的点都置零, (M, 32, 10)
        for pfn in self.pfn_layers:
            features = pfn(features)
        features = features.squeeze()  # (M, 64)
        batch_dict['pillar_features'] = features

        return batch_dict

opencood/models/sub_modules/pillar_vfe.py
MediaBrain-SJTU/CoBEVFlow
import torch
import torch.nn as nn


class PointPillarScatter(nn.Module):
    def __init__(self, model_cfg):
        super().__init__()

        self.model_cfg = model_cfg
        self.num_bev_features = self.model_cfg['num_features']
        self.nx, self.ny, self.nz = model_cfg['grid_size']  # [504, 200, 1] 
        # 这部分的参数参考 opencood/hypes_yaml/yaml_utils.py 中的 load_point_pillar_params 函数
        assert self.nz == 1

    def forward(self, batch_dict):
        """ 将生成的pillar按照坐标索引还原到原空间中
        Args:
            pillar_features:(M, 64)
            coords:(M, 4) 第一维是batch_index

        Returns:
            batch_spatial_features:(4, 64, H, W)
            
            |-------|
            |       |             |-------------|
            |       |     ->      |  *          |
            |       |             |             |
            | *     |             |-------------|
            |-------|

            Lidar Point Cloud        Feature Map
            x-axis up                Along with W 
            y-axis right             Along with H

            Something like clockwise rotation of 90 degree.

        """
        pillar_features, coords = batch_dict['pillar_features'], batch_dict[
            'voxel_coords']
        batch_spatial_features = []
        batch_size = coords[:, 0].max().int().item() + 1 #TODO: 这个应该等于 Nxk 

        for batch_idx in range(batch_size):
            spatial_feature = torch.zeros(
                self.num_bev_features,
                self.nz * self.nx * self.ny,
                dtype=pillar_features.dtype,
                device=pillar_features.device) #(C, x_grid * y_grid * z_grid)
            # batch_index的mask, (M, 1)
            batch_mask = coords[:, 0] == batch_idx
            # 根据mask提取坐标
            this_coords = coords[batch_mask, :] # (batch_idx_voxel,4)  # zyx order, x in [0,706], y in [0,200]
            # 这里的坐标是b,z,y和x的形式,且只有一层，因此计算索引的方式如下
            indices = this_coords[:, 1] + this_coords[:, 2] * self.nx + this_coords[:, 3]  # idz*ny*nx + idy*nx + idx
            # 转换数据类型
            indices = indices.type(torch.long)
            # 根据mask提取pillar_features
            pillars = pillar_features[batch_mask, :] # (batch_idx_voxel,64)
            pillars = pillars.t() # (64,batch_idx_voxel)
            # 在索引位置填充pillars
            spatial_feature[:, indices] = pillars
            # 将空间特征加入list,每个元素为(64, self.nz * self.nx * self.ny)
            batch_spatial_features.append(spatial_feature) 

        batch_spatial_features = \
            torch.stack(batch_spatial_features, 0) #(batch_size, 64, self.nz * self.nx * self.ny)
        batch_spatial_features = \
            batch_spatial_features.view(batch_size, self.num_bev_features *
                                        self.nz, self.ny, self.nx) # It put y axis(in lidar frame) as image height. [..., 200, 504]
        batch_dict['spatial_features'] = batch_spatial_features

        return batch_dict


opencood/models/sub_modules/point_pillar_scatter.py
MediaBrain-SJTU/CoBEVFlow
"""
This is pose graph optimizer, using g2o (bind to C++)
"""
import g2o
import numpy as np

class PoseGraphOptimization2D(g2o.SparseOptimizer):
    def __init__(self, verbose=False):
        super().__init__()
        # solver = g2o.BlockSolverSE2(g2o.LinearSolverCholmodSE2())
        solver = g2o.BlockSolverSE2(g2o.LinearSolverDenseSE2())
        solver = g2o.OptimizationAlgorithmLevenberg(solver)
        super().set_algorithm(solver)
        super().set_verbose(verbose)

    def optimize(self, max_iterations=1000):
        super().initialize_optimization()
        super().optimize(max_iterations)


    def add_vertex(self, id, pose, fixed=False, SE2=True):
        if SE2:
            v = g2o.VertexSE2()
        else:
            v = g2o.VertexPointXY()
        v.set_estimate(pose)
        v.set_id(id)
        v.set_fixed(fixed)
        super().add_vertex(v)


    def add_edge(self, vertices, measurement, 
            information=np.identity(3),
            robust_kernel=None, SE2 = True):
        """
        Args:
            measurement: g2o.SE2
        """
        if SE2:
            edge = g2o.EdgeSE2()
        else:
            edge = g2o.EdgeSE2PointXY()

        for i, v in enumerate(vertices):
            if isinstance(v, int):
                v = self.vertex(v)
            edge.set_vertex(i, v)

        edge.set_measurement(measurement)  # relative pose shape [3, 1] / [2, 1]
        edge.set_information(information)  # importance of each component shape [3, 3] / [2, 2]
        if robust_kernel is not None:
            edge.set_robust_kernel(robust_kernel)
        super().add_edge(edge)

    def get_pose(self, id):
        return self.vertex(id).estimate()


class PoseGraphOptimization(g2o.SparseOptimizer):
    def __init__(self):
        super().__init__()
        solver = g2o.BlockSolverSE3(g2o.LinearSolverCholmodSE3())
        solver = g2o.OptimizationAlgorithmLevenberg(solver)
        super().set_algorithm(solver)
        super().set_verbose(True)

    def optimize(self, max_iterations=50):
        super().initialize_optimization()
        super().optimize(max_iterations)

    def add_vertex(self, id, pose, fixed=False):
        v_se3 = g2o.VertexSE3()
        v_se3.set_estimate(pose)
        v_se3.set_id(id)
        v_se3.set_fixed(fixed)
        super().add_vertex(v_se3)

    def add_edge(self, vertices, measurement, 
            information=np.identity(6),
            robust_kernel=None):

        edge = g2o.EdgeSE3()
        for i, v in enumerate(vertices):
            if isinstance(v, int):
                v = self.vertex(v)
            edge.set_vertex(i, v)

        edge.set_measurement(measurement)  # relative pose, shape [4, 4]
        edge.set_information(information)  # importance of each component, shape [6, 6]
        if robust_kernel is not None:
            edge.set_robust_kernel(robust_kernel)
        super().add_edge(edge)

    def get_pose(self, id):
        return self.vertex(id).estimate()


if __name__ == "__main__":
    pgo = PoseGraphOptimization()

    with open("/GPFS/rhome/yifanlu/workspace/g2o_test/noise.g2o","r") as f:
        for line in f:
            if line.startswith("VERTEX_SE3:QUAT"):
                vertex_content = line.split(" ",1)[1]
                vertex_content_array = np.fromstring(vertex_content, dtype=float, sep=" ")
                ids = int(vertex_content_array[0])
                index = [0,1,2,6,3,4,5]
                pose_array = vertex_content_array[1:][index]

                pose = np.eye(4)
                pose[:3,3] = pose_array[:3]
                pose[:3,:3] = g2o.Quaternion(pose_array[3:]).matrix()
                pose = g2o.Isometry3d(pose)

                fixed = True if ids==6 else False
                # fixed = False
                pgo.add_vertex(id=ids, pose=pose, fixed=fixed)

            elif line.startswith("EDGE_SE3:QUAT"):
                edge_content = line.split(" ", 1)[1]
                edge_content_array = np.fromstring(edge_content, dtype=float, sep=" ")
                
                edge = [int(v) for v in edge_content_array[:2]]
                index = [0,1,2,6,3,4,5]
                pose_array = edge_content_array[2:2+7][index]
                information_array = edge_content_array[2+7:]

                pose = np.eye(4)
                pose[:3,3] = pose_array[:3]
                pose[:3,:3] = g2o.Quaternion(pose_array[3:]).matrix()
                pose = g2o.Isometry3d(pose)

                information = np.eye(6)
                information[0,0] = information_array[0]
                information[1,1] = information_array[6]
                information[2,2] = information_array[11]
                information[3,3] = information_array[15]
                information[4,4] = information_array[18]
                information[5,5] = information_array[20]

                pgo.add_edge(edge, pose, information)


    print('num vertices:', len(pgo.vertices()))
    print('num edges:', len(pgo.edges()), end='\n\n')
    pgo.optimize()

    # pgo.save("out_pose_graph2.g2o")

opencood/models/sub_modules/pose_graph_optim.py
MediaBrain-SJTU/CoBEVFlow
import torch
from torch import Tensor
import torch.nn as nn
from typing import Type, Any, Callable, Union, List, Optional


__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
           'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
           'wide_resnet50_2', 'wide_resnet101_2']


def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False, dilation=dilation)


def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class BasicBlock(nn.Module):
    expansion: int = 1

    def __init__(
        self,
        inplanes: int,
        planes: int,
        stride: int = 1,
        downsample: Optional[nn.Module] = None,
        groups: int = 1,
        base_width: int = 64,
        dilation: int = 1,
        norm_layer: Optional[Callable[..., nn.Module]] = None
    ) -> None:
        super(BasicBlock, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if groups != 1 or base_width != 64:
            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
        if dilation > 1:
            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) -> Tensor:
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
    # This variant is also known as ResNet V1.5 and improves accuracy according to
    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.

    expansion: int = 4 # original 4

    def __init__(
        self,
        inplanes: int,
        planes: int,
        stride: int = 1,
        downsample: Optional[nn.Module] = None,
        groups: int = 1,
        base_width: int = 64,
        dilation: int = 1,
        norm_layer: Optional[Callable[..., nn.Module]] = None
    ) -> None:
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.)) * groups
        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) -> Tensor:
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class ResNetModified(nn.Module):

    def __init__(
        self,
        block: Type[Union[BasicBlock, Bottleneck]],
        layers: List[int],  # number of block in one layer
        layer_strides: List[int],  #  stride after one layer
        num_filters: List[int],  # feature dim
        zero_init_residual: bool = False,
        groups: int = 1,
        width_per_group: int = 64,
        replace_stride_with_dilation: Optional[List[bool]] = None,
        norm_layer: Optional[Callable[..., nn.Module]] = None
    ) -> None:
        super(ResNetModified, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer

        self.inplanes = 64
        self.dilation = 1
        if replace_stride_with_dilation is None:
            # each element in the tuple indicates if we should replace
            # the 2x2 stride with a dilated convolution instead
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError("replace_stride_with_dilation should be None "
                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
        self.groups = groups
        self.base_width = width_per_group

        self.layer1 = self._make_layer(block, num_filters[0], layers[0], stride=layer_strides[0])
        self.layer2 = self._make_layer(block, num_filters[1], layers[1], stride=layer_strides[1],
                                       dilate=replace_stride_with_dilation[0])
        self.layer3 = self._make_layer(block, num_filters[2], layers[2], stride=layer_strides[2],
                                       dilate=replace_stride_with_dilation[1])


        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)  # type: ignore[arg-type]
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)  # type: ignore[arg-type]

    def _make_layer(self, block: Type[Union[BasicBlock, Bottleneck]], planes: int, blocks: int,
                    stride: int = 1, dilate: bool = False) -> nn.Sequential:
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                norm_layer(planes * block.expansion),
            )

        layers = []
        # if stride != 1, the first block will downsample the feature map
        # plane is the feature dim
        # if Bottleneck, then the output dim is planes * block.expansion(4)
        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
                            self.base_width, previous_dilation, norm_layer))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes, groups=self.groups,
                                base_width=self.base_width, dilation=self.dilation,
                                norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def _forward_impl(self, x: Tensor, return_interm: bool = True):
        # See note [TorchScript super()]

        x1 = self.layer1(x)
        x2 = self.layer2(x1)
        x3 = self.layer3(x2)

        if return_interm:
            return (x1,x2,x3)
        return x3

    def forward(self, x: Tensor):
        return self._forward_impl(x)


def _resnet(
    arch: str,
    block: Type[Union[BasicBlock, Bottleneck]],
    layers: List[int],
    pretrained: bool,
    progress: bool,
    **kwargs: Any
) -> ResNetModified:
    model = ResNetModified(block, layers, **kwargs)

    return model


def resnet18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNet-18 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
                   **kwargs)


def resnet34(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNet-34 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
                   **kwargs)


def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNet-50 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
                   **kwargs)


def resnet101(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNet-101 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,
                   **kwargs)


def resnet152(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNet-152 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,
                   **kwargs)


def resnext50_32x4d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNeXt-50 32x4d model from
    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['groups'] = 32
    kwargs['width_per_group'] = 4
    return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
                   pretrained, progress, **kwargs)


def resnext101_32x8d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNeXt-101 32x8d model from
    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['groups'] = 32
    kwargs['width_per_group'] = 8
    return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
                   pretrained, progress, **kwargs)


def wide_resnet50_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""Wide ResNet-50-2 model from
    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.

    The model is the same as ResNet except for the bottleneck number of channels
    which is twice larger in every block. The number of channels in outer 1x1
    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
    channels, and in Wide ResNet-50-2 has 2048-1024-2048.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['width_per_group'] = 64 * 2
    return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
                   pretrained, progress, **kwargs)


def wide_resnet101_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""Wide ResNet-101-2 model from
    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.

    The model is the same as ResNet except for the bottleneck number of channels
    which is twice larger in every block. The number of channels in outer 1x1
    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
    channels, and in Wide ResNet-50-2 has 2048-1024-2048.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['width_per_group'] = 64 * 2
    return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
                   pretrained, progress, **kwargs)




if __name__=="__main__":
    model = ResNetModified(BasicBlock, [3,4,5])
    input = torch.randn(4,64,200,704)
    output = model(input)
    from icecream import ic
    for out in output:
        ic(out.shape)

opencood/models/sub_modules/resblock.py
MediaBrain-SJTU/CoBEVFlow
import torch
from torch import Tensor
import torch.nn as nn
from typing import Type, Any, Callable, Union, List, Optional


__all__ = ['ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101',
           'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
           'wide_resnet50_2', 'wide_resnet101_2']


def conv3x3(in_planes: int, out_planes: int, stride: int = 1, groups: int = 1, dilation: int = 1) -> nn.Conv2d:
    """3x3 convolution with padding"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
                     padding=dilation, groups=groups, bias=False, dilation=dilation)


def conv1x1(in_planes: int, out_planes: int, stride: int = 1) -> nn.Conv2d:
    """1x1 convolution"""
    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)


class BasicBlock(nn.Module):
    expansion: int = 1

    def __init__(
        self,
        inplanes: int,
        planes: int,
        stride: int = 1,
        downsample: Optional[nn.Module] = None,
        groups: int = 1,
        base_width: int = 64,
        dilation: int = 1,
        norm_layer: Optional[Callable[..., nn.Module]] = None
    ) -> None:
        super(BasicBlock, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        if groups != 1 or base_width != 64:
            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
        if dilation > 1:
            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv3x3(inplanes, planes, stride)
        self.bn1 = norm_layer(planes)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(planes, planes)
        self.bn2 = norm_layer(planes)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) -> Tensor:
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class Bottleneck(nn.Module):
    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
    # This variant is also known as ResNet V1.5 and improves accuracy according to
    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.

    expansion: int = 4 # original 4

    def __init__(
        self,
        inplanes: int,
        planes: int,
        stride: int = 1,
        downsample: Optional[nn.Module] = None,
        groups: int = 1,
        base_width: int = 64,
        dilation: int = 1,
        norm_layer: Optional[Callable[..., nn.Module]] = None
    ) -> None:
        super(Bottleneck, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        width = int(planes * (base_width / 64.)) * groups
        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
        self.conv1 = conv1x1(inplanes, width)
        self.bn1 = norm_layer(width)
        self.conv2 = conv3x3(width, width, stride, groups, dilation)
        self.bn2 = norm_layer(width)
        self.conv3 = conv1x1(width, planes * self.expansion)
        self.bn3 = norm_layer(planes * self.expansion)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x: Tensor) -> Tensor:
        identity = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            identity = self.downsample(x)

        out += identity
        out = self.relu(out)

        return out


class ResNetModified(nn.Module):

    def __init__(
        self,
        block: Type[Union[BasicBlock, Bottleneck]],
        layers: List[int],  # number of block in one layer
        layer_strides: List[int],  #  stride after one layer
        num_filters: List[int],  # feature dim
        zero_init_residual: bool = False,
        groups: int = 1,
        width_per_group: int = 64,
        replace_stride_with_dilation: Optional[List[bool]] = None,
        norm_layer: Optional[Callable[..., nn.Module]] = None
    ) -> None:
        super(ResNetModified, self).__init__()
        if norm_layer is None:
            norm_layer = nn.BatchNorm2d
        self._norm_layer = norm_layer

        self.inplanes = 64
        self.dilation = 1
        if replace_stride_with_dilation is None:
            # each element in the tuple indicates if we should replace
            # the 2x2 stride with a dilated convolution instead
            replace_stride_with_dilation = [False, False, False]
        if len(replace_stride_with_dilation) != 3:
            raise ValueError("replace_stride_with_dilation should be None "
                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
        self.groups = groups
        self.base_width = width_per_group

        self.layernum = len(num_filters)
        for i in range(self.layernum):
            self.__setattr__(f"layer{i}", self._make_layer(block, num_filters[i], layers[i], stride=layer_strides[i]))


        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)

        # Zero-initialize the last BN in each residual branch,
        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
        if zero_init_residual:
            for m in self.modules():
                if isinstance(m, Bottleneck):
                    nn.init.constant_(m.bn3.weight, 0)  # type: ignore[arg-type]
                elif isinstance(m, BasicBlock):
                    nn.init.constant_(m.bn2.weight, 0)  # type: ignore[arg-type]

    def _make_layer(self, block: Type[Union[BasicBlock, Bottleneck]], planes: int, blocks: int,
                    stride: int = 1, dilate: bool = False) -> nn.Sequential:
        norm_layer = self._norm_layer
        downsample = None
        previous_dilation = self.dilation
        if dilate:
            self.dilation *= stride
            stride = 1
        if stride != 1 or self.inplanes != planes * block.expansion:
            downsample = nn.Sequential(
                conv1x1(self.inplanes, planes * block.expansion, stride),
                norm_layer(planes * block.expansion),
            )

        layers = []
        # if stride != 1, the first block will downsample the feature map
        # plane is the feature dim
        # if Bottleneck, then the output dim is planes * block.expansion(4)
        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
                            self.base_width, previous_dilation, norm_layer))
        self.inplanes = planes * block.expansion
        for _ in range(1, blocks):
            layers.append(block(self.inplanes, planes, groups=self.groups,
                                base_width=self.base_width, dilation=self.dilation,
                                norm_layer=norm_layer))

        return nn.Sequential(*layers)

    def _forward_impl(self, x: Tensor, return_interm: bool = True):
        # See note [TorchScript super()]
        interm_features = []
        for i in range(self.layernum):
            x = eval(f"self.layer{i}")(x)
            interm_features.append(x)

        if return_interm:
            return interm_features
        return x

    def forward(self, x: Tensor):
        return self._forward_impl(x)


def _resnet(
    arch: str,
    block: Type[Union[BasicBlock, Bottleneck]],
    layers: List[int],
    pretrained: bool,
    progress: bool,
    **kwargs: Any
) -> ResNetModified:
    model = ResNetModified(block, layers, **kwargs)

    return model


def resnet18(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNet-18 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress,
                   **kwargs)


def resnet34(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNet-34 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress,
                   **kwargs)


def resnet50(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNet-50 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress,
                   **kwargs)


def resnet101(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNet-101 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress,
                   **kwargs)


def resnet152(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNet-152 model from
    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress,
                   **kwargs)


def resnext50_32x4d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNeXt-50 32x4d model from
    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['groups'] = 32
    kwargs['width_per_group'] = 4
    return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3],
                   pretrained, progress, **kwargs)


def resnext101_32x8d(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""ResNeXt-101 32x8d model from
    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['groups'] = 32
    kwargs['width_per_group'] = 8
    return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3],
                   pretrained, progress, **kwargs)


def wide_resnet50_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""Wide ResNet-50-2 model from
    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.

    The model is the same as ResNet except for the bottleneck number of channels
    which is twice larger in every block. The number of channels in outer 1x1
    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
    channels, and in Wide ResNet-50-2 has 2048-1024-2048.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['width_per_group'] = 64 * 2
    return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3],
                   pretrained, progress, **kwargs)


def wide_resnet101_2(pretrained: bool = False, progress: bool = True, **kwargs: Any) -> ResNetModified:
    r"""Wide ResNet-101-2 model from
    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.

    The model is the same as ResNet except for the bottleneck number of channels
    which is twice larger in every block. The number of channels in outer 1x1
    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
    channels, and in Wide ResNet-50-2 has 2048-1024-2048.

    Args:
        pretrained (bool): If True, returns a model pre-trained on ImageNet
        progress (bool): If True, displays a progress bar of the download to stderr
    """
    kwargs['width_per_group'] = 64 * 2
    return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3],
                   pretrained, progress, **kwargs)




if __name__=="__main__":
    model = ResNetModified(BasicBlock, [3,4,5])
    input = torch.randn(4,64,200,704)
    output = model(input)
    from icecream import ic
    for out in output:
        ic(out.shape)

opencood/models/sub_modules/resblock2.py
MediaBrain-SJTU/CoBEVFlow
import copy

import torch.nn as nn
import torch
import numpy as np
from opencood.pcdet_utils.pointnet2.pointnet2_stack import \
    pointnet2_modules as pointnet2_stack_modules
from opencood.utils import common_utils
from opencood.pcdet_utils.iou3d_nms.iou3d_nms_utils import boxes_iou3d_gpu
from opencood.utils import box_utils


class RoIHead(nn.Module):
    def __init__(self, model_cfg):
        super().__init__()
        self.model_cfg = model_cfg
        input_channels = model_cfg['in_channels']
        self.code_size = 7

        mlps = copy.copy(self.model_cfg['roi_grid_pool']['mlps'])
        for k in range(len(mlps)):
            mlps[k] = [input_channels] + mlps[k]

        self.roi_grid_pool_layer = pointnet2_stack_modules.StackSAModuleMSG(
            radii=self.model_cfg['roi_grid_pool']['pool_radius'],
            nsamples=self.model_cfg['roi_grid_pool']['n_sample'],
            mlps=mlps,
            use_xyz=True,
            pool_method=self.model_cfg['roi_grid_pool']['pool_method'],
        )

        grid_size = self.model_cfg['roi_grid_pool']['grid_size']
        self.grid_size = grid_size
        c_out = sum([x[-1] for x in mlps])
        pre_channel = grid_size * grid_size * grid_size * c_out
        fc_layers = [self.model_cfg['n_fc_neurons']] * 2
        self.shared_fc_layers, pre_channel = self._make_fc_layers(pre_channel,
                                                                  fc_layers)

        self.cls_layers, pre_channel = self._make_fc_layers(pre_channel,
                                                            fc_layers,
                                                            output_channels=
                                                            self.model_cfg[
                                                                'num_cls'])
        self.iou_layers, _ = self._make_fc_layers(pre_channel, fc_layers,
                                                  output_channels=
                                                  self.model_cfg['num_cls'])
        self.reg_layers, _ = self._make_fc_layers(pre_channel, fc_layers,
                                                  output_channels=
                                                  self.model_cfg[
                                                      'num_cls'] * 7)

        self._init_weights(weight_init='xavier')

    def _init_weights(self, weight_init='xavier'):
        if weight_init == 'kaiming':
            init_func = nn.init.kaiming_normal_
        elif weight_init == 'xavier':
            init_func = nn.init.xavier_normal_
        elif weight_init == 'normal':
            init_func = nn.init.normal_
        else:
            raise NotImplementedError

        for m in self.modules():
            if isinstance(m, nn.Conv2d) or isinstance(m, nn.Conv1d):
                if weight_init == 'normal':
                    init_func(m.weight, mean=0, std=0.001)
                else:
                    init_func(m.weight)
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
        nn.init.normal_(self.reg_layers[-1].weight, mean=0, std=0.001)

    def _make_fc_layers(self, input_channels, fc_list, output_channels=None):
        fc_layers = []
        pre_channel = input_channels
        for k in range(len(fc_list)):
            fc_layers.extend([
                nn.Conv1d(pre_channel, fc_list[k], kernel_size=1, bias=False),
                # nn.BatchNorm1d(fc_list[k]),
                nn.ReLU()
            ])
            pre_channel = fc_list[k]
            if self.model_cfg['dp_ratio'] > 0:
                fc_layers.append(nn.Dropout(self.model_cfg['dp_ratio']))
        if output_channels is not None:
            fc_layers.append(
                nn.Conv1d(pre_channel, output_channels, kernel_size=1,
                          bias=True))
        fc_layers = nn.Sequential(*fc_layers)
        return fc_layers, pre_channel

    def get_global_grid_points_of_roi(self, rois):
        rois = rois.view(-1, rois.shape[-1])
        batch_size_rcnn = rois.shape[0]

        # (B, 6x6x6, 3)
        local_roi_grid_points = self.get_dense_grid_points(rois,
                                                           batch_size_rcnn,
                                                           self.grid_size)
        global_roi_grid_points = common_utils.rotate_points_along_z(
            local_roi_grid_points.clone(), rois[:, 6]
        ).squeeze(dim=1)
        global_center = rois[:, 0:3].clone()
        global_roi_grid_points += global_center.unsqueeze(dim=1)
        return global_roi_grid_points, local_roi_grid_points

    @staticmethod
    def get_dense_grid_points(rois, batch_size_rcnn, grid_size):
        """
        Get the local coordinates of each grid point of a roi in the coordinate
        system of the roi(origin lies in the center of this roi.
        """
        faked_features = rois.new_ones((grid_size, grid_size, grid_size))
        dense_idx = torch.stack(torch.where(faked_features),
                                dim=1)  # (N, 3) [x_idx, y_idx, z_idx]
        dense_idx = dense_idx.repeat(batch_size_rcnn, 1,
                                     1).float()  # (B, 6x6x6, 3)

        local_roi_size = rois.view(batch_size_rcnn, -1)[:, 3:6]
        roi_grid_points = (
                                  dense_idx + 0.5) / grid_size * local_roi_size.unsqueeze(
            dim=1) \
                          - (local_roi_size.unsqueeze(
            dim=1) / 2)  # (B, 6x6x6, 3)
        return roi_grid_points

    def assign_targets(self, batch_dict):
        batch_dict['rcnn_label_dict'] = {
            'rois': [],
            'gt_of_rois': [],
            'gt_of_rois_src': [],
            'cls_tgt': [],
            'reg_tgt': [],
            'iou_tgt': [],
            'rois_anchor': [],
            'record_len': []
        }
        # pred_boxes = [boxes[:, [0, 1, 2, 5, 4, 3, 6]] for boxes in batch_dict['boxes_fused']]
        pred_boxes = batch_dict['boxes_fused']  # hwl
        # print("batch_dict['object_bbx_center']:",batch_dict['object_bbx_center'])   ##### hwl
        gt_boxes = [b[m][:, [0, 1, 2, 5, 4, 3, 6]].float() for b, m in
                    zip(batch_dict['object_bbx_center'],
                        batch_dict['object_bbx_mask'].bool())]  # lwh order
        for rois, gts in zip(pred_boxes, gt_boxes): # each frame
            # print("rois:",rois)  # hwl
            # print('gts',gts) # lwh
            if gts.shape[0] == 0:
                batch_dict['rcnn_label_dict']['record_len'].append(0)
                continue

            gts[:, -1] *= 1
            ious = boxes_iou3d_gpu(rois, gts)
            max_ious, gt_inds = ious.max(dim=1)
            gt_of_rois = gts[gt_inds]
            rcnn_labels = (max_ious > 0.3).float()
            mask = torch.logical_not(rcnn_labels.bool())

            # set negative samples back to rois, no correction in stage2 for them
            gt_of_rois[mask] = rois[mask]
            gt_of_rois_src = gt_of_rois.clone().detach()

            # canoical transformation
            roi_center = rois[:, 0:3]
            # TODO: roi_ry > 0 in pcdet
            roi_ry = rois[:, 6] % (2 * np.pi)
            gt_of_rois[:, 0:3] = gt_of_rois[:, 0:3] - roi_center
            gt_of_rois[:, 6] = gt_of_rois[:, 6] - roi_ry

            # transfer LiDAR coords to local coords
            gt_of_rois = common_utils.rotate_points_along_z(
                points=gt_of_rois.view(-1, 1, gt_of_rois.shape[-1]),
                angle=-roi_ry.view(-1)
            ).view(-1, gt_of_rois.shape[-1])

            # flip orientation if rois have opposite orientation
            heading_label = (gt_of_rois[:, 6] + (
                    torch.div(torch.abs(gt_of_rois[:, 6].min()),
                              (2 * np.pi), rounding_mode='trunc')
                    + 1) * 2 * np.pi) % (2 * np.pi)  # 0 ~ 2pi
            opposite_flag = (heading_label > np.pi * 0.5) & (
                    heading_label < np.pi * 1.5)

            # (0 ~ pi/2, 3pi/2 ~ 2pi)
            heading_label[opposite_flag] = (heading_label[
                                                opposite_flag] + np.pi) % (
                                                   2 * np.pi)
            flag = heading_label > np.pi
            heading_label[flag] = heading_label[
                                      flag] - np.pi * 2  # (-pi/2, pi/2)
            heading_label = torch.clamp(heading_label, min=-np.pi / 2,
                                        max=np.pi / 2)
            gt_of_rois[:, 6] = heading_label

            # generate regression target
            rois_anchor = rois.clone().detach().view(-1, self.code_size)
            rois_anchor[:, 0:3] = 0
            rois_anchor[:, 6] = 0

            reg_targets = box_utils.box_encode(
                gt_of_rois.view(-1, self.code_size), rois_anchor
            )

            batch_dict['rcnn_label_dict']['rois'].append(rois)
            batch_dict['rcnn_label_dict']['gt_of_rois'].append(gt_of_rois)
            batch_dict['rcnn_label_dict']['gt_of_rois_src'].append(
                gt_of_rois_src)
            batch_dict['rcnn_label_dict']['cls_tgt'].append(rcnn_labels)
            batch_dict['rcnn_label_dict']['reg_tgt'].append(reg_targets)
            batch_dict['rcnn_label_dict']['iou_tgt'].append(max_ious)
            batch_dict['rcnn_label_dict']['rois_anchor'].append(rois_anchor)
            batch_dict['rcnn_label_dict']['record_len'].append(rois.shape[0])
            

        # cat list to tensor
        for k, v in batch_dict['rcnn_label_dict'].items():
            if k == 'record_len':
                continue
            batch_dict['rcnn_label_dict'][k] = torch.cat(v, dim=0)

        return batch_dict

    def roi_grid_pool(self, batch_dict):
        batch_size = len(batch_dict['record_len'])
        rois = batch_dict['rcnn_label_dict']['rois']
        point_coords = batch_dict['point_coords']
        point_features = batch_dict['point_features']
        label_record_len = batch_dict['rcnn_label_dict']['record_len']

        point_features = torch.cat(point_features, dim=0)
        # (BxN, 6x6x6, 3)
        global_roi_grid_points, local_roi_grid_points = \
            self.get_global_grid_points_of_roi(rois)
        # (B, Nx6x6x6, 3)
        global_roi_grid_points = global_roi_grid_points.view(batch_size, -1, 3)

        xyz = torch.cat(point_coords, dim=0)
        xyz_batch_cnt = xyz.new_zeros(batch_size).int()
        for bs_idx in range(batch_size):
            xyz_batch_cnt[bs_idx] = len(point_coords[bs_idx])
        new_xyz = global_roi_grid_points.view(-1, 3)
        new_xyz_batch_cnt = xyz.new_zeros(batch_size).int()
        for bs_idx in range(batch_size):
            new_xyz_batch_cnt[bs_idx] = label_record_len[
                                            bs_idx] * self.grid_size ** 3

        pooled_points, pooled_features = self.roi_grid_pool_layer(
            xyz=xyz[:, :3].contiguous(),
            xyz_batch_cnt=xyz_batch_cnt,
            new_xyz=new_xyz[:, :3].contiguous(),
            new_xyz_batch_cnt=new_xyz_batch_cnt,
            features=point_features.contiguous(),  # weighted point features
        )  # (M1 + M2 ..., C)
        # (BxN, 6x6x6, C)
        pooled_features = pooled_features.view(-1, self.grid_size ** 3,
                                               pooled_features.shape[-1])

        return pooled_features

    def forward(self, batch_dict):
        batch_dict = self.assign_targets(batch_dict)
        # RoI aware pooling
        pooled_features = self.roi_grid_pool(batch_dict)  # (BxN, 6x6x6, C)

        batch_size_rcnn = pooled_features.shape[0]
        pooled_features = pooled_features.permute(0, 2, 1). \
            contiguous().view(batch_size_rcnn, -1, self.grid_size,
                              self.grid_size,
                              self.grid_size)  # (BxN, C, 6, 6, 6)
        shared_features = self.shared_fc_layers(
            pooled_features.view(batch_size_rcnn, -1, 1))
        rcnn_cls = self.cls_layers(shared_features).transpose(1,
                                                              2).contiguous().squeeze(
            dim=1)  # (B, 1 or 2)
        rcnn_iou = self.iou_layers(shared_features).transpose(1,
                                                              2).contiguous().squeeze(
            dim=1)  # (B, 1)
        rcnn_reg = self.reg_layers(shared_features).transpose(1,
                                                              2).contiguous().squeeze(
            dim=1)  # (B, C)

        batch_dict['fpvrcnn_out'] = {
            'rcnn_cls': rcnn_cls,
            'rcnn_iou': rcnn_iou,
            'rcnn_reg': rcnn_reg,
        }
        return batch_dict
opencood/models/sub_modules/roi_head.py
MediaBrain-SJTU/CoBEVFlow
from functools import partial

import spconv
import torch.nn as nn


def post_act_block(in_channels, out_channels, kernel_size, indice_key=None, stride=1, padding=0,
                   conv_type='subm', norm_fn=None):

    if conv_type == 'subm':
        conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, bias=False, indice_key=indice_key)
    elif conv_type == 'spconv':
        conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, padding=padding,
                                   bias=False, indice_key=indice_key)
    elif conv_type == 'inverseconv':
        conv = spconv.SparseInverseConv3d(in_channels, out_channels, kernel_size, indice_key=indice_key, bias=False)
    else:
        raise NotImplementedError

    m = spconv.SparseSequential(
        conv,
        norm_fn(out_channels),
        nn.ReLU(),
    )

    return m


class VoxelBackBone8x(nn.Module):
    def __init__(self, model_cfg, input_channels, grid_size, **kwargs):
        super().__init__()
        self.model_cfg = model_cfg
        norm_fn = partial(nn.BatchNorm1d, eps=1e-3, momentum=0.01)

        self.sparse_shape = grid_size[::-1] + [1, 0, 0]

        self.conv_input = spconv.SparseSequential(
            spconv.SubMConv3d(input_channels, 16, 3, padding=1, bias=False, indice_key='subm1'),
            norm_fn(16),
            nn.ReLU(),
        )
        block = post_act_block

        self.conv1 = spconv.SparseSequential(
            block(16, 16, 3, norm_fn=norm_fn, padding=1, indice_key='subm1'),
        )

        self.conv2 = spconv.SparseSequential(
            # [1600, 1408, 41] <- [800, 704, 21]
            block(16, 32, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv2', conv_type='spconv'),
            block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
            block(32, 32, 3, norm_fn=norm_fn, padding=1, indice_key='subm2'),
        )

        self.conv3 = spconv.SparseSequential(
            # [800, 704, 21] <- [400, 352, 11]
            block(32, 64, 3, norm_fn=norm_fn, stride=2, padding=1, indice_key='spconv3', conv_type='spconv'),
            block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
            block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm3'),
        )

        self.conv4 = spconv.SparseSequential(
            # [400, 352, 11] <- [200, 176, 5]
            block(64, 64, 3, norm_fn=norm_fn, stride=2, padding=(0, 1, 1), indice_key='spconv4', conv_type='spconv'),
            block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
            block(64, 64, 3, norm_fn=norm_fn, padding=1, indice_key='subm4'),
        )

        last_pad = 0
        if 'num_features_out' in self.model_cfg:
            self.num_point_features = self.model_cfg['num_features_out']
        else:
            self.num_point_features = 128
        self.conv_out = spconv.SparseSequential(
            # [200, 150, 5] -> [200, 150, 2]
            spconv.SparseConv3d(64, self.num_point_features, (3, 1, 1), stride=(2, 1, 1), padding=last_pad,
                                bias=False, indice_key='spconv_down2'),
            norm_fn(self.num_point_features),
            nn.ReLU(),
        )

        self.backbone_channels = {
            'x_conv1': 16,
            'x_conv2': 32,
            'x_conv3': 64,
            'x_conv4': 64
        }

    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size: int
                vfe_features: (num_voxels, C)
                voxel_coords: (num_voxels, 4), [batch_idx, z_idx, y_idx, x_idx]
        Returns:
            batch_dict:
                encoded_spconv_tensor: sparse tensor
        """
        voxel_features, voxel_coords = batch_dict['voxel_features'], \
                                       batch_dict['voxel_coords']
        batch_size = batch_dict['batch_size']
        input_sp_tensor = spconv.SparseConvTensor(
            features=voxel_features,
            indices=voxel_coords.int(),
            spatial_shape=self.sparse_shape,
            batch_size=batch_size
        )

        x = self.conv_input(input_sp_tensor)

        x_conv1 = self.conv1(x)
        x_conv2 = self.conv2(x_conv1)
        x_conv3 = self.conv3(x_conv2)
        x_conv4 = self.conv4(x_conv3)

        # for detection head
        # [200, 176, 5] -> [200, 176, 2]
        out = self.conv_out(x_conv4)

        batch_dict.update({
            'encoded_spconv_tensor': out,
            'encoded_spconv_tensor_stride': 8
        })
        batch_dict.update({
            'multi_scale_3d_features': {
                'x_conv1': x_conv1,
                'x_conv2': x_conv2,
                'x_conv3': x_conv3,
                'x_conv4': x_conv4,
            }
        })
        batch_dict.update({
            'multi_scale_3d_strides': {
                'x_conv1': 1,
                'x_conv2': 2,
                'x_conv3': 4,
                'x_conv4': 8,
            }
        })

        return batch_dict
opencood/models/sub_modules/sparse_backbone_3d.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Sizhe Wei @ 2022/11/06

'''
instance of SyncNet
'''

import torch
import torch.nn as nn
import torch.nn.functional as F
import math

class SyncLSTM(nn.Module):
    def __init__(self, channel_size = 256, spatial_size = 32, k = 2, TM_Flag = False, compressed_size = 64):
        super(SyncLSTM, self).__init__()
        self.k = k
        self.spatial_size = spatial_size
        self.channel_size = channel_size
        self.compressed_size = compressed_size
        self.time_weight = Time_Modulation(input_channel = 2*self.compressed_size)
        self.lstmcell = MotionLSTM(32, self.compressed_size)
        self.init_c = nn.parameter.Parameter(torch.rand(self.compressed_size, spatial_size, spatial_size))
        self.TM_Flag = TM_Flag

        self.ratio = int(math.sqrt(channel_size / compressed_size))
        self.conv_pre_1 = nn.Conv2d(self.channel_size, self.ratio * self.compressed_size, kernel_size=3, stride=1, padding=1)
        self.conv_pre_2 = nn.Conv2d(self.ratio * self.compressed_size, self.compressed_size, kernel_size=3, stride=1, padding=1)
        self.bn_pre_1 = nn.BatchNorm2d(self.ratio * self.compressed_size)
        self.bn_pre_2 = nn.BatchNorm2d(self.compressed_size)
        self.conv_after_1 = nn.Conv2d(self.compressed_size, self.ratio * self.compressed_size, kernel_size=3, stride=1, padding=1)
        self.conv_after_2 = nn.Conv2d(self.ratio * self.compressed_size, self.channel_size, kernel_size=3, stride=1, padding=1)
        self.bn_after_1 = nn.BatchNorm2d(self.ratio * self.compressed_size)
        self.bn_after_2 = nn.BatchNorm2d(self.channel_size)
        
    def forward(self, x_raw, delta_t):
        batch, seq, channel, h, w = x_raw.shape
        if self.compressed_size != self.channel_size:
            x = F.relu(self.bn_pre_1(self.conv_pre_1(x_raw.view(-1,channel,h,w))))
            x = F.relu(self.bn_pre_2(self.conv_pre_2(x)))
            x = x.view(batch, seq, self.compressed_size, h, w)
        else:
            x = x_raw
        
        if delta_t[0] > 0:
            self.delta_t = delta_t[0]
            h = x[:,0]
            c = self.init_c
            for i in range(1, self.k):
                h,c = self.lstmcell(x[:,i], (h,c))
            for t in range(int(self.delta_t - 1)):
                h,c = self.lstmcell(h, (h,c))
            if self.TM_Flag:
                w = self.time_weight(torch.cat([x[:,-1], h],1), delta_t[0])
                w = torch.tanh(0.1 * int(delta_t[0] - 1) * w)
                res = w * h + (1-w) * x[:,-1]
            else:
                res = h
        else:
            res = x[:,-1]
        if self.compressed_size != self.channel_size:
            res = F.relu(self.bn_after_1(self.conv_after_1(res)))
            res = F.relu(self.bn_after_2(self.conv_after_2(res)))
            # res = res.view(batch,channel,h,w)
        else:
            res = res
        return res.unsqueeze(1)

class MotionLSTM(nn.Module):
    def __init__(self, spatial_size, input_channel_size, hidden_size = 0):
        super().__init__()
        self.input_channel_size = input_channel_size  # channel size
        self.hidden_size = hidden_size
        self.spatial_size = spatial_size

        #i_t 
        # self.U_i = nn.Parameter(torch.Tensor(input_channel_size, hidden_size)) 
        # self.V_i = nn.Parameter(torch.Tensor(hidden_size, hidden_size)) 
        
        self.U_i = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.V_i = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.b_i = nn.Parameter(torch.Tensor(1, self.input_channel_size, self.spatial_size, self.spatial_size))
        
        # #f_t 
        # self.U_f = nn.Parameter(torch.Tensor(input_channel_size, hidden_size)) 
        # self.V_f = nn.Parameter(torch.Tensor(hidden_size, hidden_size)) 
        # self.b_f = nn.Parameter(torch.Tensor(hidden_size)) 
        self.U_f = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.V_f = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.b_f = nn.Parameter(torch.Tensor(1, self.input_channel_size, self.spatial_size, self.spatial_size))

        # #c_t 
        # self.U_c = nn.Parameter(torch.Tensor(input_channel_size, hidden_size)) 
        # self.V_c = nn.Parameter(torch.Tensor(hidden_size, hidden_size)) 
        # self.b_c = nn.Parameter(torch.Tensor(hidden_size)) 
        self.U_c = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.V_c = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.b_c = nn.Parameter(torch.Tensor(1, self.input_channel_size, self.spatial_size, self.spatial_size))

        # #o_t 
        # self.U_o = nn.Parameter(torch.Tensor(input_channel_size, hidden_size)) 
        # self.V_o = nn.Parameter(torch.Tensor(hidden_size, hidden_size)) 
        # self.b_o = nn.Parameter(torch.Tensor(hidden_size)) 
        self.U_o = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.V_o = STPN_MotionLSTM(height_feat_size = self.input_channel_size)
        self.b_o = nn.Parameter(torch.Tensor(1, self.input_channel_size, self.spatial_size, self.spatial_size))

        # self.init_weights()

    # def init_weights(self):
    #     stdv = 1.0 / math.sqrt(self.hidden_size)
    #     for weight in self.parameters():
    #         weight.data.uniform_(-stdv, stdv)

    def forward(self,x,init_states=None): 
        """ 
        assumes x.shape represents (batch_size, sequence_size, input_channel_size) 
        """ 
        h, c = init_states 
        i = torch.sigmoid(self.U_i(x) + self.V_i(h) + self.b_i) 
        f = torch.sigmoid(self.U_f(x) + self.V_f(h) + self.b_f) 
        g = torch.tanh(self.U_c(x) + self.V_c(h) + self.b_c) 
        o = torch.sigmoid(self.U_o(x) + self.V_o(x) + self.b_o) 
        c_out = f * c + i * g 
        h_out = o *  torch.tanh(c_out) 

        # hidden_seq.append(h_t.unsqueeze(0)) 

        # #reshape hidden_seq p/ retornar 
        # hidden_seq = torch.cat(hidden_seq, dim=0) 
        # hidden_seq = hidden_seq.transpose(0, 1).contiguous() 
        return (h_out, c_out)



class STPN_MotionLSTM(nn.Module):
    def __init__(self, height_feat_size = 16):
        super(STPN_MotionLSTM, self).__init__()



        # self.conv3d_1 = Conv3D(4, 8, kernel_size=(3, 1, 1), stride=1, padding=(0, 0, 0))
        # self.conv3d_2 = Conv3D(8, 8, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
        # self.conv3d_1 = Conv3D(64, 64, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))
        # self.conv3d_2 = Conv3D(128, 128, kernel_size=(1, 1, 1), stride=1, padding=(0, 0, 0))

        self.conv1_1 = nn.Conv2d(height_feat_size, 2*height_feat_size, kernel_size=3, stride=2, padding=1)
        self.conv1_2 = nn.Conv2d(2*height_feat_size, 2*height_feat_size, kernel_size=3, stride=1, padding=1)

        self.conv2_1 = nn.Conv2d(2*height_feat_size, 4*height_feat_size, kernel_size=3, stride=2, padding=1)
        self.conv2_2 = nn.Conv2d(4*height_feat_size, 4*height_feat_size, kernel_size=3, stride=1, padding=1)

        self.conv7_1 = nn.Conv2d(6*height_feat_size, 2*height_feat_size, kernel_size=3, stride=1, padding=1)
        self.conv7_2 = nn.Conv2d(2*height_feat_size, 2*height_feat_size, kernel_size=3, stride=1, padding=1)

        self.conv8_1 = nn.Conv2d(3*height_feat_size , height_feat_size, kernel_size=3, stride=1, padding=1)
        self.conv8_2 = nn.Conv2d(height_feat_size, height_feat_size, kernel_size=3, stride=1, padding=1)

        self.bn1_1 = nn.BatchNorm2d(2*height_feat_size)
        self.bn1_2 = nn.BatchNorm2d(2*height_feat_size)

        self.bn2_1 = nn.BatchNorm2d(4*height_feat_size)
        self.bn2_2 = nn.BatchNorm2d(4*height_feat_size)

        self.bn7_1 = nn.BatchNorm2d(2*height_feat_size)
        self.bn7_2 = nn.BatchNorm2d(2*height_feat_size)

        self.bn8_1 = nn.BatchNorm2d(1*height_feat_size)
        self.bn8_2 = nn.BatchNorm2d(1*height_feat_size)

    def forward(self, x):

        x = x.view(-1, x.size(-3), x.size(-2), x.size(-1))
        x_1 = F.relu(self.bn1_1(self.conv1_1(x)))
        x_1 = F.relu(self.bn1_2(self.conv1_2(x_1)))
        x_2 = F.relu(self.bn2_1(self.conv2_1(x_1)))
        x_2 = F.relu(self.bn2_2(self.conv2_2(x_2)))
        x_7 = F.relu(self.bn7_1(self.conv7_1(torch.cat((F.interpolate(x_2, scale_factor=(2, 2)), x_1), dim=1))))
        x_7 = F.relu(self.bn7_2(self.conv7_2(x_7)))
        x_8 = F.relu(self.bn8_1(self.conv8_1(torch.cat((F.interpolate(x_7, scale_factor=(2, 2)), x), dim=1))))
        res_x = F.relu(self.bn8_2(self.conv8_2(x_8)))

        return res_x


class Time_Modulation(nn.Module):
    def __init__(self, input_channel = 128):
        super(Time_Modulation, self).__init__()
        self.input_channel = input_channel
        self.conv1_channel = int(self.input_channel / 2)
        self.conv2_channel = int(self.conv1_channel / 2)
        # self.ratio = math.sqrt()
        self.conv1 = nn.Conv2d(self.input_channel, self.conv1_channel, kernel_size=3, stride=1, padding=1)
        self.conv2 = nn.Conv2d(self.conv1_channel, self.conv2_channel, kernel_size=3, stride=1, padding=1)
        self.convl1 = nn.Conv2d(1, 8, kernel_size=3, stride=1, padding=1)
        self.convl2 = nn.Conv2d(8,self.conv2_channel, kernel_size=3, stride=1, padding=1)
        self.conv3 = nn.Conv2d(int(2 * self.conv2_channel), 8, kernel_size = 3, stride=1, padding = 1) 
        # self.linear3 = nn.Linear(16,8)
        self.conv4 = nn.Conv2d(8, 1, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(self.conv1_channel)
        self.bn2 = nn.BatchNorm2d(self.conv2_channel)
        self.bn3 = nn.BatchNorm2d(8)
        self.bn4 = nn.BatchNorm2d(1)
        self.bnl1 = nn.BatchNorm2d(8)
        self.bnl2 = nn.BatchNorm2d(self.conv2_channel)
        self.linear3 = nn.Linear(int(2 * self.conv2_channel), 8)
        self.linear4 = nn.Linear(8, 1)
    def forward(self, x, delta_t):
        a,b,c,d = x.size()
        y = delta_t * torch.ones((a,1,c,d)).to(x.device)
        t_y = F.relu(self.bnl1(self.convl1(y)))
        t_y = F.relu(self.bnl2(self.convl2(t_y)))
        t_x = F.relu(self.bn1(self.conv1(x)))
        t_x = F.relu(self.bn2(self.conv2(t_x)))
        t_xy = torch.cat([t_x, t_y], 1)
        t_xy = F.relu(self.bn3(self.conv3(t_xy)))
        t_xy = torch.sigmoid(self.bn4(self.conv4(t_xy)))
        # t_xy = F.relu(self.bn3(self.linear3(t_xy)))
        # t_xy = torch.sigmoid(self.bn4(self.linear4(t_xy)))
        return t_xy


opencood/models/sub_modules/syncnet.py
MediaBrain-SJTU/CoBEVFlow
"""
torch_transformation_utils.py
"""
import os

import torch
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from icecream import ic

def get_roi_and_cav_mask(shape, cav_mask, spatial_correction_matrix,
                         discrete_ratio, downsample_rate):
    """
    Get mask for the combination of cav_mask and rorated ROI mask.
    Parameters
    ----------
    shape : tuple
        Shape of (B, L, H, W, C).
    cav_mask : torch.Tensor
        Shape of (B, L).
    spatial_correction_matrix : torch.Tensor
        Shape of (B, L, 4, 4)
    discrete_ratio : float
        Discrete ratio.
    downsample_rate : float
        Downsample rate.

    Returns
    -------
    com_mask : torch.Tensor
        Combined mask with shape (B, H, W, L, 1).

    """
    B, L, H, W, C = shape
    C = 1
    # (B,L,4,4)
    dist_correction_matrix = get_discretized_transformation_matrix(
        spatial_correction_matrix, discrete_ratio,
        downsample_rate)
    # (B*L,2,3)
    T = get_transformation_matrix(
        dist_correction_matrix.reshape(-1, 2, 3), (H, W))
    # (B,L,1,H,W)
    roi_mask = get_rotated_roi((B, L, C, H, W), T)
    # (B,L,1,H,W)
    com_mask = combine_roi_and_cav_mask(roi_mask, cav_mask)
    # (B,H,W,1,L)
    com_mask = com_mask.permute(0,3,4,2,1)
    return com_mask


def combine_roi_and_cav_mask(roi_mask, cav_mask):
    """
    Combine ROI mask and CAV mask

    Parameters
    ----------
    roi_mask : torch.Tensor
        Mask for ROI region after considering the spatial transformation/correction.
    cav_mask : torch.Tensor
        Mask for CAV to remove padded 0.

    Returns
    -------
    com_mask : torch.Tensor
        Combined mask.
    """
    # (B, L, 1, 1, 1)
    cav_mask = cav_mask.unsqueeze(2).unsqueeze(3).unsqueeze(4)
    # (B, L, C, H, W)
    cav_mask = cav_mask.expand(roi_mask.shape)
    # (B, L, C, H, W)
    com_mask = roi_mask * cav_mask
    return com_mask


def get_rotated_roi(shape, correction_matrix):
    """
    Get rorated ROI mask.

    Parameters
    ----------
    shape : tuple
        Shape of (B,L,C,H,W).
    correction_matrix : torch.Tensor
        Correction matrix with shape (N,2,3).

    Returns
    -------
    roi_mask : torch.Tensor
        Roated ROI mask with shape (N,2,3).

    """
    B, L, C, H, W = shape
    # To reduce the computation, we only need to calculate the
    # mask for the first channel.
    # (B,L,1,H,W)
    x = torch.ones((B, L, 1, H, W)).to(correction_matrix.dtype).to(
        correction_matrix.device)
    # (B*L,1,H,W)
    roi_mask = warp_affine(x.reshape(-1, 1, H, W), correction_matrix,
                           dsize=(H, W), mode="nearest")
    # (B,L,C,H,W)
    roi_mask = torch.repeat_interleave(roi_mask, C, dim=1).reshape(B, L, C, H,
                                                                   W)
    return roi_mask


def get_discretized_transformation_matrix(matrix, discrete_ratio,
                                          downsample_rate):
    """
    Get disretized transformation matrix.
    Parameters
    ----------
    matrix : torch.Tensor
        Shape -- (B, L, 4, 4) where B is the batch size, L is the max cav
        number.
    discrete_ratio : float
        Discrete ratio.
    downsample_rate : float or int
        downsample_rate

    discrete_ratio * downsample_rate = ___ meter one pixel, in the current feature map.

    Returns
    -------
    matrix : torch.Tensor
        Output transformation matrix in 2D with shape (B, L, 2, 3),
        including 2D transformation and 2D rotation. 
        transformation is pixel level

    """
    matrix = matrix[:, :, [0, 1], :][:, :, :, [0, 1, 3]]
    # normalize the x,y transformation
    matrix[:, :, :, -1] = matrix[:, :, :, -1] \
                          / (discrete_ratio * downsample_rate)

    return matrix.type(dtype=torch.float)


def _torch_inverse_cast(input):
    r"""
    Helper function to make torch.inverse work with other than fp32/64.
    The function torch.inverse is only implemented for fp32/64 which makes
    impossible to be used by fp16 or others. What this function does,
    is cast input data type to fp32, apply torch.inverse,
    and cast back to the input dtype.
    Args:
        input : torch.Tensor
            Tensor to be inversed.

    Returns:
        out : torch.Tensor
            Inversed Tensor.

    """
    dtype = input.dtype
    if dtype not in (torch.float32, torch.float64):
        dtype = torch.float32
    out = torch.inverse(input.to(dtype)).to(input.dtype)
    return out


def normal_transform_pixel(
        height, width, device, dtype, eps=1e-14):
    r"""
    Compute the normalization matrix from image size in pixels to [-1, 1].
    Args:
        height : int
            Image height.
        width : int
            Image width.
        device : torch.device
            Output tensor devices.
        dtype : torch.dtype
            Output tensor data type.
        eps : float
            Epsilon to prevent divide-by-zero errors.

    Returns:
        tr_mat : torch.Tensor
            Normalized transform with shape :math:`(1, 3, 3)`.
    """
    tr_mat = torch.tensor(
        [[1.0, 0.0, -1.0], [0.0, 1.0, -1.0], [0.0, 0.0, 1.0]], device=device,
        dtype=dtype)  # 3x3

    # prevent divide by zero bugs
    width_denom = eps if width == 1 else width - 1.0
    height_denom = eps if height == 1 else height - 1.0

    tr_mat[0, 0] = tr_mat[0, 0] * 2.0 / width_denom
    tr_mat[1, 1] = tr_mat[1, 1] * 2.0 / height_denom

    return tr_mat.unsqueeze(0)  # 1x3x3


def eye_like(n, B, device, dtype):
    r"""
    Return a 2-D tensor with ones on the diagonal and
    zeros elsewhere with the same batch size as the input.
    Args:
        n : int
            The number of rows :math:`(n)`.
        B : int
            Btach size.
        device : torch.device
            Devices of the output tensor.
        dtype : torch.dtype
            Data type of the output tensor.

    Returns:
       The identity matrix with the shape :math:`(B, n, n)`.
    """

    identity = torch.eye(n, device=device, dtype=dtype)
    return identity[None].repeat(B, 1, 1)


def normalize_homography(dst_pix_trans_src_pix, dsize_src, dsize_dst=None):
    r"""
    Normalize a given homography in pixels to [-1, 1].
    Args:
        dst_pix_trans_src_pix : torch.Tensor
            Homography/ies from source to destination to be normalized with
            shape :math:`(B, 3, 3)`.
        dsize_src : Tuple[int, int]
            Size of the source image (height, width).
        dsize_dst : Tuple[int, int]
            Size of the destination image (height, width).

    Returns:
        dst_norm_trans_src_norm : torch.Tensor
            The normalized homography of shape :math:`(B, 3, 3)`.
    """
    if dsize_dst is None:
        dsize_dst = dsize_src
    # source and destination sizes
    src_h, src_w = dsize_src
    dst_h, dst_w = dsize_dst
    device = dst_pix_trans_src_pix.device
    dtype = dst_pix_trans_src_pix.dtype
    # compute the transformation pixel/norm for src/dst
    src_norm_trans_src_pix = normal_transform_pixel(src_h, src_w, device,
                                                    dtype).to(
        dst_pix_trans_src_pix)

    src_pix_trans_src_norm = _torch_inverse_cast(src_norm_trans_src_pix)
    dst_norm_trans_dst_pix = normal_transform_pixel(dst_h, dst_w, device,
                                                    dtype).to(
        dst_pix_trans_src_pix)
    # compute chain transformations
    dst_norm_trans_src_norm: torch.Tensor = dst_norm_trans_dst_pix @ (
            dst_pix_trans_src_pix @ src_pix_trans_src_norm)
    return dst_norm_trans_src_norm


def get_rotation_matrix2d(M, dsize):
    r"""
    Return rotation matrix for torch.affine_grid based on transformation matrix.
    Args:
        M : torch.Tensor
            Transformation matrix with shape :math:`(B, 2, 3)`.
        dsize : Tuple[int, int]
            Size of the source image (height, width).

    Returns:
        R : torch.Tensor
            Rotation matrix with shape :math:`(B, 2, 3)`.
    """
    H, W = dsize
    B = M.shape[0]
    center = torch.Tensor([W / 2, H / 2]).to(M.dtype).to(M.device).unsqueeze(0)
    shift_m = eye_like(3, B, M.device, M.dtype)
    shift_m[:, :2, 2] = center

    shift_m_inv = eye_like(3, B, M.device, M.dtype)
    shift_m_inv[:, :2, 2] = -center

    rotat_m = eye_like(3, B, M.device, M.dtype)
    rotat_m[:, :2, :2] = M[:, :2, :2]
    affine_m = shift_m @ rotat_m @ shift_m_inv
    return affine_m[:, :2, :]  # Bx2x3


def get_transformation_matrix(M, dsize):
    r"""
    Return transformation matrix for torch.affine_grid.
    Args:
        M : torch.Tensor
            Transformation matrix with shape :math:`(N, 2, 3)`.
        dsize : Tuple[int, int]
            Size of the source image (height, width).

    Returns:
        T : torch.Tensor
            Transformation matrix with shape :math:`(N, 2, 3)`.
    """
    T = get_rotation_matrix2d(M, dsize)
    T[..., 2] += M[..., 2]
    return T


def convert_affinematrix_to_homography(A):
    r"""
    Convert to homography coordinates
    Args:
        A : torch.Tensor
            The affine matrix with shape :math:`(B,2,3)`.

    Returns:
        H : torch.Tensor
            The homography matrix with shape of :math:`(B,3,3)`.
    """
    H: torch.Tensor = torch.nn.functional.pad(A, [0, 0, 0, 1], "constant",
                                              value=0.0)
    H[..., -1, -1] += 1.0
    return H


def warp_affine_simple(src, M, dsize,
        mode='bilinear',
        padding_mode='zeros',
        align_corners=False):

    B, C, H, W = src.size()
    grid = F.affine_grid(M,
                         [B, C, dsize[0], dsize[1]],
                         align_corners=align_corners).to(src) # .to() 统一数据格式 float32
    return F.grid_sample(src, grid, align_corners=align_corners)

def warp_affine(
        src, M, dsize,
        mode='bilinear',
        padding_mode='zeros',
        align_corners=True):
    r"""
    Transform the src based on transformation matrix M.
    Args:
        src : torch.Tensor
            Input feature map with shape :math:`(B,C,H,W)`.
        M : torch.Tensor
            Transformation matrix with shape :math:`(B,2,3)`.
        dsize : tuple
            Tuple of output image H_out and W_out.
        mode : str
            Interpolation methods for F.grid_sample.
        padding_mode : str
            Padding methods for F.grid_sample.
        align_corners : boolean
            Parameter of F.affine_grid.

    Returns:
        Transformed features with shape :math:`(B,C,H,W)`.
    """

    B, C, H, W = src.size()

    # we generate a 3x3 transformation matrix from 2x3 affine
    M_3x3 = convert_affinematrix_to_homography(M)
    dst_norm_trans_src_norm = normalize_homography(M_3x3, (H, W), dsize)

    # src_norm_trans_dst_norm = torch.inverse(dst_norm_trans_src_norm)
    src_norm_trans_dst_norm = _torch_inverse_cast(dst_norm_trans_src_norm)
    
    grid = F.affine_grid(src_norm_trans_dst_norm[:, :2, :],
                         [B, C, dsize[0], dsize[1]],
                         align_corners=align_corners)

    return F.grid_sample(src.half() if grid.dtype==torch.half else src, 
                         grid, align_corners=align_corners, mode=mode,
                         padding_mode=padding_mode)


class Test:
    """
    Test the transformation in this file.
    The methods in this class are not supposed to be used outside of this file.
    """

    def __init__(self):
        pass

    @staticmethod
    def load_img():
        torch.manual_seed(0)
        x = torch.randn(1, 5, 16, 400, 200) * 100
        # x = torch.ones(1, 5, 16, 400, 200)
        return x

    @staticmethod
    def load_raw_transformation_matrix(N):
        a = 90 / 180 * np.pi
        matrix = torch.Tensor([[np.cos(a), -np.sin(a), 10],
                               [np.sin(a), np.cos(a), 10]])
        matrix = torch.repeat_interleave(matrix.unsqueeze(0).unsqueeze(0), N,
                                         dim=1)
        return matrix

    @staticmethod
    def load_raw_transformation_matrix2(N, alpha):
        a = alpha / 180 * np.pi
        matrix = torch.Tensor([[np.cos(a), -np.sin(a), 0, 0],
                               [np.sin(a), np.cos(a), 0, 0]])
        matrix = torch.repeat_interleave(matrix.unsqueeze(0).unsqueeze(0), N,
                                         dim=1)
        return matrix

    @staticmethod
    def test():
        img = Test.load_img()
        B, L, C, H, W = img.shape
        raw_T = Test.load_raw_transformation_matrix(5)
        T = get_transformation_matrix(raw_T.reshape(-1, 2, 3), (H, W))
        img_rot = warp_affine(img.reshape(-1, C, H, W), T, (H, W))
        print(img_rot[0, 0, :, :])
        plt.matshow(img_rot[0, 0, :, :])
        plt.show()

    @staticmethod
    def test_combine_roi_and_cav_mask():
        B = 2
        L = 5
        C = 16
        H = 300
        W = 400
        # 2, 5
        cav_mask = torch.Tensor([[1, 1, 1, 0, 0], [1, 0, 0, 0, 0]])
        x = torch.zeros(B, L, C, H, W)
        correction_matrix = Test.load_raw_transformation_matrix2(5, 10)
        correction_matrix = torch.cat([correction_matrix, correction_matrix],
                                      dim=0)
        mask = get_roi_and_cav_mask((B, L, H, W, C), cav_mask, 
                                    correction_matrix, 0.4, 4)
        plt.matshow(mask[0, :, :, 0, 0])
        plt.show()



if __name__ == "__main__":
    os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
    Test.test_combine_roi_and_cav_mask()

opencood/models/sub_modules/torch_transformation_utils.py
MediaBrain-SJTU/CoBEVFlow
from icecream import ic
import torch
import math
import torch.nn as nn
from opencood.models.sub_modules.torch_transformation_utils import warp_affine_simple
from opencood.utils.transformation_utils import pose_to_tfm, tfm_to_pose_torch, tfm_to_xycs_torch, xycs_to_tfm_torch

def regroup(x, record_len):
    cum_sum_len = torch.cumsum(record_len, dim=0)
    split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
    return split_x

# Part1, Pose Regression Module

class PoseRegression(nn.Module):
    """
    Args:
        in_ch: 2*C

    forward:
        x: [N,2C,H,W] concatenated feature

    Returns:
        [N, 3]: x, y, yaw
    
    """
    def __init__(self, in_ch=512, hidden_ch=256):
        super(PoseRegression, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_ch, hidden_ch, kernel_size=(3, 3), padding=1),
            nn.LeakyReLU(negative_slope=0.01),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
            nn.Conv2d(hidden_ch, hidden_ch, kernel_size=(3, 3), padding=1),
            nn.LeakyReLU(negative_slope=0.01),
            nn.MaxPool2d(kernel_size=2, stride=2, padding=0),
            nn.Conv2d(hidden_ch, hidden_ch, kernel_size=(3, 3), padding=1),
            nn.LeakyReLU(negative_slope=0.01),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(
                hidden_ch, hidden_ch, kernel_size=(3, 3), stride=(2, 2), padding=1
            ),
            nn.LeakyReLU(negative_slope=0.01),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.AdaptiveAvgPool2d(output_size=1),
            nn.Flatten(),
            nn.Linear(in_features=hidden_ch, out_features=hidden_ch, bias=True),
            nn.LeakyReLU(negative_slope=0.01),
            nn.Linear(in_features=hidden_ch, out_features=hidden_ch, bias=True),
            nn.LeakyReLU(negative_slope=0.01),
            nn.Linear(in_features=hidden_ch, out_features=3, bias=True),
        )

    def forward(self, x):

        pose_reg = self.model(x)
        return pose_reg



class PoseRegressionWraper(nn.Module):
    """
    Args:
        features: [sum(cav), C, H, W], 
        record_len: list
        pairwise_t_matrix: [B, L, L, 4, 4], original pairwise_t_matrix, noise contains
    Retuens:
        pairwise_t_matrix_new: [B, L, L, 4, 4], the relative pose after correction.
    """
    def __init__(self, in_ch, hidden_ch, affine_parameter):
        super(PoseRegressionWraper, self).__init__()
        self.pose_regression = PoseRegression(
            in_ch=in_ch, hidden_ch=hidden_ch
        )
        self.H = affine_parameter['H']
        self.W = affine_parameter['W']
        self.downsample_rate = affine_parameter['downsample_rate']
        self.discrete_ratio = affine_parameter['discrete_ratio']

    def forward(self, features, record_len, pairwise_t_matrix):
        _, C, H, W = features.shape
        B, L = pairwise_t_matrix.shape[:2]
        split_x = regroup(features, record_len)
        pairwise_t_matrix_new = torch.eye(4, device=pairwise_t_matrix.device).view(1,1,1,4,4).repeat(B,L,L,1,1)
        pose_corr_matrix = torch.zeros((B,L,L,3),device=pairwise_t_matrix.device)
        for b in range(B):
            N = record_len[b]
            agent_features = split_x[b]
            for i in range(N):
                t_matrix = pairwise_t_matrix[b]
                t_matrix = t_matrix[:,:,[0, 1],:][:,:,:,[0, 1, 3]] # [L, L, 2, 3]
                t_matrix[...,0,1] = t_matrix[...,0,1] * H / W
                t_matrix[...,1,0] = t_matrix[...,1,0] * W / H
                t_matrix[...,0,2] = t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
                t_matrix[...,1,2] = t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2
                # (N,C,H,W)
                neighbors = warp_affine_simple(agent_features, t_matrix[i, :N, :, :], (self.H, self.W))
                # (N,C,H,W)
                ego_agent_feature = agent_features[i].unsqueeze(0).repeat(N, 1, 1, 1)
                # (N,2C,H,W)
                neighbor_feature = torch.cat(
                        [neighbors, ego_agent_feature], dim=1)
                # (N,3)
                pose_corr = self.pose_regression(neighbor_feature)
                pose_corr_matrix[b,i,:N] = pose_corr

                # (N, 4, 4)
                pose_corr_tfm = pose_to_tfm(pose_corr)
                pairwise_t_matrix_new[b,i,:N] = pose_corr_tfm @ pairwise_t_matrix[b,i,:N]

        return pose_corr_matrix, pairwise_t_matrix_new



# Part 2, Global Consistent Module
def get_intersection(pairwise_t_matrix, affine_parameter):
    """ get intersection from pairwise_t_matrix

    Args: 
        pairwise_t_matrix: torch.Tensor, shape [L, L, 4, 4]
            pairwise transformation matrix for one frame.
            pairwise_t_matrix[i,j] = Tji, i is ego
        affine_parameter: dict
            H, W, etc.


    Returns:
        intersection: torch.Tensor, shape [L, L]
    """
    H = affine_parameter['H']
    W = affine_parameter['W']
    downsample_rate = affine_parameter['downsample_rate']
    discrete_ratio = affine_parameter['discrete_ratio']
    intersections = []

    L = pairwise_t_matrix.shape[0]
    one_tensor = torch.zeros((L,1,H,W), device=pairwise_t_matrix.device)
    for i in range(L):
        t_matrix = pairwise_t_matrix[:,:,[0, 1],:][:,:,:,[0, 1, 3]] # [L, L, 2, 3]
        t_matrix[...,0,1] = t_matrix[...,0,1] * H / W
        t_matrix[...,1,0] = t_matrix[...,1,0] * W / H
        t_matrix[...,0,2] = t_matrix[...,0,2] / (downsample_rate * discrete_ratio * W) * 2
        t_matrix[...,1,2] = t_matrix[...,1,2] / (downsample_rate * discrete_ratio * H) * 2

        # [L,1,H,W]
        neighbors = warp_affine_simple(one_tensor, t_matrix[i, :L, :, :], (H, W))
        intersection = torch.sum(neighbors, dim=[1,2,3]) / (H * W)  # [L,]
        intersections.append(intersection)

    # [L, L], intersections[i,:], ego is i
    intersections = torch.stack(intersections)

    # if intersection is zero, may meet nan later.
    eps = 0.01
    intersections += eps
    
    return intersections




def WeightedMLE(pose, pairwise_t_matrix, weight):
    """ Weighted MLE for estimate mu and sigma of multivariate student t distribution.
        simutanously for all nodes
    Args:
        pose: [N,3]
        pairwise_t_matrix: [L, L, 4, 4]
        weight: [L, L]

    Returns:
        pose_mu: [N, 3] , but [N, 4] now
        pose_sigma: [N, 3, 3], but [N, 4] now
    """

    N = pose.shape[0]
    mu_list = []
    sigma_list = []

    for i in range(N):

        neighbor_ids = list(range(N))
        neighbor_ids.remove(i)

        weights = weight[i,neighbor_ids].repeat(2) # [2(N-1)]
        relative_pose1 = pairwise_t_matrix[i,neighbor_ids] # [N-1, 4, 4]  Tji
        relative_pose2 = pairwise_t_matrix[neighbor_ids,i] # [N-1, 4, 4]  Tij
        relative_pose2 = torch.inverse(relative_pose2)
        relative_pose = torch.cat([relative_pose1,relative_pose2], dim=0) # [2(N-1), 4, 4]

        tfm = pose_to_tfm(pose[neighbor_ids]).repeat(2,1,1) # [2(N-1), 4, 4]
        samples = tfm @ relative_pose # [2(N-1), 4, 4]
        # here is one problem, -179 and +179 degree. They are close actually.
        # so we use cos and sin to replace angle
        samples = tfm_to_xycs_torch(samples).to(torch.float64) # [N, 4]


        mu = samples.median(0).values
        Sigma = torch.eye(4, device=pose.device, dtype=torch.float64)
        small_identity = torch.eye(4, device=pose.device, dtype=torch.float64) * 0.05

        diff = mu[None] - samples

        v = 2
        for _ in range(15):
            eta = (v + mu.size(0)) / (
                v + torch.einsum("ni,ij,nj->n", diff, Sigma.inverse(), diff)
            )
            mu = torch.einsum("n,n,ni->i", weights, eta, samples) / (weights * eta).sum()
            diff = mu[None] - samples
            # Sigma = torch.einsum('n,n,ni,nj->ij', weights, w, diff, diff) / weights.sum()
            Sigma = (
                torch.einsum("n,ni,nj->ij", eta, diff, diff) / diff.size(0) + small_identity
            )

        mu_list.append(mu.to(torch.float32))
        sigma_list.append(Sigma.to(torch.float32))

    pose_mu = torch.stack(mu_list)
    pose_sigma = torch.stack(sigma_list)

    return pose_mu, pose_sigma


def WeightedEM(lidar_pose, pairwise_t_matrix, intersection):
    """Weighted EM algorithm, for a single frame, not batch data
    Args:
        lidar_pose : torch.Tenosr
            shape [N, 3]
        pairwise_t_matrix: torch.Tensor
            shape [L, L, 4, 4]
        intersection: torch.Tensor
            shape [L, L]

    Returns:
        pose_mu : torch.Tensor
            new lidar pose after correction. shape [N, 3]
    """
    num_iters = 10
    pose = lidar_pose
    weight = torch.ones_like(intersection, device=intersection.device)
    
    for k in range(num_iters):
        pose_mu, pose_sigma = WeightedMLE(pose, pairwise_t_matrix, weight) # [N, 4], [N, 4, 4]
        weight = update_weight(pose_mu, pose_sigma, pairwise_t_matrix, intersection)
    
    N = lidar_pose.shape[0]
    pose_new = torch.zeros((N,3), device=lidar_pose.device, dtype=lidar_pose.dtype)
    pose_new[:,:2] = pose_mu[:,:2]
    pose_new[:,2] = torch.rad2deg(torch.atan2(pose_mu[:,3], pose_mu[:,2])) # sin, cos

    return pose_new

def update_weight(pose_mu, pose_sigma, pairwise_t_matrix, intersection):
    """ using the close form to update weight w.
    Args:
        pose_mu: [N,3], but [N, 4] now 
        pose_sigma: [N, 3, 3], but [N, 4, 4] now 
        pairwise_t_matrix: [L,L,4,4]
        interesection: [L, L]
    """
    k = 120
    df = 2 # degree of freedom
    L = intersection.shape[0]
    N = pose_mu.shape[0]
    weight = torch.zeros_like(intersection, device=intersection.device)
    for i in range(N):
        for j in range(N):
            if i!=j:
                pose_estimate1 = xycs_to_tfm_torch(pose_mu[[j]])[0] @ pairwise_t_matrix[i,j] # [4,4]
                pose_estimate2 = xycs_to_tfm_torch(pose_mu[[i]])[0] @ torch.inverse(pairwise_t_matrix[i,j]) # [4,4]
                pose_estimate = torch.stack([pose_estimate1, pose_estimate2]) # [2, 4, 4]
                pose_estimate = tfm_to_xycs_torch(pose_estimate) # [2, 4]
                weight[i,j] = k * intersection[i,j] / (k - log_t(pose_estimate, pose_mu[i], pose_sigma[i], df).sum())

    return weight



def log_t(x, mu, Sigma, df): 
    """ log pdf of t distribution
    Args:
        x: [N, 3]
        mu: [3,]
        Sigma: [3,3]
        df: int, degree of freedom

    Returns:
        log_pdf: log of the pdf
    """

    assert len(x.shape) == 2
    n, p = x.shape
    # assert mu.shape[0] == p   # for now, allow multiple mu
    assert Sigma.shape == (p, p)

    v = torch.as_tensor(df, dtype=x.dtype, device=x.device)
    p = torch.as_tensor(p, dtype=x.dtype, device=x.device)
    pi = torch.tensor(math.pi, dtype=x.dtype, device=x.device)
    half_v = v / 2.0
    half_p = p / 2.0

    log_num = (half_v + half_p).lgamma()
    log_denom = half_v.lgamma() + half_p * (v.log() + pi.log()) + 0.5 * Sigma.logdet()

    d = x - mu
    log_val = -(half_p + half_v) * torch.log(
        1 + torch.einsum("ni,ij,nj->n", d, Sigma.inverse(), d) / v
    )

    log_pdf = log_num - log_denom + log_val

    return log_pdf


# Part 3, Attention Module

class Attention(nn.Module):
    """
    Args:
        in_ch: 2*C

    forward:
        x: [N,2C,H,W] concatenated feature
    
    """
    def __init__(self, in_ch, hidden_ch=160):
        super(Attention, self).__init__()
        self.model = nn.Sequential(
            nn.Conv2d(in_ch, hidden_ch, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.01),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.Conv2d(hidden_ch, hidden_ch, 3, 1, 1),
            nn.LeakyReLU(negative_slope=0.01),
            nn.MaxPool2d(kernel_size=2, stride=2),
            nn.AdaptiveMaxPool2d(output_size=1),
            nn.Flatten(),
            nn.Linear(in_features=hidden_ch, out_features=1, bias=True),
            nn.Sigmoid(),
        )

    def forward(self, x):
        out = self.model(x)
        return out

class AttentionWrapper(nn.Module):
    """ wrapper of attention scoring
    Args:
        features: [sum(cav), C, H, W], 
        record_len: list
        pairwise_t_matrix: [B, L, L, 4, 4], original pairwise_t_matrix, noise contains
    Retuens:
        pairwise_score: [B, L, L]
            pairwise_score[i,j], ego is i.
    """
    def __init__(self, in_ch, hidden_ch, affine_parameter, learnable_alpha=True):
        super(AttentionWrapper, self).__init__()
        self.attention_net = Attention(in_ch, hidden_ch)
        self.H = affine_parameter['H']
        self.W = affine_parameter['W']
        self.downsample_rate = affine_parameter['downsample_rate']
        self.discrete_ratio = affine_parameter['discrete_ratio']
        if learnable_alpha:
            self.alpha = nn.Parameter(torch.Tensor([0.15]))
        else:
            self.alpha = 0.35
    
    def forward(self, features, record_len, pairwise_t_matrix):
        _, C, H, W = features.shape
        B, L = pairwise_t_matrix.shape[:2]
        split_x = regroup(features, record_len)
        pairwise_score = torch.zeros((B, L, L), device=features.device)
        # mask = torch.eye(L, device=features.device).expand(B,L,L)


        for b in range(B):
            N = record_len[b]
            agent_features = split_x[b]
            for i in range(N):
                t_matrix = pairwise_t_matrix[b]
                t_matrix = t_matrix[:,:,[0, 1],:][:,:,:,[0, 1, 3]] # [L, L, 2, 3]
                t_matrix[...,0,1] = t_matrix[...,0,1] * self.H / self.W
                t_matrix[...,1,0] = t_matrix[...,1,0] * self.W / self.H
                t_matrix[...,0,2] = t_matrix[...,0,2] / (self.downsample_rate * self.discrete_ratio * W) * 2
                t_matrix[...,1,2] = t_matrix[...,1,2] / (self.downsample_rate * self.discrete_ratio * H) * 2

                # (N,C,H,W)
                neighbors = warp_affine_simple(agent_features, t_matrix[i, :N, :, :], (self.H, self.W))
                # (N,C,H,W)
                ego_agent_feature = agent_features[i].unsqueeze(0).repeat(N, 1, 1, 1)
                # (N,2C,H,W)
                neighbor_feature = torch.cat(
                        [neighbors, ego_agent_feature], dim=1)
                # (N,1)
                pairwise_score[b,i,:N] = self.attention_net(neighbor_feature).flatten()
        
        # pairwise_score *= mask

        scores = pairwise_score
        eps = 1e-4
        # pairwise_score (B, L, L). pairwise_score[b,i,j] is agent j' feature warping to agent i's coordinate
        # weight (B, L, L), normalized at dim=2
        weight = pairwise_score / (torch.sum(pairwise_score, dim=2, keepdim=True) + self.alpha + eps)

        return scores, weight

opencood/models/sub_modules/v2v_robust_module.py
MediaBrain-SJTU/CoBEVFlow
import math
import torch
import torch.nn as nn

from opencood.models.sub_modules.base_transformer import *
from opencood.models.sub_modules.hmsa import *
from opencood.models.sub_modules.mswin import *
from opencood.models.sub_modules.torch_transformation_utils import \
    get_transformation_matrix, warp_affine, get_roi_and_cav_mask, \
    get_discretized_transformation_matrix


class STTF(nn.Module):
    def __init__(self, args):
        super(STTF, self).__init__()
        self.discrete_ratio = args['voxel_size'][0]
        self.downsample_rate = args['downsample_rate']

    def forward(self, x, mask, spatial_correction_matrix):
        x = x.permute(0, 1, 4, 2, 3)
        dist_correction_matrix = get_discretized_transformation_matrix(
            spatial_correction_matrix, self.discrete_ratio,
            self.downsample_rate)
        # Only compensate non-ego vehicles
        B, L, C, H, W = x.shape

        T = get_transformation_matrix(
            dist_correction_matrix[:, 1:, :, :].reshape(-1, 2, 3), (H, W))
        cav_features = warp_affine(x[:, 1:, :, :, :].reshape(-1, C, H, W), T,
                                   (H, W))
        cav_features = cav_features.reshape(B, -1, C, H, W)
        x = torch.cat([x[:, 0, :, :, :].unsqueeze(1), cav_features], dim=1)
        x = x.permute(0, 1, 3, 4, 2)
        return x


class RelTemporalEncoding(nn.Module):
    """
    Implement the Temporal Encoding (Sinusoid) function.
    """

    def __init__(self, n_hid, RTE_ratio, max_len=100, dropout=0.2):
        super(RelTemporalEncoding, self).__init__()
        position = torch.arange(0., max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, n_hid, 2) *
                             -(math.log(10000.0) / n_hid))
        emb = nn.Embedding(max_len, n_hid)
        emb.weight.data[:, 0::2] = torch.sin(position * div_term) / math.sqrt(
            n_hid)
        emb.weight.data[:, 1::2] = torch.cos(position * div_term) / math.sqrt(
            n_hid)
        emb.requires_grad = False
        self.RTE_ratio = RTE_ratio
        self.emb = emb
        self.lin = nn.Linear(n_hid, n_hid)

    def forward(self, x, t):
        # When t has unit of 50ms, rte_ratio=1.
        # So we can train on 100ms but test on 50ms
        return x + self.lin(self.emb(t * self.RTE_ratio)).unsqueeze(
            0).unsqueeze(1)


class RTE(nn.Module):
    def __init__(self, dim, RTE_ratio=2):
        super(RTE, self).__init__()
        self.RTE_ratio = RTE_ratio

        self.emb = RelTemporalEncoding(dim, RTE_ratio=self.RTE_ratio)

    def forward(self, x, dts):
        # x: (B,L,H,W,C)
        # dts: (B,L)
        rte_batch = []
        for b in range(x.shape[0]):
            rte_list = []
            for i in range(x.shape[1]):
                rte_list.append(
                    self.emb(x[b, i, :, :, :], dts[b, i]).unsqueeze(0))
            rte_batch.append(torch.cat(rte_list, dim=0).unsqueeze(0))
        return torch.cat(rte_batch, dim=0)


class V2XFusionBlock(nn.Module):
    def __init__(self, num_blocks, cav_att_config, pwindow_config):
        super().__init__()
        # first multi-agent attention and then multi-window attention
        self.layers = nn.ModuleList([])
        self.num_blocks = num_blocks

        for _ in range(num_blocks):
            att = HGTCavAttention(cav_att_config['dim'],
                                  heads=cav_att_config['heads'],
                                  dim_head=cav_att_config['dim_head'],
                                  dropout=cav_att_config['dropout']) if \
                cav_att_config['use_hetero'] else \
                CavAttention(cav_att_config['dim'],
                             heads=cav_att_config['heads'],
                             dim_head=cav_att_config['dim_head'],
                             dropout=cav_att_config['dropout'])
            self.layers.append(nn.ModuleList([
                PreNorm(cav_att_config['dim'], att),
                PreNorm(cav_att_config['dim'],
                        PyramidWindowAttention(pwindow_config['dim'],
                                               heads=pwindow_config['heads'],
                                               dim_heads=pwindow_config[
                                                   'dim_head'],
                                               drop_out=pwindow_config[
                                                   'dropout'],
                                               window_size=pwindow_config[
                                                   'window_size'],
                                               relative_pos_embedding=
                                               pwindow_config[
                                                   'relative_pos_embedding'],
                                               fuse_method=pwindow_config[
                                                   'fusion_method']))]))

    def forward(self, x, mask, prior_encoding):
        for cav_attn, pwindow_attn in self.layers:
            x = cav_attn(x, mask=mask, prior_encoding=prior_encoding) + x
            x = pwindow_attn(x) + x
        return x


class V2XTEncoder(nn.Module):
    def __init__(self, args):
        super().__init__()

        cav_att_config = args['cav_att_config']
        pwindow_att_config = args['pwindow_att_config']
        feed_config = args['feed_forward']

        num_blocks = args['num_blocks']
        depth = args['depth']
        mlp_dim = feed_config['mlp_dim']
        dropout = feed_config['dropout']

        self.downsample_rate = args['sttf']['downsample_rate']
        self.discrete_ratio = args['sttf']['voxel_size'][0]
        self.use_roi_mask = args['use_roi_mask']
        self.use_RTE = cav_att_config['use_RTE']
        self.RTE_ratio = cav_att_config['RTE_ratio']
        self.sttf = STTF(args['sttf'])
        # adjust the channel numbers from 256+3 -> 256
        self.prior_feed = nn.Linear(cav_att_config['dim'] + 3,
                                    cav_att_config['dim'])
        self.layers = nn.ModuleList([])
        if self.use_RTE:
            self.rte = RTE(cav_att_config['dim'], self.RTE_ratio)
        for _ in range(depth):
            self.layers.append(nn.ModuleList([
                V2XFusionBlock(num_blocks, cav_att_config, pwindow_att_config),
                PreNorm(cav_att_config['dim'],
                        FeedForward(cav_att_config['dim'], mlp_dim,
                                    dropout=dropout))
            ]))

    def forward(self, x, mask, spatial_correction_matrix):

        # transform the features to the current timestamp
        # velocity, time_delay, infra
        # (B,L,H,W,3)
        prior_encoding = x[..., -3:]
        # (B,L,H,W,C)
        x = x[..., :-3]
        if self.use_RTE:
            # dt: (B,L)
            dt = prior_encoding[:, :, 0, 0, 1].to(torch.int)
            x = self.rte(x, dt)
        x = self.sttf(x, mask, spatial_correction_matrix)
        com_mask = mask.unsqueeze(1).unsqueeze(2).unsqueeze(
            3) if not self.use_roi_mask else get_roi_and_cav_mask(x.shape,
                                                                  mask,
                                                                  spatial_correction_matrix,
                                                                  self.discrete_ratio,
                                                                  self.downsample_rate)
        for attn, ff in self.layers:
            x = attn(x, mask=com_mask, prior_encoding=prior_encoding)
            x = ff(x) + x
        return x


class V2XTransformer(nn.Module):
    def __init__(self, args):
        super(V2XTransformer, self).__init__()

        encoder_args = args['encoder']
        self.encoder = V2XTEncoder(encoder_args)

    def forward(self, x, mask, spatial_correction_matrix):
        output = self.encoder(x, mask, spatial_correction_matrix)
        output = output[:, 0]
        return output
opencood/models/sub_modules/v2xvit_basic.py
MediaBrain-SJTU/CoBEVFlow
import copy
import random

import torch
import torch.nn as nn

from opencood.pcdet_utils.pointnet2.pointnet2_stack import pointnet2_modules as pointnet2_stack_modules
from opencood.pcdet_utils.pointnet2.pointnet2_stack import pointnet2_utils as pointnet2_stack_utils
from opencood.pcdet_utils.roiaware_pool3d.roiaware_pool3d_utils import points_in_boxes_gpu
from opencood.utils import common_utils


def bilinear_interpolate_torch(im, x, y):
    """
    Args:
        im: (H, W, C) [y, x]
        x: (N)
        y: (N)
    Returns:
    """
    x0 = torch.floor(x).long()
    x1 = x0 + 1

    y0 = torch.floor(y).long()
    y1 = y0 + 1

    x0 = torch.clamp(x0, 0, im.shape[1] - 1)
    x1 = torch.clamp(x1, 0, im.shape[1] - 1)
    y0 = torch.clamp(y0, 0, im.shape[0] - 1)
    y1 = torch.clamp(y1, 0, im.shape[0] - 1)

    Ia = im[y0, x0]
    Ib = im[y1, x0]
    Ic = im[y0, x1]
    Id = im[y1, x1]

    wa = (x1.type_as(x) - x) * (y1.type_as(y) - y)
    wb = (x1.type_as(x) - x) * (y - y0.type_as(y))
    wc = (x - x0.type_as(x)) * (y1.type_as(y) - y)
    wd = (x - x0.type_as(x)) * (y - y0.type_as(y))
    ans = torch.t((torch.t(Ia) * wa)) + torch.t(torch.t(Ib) * wb) + torch.t(torch.t(Ic) * wc) + torch.t(torch.t(Id) * wd)
    return ans


class VoxelSetAbstraction(nn.Module):
    def __init__(self, model_cfg, voxel_size, point_cloud_range, num_bev_features=None,
                 num_rawpoint_features=None, **kwargs):
        super().__init__()
        self.model_cfg = model_cfg
        self.voxel_size = voxel_size
        self.point_cloud_range = point_cloud_range

        SA_cfg = self.model_cfg['sa_layer']

        self.SA_layers = nn.ModuleList()
        self.SA_layer_names = []
        self.downsample_times_map = {}
        c_in = 0
        for src_name in self.model_cfg['features_source']:
            if src_name in ['bev', 'raw_points']:
                continue
            self.downsample_times_map[src_name] = SA_cfg[src_name]['downsample_factor']
            mlps = copy.copy(SA_cfg[src_name]['mlps'])
            for k in range(len(mlps)):
                mlps[k] = [mlps[k][0]] + mlps[k]
            cur_layer = pointnet2_stack_modules.StackSAModuleMSG(
                radii=SA_cfg[src_name]['pool_radius'],
                nsamples=SA_cfg[src_name]['n_sample'],
                mlps=mlps,
                use_xyz=True,
                pool_method='max_pool',
            )
            self.SA_layers.append(cur_layer)
            self.SA_layer_names.append(src_name)

            c_in += sum([x[-1] for x in mlps])

        if 'bev' in self.model_cfg['features_source']:
            c_bev = num_bev_features
            c_in += c_bev

        if 'raw_points' in self.model_cfg['features_source']:
            mlps = copy.copy(SA_cfg['raw_points']['mlps'])
            for k in range(len(mlps)):
                mlps[k] = [num_rawpoint_features - 3] + mlps[k]

            self.SA_rawpoints = pointnet2_stack_modules.StackSAModuleMSG(
                radii=SA_cfg['raw_points']['pool_radius'],
                nsamples=SA_cfg['raw_points']['n_sample'],
                mlps=mlps,
                use_xyz=True,
                pool_method='max_pool'
            )
            c_in += sum([x[-1] for x in mlps])

        self.vsa_point_feature_fusion = nn.Sequential(
            nn.Linear(c_in, self.model_cfg['num_out_features'], bias=False),
            nn.BatchNorm1d(self.model_cfg['num_out_features']),
            nn.ReLU(),
        )
        self.num_point_features = self.model_cfg['num_out_features']
        self.num_point_features_before_fusion = c_in

    def interpolate_from_bev_features(self, keypoints, bev_features, batch_size, bev_stride):
        x_idxs = (keypoints[:, :, 0] - self.point_cloud_range[0]) / self.voxel_size[0]
        y_idxs = (keypoints[:, :, 1] - self.point_cloud_range[1]) / self.voxel_size[1]
        x_idxs = x_idxs / bev_stride
        y_idxs = y_idxs / bev_stride

        point_bev_features_list = []
        for k in range(batch_size):
            cur_x_idxs = x_idxs[k]
            cur_y_idxs = y_idxs[k]
            cur_bev_features = bev_features[k].permute(1, 2, 0)  # (H, W, C)
            point_bev_features = bilinear_interpolate_torch(cur_bev_features, cur_x_idxs, cur_y_idxs)
            point_bev_features_list.append(point_bev_features.unsqueeze(dim=0))

        point_bev_features = torch.cat(point_bev_features_list, dim=0)  # (B, N, C0)
        return point_bev_features

    def get_sampled_points(self, batch_dict):
        batch_size = batch_dict['batch_size']
        if self.model_cfg['point_source'] == 'raw_points':
            src_points = batch_dict['origin_lidar_for_vsa'][:, 1:]
            batch_indices = batch_dict['origin_lidar_for_vsa'][:, 0].long()
        elif self.model_cfg['point_source'] == 'voxel_centers':
            src_points = common_utils.get_voxel_centers(
                batch_dict['voxel_coords'][:, 1:4],
                downsample_times=1,
                voxel_size=self.voxel_size,
                point_cloud_range=self.point_cloud_range
            )
            batch_indices = batch_dict['voxel_coords'][:, 0].long()
        else:
            raise NotImplementedError

        keypoints_batch = torch.randn((batch_size, self.model_cfg['num_keypoints'], 4), device=src_points.device)
        keypoints_batch[..., 0] = keypoints_batch[..., 0] * 140
        keypoints_batch[..., 1] = keypoints_batch[..., 0] * 40
        # points with height flag 10 are padding/invalid, for later filtering
        keypoints_batch[..., 2] = 10.0
        for bs_idx in range(batch_size):
            bs_mask = (batch_indices == bs_idx)
            sampled_points = src_points[bs_mask].unsqueeze(dim=0)  # (1, N, 3)
            # sample points with FPS
            # some cropped pcd may have very few points, select various number
            # of points to ensure similar sample density
            # 50000 is approximately the number of points in one full pcd
            num_kpts = int(self.model_cfg['num_keypoints'] * sampled_points.shape[1] / 50000) + 1
            num_kpts = min(num_kpts, self.model_cfg['num_keypoints'])
            cur_pt_idxs = pointnet2_stack_utils.furthest_point_sample(
                sampled_points[:, :, 0:3].contiguous(), num_kpts
            ).long()

            if sampled_points.shape[1] < num_kpts:
                empty_num = num_kpts - sampled_points.shape[1]
                cur_pt_idxs[0, -empty_num:] = cur_pt_idxs[0, :empty_num]

            keypoints = sampled_points[0][cur_pt_idxs[0]].unsqueeze(dim=0)

            keypoints_batch[bs_idx, :len(keypoints[0]), :] = keypoints

        # keypoints = torch.cat(keypoints_list, dim=0)  # (B, M, 3)
        return keypoints_batch

    def forward(self, batch_dict):
        """
        Args:
            batch_dict:
                batch_size:
                keypoints: (B, num_keypoints, 3)
                multi_scale_3d_features: {
                        'x_conv4': ...
                    }
                points: optional (N, 1 + 3 + C) [bs_idx, x, y, z, ...]
                spatial_features: optional
                spatial_features_stride: optional
        Returns:
            point_features: (N, C)
            point_coords: (N, 4)
        """
        keypoints = self.get_sampled_points(batch_dict) # BxNx4
        kpt_mask1 = torch.logical_and(keypoints[..., 2] > -2.8, keypoints[..., 2] < 1.0)
        kpt_mask2 = None
        # Only select the points that are in the predicted bounding boxes
        if 'det_boxes' in batch_dict:
            dets_list = batch_dict['det_boxes']
            max_len = max([len(dets) for dets in dets_list])
            boxes = torch.zeros((len(dets_list), max_len, 7), dtype=dets_list[0].dtype,
                                device=dets_list[0].device)
            for i, dets in enumerate(dets_list):
                if len(dets)==0:
                    continue
                cur_dets = dets.clone()
                if self.model_cfg['enlarge_selection_boxes']:
                    cur_dets[:, 3:6] += 0.5
                boxes[i, :len(dets)] = cur_dets
            # mask out some keypoints to spare the GPU storage
            kpt_mask2 = points_in_boxes_gpu(keypoints[..., :3], boxes) >= 0

        kpt_mask = torch.logical_and(kpt_mask1, kpt_mask2) if kpt_mask2 is not None else kpt_mask1
        # Ensure there are more than 2 points are selected to satisfy the
        # condition of batch norm in the FC layers of feature fusion module
        if (kpt_mask).sum() < 2:
            kpt_mask[0, torch.randint(0, 1024, (2,))] = True


        point_features_list = []
        if 'bev' in self.model_cfg['features_source']:
            point_bev_features = self.interpolate_from_bev_features(
                keypoints[..., :3], batch_dict['spatial_features'], batch_dict['batch_size'],
                bev_stride=batch_dict['spatial_features_stride']
            )
            point_features_list.append(point_bev_features[kpt_mask])

        batch_size, num_keypoints, _ = keypoints.shape

        new_xyz = keypoints[kpt_mask]
        new_xyz_batch_cnt = torch.tensor([(mask).sum() for mask in kpt_mask], device=new_xyz.device).int()

        if 'raw_points' in self.model_cfg['features_source']:
            raw_points = batch_dict['origin_lidar_for_vsa']
            xyz = raw_points[:, 1:4]
            xyz_batch_cnt = xyz.new_zeros(batch_size).int()
            indices = raw_points[:, 0].long()
            for bs_idx in range(batch_size):
                xyz_batch_cnt[bs_idx] = (indices == bs_idx).sum()
            point_features = None

            pooled_points, pooled_features = self.SA_rawpoints(
                xyz=xyz.contiguous(),
                xyz_batch_cnt=xyz_batch_cnt,
                new_xyz=new_xyz[:, :3].contiguous(),
                new_xyz_batch_cnt=new_xyz_batch_cnt,
                features=point_features,
            )
            point_features_list.append(pooled_features)

        for k, src_name in enumerate(self.SA_layer_names):
            cur_coords = batch_dict['multi_scale_3d_features'][src_name].indices
            xyz = common_utils.get_voxel_centers(
                cur_coords[:, 1:4],
                downsample_times=self.downsample_times_map[src_name],
                voxel_size=self.voxel_size,
                point_cloud_range=self.point_cloud_range
            )
            xyz_batch_cnt = xyz.new_zeros(batch_size).int()
            for bs_idx in range(batch_size):
                xyz_batch_cnt[bs_idx] = (cur_coords[:, 0] == bs_idx).sum()

            pooled_points, pooled_features = self.SA_layers[k](
                xyz=xyz.contiguous(),
                xyz_batch_cnt=xyz_batch_cnt,
                new_xyz=new_xyz[:, :3].contiguous(),
                new_xyz_batch_cnt=new_xyz_batch_cnt,
                features=batch_dict['multi_scale_3d_features'][src_name].features.contiguous(),
            )

            point_features_list.append(pooled_features)

        point_features = torch.cat(point_features_list, dim=1)
        batch_dict['point_features_before_fusion'] = point_features.view(-1, point_features.shape[-1])  # torch.Size([373, 640])
        point_features = self.vsa_point_feature_fusion(point_features.view(-1, point_features.shape[-1]))  # (0): Linear(in_features=512, out_features=32, bias=False)

        cur_idx = 0
        batch_dict['point_features'] = []
        batch_dict['point_coords'] = []
        for num in new_xyz_batch_cnt:
            batch_dict['point_features'].append(point_features[cur_idx:cur_idx + num])
            batch_dict['point_coords'].append(new_xyz[cur_idx:cur_idx + num])
            cur_idx += num

        return batch_dict
opencood/models/sub_modules/vsa.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
VoxelNet for intermediate fusion
"""
import torch.nn as nn
import torch.nn.functional as F
import torch
import numpy as np
from torch.autograd import Variable

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


# conv2d + bn + relu
class Conv2d(nn.Module):

    def __init__(self, in_channels, out_channels, k, s, p, activation=True,
                 batch_norm=True, bias=True):
        super(Conv2d, self).__init__()
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=k,
                              stride=s, padding=p, bias=bias)
        if batch_norm:
            self.bn = nn.BatchNorm2d(out_channels)
        else:
            self.bn = None
        self.activation = activation

    def forward(self, x):
        x = self.conv(x)
        if self.bn is not None:
            x = self.bn(x)
        if self.activation:
            return F.relu(x, inplace=True)
        else:
            return x


class NaiveFusion(nn.Module):

    def __init__(self):
        super(NaiveFusion, self).__init__()
        self.conv1 = Conv2d(128 * 5, 256, 3, 1, 1,
                            batch_norm=False, bias=False)
        self.conv2 = Conv2d(256, 128, 3, 1, 1)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)

        return x


class VoxelNetIntermediate(nn.Module):
    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'])

        self.proj_first = True
        if ('proj_first' in args) and (args['proj_first'] is False):
            self.proj_first = False

        self.cml = CML()
        self.fusion_net = AttFusion(128)
        self.rpn = RPN(args['anchor_num'])

        self.N = args['N']
        self.D = args['D']
        self.H = args['H']
        self.W = args['W']
        self.T = args['T']
        self.anchor_num = args['anchor_num']

        self.compression = False
        if 'compression' in args and args['compression'] > 0:
            self.compression = True
            self.compression_layer = AutoEncoder(128, args['compression'])

    def voxel_indexing(self, sparse_features, coords):
        dim = sparse_features.shape[-1]

        dense_feature = Variable(
            torch.zeros(dim, self.N, self.D, self.H, self.W).cuda())

        dense_feature[:, coords[:, 0], coords[:, 1], coords[:, 2],
        coords[:, 3]] = sparse_features.transpose(0, 1)

        return dense_feature.transpose(0, 1)

    def regroup(self, dense_feature, record_len):
        """
        Regroup the data based on the record_len.

        Parameters
        ----------
        dense_feature : torch.Tensor
            N, C, H, W
        record_len : list
            [sample1_len, sample2_len, ...]

        Returns
        -------
        regroup_feature : torch.Tensor
            B, 5C, H, W
        """
        cum_sum_len = list(np.cumsum(record_len))
        split_features = torch.tensor_split(dense_feature,
                                            cum_sum_len[:-1])
        regroup_features = []

        for split_feature in split_features:
            # M, C, H, W
            feature_shape = split_feature.shape

            # the maximum M is 5 as most 5 cavs
            padding_len = 5 - feature_shape[0]
            padding_tensor = torch.zeros(padding_len, feature_shape[1],
                                         feature_shape[2], feature_shape[3])
            padding_tensor = padding_tensor.to(split_feature.device)

            split_feature = torch.cat([split_feature, padding_tensor],
                                      dim=0)

            # 1, 5C, H, W
            split_feature = split_feature.view(-1,
                                               feature_shape[2],
                                               feature_shape[3]).unsqueeze(0)
            regroup_features.append(split_feature)

        # B, 5C, H, W
        regroup_features = torch.cat(regroup_features, dim=0)

        return regroup_features

    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']
        record_len = data_dict['record_len']
        pairwise_t_matrix = data_dict['pairwise_t_matrix']


        batch_dict = {'voxel_features': voxel_features,
                      'voxel_coords': voxel_coords,
                      'voxel_num_points': voxel_num_points,
                      'pairwise_t_matrix': pairwise_t_matrix}

        if voxel_coords.is_cuda:
            record_len_tmp = record_len.cpu()

        record_len_tmp = list(record_len_tmp.numpy())

        self.N = sum(record_len_tmp)

        # 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)

        # convolutional middle network
        vwfs = self.cml(vwfs)
        # convert from 3d to 2d N C H W
        vmfs = vwfs.view(self.N, -1, self.H, self.W)

        # compression layer
        if self.compression:
            vmfs = self.compression_layer(vmfs)

        # pairwise_t_matrix
        # project_first must be right
        # have not check project_first=False
        pairwise_t_matrix = pairwise_t_matrix[:,:,:,[0, 1],:][:,:,:,:,[0, 1, 3]] # [B, L, L, 2, 3]

        if not self.proj_first:
            pairwise_t_matrix[...,0,1] = pairwise_t_matrix[...,0,1] * self.H / self.W
            pairwise_t_matrix[...,1,0] = pairwise_t_matrix[...,1,0] * self.W / self.H
            pairwise_t_matrix[...,0,2] = pairwise_t_matrix[...,0,2] / self.W * 2
            pairwise_t_matrix[...,1,2] = pairwise_t_matrix[...,1,2] / self.H * 2


        # information naive fusion
        vmfs_fusion = self.fusion_net(vmfs, record_len, pairwise_t_matrix)

        # region proposal network
        # merge the depth and feature dim into one, output probability score
        # map and regression map
        psm, rm = self.rpn(vmfs_fusion)

        output_dict = {'psm': psm,
                       'rm': rm}

        return output_dict

opencood/models/voxel_net_intermediate.py
MediaBrain-SJTU/CoBEVFlow
# 验证 timestamp 的数量和regular 的关系
import os
import sys
import numpy as np 
import time
from tqdm import tqdm
from tqdm.contrib import tenumerate
from tqdm.auto import trange

def extract_timestamps(yaml_files):
    """
    Given the list of the yaml files, extract the mocked timestamps.

    Parameters
    ----------
    yaml_files : list
        The full path of all yaml files of ego vehicle

    Returns
    -------
    timestamps : list
        The list containing timestamps only.
    """
    timestamps = []

    for file in yaml_files:
        res = file.split('/')[-1]
        if res.endswith('.yaml'):
            timestamp = res.replace('.yaml', '')
        elif res.endswith('.json'):
            timestamp = res.replace('.json', '')
        else:
            print("Woops! There is no processing method for file {}".format(res))
            sys.exit(1)
        timestamps.append(timestamp)

    return timestamps

# 处理的数据集
data_seperate_name = "train"
# pcd 文件路径
root_dir = "/GPFS/public/OPV2V_irregular_npy/" + data_seperate_name

# first load all paths of different scenarios
scenario_folders = sorted([os.path.join(root_dir, x)
                            for x in os.listdir(root_dir) if
                            os.path.isdir(os.path.join(root_dir, x))])
scenario_folders_name = sorted([x
                            for x in os.listdir(root_dir) if
                            os.path.isdir(os.path.join(root_dir, x))])

# loop over all scenarios
for (i, scenario_folder) in enumerate(scenario_folders):
    # if i!=43:
    #     continue
    print("==== scenario %s ====" % (scenario_folders_name[i]))
    # start_time = time.time()
    
    # copy timestamps npy file
    timestamps_file = os.path.join(scenario_folder, 'timestamps.npy')
    time_annotations = np.load(timestamps_file)

    # iterate all cav in this scenario
    cav_list = sorted([x 
                        for x in os.listdir(scenario_folder) if 
                        os.path.isdir(os.path.join(scenario_folder, x))], 
                        key=lambda y: int(y))
    assert len(cav_list) > 0

    yaml_files = sorted([x
                        for x in os.listdir(os.path.join(scenario_folder, cav_list[0])) if
                        x.endswith(".json")], 
                        key=lambda y:float((y.split('/')[-1]).split('.json')[0]))
    start_timestamp = int(float(extract_timestamps(yaml_files)[0]))
    while(1):
        time_id_json = ("%.3f" % float(start_timestamp)) + ".json"
        time_id_yaml = ("%.3f" % float(start_timestamp)) + ".yaml"
        if not (time_id_json in yaml_files or time_id_yaml in yaml_files):
            start_timestamp += 1
        else:
            break

    end_timestamp = int(float(extract_timestamps(yaml_files)[-1]))
    if start_timestamp%2 == 0:
        # even
        end_timestamp = end_timestamp-1 if end_timestamp%2==1 else end_timestamp
    else:
        end_timestamp = end_timestamp-1 if end_timestamp%2==0 else end_timestamp
    num_timestamps = int((end_timestamp - start_timestamp)/2 + 1)

    all_timestamps = []
    for (j, cav_id) in enumerate(cav_list):  
        if j==0: # ego
            timestamps = [start_timestamp+2*i for i in range(num_timestamps)]
        else:
            timestamps = list(time_annotations[j-1, :])
        all_timestamps.append(timestamps)
        print("cav %s has %d timestamps" % (cav_id, len(timestamps)))
    all_timestamps = np.array(all_timestamps)
    order_flag = (all_timestamps[1:, :] - all_timestamps[0, :]) < 0
    if int(order_flag.max()) > 0:
        print("Fuck!")
        

        # # use the frame number as key, the full path as the values
        # yaml_files = \
        #     sorted([os.path.join(cav_path_yaml, x)
        #             for x in os.listdir(cav_path_yaml) if
        #             x.endswith('.yaml')])
        # timestamps = extract_timestamps(yaml_files)	
        
    """
        for k, timestamp in tenumerate(timestamps):
            # if k > 30:
            #     break

            # new save dir
            save_path = os.path.join(save_dir, scenario_folders_name[i], cav_id)
            if not os.path.exists(save_path):
                os.makedirs(save_path)
                # print("======Path {} created successfully!======".format(save_path))

            # change pcd to npy 
            lidar_file = os.path.join(cav_path_pcd, timestamp + '.pcd')
            new_lidar_file = os.path.join(save_path, timestamp + '.npy')
            if not os.path.exists(new_lidar_file):
                lidar = pcd_to_np(lidar_file)
                np.save(new_lidar_file, lidar)

            # change yaml to json
            yaml_file = os.path.join(cav_path_yaml, timestamp + '.yaml')
            json_file = os.path.join(cav_path_yaml, timestamp + '.json')
            new_json_file = os.path.join(save_path, timestamp + '.json')
            if os.path.exists(json_file): # 原数据集有json文件
                if not os.path.exists(new_json_file):
                    shutil.copyfile(json_file, new_json_file)
            else:
                if not os.path.exists(new_json_file):
                    params = load_yaml(yaml_file)
                    json_str = json.dumps(params, indent=4)
                    with open(new_json_file, 'w') as json_file:
                        json_file.write(json_str)
            
            # new_yaml_file = os.path.join(save_path, timestamp + '.yaml')
            # if not os.path.exists(new_yaml_file):
            #     shutil.copyfile(yaml_file, new_yaml_file)

            # print(timestamp)
    end_time = time.time()
    print("=== %d-th scenario (%s) finished, consumed %.3f mins. ===" \
        % (i, scenario_folders_name[i], (end_time-start_time)/60.0))
    """
opencood/test.py
MediaBrain-SJTU/CoBEVFlow
import torch
import numpy as np
import opencood.visualization.simple_plot3d.canvas_bev as canvas_bev
import sys
import os
import glob
from opencood.utils.transformation_utils import x1_to_x2
from opencood.utils.box_utils import project_box3d
from collections import OrderedDict
from icecream import ic
from matplotlib import pyplot as plt

def draw_trajectory(pt_file, save_path):
    '''
    box_dict : {
        'single_detection_bbx': # dict, [0, 1, ... , N-1],  
            [i]: dict:{
                [0]/[1]/[2]: past 3 frames detection results:{ # 都是在[0]的各自坐标系下的检测结果。
                    pred_box_3dcorner_tensor
                    pred_box_center_tensor
                    scores
                }
            }
        }
        'lidar_pose_0': 过去第一帧, 所有车都的pose
        'lidar_pose_current': 当前帧, 所有车都的pose
        'compensated_results_list': # len=N-1, 每个non-ego补偿出来的box, each of which is [N_obj, 4, 3], 注意这里用的是4个点, 最后一个维度上是 xyz, z是多余的
        'matched_idx_list': matched_idx_list, # len=N-1, 表示每个non-ego的过去2/3帧的匹配id, each of which is [N_obj, 2or3], 比如 ['matched_idx_list'][0] shape(22,2) 表示过去第一帧的22个框与第二帧的22个框的索引的匹配情况
        'single_gt_box_tensor': # list, len=N, 表示ego与non-ego每辆车在current时刻的检测框结果, 例如box_dict['single_gt_box_tensor'][1]大小为[N_obj, 8, 3] 表示第二辆车在current时刻的检测框结果。在各自current时刻坐标系下
        'single_lidar': # list, len=N, 表示ego与non-ego每辆车在current时刻的lidar np
        'gt_range': # [-140.8, -40, -3, 140.8, 40, 1], 表示lidar的范围
    }

    TODO
    1. transform detection results (from past 1 frame coordinate) into current frame
    2. visualzie GT box
    3. visualize detection results
    4. draw lidar point cloud
    5. clustering the same object (matched_idx_list)
    '''
    data = torch.load(pt_file)

    agent_num = len(data["single_detection_bbx"])
    left_hand = True
    pc_range = data['gt_range']

    for agent_id in range(1, agent_num):
        # step0: initailize canvas
        canvas = canvas_bev.Canvas_BEV_heading_right(canvas_shape=(int((pc_range[4]-pc_range[1])*10), int((pc_range[3]-pc_range[0])*10)),
                                                canvas_x_range=(pc_range[0], pc_range[3]), 
                                                canvas_y_range=(pc_range[1], pc_range[4]),
                                                left_hand=left_hand) 
        # step 1: transform detection results
        lidar_pose_past = data['lidar_pose_0'][agent_id].cpu().numpy()
        lidar_pose_current = data['lidar_pose_current'][agent_id].cpu().numpy()

        past_to_current = x1_to_x2(lidar_pose_past, lidar_pose_current)
        past_to_current = torch.from_numpy(past_to_current).to('cuda').float()
        boxes_past_in_past_coor_dict = data['single_detection_bbx'][agent_id] 
        boxes_past_in_current_coor_dict = OrderedDict()

        for past_frame, detections in boxes_past_in_past_coor_dict.items():
            if past_frame == 'past_k_time_diff':
                continue
            boxes_past_in_current_coor_dict[past_frame] = OrderedDict()
            boxes_past_in_current_coor_dict[past_frame]['pred_box_3dcorner_tensor'] = \
                project_box3d(detections['pred_box_3dcorner_tensor'].float(), past_to_current)
            boxes_past_in_current_coor_dict[past_frame]['scores'] = detections['scores']


        # step 2: draw GT box
        gt_box_np = data['single_gt_box_tensor'][agent_id].cpu().numpy()
        canvas.draw_boxes(gt_box_np,colors=(0,255,0), texts=['gt']*gt_box_np.shape[0])

        # step 3: visualze detection result, compensation, Establish inverted indices
        matched_idx_list = data['matched_idx_list']
        match_pairs = matched_idx_list[agent_id - 1]
        label_1 = [''] * boxes_past_in_current_coor_dict[1]['pred_box_3dcorner_tensor'].shape[0]
        label_0 = [''] * boxes_past_in_current_coor_dict[0]['pred_box_3dcorner_tensor'].shape[0]
        
        for pair_idx, pair in enumerate(match_pairs):
            idx_in_0 = pair[0]
            idx_in_1 = pair[1]
            label_1[idx_in_1] = pair_idx
            label_0[idx_in_0] = pair_idx

        ## past 2 frame (key is 1)
        pred_box_np_1 = boxes_past_in_current_coor_dict[1]['pred_box_3dcorner_tensor'].cpu().numpy()
        canvas.draw_boxes(pred_box_np_1,colors=(int(255*0.5),0,0), texts=label_1)

        ## past 1 frame (key is 0)
        pred_box_np_0 = boxes_past_in_current_coor_dict[0]['pred_box_3dcorner_tensor'].cpu().numpy()
        canvas.draw_boxes(pred_box_np_0,colors=(int(255*0.8),0,0), texts=label_0)

        ## compensate box
        compensate_box_np = data['compensated_results_list'][agent_id-1].cpu().numpy()

        if compensate_box_np.shape[0] == 0:
            print("Error: no compensate box: ", pt_file, " ", agent_id)
            return
        compensate_box_np = np.concatenate([compensate_box_np, compensate_box_np], axis=1)
        compensate_box_np = project_box3d(compensate_box_np, past_to_current.cpu().numpy())
        canvas.draw_boxes(compensate_box_np,colors=(0, int(255*0.83), int(255*0.83)), texts=['cp']*compensate_box_np.shape[0])




        # step 4: draw lidar point cloud
        pcd_np = data['single_lidar'][agent_id].cpu().numpy()
        canvas_xy, valid_mask = canvas.get_canvas_coords(pcd_np) # Get Canvas Coords
        canvas.draw_canvas_points(canvas_xy[valid_mask]) # Only draw valid points


        plt.axis("off")
        plt.imshow(canvas.canvas)
        plt.tight_layout()
        save_name = pt_file.split("/")[-1].rstrip(".pt")
        plt.savefig(f"{save_path}{save_name}_{agent_id}.png", dpi=400)
        plt.close()


def main():
    data_dir = '/remote-home/share/sizhewei/logs/irv2v_where2comm_cobevflow_w_dir_finetune/vis_cobevflow_reverse_0.50/bbx_folder' #'/remote-home/share/sizhewei/logs/irv2v_where2comm_cobevflow_w_dir_finetune/vis_cobevflow_521_trick_reverse_new_0.50_0'

    files = glob.glob(os.path.join(data_dir, 'bbx_folder', '*.pt'))
    files.sort()

    save_path = os.path.join(data_dir, 'multi_results')
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    for file in files:
        draw_trajectory(file, save_path)

if __name__ == "__main__":
    main()


opencood/tools/box_drawing_multi_frame.py
MediaBrain-SJTU/CoBEVFlow
import argparse
import yaml
import opencood.hypes_yaml.yaml_utils as yaml_utils

def get_parser():
    parser = argparse.ArgumentParser(description="synthetic data generation")
    parser.add_argument("--hypes_yaml", "-y", type=str, required=True,
                            help='data generation yaml file needed ')
    parser.add_argument('--model_dir', default='',
                        help='Continued training path')
    parser.add_argument("--save_name", "-s", type=str, required=True,
                            help='yaml save path ')
    opt = parser.parse_args()
    return opt
if __name__ == "__main__":
    opt = get_parser()
    hypes = yaml_utils.load_yaml(opt.hypes_yaml, opt)
    with open(opt.save_name, 'w') as outfile:
        yaml.dump(hypes, outfile)
opencood/tools/config_generate.py
MediaBrain-SJTU/CoBEVFlow
'''
Prediction model input data generation.
SizheWei 2023.4.23
'''

import os
import torch
from scipy import stats
from collections import OrderedDict

def retrieve_base_data(scenario_database, len_record, idx, binomial_n=10, binomial_p=0.1, k=3, is_no_shift=False, is_same_sample_interval=False):
    """
    Given the index, return the corresponding data.

    Parameters
    ----------
    idx : int
        Index given by dataloader.

    Returns
    -------
    data : dict
        The dictionary contains loaded yaml params and lidar data for
        each cav.
        Structure: 
        {
            cav_id_1 : {
                'ego' : true,
                curr : {
                    'params': json_path,
                    'timestamp': string
                },
                past_k : {		                # (k) totally
                    [0]:{
                        'params': json_path,
                        'timestamp': string,
                        'time_diff': float,
                        'sample_interval': int
                    },
                    [1] : {},	(id-1)
                    ...,		
                    [k-1] : {} (id-(k-1))
                },
            }, 
            cav_id_2 : {		                # (k) totally
                'ego': false, 
                ...
            }, 
            ...
        }
    """
    sample_interval_exp = int(binomial_n * binomial_p)
    # we loop the accumulated length list to get the scenario index
    scenario_index = 0
    for i, ele in enumerate(len_record):
        if idx < ele:
            scenario_index = i
            break
    scenario_database = scenario_database[scenario_index]
    
    # 生成冻结分布函数
    bernoulliDist = stats.bernoulli(binomial_p) 

    data = OrderedDict()
    # 找到 current 时刻的 timestamp_index 这对于每辆车来讲都一样
    curr_timestamp_idx = idx if scenario_index == 0 else \
                    idx - len_record[scenario_index - 1]
    curr_timestamp_idx = curr_timestamp_idx + binomial_n * k
    
    # load files for all CAVs
    for cav_id, cav_content in scenario_database.items():
        '''
        cav_content 
        {
            timestamp_1 : {
                yaml: path,
                lidar: path, 
                cameras: list of path
            },
            ...
            timestamp_n : {
                yaml: path,
                lidar: path, 
                cameras: list of path
            },
            'regular' : {
                timestamp_1 : {},
                ...
                timestamp_n : {}
            },
            'ego' : true / false , 
        },
        '''
        data[cav_id] = OrderedDict()
        
        # 1. set 'ego' flag
        data[cav_id]['ego'] = cav_content['ego']
        
        # 2. current frame, for co-perception lable use
        data[cav_id]['curr'] = {}

        timestamp_key = list(cav_content['regular'].items())[curr_timestamp_idx][0]
        
        # 2.1 load curr params
        # json is faster than yaml
        json_file = cav_content['regular'][timestamp_key]['yaml'].replace("yaml", "json")
        json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")
        data[cav_id]['curr']['params'] = json_file

        # 2.3 store curr timestamp and time_diff
        data[cav_id]['curr']['timestamp'] = timestamp_key
        data[cav_id]['curr']['time_diff'] = 0.0
        data[cav_id]['curr']['sample_interval'] = 0

        # 3. past frames, for model input
        data[cav_id]['past_k'] = OrderedDict()
        latest_sample_stamp_idx = curr_timestamp_idx
        # past k frames, pose | lidar | label(for single view confidence map generator use)
        for i in range(k):
            # sample_interval
            if data[cav_id]['ego']:             # ego sample_interval = E(B(n, p))
                if i == 0: # ego-past-0 与 ego-curr 是一样的
                    data[cav_id]['past_k'][i] = data[cav_id]['curr']
                    continue
                sample_interval = sample_interval_exp
                if sample_interval == 0:
                    sample_interval = 1
            else:                               # non-ego sample_interval ~ B(n, p)
                if sample_interval_exp==0 \
                    and is_no_shift \
                        and i == 0:
                    data[cav_id]['past_k'][i] = data[cav_id]['curr']
                    continue
                if is_same_sample_interval:
                    sample_interval = sample_interval_exp
                else:
                    # B(n, p)
                    trails = bernoulliDist.rvs(binomial_n)
                    sample_interval = sum(trails)
                if sample_interval==0:
                    if i==0: # 检查past 0 的实际时间是否在curr 的后面
                        tmp_time_key = list(cav_content.items())[latest_sample_stamp_idx][0]
                        if dist_time(tmp_time_key, data[cav_id]['curr']['timestamp'])>0:
                            sample_interval = 1
                    if i>0: # 过去的几帧不要重复
                        sample_interval = 1                

            # check the timestamp index
            data[cav_id]['past_k'][i] = {}
            latest_sample_stamp_idx -= sample_interval
            timestamp_key = list(cav_content.items())[latest_sample_stamp_idx][0]
            # load the corresponding data into the dictionary
            # load param file: json is faster than yaml
            json_file = cav_content[timestamp_key]['yaml'].replace("yaml", "json")
            json_file = json_file.replace("OPV2V_irregular_npy", "OPV2V_irregular_npy_updated")
            data[cav_id]['past_k'][i]['params'] = json_file

            data[cav_id]['past_k'][i]['timestamp'] = timestamp_key
            data[cav_id]['past_k'][i]['sample_interval'] = sample_interval
            data[cav_id]['past_k'][i]['time_diff'] = \
                dist_time(timestamp_key, data[cav_id]['curr']['timestamp'])

    return data

def dist_time(ts1, ts2, i = -1):
    """caculate the time interval between two timestamps

    Args:
        ts1 (string): time stamp at some time
        ts2 (string): current time stamp
        i (int, optional): past frame id, for debug use. Defaults to -1.
    
    Returns:
        time_diff (float): time interval (ts1 - ts2)
    """
    if not i==-1:
        return -i
    else:
        return (float(ts1) - float(ts2))


# debug_path = '/remote-home/share/OPV2V_irregular_npy'
debug_path = '/dssg/home/acct-seecsh/seecsh/sizhewei/data_sftp'
scenario_database = torch.load(os.path.join(debug_path, 'scenario_database.pt'))
len_record = torch.load(os.path.join(debug_path, 'len_record.pt'))


unit_data = retrieve_base_data(scenario_database, len_record, 1)
opencood/tools/fornewdataset.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Sizhe Wei <sizhewei@sjtu.edu.cn>

import argparse
import os
import time
from typing import OrderedDict
import sys
sys.path.append(os.getcwd())

import torch
import open3d as o3d
from torch.utils.data import DataLoader
import numpy as np

import opencood.hypes_yaml.yaml_utils as yaml_utils
from opencood.tools import train_utils, inference_utils
from opencood.data_utils.datasets import build_dataset
from opencood.utils import eval_utils, box_utils
from opencood.visualization import vis_utils, my_vis, simple_vis

from tqdm import tqdm
from tqdm.contrib import tenumerate
from tqdm.auto import trange

def test_parser():
    parser = argparse.ArgumentParser(description="synthetic data generation")
    parser.add_argument('--model_dir', type=str, required=True,
                        help='Continued training path')
    parser.add_argument('--fusion_method', type=str,
                        default='intermediate',
                        help='no, no_w_uncertainty, late, early or intermediate')
    parser.add_argument('--save_vis_interval', type=int, default=40,
                        help='interval of saving visualization')
    parser.add_argument('--save_npy', action='store_true',
                        help='whether to save prediction and gt result'
                             'in npy file')
    parser.add_argument('--pretrained_path', default='',
                        help='The path of the model need to be fine tuned.')
    parser.add_argument('--note', default="ir_thre_0_d_20", type=str, help='save folder name')
    parser.add_argument('--p', default=None, type=float, help='binomial probability')
    parser.add_argument('--ir_range', default=None, type=float, help='binomial probability')
    parser.add_argument('--two_stage', help='whether to use two stage training', default=0, type=int)
    parser.add_argument('--config_suffix', default='', type=str, help='config suffix')
    parser.add_argument('--dataset', default='o', type=str, choices=['o', 'd'], help='which dataset will be used, o is for OPV2V/IRV2V, d is for DAIR-V2X.')
    opt = parser.parse_args()
    return opt


def main():
    opt = test_parser()

    assert opt.fusion_method in ['late', 'early', 'intermediate', 'no', 'no_w_uncertainty'] 

    hypes = yaml_utils.load_yaml(None, opt, config_suffix=opt.config_suffix)
    
    hypes['validate_dir'] = hypes['test_dir']
    if "OPV2V" in hypes['test_dir'] or "v2xsim" in hypes['test_dir']:
        assert "test" in hypes['validate_dir']
    
    # update binomial prob
    if 'binomial_p' in hypes or opt.p is not None:
        hypes['binomial_p'] = opt.p
    if 'binomial_p' not in hypes and opt.p is None:
        hypes['binomial_p'] = 0.0
        print(f"!!! No binomial probability is set, using default value: {hypes['binomial_p']} !!!")

    if 'binomial_n' not in hypes:
        hypes['binomial_n'] = 1
        print(f"!!! No binomial n is set, using default value: {hypes['binomial_n']} !!!")

    hypes['time_delay'] = int(hypes['binomial_p'] * hypes['binomial_n'])
    
    if 'ir_range' in hypes and opt.ir_range is not None:
        hypes['ir_range'] = opt.ir_range
    if 'ir_range' not in hypes and opt.ir_range is None:
        hypes['ir_range'] = 0
        print(f"!!! No ir range is set, using default value: {hypes['ir_range']} !!!")
    
    viz_bbx_flag = False
    if 'viz_bbx_flag' in hypes:
        viz_bbx_flag = hypes['viz_bbx_flag']

    # This is used in visualization
    # left hand: IRV2V, OPV2V, V2XSET
    # right hand: V2X-Sim 2.0 and DAIR-V2X
    left_hand = True if "OPV2V" in hypes['test_dir'] or "V2XSET" in hypes['test_dir'] else False

    # control comm volume
    num_roi_thres = -1
    if 'num_roi_thres' in hypes:
        num_roi_thres = hypes['num_roi_thres']

    # add pose noise:
    noise = {'pos_std':0, 'rot_std':0, 'pos_mean':0, 'rot_mean':0}
    if 'noise' in hypes:
        noise = hypes['noise']
    noise_note = f"noise_{noise['pos_std']}_{noise['rot_std']}_{noise['pos_mean']}_{noise['rot_mean']}"

    print(f"Left hand visualizing: {left_hand}")

    if 'box_align' in hypes.keys():
        hypes['box_align']['val_result'] = hypes['box_align']['test_result']
        
    print('Creating Model')
    model = train_utils.create_model(hypes)
    # we assume gpu is necessary
    if torch.cuda.is_available():
        model.cuda()
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    print('Loading Model from checkpoint')
    saved_path = opt.model_dir
    _, model = train_utils.load_saved_model_diff(saved_path, model)
    
    if opt.pretrained_path: # load traj pred model
        saved_path = opt.pretrained_path
        pretrained_model_dict = torch.load(saved_path, map_location='cpu')
        diff_keys = {k:v for k, v in pretrained_model_dict.items() if k not in model.state_dict()}
        modified_pretrained_model_dict = OrderedDict()
        for k, v in pretrained_model_dict['model_dict'].items():
            modified_pretrained_model_dict.update({'matcher.compensate_motion.'+k: v})
        diff_keys = {k:v for k, v in modified_pretrained_model_dict.items() if k not in model.state_dict()}
        if diff_keys:
            print(f"!!! PreTrained model has keys: {diff_keys.keys()}, \
                which are not in the model you have created!!!")
        model.load_state_dict(modified_pretrained_model_dict, strict=False)
    
    model.eval()

    # setting noise
    np.random.seed(303)
    noise_setting = OrderedDict()
    noise_setting['add_noise'] = False
    
    # build dataset for each noise setting
    print('Dataset Building')
    print(f"No Noise Added.")
    hypes.update({"noise_setting": noise_setting})
    opencood_dataset = build_dataset(hypes, visualize=True, train=False)
    data_loader = DataLoader(opencood_dataset,
                            batch_size=1,
                            num_workers=4,
                            collate_fn=opencood_dataset.collate_batch_test,
                            shuffle=False,
                            pin_memory=False,
                            drop_last=False)
    
    # Create the dictionary for evaluation
    result_stat = {0.3: {'tp': [], 'fp': [], 'gt': 0, 'score': []},
                0.5: {'tp': [], 'fp': [], 'gt': 0, 'score': []},
                0.7: {'tp': [], 'fp': [], 'gt': 0, 'score': []}}
    
    noise_level = "no_noise"

    start_time = time.time()
    # infer_info = opt.fusion_method + f"{opt.cavnum}agent" + opt.note
    i = -1
    avg_time_delay = 0.0
    avg_sample_interval = 0.0
    avg_time_var = 0.0
    if opt.dataset == 'd':
        avg_cp_rate = 0.0
    for i, batch_data in tenumerate(data_loader):
        if batch_data is None:
            continue
        with torch.no_grad():
            if opt.fusion_method == 'late':
                if opt.dataset == 'o':
                    unit_time_delay = []
                    unit_sample_interval = []
                    for cav_id, cav_content in batch_data.items():
                        unit_time_delay.append(cav_content['debug']['time_diff'])
                        unit_sample_interval.append(hypes['time_delay']*100)
                        # unit_sample_interval.append(cav_content['debug']['sample_interval'])
                    avg_time_delay += (sum(unit_time_delay[1:])/len(unit_time_delay[1:]))
                    avg_sample_interval += (float(sum(unit_sample_interval[1:]))/len(unit_sample_interval[1:]))
            if opt.fusion_method == 'intermediate':
                if opt.dataset == 'o':
                    try:
                        avg_time_delay += batch_data['ego']['avg_time_delay']
                        avg_sample_interval += batch_data['ego']['avg_sample_interval']
                    except KeyError:
                        avg_sample_interval += 0
                        avg_time_delay += 0
                    try: 
                        avg_time_var += batch_data['ego']['avg_time_var']
                    except:
                        avg_time_var += -1
                elif opt.dataset == 'd':
                    try:
                        avg_time_delay += (-batch_data['ego']['avg_time_delay'])
                    except KeyError:
                        avg_time_delay += 0
                    try:
                        avg_sample_interval += batch_data['ego']['avg_time_delay']
                    except KeyError:
                        avg_sample_interval += 0
                    try:
                        avg_cp_rate += (batch_data['ego']['cp_rate'])
                    except KeyError:
                        avg_cp_rate += 1
                        
            batch_data = train_utils.to_device(batch_data, device)
            uncertainty_tensor = None
            if opt.fusion_method == 'late':
                pred_box_tensor, pred_score, gt_box_tensor = \
                    inference_utils.inference_late_fusion(batch_data,
                                                        model,
                                                        opencood_dataset)
            elif opt.fusion_method == 'early':
                pred_box_tensor, pred_score, gt_box_tensor = \
                    inference_utils.inference_early_fusion(batch_data,
                                                        model,
                                                        opencood_dataset)
            elif opt.fusion_method == 'intermediate':
                if opt.two_stage == 1:
                    if viz_bbx_flag:
                        pred_box_tensor, pred_score, gt_box_tensor, single_detection_bbx, matched_idx_list, compensated_results_list, single_updated_feature, single_original_feature, single_flow_map, single_reserved_mask, single_original_reserved_mask = \
                            inference_utils.inference_intermediate_fusion_flow_module(batch_data,
                                                                        model,
                                                                        opencood_dataset, True)

                    else:
                        pred_box_tensor, pred_score, gt_box_tensor = \
                            inference_utils.inference_intermediate_fusion_flow_module(batch_data,
                                                                        model,
                                                                        opencood_dataset)
                else:
                    pred_box_tensor, pred_score, gt_box_tensor = \
                        inference_utils.inference_intermediate_fusion(batch_data,
                                                                    model,
                                                                    opencood_dataset)
            elif opt.fusion_method == 'no':
                pred_box_tensor, pred_score, gt_box_tensor = \
                    inference_utils.inference_no_fusion(batch_data,
                                                                model,
                                                                opencood_dataset)
            elif opt.fusion_method == 'no_w_uncertainty':
                pred_box_tensor, pred_score, gt_box_tensor, uncertainty_tensor = \
                    inference_utils.inference_no_fusion_w_uncertainty(batch_data,
                                                                model,
                                                                opencood_dataset)
            else:
                raise NotImplementedError('Only no, no_w_uncertainty, early, late and intermediate'
                                        'fusion is supported.')
            
            
            eval_utils.caluclate_tp_fp(pred_box_tensor,
                                    pred_score,
                                    gt_box_tensor,
                                    result_stat,
                                    0.3)
            eval_utils.caluclate_tp_fp(pred_box_tensor,
                                    pred_score,
                                    gt_box_tensor,
                                    result_stat,
                                    0.5)
            eval_utils.caluclate_tp_fp(pred_box_tensor,
                                    pred_score,
                                    gt_box_tensor,
                                    result_stat,
                                    0.7)
            if opt.save_npy:
                npy_save_path = os.path.join(opt.model_dir, 'npy')
                if not os.path.exists(npy_save_path):
                    os.makedirs(npy_save_path)
                inference_utils.save_prediction_gt(pred_box_tensor,
                                                gt_box_tensor,
                                                batch_data['ego']['origin_lidar'][0],
                                                i,
                                                npy_save_path)

            if (i % opt.save_vis_interval == 0) and (pred_box_tensor is not None):
                vis_save_path_root = os.path.join(opt.model_dir, f'vis_{opt.note}_%.2f_{noise_note}_roi_{num_roi_thres}'%(hypes['binomial_p']))
                if not os.path.exists(vis_save_path_root):
                    os.makedirs(vis_save_path_root)

                try:
                    debug_path = batch_data['ego']['debug']['scene_name'] + '_' + batch_data['ego']['debug']['cav_id'] + '_' + batch_data['ego']['debug']['timestamp']
                    # print(debug_path)
                except:
                    debug_path = 'path'
                vis_save_path = os.path.join(vis_save_path_root, 'bev_%05d_%s.png' % (i, debug_path))
                simple_vis.visualize(pred_box_tensor,
                                    gt_box_tensor,
                                    batch_data['ego']['origin_lidar'][0],
                                    hypes['postprocess']['gt_range'],
                                    vis_save_path,
                                    method='bev',
                                    left_hand=left_hand,
                                    uncertainty=uncertainty_tensor)
                if viz_bbx_flag:
                    '''
                    box_dict : {
                        'single_detection_bbx': # dict, [0, 1, ... , N-1], 
                            [i]: dict:{
                                [0]/[1]/[2]: past 3 frames detection results:{  past 3 frames detection results:{ # 都是在[0]的各自坐标系下的检测结果。
                                    pred_box_3dcorner_tensor
                                    pred_box_center_tensor
                                    scores
                                }
                            }
                        }
                        'lidar_pose_0': 过去第一帧 所有车的pose [N, 6]
                        'lidar_pose_current': 当前帧 所有车的pose [N, 6]
                        'matched_idx_list': matched_idx_list, # len=N-1, 表示每个non-ego的过去两帧的匹配id, each of which is [N_obj, 2], 比如 ['matched_idx_list'][0] shape(22,2) 表示过去第一帧的22个框与第二帧的22个框的索引的匹配情况
                        'compensated_results_list': # len=N-1, 每个non-ego补偿出来的box, each of which is [N_obj, 4, 3], 注意这里用的是4个点, 最后一个维度上是 xyz, z是多余的
                        'single_gt_box_tensor': # list, len=N, 表示ego与non-ego每辆车在current时刻的检测框结果, 例如box_dict['single_gt_box_tensor'][1]大小为[N_obj, 8, 3] 表示第二辆车在current时刻的检测框结果
                        'single_lidar': # list, len=N, 表示ego与non-ego每辆车在current时刻的lidar np
                        'gt_range': # [-140.8, -40, -3, 140.8, 40, 1], 表示lidar的范围
                        'single_updated_feature': tensor, [N, C, H, W], 表示ego与non-ego每辆车在past-0时刻的更新后的feature
                        'single_original_feature': tensor, [N, C, H, W], 表示ego与non-ego每辆车在past-0时刻的原始feature
                        'single_flow_map': tensor, [N, H, W, 2], 表示ego与non-ego每辆车在past-0时刻的flow map
                        'single_reserved_mask': tensor, [N, C, H, W], 表示ego与non-ego每辆车在past-0时刻的reserved mask
                        'single_original_reserved_mask': tensor, [N, C, H, W], 表示ego与non-ego每辆车在past-0时刻的原始reserved mask
                    }
                    '''
                    box_dict = {}
                    box_dict.update({
                        'single_lidar': batch_data['ego']['single_lidar_list'], # len = N, (n_lidar, 3)
                        'single_past_lidar': batch_data['ego']['single_past_lidar_list'], 
                        'gt_range': hypes['postprocess']['gt_range'],
                        'lidar_pose_current': batch_data['ego']['curr_lidar_pose'], #[N, 6] 
                        'lidar_pose_0': batch_data['ego']['past_lidar_pose'][:, 0, :], # (N, 6)
                        'pred_box_tensor': pred_box_tensor,
                        'gt_box_tensor': gt_box_tensor
                    })
                    '''
                    box_dict : {
                        'lidar_pose_0': 过去第一帧 所有车的pose [N, 6]
                        'lidar_pose_current': 当前帧 所有车的pose [N, 6]
                        'single_lidar': # list, len=N, 表示ego与non-ego每辆车在current时刻的lidar np, 每一个都是 (n_lidar, 3)
                        'single_past_lidar': # list, len=N, 表示ego与non-ego每辆车在past-0时刻的lidar np, 每一个都是 (n_lidar, 3)
                        'pred_box_tensor': # tensor, [N, 8, 3], 表示ego与non-ego每辆车在current时刻的检测框结果
                        'gt_box_tensor': # tensor, [N, 8, 3], 表示ego与non-ego每辆车在current时刻的检测框结果
                        'gt_range': # [-140.8, -40, -3, 140.8, 40, 1], 表示lidar的范围
                    }
                    '''
                    box_save_folder = os.path.join(vis_save_path_root, 'bbx_folder')
                    if not os.path.exists(box_save_folder):
                        os.mkdir(box_save_folder)
                    box_save_path = os.path.join(box_save_folder, 'bbx_%05d_%s.pt' % (i, debug_path))
                    torch.save(box_dict, box_save_path)

                    # box_dict = {}
                    # single_gt_box_tensor = []
                    # for cav in range(len(batch_data['ego']['single_object_bbx_center'])):
                    #     object_bbx_center = batch_data['ego']['single_object_bbx_center'][cav]
                    #     object_bbx_corner = box_utils.boxes_to_corners_3d(object_bbx_center, 'hwl')
                    #     single_gt_box_tensor.append(object_bbx_corner)
                    # box_dict.update({
                    #     'single_detection_bbx': single_detection_bbx, # dict, [0, 1, 2] 
                    #     'compensated_results_list': compensated_results_list, # len=N-1, each of which is [N_obj, 8, 3]
                    #     'matched_idx_list': matched_idx_list, # len=N-1, each of which is [N_obj, 2]
                    #     'single_gt_box_tensor': single_gt_box_tensor, # tensor, [N_obj, 8, 3]
                    #     'single_lidar': batch_data['ego']['single_lidar_list'], # len = N, (n_lidar, 3)
                    #     'single_past_lidar': batch_data['ego']['single_past_lidar_list'], 
                    #     'gt_range': hypes['postprocess']['gt_range'],
                    #     'lidar_pose_current': batch_data['ego']['curr_lidar_pose'], #[N, 6] 
                    #     'lidar_pose_0': batch_data['ego']['past_lidar_pose'][:, 0, :], # (N, 6)
                    #     'single_updated_feature': single_updated_feature, # tensor, [N, C, H, W], 表示ego与non-ego每辆车在current时刻的更新后的feature
                    #     'single_original_feature': single_original_feature, # tensor, [N, C, H, W], 表示ego与non-ego每辆车在current时刻的原始feature
                    #     'single_flow_map': single_flow_map, # tensor, [N, H, W, 2], 表示ego与non-ego每辆车在past-0时刻的flow map
                    #     'single_reserved_mask': single_reserved_mask, # tensor, [N, C, H, W], 表示ego与non-ego每辆车在past-0时刻的reserved mask
                    #     'single_original_reserved_mask': single_original_reserved_mask, # tensor, [N, C, H, W], 表示ego与non-ego每辆车在past-0时刻的原始reserved mask
                    # })
                    # box_save_folder = os.path.join(vis_save_path_root, 'bbx_folder')
                    # if not os.path.exists(box_save_folder):
                    #     os.mkdir(box_save_folder)
                    # box_save_path = os.path.join(box_save_folder, 'bbx_%05d_%s.pt' % (i, debug_path))
                    # torch.save(box_dict, box_save_path)
        torch.cuda.empty_cache()
    end_time = time.time()
    print("Time Consumed: %.2f minutes" % ((end_time - start_time)/60))
        
    if opt.dataset == 'o':
        avg_time_delay = (avg_time_delay/i) * 50 # unit is ms
        avg_sample_interval /= i
        avg_time_var /= i
        ap30, ap50, ap70 = eval_utils.eval_final_results(result_stat,
                                    opt.model_dir, noise_level, avg_time_delay, avg_sample_interval, avg_time_var, opt.note+'_'+str("%.2f"%hypes['binomial_p'])+f'_{noise_note}_roi_{num_roi_thres}')
    elif opt.dataset == 'd':
        avg_sample_interval /= i
        avg_cp_rate /= i
        ap30, ap50, ap70 = eval_utils.eval_final_results(result_stat,
                                    opt.model_dir, noise_level=noise_level, avg_time_delay=avg_cp_rate, avg_sample_interval=avg_sample_interval, note=opt.note+'_'+str("%.2f"%hypes['binomial_p'])+f'_{noise_note}_roi_{num_roi_thres}', dataset=opt.dataset)
    print("Module with sample interval expection: {}".format(hypes['binomial_n']*hypes['binomial_p']))
    if opt.dataset == 'o':
        print(f"IR sample range is {hypes['ir_range']}")

if __name__ == '__main__':
    main()

opencood/tools/inference.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import os
from collections import OrderedDict

import numpy as np
import torch

from opencood.utils.common_utils import torch_tensor_to_numpy
from opencood.visualization import vis_utils, my_vis, simple_vis
from opencood.tools.debug_tools import viz_compensation_latefusion_flow
def inference_late_fusion(batch_data, model, dataset):
    """
    Model inference for late fusion.

    Parameters
    ----------
    batch_data : dict
    model : opencood.object
    dataset : opencood.LateFusionDataset

    Returns
    -------
    pred_box_tensor : torch.Tensor
        The tensor of prediction bounding box after NMS.
    gt_box_tensor : torch.Tensor
        The tensor of gt bounding box.
    """
    output_dict = OrderedDict()

    for cav_id, cav_content in batch_data.items():
        output_dict[cav_id] = model(cav_content)

    pred_box_tensor, pred_score, gt_box_tensor = \
        dataset.post_process(batch_data,
                             output_dict)

    return pred_box_tensor, pred_score, gt_box_tensor

def inference_late_fusion_flow(batch_data, model, dataset, batch_id = 0):
    """
    Model inference for late fusion.

    Parameters
    ----------
    batch_data : dict
    model : opencood.object
    dataset : opencood.LateFusionDataset

    Returns
    -------
    pred_box_tensor : torch.Tensor
        The tensor of prediction bounding box after NMS.
    gt_box_tensor : torch.Tensor
        The tensor of gt bounding box.
    """
    output_dict = OrderedDict()
    thre = 20
    delay = 300

    folder_name=f'viz_wo_norm_thre_{thre}_delay_{delay}'

    for cav_id, cav_content in batch_data.items():
        output_dict[cav_id] = model(cav_content)

    box_results = dataset.generate_pred_bbx_frames(batch_data, output_dict)
    # if batch_id%10 == 0:
    #     viz_compensation_latefusion_flow(dataset, batch_data, box_results, file_name=folder_name, save_notes='no_comp', vis_comp_box=False, batch_id=batch_id)

    ''' 
    box_result : dict for each cav at each time
        {
            'ego' / cav_id : {
                'past_k_time_diff' : 
                [0] ... [k-1] : {
                        pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                        pred_box_center_tensor : n, 7
                        scores: (n, )
                }
            }
        }
    ''' 
    
    updated_box = box_flow_update(box_results, batch_id)
    '''
    updated_box: dict {
        'past_k_time_diff' : 
        [0], [1], ... , [k-1]
        'comp' : {
            pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3  (cav-past0)
            pred_box_center_tensor : n, 7  (cav-past0)
            scores: (n, )  
        }
    } 
    '''

    # 可视化每一帧的 box frame 与 final GT 进行对比
    # if batch_id%10 == 0:
    #     viz_compensation_latefusion_flow(dataset, batch_data, updated_box, file_name=folder_name, save_notes='w_comp', vis_comp_box=True, batch_id=batch_id)

    # return batch_id, batch_id, batch_id

    # TODO: warp each cavs' estimated frame to ego (cav-past-0 to ego-curr)
    pred_box_tensor, pred_score, delay_box_tensor = dataset.post_process_updated(batch_data, updated_box)
    gt_box_tensor = dataset.post_processor.generate_gt_bbx(batch_data)
    
    return pred_box_tensor, pred_score, gt_box_tensor, delay_box_tensor

def box_flow_update(box_results, batch_id=0):
    """
    Calculate flow using the detection results of two adjacent frames, 
    and update the detection results using the flow.

    Parameters
    ----------
    box_results: dict

    Return 
    pred_box_tensor: torch.Tensor, boxes after late fusion. N_pred, 8, 3 
    pred_score: torch.Tensor, scores after NMS filter
        
    """
    from opencood.tools.matcher import Matcher
    matcher = Matcher('box', 1, 1)
    updated_box = matcher(box_results, batch_id=batch_id)

    return updated_box

def inference_intermediate_fusion_flow_module(batch_data, model, dataset, viz_bbx_flag=False):
    """
    Model inference for flow module
    利用box预测结果生成flow 根据该flow更新feature

    Parameters
    ----------
    batch_data : dict
    model : opencood.object
    dataset : opencood.EarlyFusionDataset

    Returns
    -------
    pred_box_tensor : torch.Tensor
        The tensor of prediction bounding box after NMS.
    gt_box_tensor : torch.Tensor
        The tensor of gt bounding box.
    """
    output_dict = OrderedDict()
    cav_content = batch_data['ego']
    output_dict['ego'] = model(cav_content, dataset)
    
    pred_box_tensor, pred_score, gt_box_tensor = \
        dataset.post_process(batch_data,
                             output_dict)
    # TODO: uncomment
    # if viz_bbx_flag:
    #     single_detection_bbx = output_dict['ego']['single_detection_bbx']
    #     matched_idx_list = output_dict['ego']['matched_idx_list']
    #     compensated_results_list = output_dict['ego']['compensated_results_list']
    #     single_updated_feature = output_dict['ego']['single_updated_feature']
    #     single_original_feature = output_dict['ego']['single_original_feature']
    #     single_flow_map = output_dict['ego']['single_flow_map']
    #     single_reserved_mask = output_dict['ego']['single_reserved_mask']
    #     single_original_reserved_mask = output_dict['ego']['single_original_reserved_mask']

    #     return pred_box_tensor, pred_score, gt_box_tensor, single_detection_bbx, matched_idx_list, compensated_results_list, single_updated_feature, single_original_feature, single_flow_map, single_reserved_mask, single_original_reserved_mask

    return pred_box_tensor, pred_score, gt_box_tensor

def inference_intermediate_fusion_flow(batch_data, model, dataset):
    """
    Model inference for early fusion.

    Parameters
    ----------
    batch_data : dict
    model : opencood.object
    dataset : opencood.EarlyFusionDataset

    Returns
    -------
    pred_box_tensor : torch.Tensor
        The tensor of prediction bounding box after NMS.
    gt_box_tensor : torch.Tensor
        The tensor of gt bounding box.
    """
    output_dict = OrderedDict()
    thre = 20
    delay = 300

    folder_name=f'viz_wo_norm_thre_{thre}_delay_{delay}_inte_flow'

    for cav_id, cav_content in batch_data.items():
        output_dict[cav_id] = model(cav_content) # point_pillar_w_flow

    # 要对 output_dict 里面的 spatial_feature_2d 进行 flow update
    # step1: 生成每一帧的 box 结果
    box_results = dataset.generate_pred_bbx_frames(batch_data, output_dict)
    # step2: a. 根据检测结果生成 flow ; b. 更新 spatial feature  :  'updated_spatial_feature_2d' # TODO: 这里面添加一个可视化
    updated_output_dict = feature_flow_update(output_dict, box_results)

    ####### debug use, viz updated feature of each cav
    # from matplotlib import pyplot as plt
    # viz_save_path = '/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/debug_4_feature_flow'
    # for cav_id, cav_content in updated_output_dict.items():
    #     viz_content = torch.max(cav_content['updated_spatial_feature_2d'], dim=0)[0].detach().cpu()
    #     plt.imshow(viz_content)
    #     plt.savefig(viz_save_path+f'/updated_feature_{cav_id}.png')
    ##############
    
    # step3: 根据 spatial feature, 融合到 ego-curr
    # step4: detection header 
    pairwise_t_matrix = batch_data['ego']['pairwise_t_matrix']
    fused_dict = model(updated_output_dict, stage=2, pairwise_t_matrix=pairwise_t_matrix)
    pred_box_tensor, pred_score, gt_box_tensor = dataset.post_process_for_intermediate(batch_data, fused_dict)
    
    return pred_box_tensor, pred_score, gt_box_tensor

    # return inference_early_fusion(batch_data, model, dataset)

def feature_flow_update(output_dict, box_results):
    """
    generate box flow, and use the flow to update the feature 
    """
    # step 1. use box_results to generate flow
    # step 2. use flow to update spatial feature
    from opencood.tools.matcher import Matcher
    matcher = Matcher(1, 1, thre=20)
    updated_output_dict = matcher(box_results, fusion='feature', feature=output_dict)
    return updated_output_dict

def inference_no_fusion(batch_data, model, dataset):
    """
    Model inference for no fusion.

    Parameters
    ----------
    batch_data : dict
    model : opencood.object
    dataset : opencood.LateFusionDataset

    Returns
    -------
    pred_box_tensor : torch.Tensor
        The tensor of prediction bounding box after NMS.
    gt_box_tensor : torch.Tensor
        The tensor of gt bounding box.
    """
    output_dict_ego = OrderedDict()

    output_dict_ego['ego'] = model(batch_data['ego'])
    # output_dict only contains ego
    # but batch_data havs all cavs, because we need the gt box inside.

    pred_box_tensor, pred_score, gt_box_tensor = \
        dataset.post_process_no_fusion(batch_data,  # only for late fusion dataset
                             output_dict_ego)

    return pred_box_tensor, pred_score, gt_box_tensor

def inference_no_fusion_w_uncertainty(batch_data, model, dataset):
    """
    Model inference for no fusion.

    Parameters
    ----------
    batch_data : dict
    model : opencood.object
    dataset : opencood.LateFusionDataset

    Returns
    -------
    pred_box_tensor : torch.Tensor
        The tensor of prediction bounding box after NMS.
    gt_box_tensor : torch.Tensor
        The tensor of gt bounding box.
    """
    output_dict_ego = OrderedDict()

    output_dict_ego['ego'] = model(batch_data['ego'])
    # output_dict only contains ego
    # but batch_data havs all cavs, because we need the gt box inside.

    pred_box_tensor, pred_score, gt_box_tensor, uncertainty_tensor = \
        dataset.post_process_no_fusion(batch_data, # only for late fusion dataset
                             output_dict_ego, return_uncertainty=True)

    return pred_box_tensor, pred_score, gt_box_tensor, uncertainty_tensor


def inference_early_fusion(batch_data, model, dataset):
    """
    Model inference for early fusion.

    Parameters
    ----------
    batch_data : dict
    model : opencood.object
    dataset : opencood.EarlyFusionDataset

    Returns
    -------
    pred_box_tensor : torch.Tensor
        The tensor of prediction bounding box after NMS.
    gt_box_tensor : torch.Tensor
        The tensor of gt bounding box.
    """
    output_dict = OrderedDict()
    cav_content = batch_data['ego']
    output_dict['ego'] = model(cav_content)
    
    pred_box_tensor, pred_score, gt_box_tensor = \
        dataset.post_process(batch_data,
                             output_dict)

    # return_dict = {"pred_box_tensor" : pred_box_tensor, \
    #                     "pred_score" : pred_score, \
    #                     "gt_box_tensor" : gt_box_tensor}

    return pred_box_tensor, pred_score, gt_box_tensor


def inference_intermediate_fusion(batch_data, model, dataset):
    """
    Model inference for early fusion.

    Parameters
    ----------
    batch_data : dict
    model : opencood.object
    dataset : opencood.EarlyFusionDataset

    Returns
    -------
    pred_box_tensor : torch.Tensor
        The tensor of prediction bounding box after NMS.
    gt_box_tensor : torch.Tensor
        The tensor of gt bounding box.
    """
    return inference_early_fusion(batch_data, model, dataset)


def save_prediction_gt(pred_tensor, gt_tensor, pcd, timestamp, save_path):
    """
    Save prediction and gt tensor to txt file.
    """
    pred_np = torch_tensor_to_numpy(pred_tensor)
    gt_np = torch_tensor_to_numpy(gt_tensor)
    pcd_np = torch_tensor_to_numpy(pcd)

    np.save(os.path.join(save_path, '%04d_pcd.npy' % timestamp), pcd_np)
    np.save(os.path.join(save_path, '%04d_pred.npy' % timestamp), pred_np)
    np.save(os.path.join(save_path, '%04d_gt.npy' % timestamp), gt_np)

opencood/tools/inference_utils.py
MediaBrain-SJTU/CoBEVFlow
def rename_model_dict_keys(pretrained_dict_path, rename_dict):
    """ load pretrained state dict, keys may not match with model

    Args:
        model: nn.Module

        pretrained_dict: collections.OrderedDict

    """
    pretrained_dict = torch.load(pretrained_dict_path)
    # 1. filter out unnecessary keys
    for oldname, newname in rename_dict.items():
        if oldname.endswith("*"):
            _oldnames = list(pretrained_dict.keys())
            _oldnames = [x for x in _oldnames if x.startswith(oldname[:-1])]
            for _oldname in _oldnames:
                if newname != "":
                    _newname = _oldname.replace(oldname[:-1], newname[:-1])
                    pretrained_dict[_newname] = pretrained_dict[_oldname]
                pretrained_dict.pop(_oldname)
        else:
            if newname != "":
                pretrained_dict[newname] = pretrained_dict[oldname]
            pretrained_dict.pop(oldname)
    torch.save(pretrained_dict, pretrained_dict_path)


if __name__ == "__main__":
    dict_path = "/GPFS/rhome/yifanlu/workspace/OpenCOODv2/opencood/logs/then_SECOND_LID_72_48/net_epoch1.pth"
    rename_dict = {"lidar_encoder.*": "",
                   "shrink_lidar.*": ""}
    rename_model_dict_keys(dict_path, rename_dict)

    rename_dict = {"lidar_encoder.*": "", 
                "shrink_lidar.*": ""} # if value is empty, pop this key, if value is non empty, replace this key with new name. * means all keys start with this prefix
opencood/tools/load_pretrain.py
MediaBrain-SJTU/CoBEVFlow
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
"""
Modules to compute the matching cost and solve the corresponding LSAP.
"""
import torch
import torch.nn.functional as F 
from scipy.optimize import linear_sum_assignment
from torch import nn
from torch.autograd import Variable
from opencood.utils import box_utils
import copy
import math

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.embedding = nn.Sequential(nn.Linear(vocab,d_model),nn.ReLU(),nn.Linear(d_model,d_model))

    def forward(self, x):
        return self.embedding(x)

class PositionalEncoding_irregular(nn.Module):

    def __init__(self, d_model, dropout, max_len=5000 ):
        super(PositionalEncoding_irregular, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.max_len = max_len
        self.d_model = d_model
        self.linear = nn.Linear(2*d_model,d_model)
         
    def forward(self, x, time):
        ### x:[batch, agent, k, d] time:[batch, agent, k]
        if len(x.shape) > 3:
            last_shape = x.shape
            x = x.reshape(-1,x.shape[-2],x.shape[-1])
            time = time.reshape(-1,time.shape[-1])

        pe = torch.zeros(x.shape[0],time.shape[1], self.d_model).cuda()
        position = time.unsqueeze(2) # 增加维度
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).cuda()#相对位置公式
         
        pe[:,:, 0::2] = torch.sin(position * div_term)   #取奇数列
        pe[:,:, 1::2] = torch.cos(position * div_term)   #取偶数列

        # x = x + Variable(pe[:, :x.size(1)], requires_grad=False) # embedding 与positional相加
        
        x = self.linear(torch.cat((x,Variable(pe[:, :x.size(1)], requires_grad=False)),dim=-1))
        x = x.reshape(last_shape)
        return self.dropout(x)

class Encoder(nn.Module):
    "Core encoder is a stack of N layers"

    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.N = N

        self.norm = LayerNorm(layer.size) #归一化

    def forward(self, x, mask):
        "Pass the input (and mask) through each layer in turn."
        for i in range(self.N):
            # print('x',x.shape)
            x = self.layers[i](x, mask) 

        return self.norm(x)

class EncoderLayer(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn #sublayer 1 
        self.feed_forward = feed_forward #sublayer 2 
        self.sublayer = clones(SublayerConnection(size, dropout), 2) #拷贝两个SublayerConnection，一个为了attention，一个是独自作为简单的神经网络
        self.size = size

    def forward(self, x, mask):
        "Follow Figure 1 (left) for connections."
        # print('x',x.shape,'attn',self.self_attn(x,x,x,mask).shape)
        
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) #attention层 对应层1
        return self.sublayer[1](x, self.feed_forward) # 对应 层2
        return x


class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
 
    def forward(self, x):
        "Norm"
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, x, sublayer):
        # print('sub',x.shape,self.dropout(sublayer(self.norm(x))).shape)
        return x + self.dropout(sublayer(self.norm(x))) #残差连接

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model,d_out, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_out)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

class DecoderLayer(nn.Module):
    "Decoder is made of self-attn, src-attn, and feed forward (defined below)"
    def __init__(self, size, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.src_attn = src_attn #解码的attention
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)

    def forward(self, x, src_key,memory, tgt_mask):
        m = memory
        # x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask)) #self-attention
        x = self.sublayer[0](x, lambda x: self.src_attn(x, src_key, m, tgt_mask)) #解码
        return self.sublayer[1](x, self.feed_forward)


class Decoder(nn.Module):
    "Generic N layer decoder with masking."
    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        # print('layersize',layer.size)
        self.norm = LayerNorm(layer.size)
        
    def forward(self, x, src_key ,memory, tgt_mask):
        for layer in self.layers:
            x = layer(x, src_key, memory, tgt_mask) #添加编码的后的结果
        return self.norm(x)

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    
    # print('score',scores.shape,'mask',mask.shape)
    if mask is not None:
        mask = mask.transpose(1,2).cuda()
        scores = scores.masked_fill(mask == 0, -1e9) #mask必须是一个ByteTensor 而且shape必须和 a一样 并且元素只能是 0或者1 ，是将 mask中为1的 元素所在的索引，在a中相同的的索引处替换为 value  ,mask value必须同为tensor
    
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, q_dim, k_dim,v_dim, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h  # d_v=d_k=d_model/h 
        self.h = h # heads 的数目文中为8
        self.q_dim = q_dim
        self.k_dim = k_dim
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.linear_q = nn.Linear(q_dim,d_model)
        self.linear_k = nn.Linear(k_dim,d_model)
        self.linear_v = nn.Linear(v_dim,d_model)
        self.linear_out = nn.Linear(d_model,q_dim)
        self.attn = None  
        self.dropout = nn.Dropout(p=dropout)


    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        
        dim_to_keep = query.size(1)

        # 1) Do all the linear projections in batch from d_model => h x d_k 


        query = self.linear_q(query).view(nbatches,dim_to_keep,-1,self.h,self.d_k).transpose(2,3)
        key = self.linear_k(key).view(nbatches,dim_to_keep,-1,self.h,self.d_k).transpose(2,3)
        value = self.linear_v(value).view(nbatches,dim_to_keep,-1,self.h,self.d_k).transpose(2,3)

        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(query, key, value, mask=mask,# 进行attention
                                 dropout=self.dropout)

        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(2, 3).contiguous() \
            .view(nbatches, dim_to_keep ,-1, self.h * self.d_k) # 还原序列[batch_size,len,d_model]

        return self.linear_out(x)

class Generator(nn.Module):
    def __init__(self, d_model, out_dim):
        
        super(Generator, self).__init__()
        
        self.proj = nn.Linear(d_model, out_dim)

    def forward(self, x):
        
        return self.proj(x)


class Motion_prediction(nn.Module):
    def __init__(self, encoder, decoder, embedding_src, embedding_tgt, position_src, position_tgt, generator, embed_dim):
        super(Motion_prediction, self).__init__()
        
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = embedding_src
        self.tgt_embed = embedding_tgt
        self.src_pe = position_src
        self.tgt_pe = position_tgt
        self.generator = generator

        

    def forward(self, input, query, time, future_time, mask=None):
        ################################ input: [batch, agent, past_frames, dim]  mask: [batch, agent, past_frames, past_frames]
        input = self.src_embed(input)
        input = self.src_pe(input, time)

        query = self.tgt_embed(query)
        query = self.tgt_pe(query, future_time)

        input_features = self.encoder(input, mask)
        
        prediction_features = self.decoder(query, input_features, input_features,mask)
        predictions = self.generator(prediction_features)
        return predictions



class Motion_interaction(nn.Module):
    def __init__(self, encoder, embedding, generator, neighbor_threshold):
        super(Motion_interaction, self).__init__()
        self.encoder = encoder
        self.src_embed = embedding
        self.generator = generator
        self.neighbor_threshold = neighbor_threshold


    def forward(self, input):
        ####input:[batch,agent,2]
        mask = torch.cdist(input,input)

        mask = torch.where(mask<self.neighbor_threshold,1,0)
        input = input.unsqueeze(1)
        mask = mask.unsqueeze(1)

        input = self.src_embed(input)

        input_features = self.encoder(input, mask)
        predictions = self.generator(input_features)

        return predictions


def make_model(input_dim, output_dim, num_layers=2, d_model=64, d_ff=128, num_heads=2, dropout=0.1,neighbor_shreshold=10):
    c = copy.deepcopy
    attn = MultiHeadedAttention(num_heads, d_model,d_model,d_model,d_model)
    attn_decoder = MultiHeadedAttention(num_heads,d_model,d_model,d_model,d_model)
    ff = PositionwiseFeedForward(d_model, d_model, d_ff,dropout)
    position = PositionalEncoding_irregular(d_model, dropout)

    model_prediction = Motion_prediction(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), num_layers),
        Decoder(DecoderLayer(d_model, c(attn), c(ff), dropout), num_layers),
        Embeddings(d_model, input_dim),
        Embeddings(d_model, output_dim), 
        c(position),
        c(position),
        Generator(d_model, output_dim),
        d_model)

    model_interaction = Motion_interaction(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), num_layers),
        Embeddings(d_model, input_dim), 
        Generator(d_model, output_dim),
        neighbor_shreshold)
    
    for p in model_prediction.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

    for p in model_interaction.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

    return model_prediction, model_interaction

# from util.box_ops import box_cxcywh_to_xyxy, generalized_box_iou

'''
class HungarianMatcher(nn.Module):
    """This class computes an assignment between the targets and the predictions of the network
    For efficiency reasons, the targets don't include the no_object. Because of this, in general,
    there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,
    while the others are un-matched (and thus treated as non-objects).
    """

    def __init__(self, cost_class: float = 1, cost_bbox: float = 1, cost_giou: float = 1):
        """Creates the matcher
        Params:
            cost_class: This is the relative weight of the classification error in the matching cost
            cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost
            cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost
        """
        super().__init__()
        self.cost_class = cost_class
        self.cost_bbox = cost_bbox
        self.cost_giou = cost_giou
        assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0, "all costs cant be 0"

    @torch.no_grad()
    def forward(self, outputs, targets):
        """ Performs the matching
        Params:
            outputs: This is a dict that contains at least these entries:
                 "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
                 "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates
            targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:
                 "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth
                           objects in the target) containing the class labels
                 "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates
        Returns:
            A list of size batch_size, containing tuples of (index_i, index_j) where:
                - index_i is the indices of the selected predictions (in order)
                - index_j is the indices of the corresponding selected targets (in order)
            For each batch element, it holds:
                len(index_i) = len(index_j) = min(num_queries, num_target_boxes)
        """
        bs, num_queries = outputs["pred_logits"].shape[:2]

        # We flatten to compute the cost matrices in a batch
        out_prob = outputs["pred_logits"].flatten(0, 1).softmax(-1)  # [batch_size * num_queries, num_classes]
        out_bbox = outputs["pred_boxes"].flatten(0, 1)  # [batch_size * num_queries, 4]

        # Also concat the target labels and boxes
        tgt_ids = torch.cat([v["labels"] for v in targets])
        tgt_bbox = torch.cat([v["boxes"] for v in targets])

        # Compute the classification cost. Contrary to the loss, we don't use the NLL,
        # but approximate it in 1 - proba[target class].
        # The 1 is a constant that doesn't change the matching, it can be ommitted.
        cost_class = -out_prob[:, tgt_ids]

        # Compute the L1 cost between boxes
        cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1)

        # Compute the giou cost betwen boxes
        cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox))

        # Final cost matrix
        C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou
        C = C.view(bs, num_queries, -1).cpu()

        sizes = [len(v["boxes"]) for v in targets]
        indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]
        return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]
'''

class Matcher(nn.Module):
    """This class computes an assignment between the targets and the predictions of the network
    For efficiency reasons, the targets don't include the no_object. Because of this, in general,
    there are more predictions than targets. In this case, we do a 1-to-1 matching of the best predictions,
    while the others are un-matched (and thus treated as non-objects).
    """

    def __init__(self, fusion, cost_dist: float = 1, cost_giou: float = 1, thre: float = 20):
        """Creates the matcher
        Params:
            cost_class: This is the relative weight of the classification error in the matching cost
            cost_bbox: This is the relative weight of the L1 error of the bounding box coordinates in the matching cost
            cost_giou: This is the relative weight of the giou loss of the bounding box in the matching cost
        """
        super().__init__()
        self.cost_dist = cost_dist
        self.cost_giou = cost_giou
        self.thre = thre

        self.fusion = fusion
        if fusion=='flow':
            m1, m2 = make_model(4,2)
            self.compensate_motion = m1

    @torch.no_grad()
    def forward(self, input_dict, feature=None, shape_list=None, batch_id=0, viz_flag=False):
        self.viz_flag = viz_flag
        if self.fusion=='box':
            # return self.forward_box(input_dict, batch_id)
            return self.forward_box_w_dir(input_dict, batch_id)
        elif self.fusion=='feature':
            return self.forward_feature(input_dict, feature)
        # elif self.fusion=='flow':
        #     return self.forward_flow_multi_frames(input_dict, shape_list)
        elif self.fusion=='linear':
            return self.forward_flow(input_dict, shape_list)
        elif self.fusion=='flow': # TODO: flow_dir
            return self.forward_flow_dir(input_dict, shape_list)
        else:
            print("Attention, fusion method must be in box or feature!")

    def forward_box(self, input_dict, batch_id):
        """ Performs the matching
        input_dict: 
            {
                'ego' : {
                    'past_k_time_diff' : 
                    [0] {
                        pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                        pred_box_center_tensor : n, 7
                        scores: (n, )
                    }
                    ... 
                    [k-1]
                    ['comp']{
                        pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                        pred_box_center_tensor : n, 7
                        scores: (n, )
                    }
                }
                cav_id {}
            }
        """

        # output_dict_past = {}
        # output_dict_current = {}
        pred_bbox_list = []
        for cav, cav_content in input_dict.items():
            if cav == 'ego':
                estimated_box_center_current = cav_content[0]['pred_box_center_tensor']
                estimated_box_3dcorner_current = cav_content[0]['pred_box_3dcorner_tensor']
            else:
                coord_past1 = cav_content[0]
                coord_past2 = cav_content[1]

                center_points_past1 = coord_past1['pred_box_center_tensor'][:,:2]
                center_points_past2 = coord_past2['pred_box_center_tensor'][:,:2]
                
                # TODO: norm
                # center_points_past1_norm = center_points_past1 - torch.mean(center_points_past1,dim=0,keepdim=True)
                # center_points_past2_norm = center_points_past2 - torch.mean(center_points_past2,dim=0,keepdim=True)
                # cost_mat_center = torch.cdist(center_points_past2_norm,center_points_past1_norm) # [num_cav_past2,num_cav_past1]

                cost_mat_center = torch.cdist(center_points_past2, center_points_past1) # [num_cav_past2,num_cav_past1]

                cost_mat_center_drop_2 = torch.sum(torch.where(cost_mat_center > self.thre, 1, 0), dim=1)
                dist_valid_past2 = torch.where(cost_mat_center_drop_2 < center_points_past1.shape[0])
                cost_mat_center_drop_1 = torch.sum(torch.where(cost_mat_center > self.thre, 1, 0), dim=0)
                dist_valid_past1 = torch.where(cost_mat_center_drop_1 < center_points_past2.shape[0])

                cost_mat_center = cost_mat_center[dist_valid_past2[0], :]
                cost_mat_center = cost_mat_center[:, dist_valid_past1[0]]
                
                # cost_mat_iou = get_ious()
                cost_mat = cost_mat_center
                past2_ids, past1_ids = linear_sum_assignment(cost_mat.cpu())
                
                past2_ids = dist_valid_past2[0][past2_ids]
                past1_ids = dist_valid_past1[0][past1_ids]
                # output_dict_past.update({car:past_ids})
                # output_dict_current.update({car:current_ids})

                matched_past2 = center_points_past2[past2_ids]
                matched_past1 = center_points_past1[past1_ids]

                time_length = cav_content['past_k_time_diff'][0] - cav_content['past_k_time_diff'][1]
                if time_length == 0:
                    time_length = 1
                flow = (matched_past1 - matched_past2) / time_length

                # if flow.shape[0] != 0:
                #     print(f"max flow is {flow.max()}")

                estimate_position = matched_past1 + flow*(0-cav_content['past_k_time_diff'][0]) 

                # from copy import deepcopy
                # estimated_box_center_current = deepcopy(coord_past1['pred_box_center_tensor'].detach())
                estimated_box_center_current = coord_past1['pred_box_center_tensor'].detach().clone()
                estimated_box_center_current[past1_ids, :2] = estimate_position  # n, 7

                # estimated_box_center_current = torch.zeros_like(coord_past1['pred_box_center_tensor']).to(estimate_position.device)
                # estimated_box_center_current[past1_ids] += torch.cat([estimate_position, coord_past1['pred_box_center_tensor'][past1_ids][:,2:]], dim=-1)
                # no_past1_ids = [x for x in range(coord_past1['pred_box_center_tensor'].shape[0]) if x not in list(past1_ids)]
                # estimated_box_center_current[no_past1_ids] += coord_past1['pred_box_center_tensor'][no_past1_ids]

                estimated_box_3dcorner_current = box_utils.boxes_to_corners_3d(estimated_box_center_current, order='hwl')

            # debug use, update input dict adding estimated frame at cav-past0
            input_dict[cav]['comp'] = {}
            input_dict[cav]['comp'].update({
                'pred_box_center_tensor': estimated_box_center_current,
                'pred_box_3dcorner_tensor': estimated_box_3dcorner_current,
                'scores': cav_content[0]['scores']
            })

        return input_dict

    def forward_box_w_dir(self, input_dict, batch_id):
        """ Performs the matching
        input_dict: 
            {
                'ego' : {
                    'past_k_time_diff' : 
                    [0] {
                        pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                        pred_box_center_tensor : n, 7
                        scores: (n, )
                    }
                    ... 
                    [k-1]
                    ['comp']{
                        pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                        pred_box_center_tensor : n, 7
                        scores: (n, )
                    }
                }
                cav_id {}
            }
        """

        # output_dict_past = {}
        # output_dict_current = {}
        pred_bbox_list = []
        for cav, cav_content in input_dict.items():
            if cav == 'ego':
                estimated_box_center_current = cav_content[0]['pred_box_center_tensor']
                estimated_box_3dcorner_current = cav_content[0]['pred_box_3dcorner_tensor']
            else:
                coord_past1 = cav_content[0]
                coord_past2 = cav_content[1]

                center_points_past1 = coord_past1['pred_box_center_tensor'][:,:2]
                center_points_past2 = coord_past2['pred_box_center_tensor'][:,:2]

                cost_mat_center = torch.zeros((center_points_past2.shape[0], center_points_past1.shape[0])).to(center_points_past1.device)

                center_points_past1_repeat = center_points_past1.unsqueeze(0).repeat(center_points_past2.shape[0], 1, 1)
                center_points_past2_repeat = center_points_past2.unsqueeze(1).repeat(1, center_points_past1.shape[0], 1)

                delta_mat = center_points_past1_repeat - center_points_past2_repeat

                angle_mat = torch.atan2(delta_mat[:,:,1], delta_mat[:,:,0]) # [num_cav_past2,num_cav_past1]

                coord_past2_angle_reverse = coord_past2['pred_box_center_tensor'][:,6].unsqueeze(1).repeat(1, center_points_past1.shape[0]).clone()

                visible_mat = torch.where((torch.abs(angle_mat-coord_past2['pred_box_center_tensor'][:,6].unsqueeze(1).repeat(1, center_points_past1.shape[0])) < 0.785) | (torch.abs(angle_mat-coord_past2['pred_box_center_tensor'][:,6].unsqueeze(1).repeat(1, center_points_past1.shape[0])) > 5.495) , 1, 0) # [num_cav_past2,num_cav_past1]

                cost_mat_center = torch.cdist(center_points_past2, center_points_past1) # [num_cav_past2,num_cav_past1]

                visible_mat = torch.where(cost_mat_center<0.5, 1, visible_mat)

                tmp_thre = torch.tensor(1000.).to(torch.float32).to(visible_mat.device)
                cost_mat_center = torch.where(visible_mat==1, cost_mat_center, tmp_thre)

                if cost_mat_center.shape[1] == 0 or cost_mat_center.shape[0] == 0:
                    estimated_box_center_current = cav_content[0]['pred_box_center_tensor']
                    estimated_box_3dcorner_current = cav_content[0]['pred_box_3dcorner_tensor']

                else:
                    match = torch.min(cost_mat_center, dim=1)
                    match_to_keep = torch.where(match[0] < 5)

                    past2_ids = match_to_keep[0]
                    past1_ids = match[1][match_to_keep[0]]

                    coord_past2_angle_reverse += 3.1415926
                    coord_past2_angle_reverse[coord_past2_angle_reverse>3.1415926] -= 6.2831852

                    left_past2_id = [n for n in range(cost_mat_center.shape[0]) if n not in past2_ids]
                    left_past1_id = [n for n in range(cost_mat_center.shape[1]) if n not in past1_ids]

                    angle_mat_left = angle_mat[left_past2_id, :][:, left_past1_id]

                    coord_past2_angle_reverse_left = coord_past2_angle_reverse[left_past2_id, :][: ,left_past1_id]

                    visible_mat_left = torch.where((torch.abs(angle_mat_left-coord_past2_angle_reverse_left) < 0.785) | (torch.abs(angle_mat_left-coord_past2_angle_reverse_left) > 5.495) , 1, 0) # [num_cav_past2,num_cav_past1]

                    cost_mat_center_left = torch.cdist(center_points_past2[left_past2_id], center_points_past1[left_past1_id]) # [num_cav_past2,num_cav_past1]
                    # cost_mat_center_left = cost_mat_center[left_past2_id, :][:, left_past1_id]

                    visible_mat_left = torch.where(cost_mat_center_left<0.5, 1, visible_mat_left)
                    cost_mat_center_left = torch.where(visible_mat_left==1, cost_mat_center_left, tmp_thre)

                    if cost_mat_center_left.shape[1] != 0 and cost_mat_center_left.shape[0] != 0:
                        match_left = torch.min(cost_mat_center_left, dim=1)
                        match_to_keep_left = torch.where(match_left[0] < 5)

                        if match_to_keep_left[0].shape[0] != 0:
                            past2_ids_left = match_to_keep_left[0]
                            past2_ids = torch.cat([past2_ids, torch.tensor(left_past2_id)[past2_ids_left].to(past2_ids.device)])
                            past1_ids_left = match_left[1][match_to_keep_left[0]]
                            past1_ids = torch.cat([past1_ids, torch.tensor(left_past1_id)[past1_ids_left].to(past1_ids.device)])

                    matched_past2 = center_points_past2[past2_ids]
                    matched_past1 = center_points_past1[past1_ids]

                    time_length = cav_content['past_k_time_diff'][0] - cav_content['past_k_time_diff'][1]
                    if time_length == 0:
                        time_length = 1
                    
                    flow = (matched_past1 - matched_past2) / time_length

                    # if flow.shape[0] != 0:
                    #     print(f"max flow is {flow.max()}")

                    estimate_position = matched_past1 + flow*(0-cav_content['past_k_time_diff'][0]) 

                    # from copy import deepcopy
                    # estimated_box_center_current = deepcopy(coord_past1['pred_box_center_tensor'].detach())
                    estimated_box_center_current = coord_past1['pred_box_center_tensor'].detach().clone()
                    estimated_box_center_current[past1_ids, :2] = estimate_position  # n, 7

                    estimated_box_3dcorner_current = box_utils.boxes_to_corners_3d(estimated_box_center_current, order='hwl')

            # debug use, update input dict adding estimated frame at cav-past0
            input_dict[cav]['comp'] = {}
            input_dict[cav]['comp'].update({
                'pred_box_center_tensor': estimated_box_center_current,
                'pred_box_3dcorner_tensor': estimated_box_3dcorner_current,
                'scores': cav_content[0]['scores']
            })

        return input_dict


    def forward_feature(self, input_dict, features_dict):
        """
        Parameters:
        -----------
        input_dict : The dictionary containing the box detections on each frame of each cav.
            dict : { 
                'ego' / cav_id : {
                    [0] / [1] : {
                        pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                        pred_box_center_tensor : n, 7
                        scores: (n, )
                    }
                }
            }
        features_dict : The dictionary containing the output of the model.
            dict : {
                'ego' / cav_id : {
                    'spatial_features_2d'
                    'spatial_features'
                    'psm'
                    'rm'
                }
            }

        Returns:
        --------
        features_dict :
            dict : {
                'ego' / cav_id : {
                    'spatial_features_2d'
                    'spatial_features'
                    'psm'
                    'rm'
                    'updated_spatial_features_2d'
                    'updated_spatial_features'
                }
            }
        """
        pred_bbox_list = []
        
        ######## TODO: debug use
        # # pop up other cavs, only keep ego:
        # debug_dict = {'ego': input_dict['ego']}
        # input_dict = debug_dict.copy()
        # debug_features_dict = {'ego': features_dict['ego']}
        # features_dict = debug_features_dict.copy()
        ##############

        for cav, cav_content in input_dict.items():
            if cav == 'ego':
                updated_spatial_features_2d = features_dict[cav]['spatial_features_2d'][0]
                updated_spatial_features = features_dict[cav]['spatial_features'][0]
            else:
                coord_past1 = cav_content[0]
                coord_past2 = cav_content[1]

                center_points_past1 = coord_past1['pred_box_center_tensor'][:,:2]
                center_points_past2 = coord_past2['pred_box_center_tensor'][:,:2]

                cost_mat_center = torch.cdist(center_points_past2, center_points_past1) # [num_cav_past2,num_cav_past1]

                cost_mat_center_drop_2 = torch.sum(torch.where(cost_mat_center > self.thre, 1, 0), dim=1)
                dist_valid_past2 = torch.where(cost_mat_center_drop_2 < center_points_past1.shape[0])
                cost_mat_center_drop_1 = torch.sum(torch.where(cost_mat_center > self.thre, 1, 0), dim=0)
                dist_valid_past1 = torch.where(cost_mat_center_drop_1 < center_points_past2.shape[0])

                cost_mat_center = cost_mat_center[dist_valid_past2[0], :]
                cost_mat_center = cost_mat_center[:, dist_valid_past1[0]]
                
                # cost_mat_iou = get_ious()
                cost_mat = cost_mat_center
                past2_ids, past1_ids = linear_sum_assignment(cost_mat.cpu())
                
                past2_ids = dist_valid_past2[0][past2_ids]
                past1_ids = dist_valid_past1[0][past1_ids]
                # output_dict_past.update({car:past_ids})
                # output_dict_current.update({car:current_ids})

                matched_past2 = center_points_past2[past2_ids]
                matched_past1 = center_points_past1[past1_ids]

                time_length = cav_content['past_k_time_diff'][0] - cav_content['past_k_time_diff'][1]
                if time_length == 0:
                    time_length = 1
                flow = (matched_past1 - matched_past2) / time_length

                # TODO: flow * (0-past_k[0])
                flow = flow*(0-cav_content['past_k_time_diff'][0])
                selected_box_3dcenter_past0 = coord_past1['pred_box_center_tensor'][past1_ids,]
                selected_box_3dcorner_past0 = box_utils.boxes_to_corners2d(selected_box_3dcenter_past0, order='hwl')
                
                # debug use
                debug_flag = False
                if debug_flag and flow.shape[0] != 0:
                    viz_save_path = '/DB/data/sizhewei/logs/where2comm_max_multiscale_resnet_32ch/vis_debug'
                    torch.save(features_dict[cav]['spatial_features_2d'][0], viz_save_path+'/features_2d.pt')
                    torch.save(features_dict[cav]['spatial_features'][0], viz_save_path+'/features.pt')
                    torch.save(selected_box_3dcorner_past0, viz_save_path+'/bbx_list.pt')
                    torch.save(flow, viz_save_path+'/flow.pt')
                    print(f"===saved, max flow is {flow.max()}===")
                ############
                
                updated_spatial_features_2d = self.feature_warp(features_dict[cav]['spatial_features_2d'][0], selected_box_3dcorner_past0, flow, scale=1.25)
                updated_spatial_features = self.feature_warp(features_dict[cav]['spatial_features'][0], selected_box_3dcorner_past0, flow, scale=2.5)

                # debug use
                if debug_flag and flow.shape[0] != 0:
                    torch.save(updated_spatial_features_2d, viz_save_path+'/updated_feature.pt')
                ############

            features_dict[cav].update({
                'updated_spatial_features_2d': updated_spatial_features_2d
            })
            features_dict[cav].update({
                'updated_spatial_features': updated_spatial_features
            })

        return features_dict

    def forward_flow_multi_frames(self, input_dict, shape_list):
        """
        Parameters:
        -----------
        input_dict : The dictionary containing the box detections on each frame of each cav.
            dict : { 
                cav_idx : {
                    'past_k_time_diff':
                    [0], [1], ..., [k-1] : {
                        pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                        pred_box_center_tensor : n, 7
                        scores: (n, )
                    }
                }
            }
        features_dict : The dictionary containing the output of the model.
            dict : {
                'ego' / cav_id : {
                    'spatial_features_2d'
                    'spatial_features'
                    'psm'
                    'rm'
                }
            }

        Returns:
        --------
        features_dict :
            dict : {
                'ego' / cav_id : {
                    'spatial_features_2d'
                    'spatial_features'
                    'psm'
                    'rm'
                    'updated_spatial_features_2d'
                    'updated_spatial_features'
                }
            }
        """
        flow_map_list = []
        reserved_mask = []
        if self.viz_flag:
            matched_idx_list = []
            compensated_results_list = []
        for cav, cav_content in input_dict.items():
            if cav == 0:
                # ego do not need warp
                C, H, W = shape_list
                basic_mat = torch.tensor([[1,0,0],[0,1,0]]).unsqueeze(0).to(torch.float32)
                basic_warp_mat = F.affine_grid(basic_mat, [1, C, H, W], align_corners=False).to(shape_list.device)
                mask = torch.ones(1, C, H, W).to(shape_list)
                flow_map_list.append(basic_warp_mat)
                reserved_mask.append(mask)
            else:
                past_k_time_diff = cav_content['past_k_time_diff']

                coord_past1 = cav_content[0]
                coord_past2 = cav_content[1]
                coord_past3 = cav_content[2]

                center_points_past1 = coord_past1['pred_box_center_tensor'][:,:2]
                center_points_past2 = coord_past2['pred_box_center_tensor'][:,:2]
                center_points_past3 = coord_past3['pred_box_center_tensor'][:,:2]

                self.thre_post_process = 10

                # past1 and past2 match
                cost_mat_center_a = torch.cdist(center_points_past2, center_points_past1) # [num_cav_past2,num_cav_past1]
                # original_cost_mat_center_a = cost_mat_center_a.clone()
                # cost_mat_center_a[cost_mat_center_a > self.thre_post_process] = 1000

                cost_mat_center_drop_2_a = torch.sum(torch.where(cost_mat_center_a > self.thre, 1, 0), dim=1)
                dist_valid_past2_a = torch.where(cost_mat_center_drop_2_a < center_points_past1.shape[0])
                cost_mat_center_drop_1 = torch.sum(torch.where(cost_mat_center_a > self.thre, 1, 0), dim=0)
                dist_valid_past1 = torch.where(cost_mat_center_drop_1 < center_points_past2.shape[0])

                cost_mat_center_a = cost_mat_center_a[dist_valid_past2_a[0], :]
                cost_mat_center_a = cost_mat_center_a[:, dist_valid_past1[0]]
                
                cost_mat = cost_mat_center_a.clone()
                past2_ids_a, past1_ids = linear_sum_assignment(cost_mat.cpu())
                
                past2_ids_a = dist_valid_past2_a[0][past2_ids_a]
                past1_ids = dist_valid_past1[0][past1_ids]

                if len(past2_ids_a)==0:
                    print('======= No matched boxes between latest 2 frames! =======')
                    C, H, W = shape_list
                    basic_mat = torch.tensor([[1,0,0],[0,1,0]]).unsqueeze(0).to(torch.float32)
                    basic_warp_mat = F.affine_grid(basic_mat, [1, C, H, W], align_corners=False).to(shape_list.device)
                    mask = torch.ones(1, C, H, W).to(shape_list) # TODO: rethink ones or zeros
                    flow_map_list.append(basic_warp_mat)
                    reserved_mask.append(mask)
                    if self.viz_flag:
                        matched_idx_list.append(torch.stack([torch.tensor([]), torch.tensor([])], dim=1).to(shape_list.device))
                        compensated_results_list.append(torch.zeros(0, 8, 3).to(shape_list.device))
                    continue

                # past2 and past3 match
                cost_mat_center_b = torch.cdist(center_points_past3, center_points_past2) # [num_cav_past3,num_cav_past2]

                # original_cost_mat_center_b = cost_mat_center_b.clone()
                # cost_mat_center_b[cost_mat_center_b > self.thre_post_process] = 1000

                cost_mat_center_drop_3 = torch.sum(torch.where(cost_mat_center_b > self.thre, 1, 0), dim=1)
                dist_valid_past3 = torch.where(cost_mat_center_drop_3 < center_points_past2.shape[0])
                cost_mat_center_drop_2_b = torch.sum(torch.where(cost_mat_center_b > self.thre, 1, 0), dim=0)
                dist_valid_past2_b = torch.where(cost_mat_center_drop_2_b < center_points_past3.shape[0])

                cost_mat_center_b = cost_mat_center_b[dist_valid_past3[0], :]
                cost_mat_center_b = cost_mat_center_b[:, dist_valid_past2_b[0]]
                
                # cost_mat_iou = get_ious()
                cost_mat = cost_mat_center_b.clone()
                past3_ids, past2_ids_b = linear_sum_assignment(cost_mat.cpu())
                
                past3_ids = dist_valid_past3[0][past3_ids]
                past2_ids_b = dist_valid_past2_b[0][past2_ids_b]
                # output_dict_past.update({car:past_ids})
                # output_dict_current.update({car:current_ids})

                # find the matched obj among three frames
                a_idx, b_idx = self.get_common_elements(past2_ids_a, past2_ids_b)

                # 最近两帧有匹配的，但是最近的三帧没有
                # there is no matched object in past frames
                if len(a_idx)==0 or len(b_idx)==0:
                    matched_past2 = center_points_past2[past2_ids_a]
                    matched_past1 = center_points_past1[past1_ids]

                    time_length = cav_content['past_k_time_diff'][0] - cav_content['past_k_time_diff'][1]
                    if time_length == 0:
                        time_length = 1
                    flow = (matched_past1 - matched_past2) / time_length

                    flow = flow*(0-cav_content['past_k_time_diff'][0])
                    selected_box_3dcenter_past0 = coord_past1['pred_box_center_tensor'][past1_ids,]
                    selected_box_3dcorner_past0 = box_utils.boxes_to_corners2d(selected_box_3dcenter_past0, order='hwl')
                    flow_map, mask = self.generate_flow_map(flow, selected_box_3dcorner_past0, scale=2.5, shape_list=shape_list)
                    flow_map_list.append(flow_map)
                    reserved_mask.append(mask)
                    if self.viz_flag:
                        matched_idx_list.append(torch.stack([past1_ids, past2_ids_a], dim=1))
                        selected_box_3dcorner_compensated = selected_box_3dcorner_past0.clone()
                        selected_box_3dcorner_compensated[:, :, :-1] += flow.unsqueeze(1).repeat(1, 4, 1) 
                        compensated_results_list.append(selected_box_3dcorner_compensated)
                    continue

                # 三帧匹配结果输入预测模块
                matched_past1 = center_points_past1[past1_ids[a_idx]].unsqueeze(0)
                matched_past2 = center_points_past2[past2_ids_a[a_idx]].unsqueeze(0)
                matched_past3 = center_points_past3[past3_ids[b_idx]].unsqueeze(0)

                obj_coords = torch.cat([matched_past3, matched_past2, matched_past1], dim=0)
                obj_coords = obj_coords.permute(1, 0, 2) # (N, k, 2)

                obj_coords_norm = obj_coords - obj_coords[:, -1:, :] # (N, k, 2)
                past_k_time_diff = torch.flip(past_k_time_diff, dims=[0]) # (k,) TODO: check if this is correct
                past_k_time_diff_norm = past_k_time_diff - past_k_time_diff[-1] # (k,)

                speed = torch.zeros_like(obj_coords_norm) # (N, k, 2)
                speed[:, 1:, :] = torch.div((obj_coords_norm[:, 1:, :] - obj_coords_norm[:, :-1, :]), ((past_k_time_diff[1:] - past_k_time_diff[:-1]).unsqueeze(-1)).unsqueeze(0)) # (N, k-1, 2) / (1, k-1, 1)    

                obj_input = torch.cat([obj_coords_norm, speed], dim=-1) # (N, k, 4)

                obj_input = obj_input.unsqueeze(0) # (1, N, k, 4)
                
                last_time_length = (past_k_time_diff_norm[-1] - past_k_time_diff_norm[-2])
                if last_time_length == 0:
                    print("==== Warning! You met repeated package! ====")
                    query = torch.zeros(obj_input.shape)[:,:,:1,:2].to(obj_input.device) # (1, N, 1, 2)
                else: 
                    query = obj_coords_norm[:, -1:, :] + \
                        (obj_coords_norm[:, -1:, :]-obj_coords_norm[:, -2:-1, :])*(0-past_k_time_diff[-1]) / \
                            last_time_length
                    query = query.unsqueeze(0) # (1, N, 1, 2)

                target_time_diff = torch.tensor([-past_k_time_diff[-1]]).to(obj_input.device) # (1,)

                # target_time_diff = torch.tensor([-past_k_time_diff[0]]).to(obj_input.device) # (1,)
                compensated_coords_norm = self.compensate_motion(obj_input, query, past_k_time_diff_norm, target_time_diff) + query

                flow = compensated_coords_norm.squeeze(0).squeeze(1) # (N, 2)

                # flow = flow*(0-cav_content['past_k_time_diff'][0])
                selected_box_3dcenter_past0 = coord_past1['pred_box_center_tensor'][past1_ids[a_idx],]
                selected_box_3dcorner_past0 = box_utils.boxes_to_corners2d(selected_box_3dcenter_past0, order='hwl')

                if self.viz_flag and not(len(a_idx) < len(past2_ids_a)):
                    unit_matched_list = torch.stack([past1_ids[a_idx], past2_ids_a[a_idx], past3_ids[b_idx]], dim=1) # (N_obj, 3)
                    selected_box_3dcorner_compensated = selected_box_3dcorner_past0.clone()
                    selected_box_3dcorner_compensated[:, :, :-1] += flow.unsqueeze(1).repeat(1, 4, 1) 

                # 两帧匹配成功 but三帧匹配失败 将两帧的结果进行插值 com 表示补集
                if len(a_idx) < len(past2_ids_a): 
                    com_past1_ids = [elem.item() for id, elem in enumerate(past1_ids) if id not in a_idx]
                    com_past2_ids = [elem.item() for id, elem in enumerate(past2_ids_a) if id not in a_idx]
                    matched_past1 = center_points_past1[com_past1_ids]
                    matched_past2 = center_points_past2[com_past2_ids]
                    
                    time_length = cav_content['past_k_time_diff'][0] - cav_content['past_k_time_diff'][1]
                    if time_length == 0:
                        time_length = 1
                    com_flow = (matched_past1 - matched_past2) / time_length

                    com_flow = com_flow*(0-cav_content['past_k_time_diff'][0])
                    com_selected_box_3dcenter_past0 = coord_past1['pred_box_center_tensor'][com_past1_ids,]
                    com_selected_box_3dcorner_past0 = box_utils.boxes_to_corners2d(com_selected_box_3dcenter_past0, order='hwl')
                    
                    flow = torch.cat([flow, com_flow], dim=0)
                    selected_box_3dcorner_past0 = torch.cat([selected_box_3dcorner_past0, com_selected_box_3dcorner_past0], dim=0)

                    # matched: 
                    # past1: past1_ids[a_idx] + com_past1_ids
                    # past2: past2_ids_a[a_idx] + com_past2_ids
                    # past3: past3_ids[b_idx]
                    if self.viz_flag:
                        tmp_past_1 = torch.cat([past1_ids[a_idx], torch.tensor(com_past1_ids).to(past1_ids)], dim=0)
                        tmp_past_2 = torch.cat([past2_ids_a[a_idx], torch.tensor(com_past2_ids).to(past1_ids)], dim=0)
                        unit_matched_list = torch.stack([tmp_past_1, tmp_past_2], dim=1)  # (N_obj, 2)
                        selected_box_3dcorner_compensated = selected_box_3dcorner_past0.clone()
                        selected_box_3dcorner_compensated[:, :, :-1] += flow.unsqueeze(1).repeat(1, 4, 1) 

                if self.viz_flag:
                    matched_idx_list.append(unit_matched_list)
                    compensated_results_list.append(selected_box_3dcorner_compensated)
                
                # debug use
                debug_flag = False
                if debug_flag and flow.shape[0] != 0:
                    viz_save_path = '/DB/data/sizhewei/logs/where2comm_max_multiscale_resnet_32ch/vis_debug'
                    torch.save(features_dict[cav]['spatial_features_2d'][0], viz_save_path+'/features_2d.pt')
                    torch.save(features_dict[cav]['spatial_features'][0], viz_save_path+'/features.pt')
                    torch.save(selected_box_3dcorner_past0, viz_save_path+'/bbx_list.pt')
                    torch.save(flow, viz_save_path+'/flow.pt')
                    print(f"===saved, max flow is {flow.max()}===")
                ############
                
                flow_map, mask = self.generate_flow_map(flow, selected_box_3dcorner_past0, scale=2.5, shape_list=shape_list)
                flow_map_list.append(flow_map)
                reserved_mask.append(mask)
        
        final_flow_map = torch.concat(flow_map_list, dim=0) # [N_b, H, W, 2]
        reserved_mask = torch.concat(reserved_mask, dim=0)  # [N_b, C, H, W]
        
        if self.viz_flag:
            return final_flow_map, reserved_mask, matched_idx_list, compensated_results_list
        
        return final_flow_map, reserved_mask

    def forward_flow(self, input_dict, shape_list):
        """
        Parameters:
        -----------
        input_dict : The dictionary containing the box detections on each frame of each cav.
            dict : { 
                cav_idx : {
                    'past_k_time_diff':
                    [0], [1], ..., [k-1] : {
                        pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                        pred_box_center_tensor : n, 7
                        scores: (n, )
                    }
                }
            }
        features_dict : The dictionary containing the output of the model.
            dict : {
                'ego' / cav_id : {
                    'spatial_features_2d'
                    'spatial_features'
                    'psm'
                    'rm'
                }
            }

        Returns:
        --------
        features_dict :
            dict : {
                'ego' / cav_id : {
                    'spatial_features_2d'
                    'spatial_features'
                    'psm'
                    'rm'
                    'updated_spatial_features_2d'
                    'updated_spatial_features'
                }
            }
        """
        flow_map_list = []
        reserved_mask = []
        if self.viz_flag:
            matched_idx_list = []
        for cav, cav_content in input_dict.items():
            if cav == 0:
                # ego do not need warp
                C, H, W = shape_list
                basic_mat = torch.tensor([[1,0,0],[0,1,0]]).unsqueeze(0).to(torch.float32)
                basic_warp_mat = F.affine_grid(basic_mat, [1, C, H, W], align_corners=False).to(shape_list.device)
                mask = torch.ones(1, C, H, W).to(shape_list)
                flow_map_list.append(basic_warp_mat)
                reserved_mask.append(mask)
            else:
                coord_past1 = cav_content[0]
                coord_past2 = cav_content[1]

                center_points_past1 = coord_past1['pred_box_center_tensor'][:,:2]
                center_points_past2 = coord_past2['pred_box_center_tensor'][:,:2]

                cost_mat_center = torch.cdist(center_points_past2, center_points_past1) # [num_cav_past2,num_cav_past1]

                self.thre_post_process = 10
                original_cost_mat_center = cost_mat_center.clone()
                cost_mat_center[cost_mat_center > self.thre_post_process] = 1000

                cost_mat_center_drop_2 = torch.sum(torch.where(cost_mat_center > self.thre, 1, 0), dim=1)
                dist_valid_past2 = torch.where(cost_mat_center_drop_2 < center_points_past1.shape[0])
                cost_mat_center_drop_1 = torch.sum(torch.where(cost_mat_center > self.thre, 1, 0), dim=0)
                dist_valid_past1 = torch.where(cost_mat_center_drop_1 < center_points_past2.shape[0])

                cost_mat_center = cost_mat_center[dist_valid_past2[0], :]
                cost_mat_center = cost_mat_center[:, dist_valid_past1[0]]
                
                # cost_mat_iou = get_ious()
                cost_mat = cost_mat_center
                past2_ids, past1_ids = linear_sum_assignment(cost_mat.cpu())
                
                past2_ids = dist_valid_past2[0][past2_ids]
                past1_ids = dist_valid_past1[0][past1_ids]
                
                ### a trick
                matched_cost = original_cost_mat_center[past2_ids, past1_ids]
                valid_mat_idx = torch.where(matched_cost < self.thre_post_process)
                past2_ids = past2_ids[valid_mat_idx[0]]
                past1_ids = past1_ids[valid_mat_idx[0]]
                ####################

                matched_past2 = center_points_past2[past2_ids]
                matched_past1 = center_points_past1[past1_ids]

                if self.viz_flag:
                    matched_idx_list.append(torch.stack([past1_ids, past2_ids], dim=1))

                time_length = cav_content['past_k_time_diff'][0] - cav_content['past_k_time_diff'][1]
                if time_length == 0:
                    time_length = 1
                flow = (matched_past1 - matched_past2) / time_length

                flow = flow*(0-cav_content['past_k_time_diff'][0])
                selected_box_3dcenter_past0 = coord_past1['pred_box_center_tensor'][past1_ids,]
                selected_box_3dcorner_past0 = box_utils.boxes_to_corners2d(selected_box_3dcenter_past0, order='hwl') # TODO: box order should be a parameter
                
                # debug use
                debug_flag = False
                if debug_flag and flow.shape[0] != 0:
                    viz_save_path = '/DB/data/sizhewei/logs/where2comm_max_multiscale_resnet_32ch/vis_debug'
                    torch.save(features_dict[cav]['spatial_features_2d'][0], viz_save_path+'/features_2d.pt')
                    torch.save(features_dict[cav]['spatial_features'][0], viz_save_path+'/features.pt')
                    torch.save(selected_box_3dcorner_past0, viz_save_path+'/bbx_list.pt')
                    torch.save(flow, viz_save_path+'/flow.pt')
                    print(f"===saved, max flow is {flow.max()}===")
                ############
                
                flow_map, mask = self.generate_flow_map(flow, selected_box_3dcorner_past0, scale=2.5, shape_list=shape_list)
                flow_map_list.append(flow_map)
                reserved_mask.append(mask)
                continue
        
        final_flow_map = torch.concat(flow_map_list, dim=0) # [N_b, H, W, 2]
        reserved_mask = torch.concat(reserved_mask, dim=0)  # [N_b, C, H, W]

        if self.viz_flag:
            return final_flow_map, reserved_mask, matched_idx_list
        return final_flow_map, reserved_mask
        '''
                updated_spatial_features_2d = self.feature_warp(features_dict[cav]['spatial_features_2d'][0], selected_box_3dcorner_past0, flow, scale=1.25)
                updated_spatial_features = self.feature_warp(features_dict[cav]['spatial_features'][0], selected_box_3dcorner_past0, flow, scale=2.5)

                # debug use
                if debug_flag and flow.shape[0] != 0:
                    torch.save(updated_spatial_features_2d, viz_save_path+'/updated_feature.pt')
                ############

            features_dict[cav].update({
                'updated_spatial_features_2d': updated_spatial_features_2d
            })
            features_dict[cav].update({
                'updated_spatial_features': updated_spatial_features
            })

        return features_dict
        '''

    def forward_flow_dir(self, input_dict, shape_list):
        """
        Parameters:
        -----------
        input_dict : The dictionary containing the box detections on each frame of each cav.
            dict : { 
                cav_idx : {
                    'past_k_time_diff':
                    [0], [1], ..., [k-1] : {
                        pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                        pred_box_center_tensor : n, 7
                        scores: (n, )
                    }
                }
            }
        features_dict : The dictionary containing the output of the model.
            dict : {
                'ego' / cav_id : {
                    'spatial_features_2d'
                    'spatial_features'
                    'psm'
                    'rm'
                }
            }

        Returns:
        --------
        features_dict :
            dict : {
                'ego' / cav_id : {
                    'spatial_features_2d'
                    'spatial_features'
                    'psm'
                    'rm'
                    'updated_spatial_features_2d'
                    'updated_spatial_features'
                }
            }
        """
        flow_map_list = []
        reserved_mask = []
        if self.viz_flag:
            original_reserved_mask = []
            matched_idx_list = []
            compensated_results_list = []
        for cav, cav_content in input_dict.items():
            if cav == 0:
                # ego do not need warp
                C, H, W = shape_list
                basic_mat = torch.tensor([[1,0,0],[0,1,0]]).unsqueeze(0).to(torch.float32)
                basic_warp_mat = F.affine_grid(basic_mat, [1, C, H, W], align_corners=False).to(shape_list.device)
                mask = torch.ones(1, C, H, W).to(shape_list)
                flow_map_list.append(basic_warp_mat)
                reserved_mask.append(mask)
                if self.viz_flag:
                    original_reserved_mask.append(mask)
            else:
                coord_past1 = cav_content[0]
                coord_past2 = cav_content[1]

                center_points_past1 = coord_past1['pred_box_center_tensor'][:,:2]
                center_points_past2 = coord_past2['pred_box_center_tensor'][:,:2]

                cost_mat_center = torch.zeros((center_points_past2.shape[0], center_points_past1.shape[0])).to(center_points_past1.device)

                center_points_past1_repeat = center_points_past1.unsqueeze(0).repeat(center_points_past2.shape[0], 1, 1)
                center_points_past2_repeat = center_points_past2.unsqueeze(1).repeat(1, center_points_past1.shape[0], 1)

                delta_mat = center_points_past1_repeat - center_points_past2_repeat

                angle_mat = torch.atan2(delta_mat[:,:,1], delta_mat[:,:,0]) # [num_cav_past2,num_cav_past1]

                coord_past2_angle_reverse = coord_past2['pred_box_center_tensor'][:,6].unsqueeze(1).repeat(1, center_points_past1.shape[0]).clone()

                visible_mat = torch.where((torch.abs(angle_mat-coord_past2['pred_box_center_tensor'][:,6].unsqueeze(1).repeat(1, center_points_past1.shape[0])) < 0.785) | (torch.abs(angle_mat-coord_past2['pred_box_center_tensor'][:,6].unsqueeze(1).repeat(1, center_points_past1.shape[0])) > 5.495) , 1, 0) # [num_cav_past2,num_cav_past1]

                cost_mat_center = torch.cdist(center_points_past2, center_points_past1) # [num_cav_past2,num_cav_past1]

                visible_mat = torch.where(cost_mat_center<0.5, 1, visible_mat)

                tmp_thre = torch.tensor(1000.).to(torch.float32).to(visible_mat.device)
                cost_mat_center = torch.where(visible_mat==1, cost_mat_center, tmp_thre)

                if cost_mat_center.shape[1] == 0 or cost_mat_center.shape[0] == 0:
                    C, H, W = shape_list
                    basic_mat = torch.tensor([[1,0,0],[0,1,0]]).unsqueeze(0).to(torch.float32)
                    basic_warp_mat = F.affine_grid(basic_mat, [1, C, H, W], align_corners=False).to(shape_list.device)
                    mask = torch.ones(1, C, H, W).to(shape_list)
                    flow_map_list.append(basic_warp_mat)
                    reserved_mask.append(mask)
                    if self.viz_flag:
                        matched_idx_list.append(torch.stack([torch.tensor([]), torch.tensor([])], dim=1).to(shape_list.device))
                        compensated_results_list.append(torch.zeros(0, 8, 3).to(shape_list.device))
                        original_reserved_mask.append(mask)
                    continue

                match = torch.min(cost_mat_center, dim=1)
                match_to_keep = torch.where(match[0] < 5)

                past2_ids = match_to_keep[0]
                past1_ids = match[1][match_to_keep[0]]
                
                coord_past2_angle_reverse += 3.1415926
                coord_past2_angle_reverse[coord_past2_angle_reverse>3.1415926] -= 6.2831852

                left_past2_id = [n for n in range(cost_mat_center.shape[0]) if n not in past2_ids]
                left_past1_id = [n for n in range(cost_mat_center.shape[1]) if n not in past1_ids]

                angle_mat_left = angle_mat[left_past2_id, :][:, left_past1_id]

                coord_past2_angle_reverse_left = coord_past2_angle_reverse[left_past2_id, :][: ,left_past1_id]

                visible_mat_left = torch.where((torch.abs(angle_mat_left-coord_past2_angle_reverse_left) < 0.785) | (torch.abs(angle_mat_left-coord_past2_angle_reverse_left) > 5.495) , 1, 0) # [num_cav_past2,num_cav_past1]
                
                cost_mat_center_left = torch.cdist(center_points_past2[left_past2_id], center_points_past1[left_past1_id]) # [num_cav_past2,num_cav_past1]
                # cost_mat_center_left = cost_mat_center[left_past2_id, :][:, left_past1_id]

                visible_mat_left = torch.where(cost_mat_center_left<0.5, 1, visible_mat_left)
                cost_mat_center_left = torch.where(visible_mat_left==1, cost_mat_center_left, tmp_thre)

                if cost_mat_center_left.shape[1] != 0 and cost_mat_center_left.shape[0] != 0:
                    match_left = torch.min(cost_mat_center_left, dim=1)
                    match_to_keep_left = torch.where(match_left[0] < 5)

                    if match_to_keep_left[0].shape[0] != 0:
                        past2_ids_left = match_to_keep_left[0]
                        past2_ids = torch.cat([past2_ids, torch.tensor(left_past2_id)[past2_ids_left].to(past2_ids.device)])
                        past1_ids_left = match_left[1][match_to_keep_left[0]]
                        past1_ids = torch.cat([past1_ids, torch.tensor(left_past1_id)[past1_ids_left].to(past1_ids.device)])

                # ##############################################
                # self.thre_post_process = 10
                # original_cost_mat_center = cost_mat_center.clone()
                # cost_mat_center[cost_mat_center > self.thre_post_process] = 1000

                # cost_mat_center_drop_2 = torch.sum(torch.where(cost_mat_center > self.thre, 1, 0), dim=1)
                # dist_valid_past2 = torch.where(cost_mat_center_drop_2 < center_points_past1.shape[0])
                # cost_mat_center_drop_1 = torch.sum(torch.where(cost_mat_center > self.thre, 1, 0), dim=0)
                # dist_valid_past1 = torch.where(cost_mat_center_drop_1 < center_points_past2.shape[0])

                # cost_mat_center = cost_mat_center[dist_valid_past2[0], :]
                # cost_mat_center = cost_mat_center[:, dist_valid_past1[0]]
                
                # # cost_mat_iou = get_ious()
                # cost_mat = cost_mat_center
                # past2_ids, past1_ids = linear_sum_assignment(cost_mat.cpu())
                
                # past2_ids = dist_valid_past2[0][past2_ids]
                # past1_ids = dist_valid_past1[0][past1_ids]
                
                # ### a trick
                # matched_cost = original_cost_mat_center[past2_ids, past1_ids]
                # valid_mat_idx = torch.where(matched_cost < self.thre_post_process)
                # past2_ids = past2_ids[valid_mat_idx[0]]
                # past1_ids = past1_ids[valid_mat_idx[0]]
                # ####################

                matched_past2 = center_points_past2[past2_ids]
                matched_past1 = center_points_past1[past1_ids]

                time_length = cav_content['past_k_time_diff'][0] - cav_content['past_k_time_diff'][1]
                if time_length == 0:
                    time_length = 1
                flow = (matched_past1 - matched_past2) / time_length

                flow = flow*(0-cav_content['past_k_time_diff'][0])
                selected_box_3dcenter_past0 = coord_past1['pred_box_center_tensor'][past1_ids,]
                selected_box_3dcorner_past0 = box_utils.boxes_to_corners2d(selected_box_3dcenter_past0, order='hwl') # TODO: box order should be a parameter

                if self.viz_flag:
                    matched_idx_list.append(torch.stack([past1_ids, past2_ids], dim=1))
                    selected_box_3dcorner_compensated = selected_box_3dcorner_past0.clone()
                    selected_box_3dcorner_compensated[:, :, :-1] += flow.unsqueeze(1).repeat(1, 4, 1) 
                    compensated_results_list.append(selected_box_3dcorner_compensated)
                
                # debug use
                debug_flag = False
                if debug_flag and flow.shape[0] != 0:
                    viz_save_path = '/DB/data/sizhewei/logs/where2comm_max_multiscale_resnet_32ch/vis_debug'
                    torch.save(features_dict[cav]['spatial_features_2d'][0], viz_save_path+'/features_2d.pt')
                    torch.save(features_dict[cav]['spatial_features'][0], viz_save_path+'/features.pt')
                    torch.save(selected_box_3dcorner_past0, viz_save_path+'/bbx_list.pt')
                    torch.save(flow, viz_save_path+'/flow.pt')
                    print(f"===saved, max flow is {flow.max()}===")
                ############
                
                if self.viz_flag:
                    flow_map, mask, single_ori_mask = self.generate_flow_map(flow, selected_box_3dcorner_past0, scale=2.5, shape_list=shape_list)
                    original_reserved_mask.append(single_ori_mask)
                else:
                    flow_map, mask = self.generate_flow_map(flow, selected_box_3dcorner_past0, scale=2.5, shape_list=shape_list)
                flow_map_list.append(flow_map)
                reserved_mask.append(mask)
                continue
        
        final_flow_map = torch.concat(flow_map_list, dim=0) # [N_b, H, W, 2]
        reserved_mask = torch.concat(reserved_mask, dim=0)  # [N_b, C, H, W]

        if self.viz_flag:
            original_reserved_mask = torch.concat(original_reserved_mask, dim=0)  # [N_b, C, H, W]
            return final_flow_map, reserved_mask, original_reserved_mask, matched_idx_list, compensated_results_list
        return final_flow_map, reserved_mask
        '''
                updated_spatial_features_2d = self.feature_warp(features_dict[cav]['spatial_features_2d'][0], selected_box_3dcorner_past0, flow, scale=1.25)
                updated_spatial_features = self.feature_warp(features_dict[cav]['spatial_features'][0], selected_box_3dcorner_past0, flow, scale=2.5)

                # debug use
                if debug_flag and flow.shape[0] != 0:
                    torch.save(updated_spatial_features_2d, viz_save_path+'/updated_feature.pt')
                ############

            features_dict[cav].update({
                'updated_spatial_features_2d': updated_spatial_features_2d
            })
            features_dict[cav].update({
                'updated_spatial_features': updated_spatial_features
            })

        return features_dict
        '''

    def generate_flow_map(self, flow, bbox_list, scale=1.25, shape_list=None, align_corners=False, file_suffix=""):
        """
        Parameters
        -----------
        feature: [C, H, W] at voxel scale
        bbox_list: [num_cav, 4, 3] at cav coodinate system and lidar scale
        flow:[num_cav, 2] at cav coodinate system and lidar scale
            bbox_list & flow : x and y are exactly image coordinate
            ------------> x
            |
            |
            |
            y
        scale: float, scale meters to voxel, feature_length / lidar_range_length = 1.25 or 2.5

        Returns
        -------
        updated_feature: feature after being warped by flow, [C, H, W]
        """
        # flow = torch.tensor([70, 0]).unsqueeze(0).to(feature)

        # only use x and y
        bbox_list = bbox_list[:, :, :2]

        # scale meters to voxel, feature_length / lidar_range_length = 1.25
        flow = flow * scale
        bbox_list = bbox_list * scale

        flag_viz = False
        #######
        # store two parts of bbx: 1. original bbx, 2. 
        if flag_viz:
            viz_bbx_list = bbox_list
            fig, ax = plt.subplots(4, 1, figsize=(5,11))
            ######## viz-0: original feature, original bbx
            canvas_ori = viz_on_canvas(feature, bbox_list, scale=scale)
            plt.sca(ax[0])
            # plt.axis("off")
            plt.imshow(canvas_ori.canvas)
            ##########
        #######

        C, H, W = shape_list
        num_cav = bbox_list.shape[0]
        basic_mat = torch.tensor([[1,0,0],[0,1,0]]).unsqueeze(0).to(torch.float32)
        basic_warp_mat = F.affine_grid(basic_mat, [1, C, H, W], align_corners=align_corners).to(shape_list.device)
        reserved_area = torch.zeros((C, H, W)).to(shape_list.device)  # C, H, W 
        if flow.shape[0] == 0 : 
            # reserved_area = torch.ones((C, H, W)).to(shape_list.device)  # C, H, W 
            return basic_warp_mat,  reserved_area.unsqueeze(0)  # 返回不变的矩阵

        '''
        create affine matrix:
        ------------
        1  0  -2*t_y/W
        0  1  -2*t_x/H
        0  0    1 
        ------------
        '''
        flow_clone = flow.detach().clone()

        affine_matrices = torch.eye(3).unsqueeze(0).repeat(flow.shape[0], 1, 1)
        flow_clone = -2 * flow_clone / torch.tensor([W, H]).to(torch.float32).to(shape_list.device)
        # flow_clone = flow_clone[:, [1, 0]]
        affine_matrices[:, :2, 2] = flow_clone 
        
        cav_t_mat = affine_matrices[:, :2, :]   # n, 2, 3
        # print("cav_t_mat", cav_t_mat)

        cav_warp_mat = F.affine_grid(cav_t_mat,
                            [num_cav, C, H, W],
                            align_corners=align_corners).to(shape_list.device) # .to() 统一数据格式 float32
        
        flowed_bbx_list = bbox_list + flow.unsqueeze(1).repeat(1,4,1)  # n, 4, 2
        ######### viz-1: original feature, original bbx and flowed bbx
        if flag_viz:
            viz_bbx_list = torch.cat((bbox_list, flowed_bbx_list), dim=0)
            canvas_hidden = viz_on_canvas(feature, viz_bbx_list, scale=scale)
            plt.sca(ax[1])
            # plt.axis("off") 
            plt.imshow(canvas_hidden.canvas)
        ##########

        x_min = torch.min(flowed_bbx_list[:,:,0],dim=1)[0] - 1
        x_max = torch.max(flowed_bbx_list[:,:,0],dim=1)[0] + 1
        y_min = torch.min(flowed_bbx_list[:,:,1],dim=1)[0] - 1
        y_max = torch.max(flowed_bbx_list[:,:,1],dim=1)[0] + 1
        x_min_fid = (x_min + int(W/2)).to(torch.int)
        x_max_fid = (x_max + int(W/2)).to(torch.int)
        y_min_fid = (y_min + int(H/2)).to(torch.int)
        y_max_fid = (y_max + int(H/2)).to(torch.int)

        for cav in range(num_cav):
            basic_warp_mat[0,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]] = cav_warp_mat[cav,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]]

        # generate mask
        x_min_ori = torch.min(bbox_list[:,:,0],dim=1)[0]
        x_max_ori = torch.max(bbox_list[:,:,0],dim=1)[0]
        y_min_ori = torch.min(bbox_list[:,:,1],dim=1)[0]
        y_max_ori = torch.max(bbox_list[:,:,1],dim=1)[0]
        x_min_fid_ori = (x_min_ori + int(W/2)).to(torch.int)
        x_max_fid_ori = (x_max_ori + int(W/2)).to(torch.int)
        y_min_fid_ori = (y_min_ori + int(H/2)).to(torch.int)
        y_max_fid_ori = (y_max_ori + int(H/2)).to(torch.int)
        # set original location as 0
        for cav in range(num_cav):
            reserved_area[:,y_min_fid_ori[cav]:y_max_fid_ori[cav],x_min_fid_ori[cav]:x_max_fid_ori[cav]] = 0
        # set warped location as 1
        for cav in range(num_cav):
            reserved_area[:,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]] = 1

        if self.viz_flag:
            single_reserved_area = torch.zeros_like(reserved_area)
            for cav in range(num_cav):
                single_reserved_area[:,y_min_fid_ori[cav]:y_max_fid_ori[cav],x_min_fid_ori[cav]:x_max_fid_ori[cav]] = 1
            return basic_warp_mat, reserved_area.unsqueeze(0), single_reserved_area.unsqueeze(0)

        return basic_warp_mat, reserved_area.unsqueeze(0)
        '''
        ##################################### below is not used
        final_feature = F.grid_sample(feature.unsqueeze(0), basic_warp_mat, align_corners=align_corners)[0]
        
        ####### viz-2: warped feature, flowed box and warped 
        if flag_viz:
            p_0 = torch.stack((x_min, y_min), dim=1).to(torch.int)
            p_1 = torch.stack((x_min, y_max), dim=1).to(torch.int)
            p_2 = torch.stack((x_max, y_max), dim=1).to(torch.int)
            p_3 = torch.stack((x_max, y_min), dim=1).to(torch.int)
            warp_area_bbox_list = torch.stack((p_0, p_1, p_2, p_3), dim=1)
            viz_bbx_list = torch.cat((flowed_bbx_list, warp_area_bbox_list), dim=0)
            canvas_new = viz_on_canvas(final_feature, viz_bbx_list, scale=scale)
            plt.sca(ax[2]) 
            # plt.axis("off") 
            plt.imshow(canvas_new.canvas)
        ############## 

        reserved_area = torch.ones_like(feature)  # C, H, W
        x_min_ori = torch.min(bbox_list[:,:,0],dim=1)[0]
        x_max_ori = torch.max(bbox_list[:,:,0],dim=1)[0]
        y_min_ori = torch.min(bbox_list[:,:,1],dim=1)[0]
        y_max_ori = torch.max(bbox_list[:,:,1],dim=1)[0]
        x_min_fid_ori = (x_min_ori + int(W/2)).to(torch.int)
        x_max_fid_ori = (x_max_ori + int(W/2)).to(torch.int)
        y_min_fid_ori = (y_min_ori + int(H/2)).to(torch.int)
        y_max_fid_ori = (y_max_ori + int(H/2)).to(torch.int)
        # set original location as 0
        for cav in range(num_cav):
            reserved_area[:,y_min_fid_ori[cav]:y_max_fid_ori[cav],x_min_fid_ori[cav]:x_max_fid_ori[cav]] = 0
        # set warped location as 1
        for cav in range(num_cav):
            reserved_area[:,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]] = 1
        final_feature = final_feature * reserved_area

        ####### viz-3: mask area out of warped bbx
        if flag_viz:
            partial_feature_one = torch.zeros_like(feature)  # C, H, W
            for cav in range(num_cav):
                partial_feature_one[:,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]] = 1
            masked_final_feature = partial_feature_one * final_feature
            canvas_hidden = viz_on_canvas(masked_final_feature, warp_area_bbox_list, scale=scale)
            plt.sca(ax[3]) 
            # plt.axis("off") 
            plt.imshow(canvas_hidden.canvas)
        ##############

        ####### viz: draw figures
        if flag_viz:
            plt.tight_layout()
            plt.savefig(f'result_canvas_{file_suffix}.jpg', transparent=False, dpi=400)
            plt.clf()

            fig, axes = plt.subplots(2, 1, figsize=(4, 4))
            major_ticks_x = np.linspace(0,350,8)
            minor_ticks_x = np.linspace(0,350,15)
            major_ticks_y = np.linspace(0,100,3)
            minor_ticks_y = np.linspace(0,100,5)
            for i, ax in enumerate(axes):
                plt.sca(ax); #plt.axis("off")
                ax.set_xticks(major_ticks_x); ax.set_xticks(minor_ticks_x, minor=True)
                ax.set_yticks(major_ticks_y); ax.set_yticks(minor_ticks_y, minor=True)
                ax.grid(which='major', color='w', linewidth=0.4)
                ax.grid(which='minor', color='w', linewidth=0.2, alpha=0.5)
                if i==0:
                    plt.imshow(torch.max(feature, dim=0)[0].cpu())
                else:
                    plt.imshow(torch.max(final_feature, dim=0)[0].cpu())
            plt.tight_layout()
            plt.savefig(f'result_features_{file_suffix}.jpg', transparent=False, dpi=400)
            plt.clf()
        #######

        return final_feature
        '''

    def feature_warp(self, feature, bbox_list, flow, scale=1.25, align_corners=False, file_suffix=""):
        """
        Parameters
        -----------
        feature: [C, H, W] at voxel scale
        bbox_list: [num_cav, 4, 3] at cav coodinate system and lidar scale
        flow:[num_cav, 2] at cav coodinate system and lidar scale
            bbox_list & flow : x and y are exactly image coordinate
            ------------> x
            |
            |
            |
            y
        scale: float, scale meters to voxel, feature_length / lidar_range_length = 1.25 or 2.5

        Returns
        -------
        updated_feature: feature after being warped by flow, [C, H, W]
        """
        # flow = torch.tensor([70, 0]).unsqueeze(0).to(feature)

        if flow.shape[0] == 0 : 
            return feature

        # only use x and y
        bbox_list = bbox_list[:, :, :2]

        # scale meters to voxel, feature_length / lidar_range_length = 1.25
        flow = flow * scale
        bbox_list = bbox_list * scale

        flag_viz = False
        #######
        # store two parts of bbx: 1. original bbx, 2. 
        if flag_viz:
            viz_bbx_list = bbox_list
            fig, ax = plt.subplots(4, 1, figsize=(5,11))
            ######## viz-0: original feature, original bbx
            canvas_ori = viz_on_canvas(feature, bbox_list, scale=scale)
            plt.sca(ax[0])
            # plt.axis("off")
            plt.imshow(canvas_ori.canvas)
            ##########
        #######

        C, H, W = feature.size()
        num_cav = bbox_list.shape[0]
        basic_mat = torch.tensor([[1,0,0],[0,1,0]]).unsqueeze(0).to(torch.float32)
        basic_warp_mat = F.affine_grid(basic_mat, [1, C, H, W], align_corners=align_corners).to(feature)

        '''
        create affine matrix:
        ------------
        1  0  -2*t_y/W
        0  1  -2*t_x/H
        0  0    1 
        ------------
        '''
        flow_clone = flow.detach().clone()

        affine_matrices = torch.eye(3).unsqueeze(0).repeat(flow.shape[0], 1, 1)
        flow_clone = -2 * flow_clone / torch.tensor([feature.shape[2], feature.shape[1]]).to(feature)
        # flow_clone = flow_clone[:, [1, 0]]
        affine_matrices[:, :2, 2] = flow_clone 
        
        cav_t_mat = affine_matrices[:, :2, :]   # n, 2, 3
        # print("cav_t_mat", cav_t_mat)

        cav_warp_mat = F.affine_grid(cav_t_mat,
                            [num_cav, C, H, W],
                            align_corners=align_corners).to(feature) # .to() 统一数据格式 float32
        
        flowed_bbx_list = bbox_list + flow.unsqueeze(1).repeat(1,4,1)  # n, 4, 2
        ######### viz-1: original feature, original bbx and flowed bbx
        if flag_viz:
            viz_bbx_list = torch.cat((bbox_list, flowed_bbx_list), dim=0)
            canvas_hidden = viz_on_canvas(feature, viz_bbx_list, scale=scale)
            plt.sca(ax[1])
            # plt.axis("off") 
            plt.imshow(canvas_hidden.canvas)
        ##########

        x_min = torch.min(flowed_bbx_list[:,:,0],dim=1)[0] - 1
        x_max = torch.max(flowed_bbx_list[:,:,0],dim=1)[0] + 1
        y_min = torch.min(flowed_bbx_list[:,:,1],dim=1)[0] - 1
        y_max = torch.max(flowed_bbx_list[:,:,1],dim=1)[0] + 1
        x_min_fid = (x_min + int(W/2)).to(torch.int)
        x_max_fid = (x_max + int(W/2)).to(torch.int)
        y_min_fid = (y_min + int(H/2)).to(torch.int)
        y_max_fid = (y_max + int(H/2)).to(torch.int)

        for cav in range(num_cav):
            basic_warp_mat[0,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]] = cav_warp_mat[cav,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]]

        final_feature = F.grid_sample(feature.unsqueeze(0), basic_warp_mat, align_corners=align_corners)[0]
        
        ####### viz-2: warped feature, flowed box and warped 
        if flag_viz:
            p_0 = torch.stack((x_min, y_min), dim=1).to(torch.int)
            p_1 = torch.stack((x_min, y_max), dim=1).to(torch.int)
            p_2 = torch.stack((x_max, y_max), dim=1).to(torch.int)
            p_3 = torch.stack((x_max, y_min), dim=1).to(torch.int)
            warp_area_bbox_list = torch.stack((p_0, p_1, p_2, p_3), dim=1)
            viz_bbx_list = torch.cat((flowed_bbx_list, warp_area_bbox_list), dim=0)
            canvas_new = viz_on_canvas(final_feature, viz_bbx_list, scale=scale)
            plt.sca(ax[2]) 
            # plt.axis("off") 
            plt.imshow(canvas_new.canvas)
        ############## 

        reserved_area = torch.ones_like(feature)  # C, H, W
        x_min_ori = torch.min(bbox_list[:,:,0],dim=1)[0]
        x_max_ori = torch.max(bbox_list[:,:,0],dim=1)[0]
        y_min_ori = torch.min(bbox_list[:,:,1],dim=1)[0]
        y_max_ori = torch.max(bbox_list[:,:,1],dim=1)[0]
        x_min_fid_ori = (x_min_ori + int(W/2)).to(torch.int)
        x_max_fid_ori = (x_max_ori + int(W/2)).to(torch.int)
        y_min_fid_ori = (y_min_ori + int(H/2)).to(torch.int)
        y_max_fid_ori = (y_max_ori + int(H/2)).to(torch.int)
        # set original location as 0
        for cav in range(num_cav):
            reserved_area[:,y_min_fid_ori[cav]:y_max_fid_ori[cav],x_min_fid_ori[cav]:x_max_fid_ori[cav]] = 0
        # set warped location as 1
        for cav in range(num_cav):
            reserved_area[:,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]] = 1
        final_feature = final_feature * reserved_area

        ####### viz-3: mask area out of warped bbx
        if flag_viz:
            partial_feature_one = torch.zeros_like(feature)  # C, H, W
            for cav in range(num_cav):
                partial_feature_one[:,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]] = 1
            masked_final_feature = partial_feature_one * final_feature
            canvas_hidden = viz_on_canvas(masked_final_feature, warp_area_bbox_list, scale=scale)
            plt.sca(ax[3]) 
            # plt.axis("off") 
            plt.imshow(canvas_hidden.canvas)
        ##############

        ####### viz: draw figures
        if flag_viz:
            plt.tight_layout()
            plt.savefig(f'result_canvas_{file_suffix}.jpg', transparent=False, dpi=400)
            plt.clf()

            fig, axes = plt.subplots(2, 1, figsize=(4, 4))
            major_ticks_x = np.linspace(0,350,8)
            minor_ticks_x = np.linspace(0,350,15)
            major_ticks_y = np.linspace(0,100,3)
            minor_ticks_y = np.linspace(0,100,5)
            for i, ax in enumerate(axes):
                plt.sca(ax); #plt.axis("off")
                ax.set_xticks(major_ticks_x); ax.set_xticks(minor_ticks_x, minor=True)
                ax.set_yticks(major_ticks_y); ax.set_yticks(minor_ticks_y, minor=True)
                ax.grid(which='major', color='w', linewidth=0.4)
                ax.grid(which='minor', color='w', linewidth=0.2, alpha=0.5)
                if i==0:
                    plt.imshow(torch.max(feature, dim=0)[0].cpu())
                else:
                    plt.imshow(torch.max(final_feature, dim=0)[0].cpu())
            plt.tight_layout()
            plt.savefig(f'result_features_{file_suffix}.jpg', transparent=False, dpi=400)
            plt.clf()
        #######

        return final_feature

    def backup_feature_warp(self, feature, bbox_list, flow, scale=1.25, align_corners=False):
        """
        Parameters
        -----------
        feature: [C, H, W] at voxel scale
        bbox_list: [num_cav, 4, 3] at cav coodinate system and lidar scale
        flow:[num_cav, 2] at cav coodinate system and lidar scale
            bbox_list & flow : x and y are exactly image coordinate
            ------------> x
            |
            |
            |
            y
        scale: float, scale meters to voxel, feature_length / lidar_range_length = 1.25 or 2.5

        Returns
        -------
        updated_feature: feature after being warped by flow, [C, H, W]
        """
        # flow = torch.tensor([70, 0]).unsqueeze(0).to(feature)

        if flow.shape[0] == 0 : 
            return feature

        # only use x and y
        bbox_list = bbox_list[:, :, :2]

        # scale meters to voxel, feature_length / lidar_range_length = 1.25
        flow = flow * scale
        bbox_list = bbox_list * scale

        # # store two parts of bbx: 1. original bbx, 2. 
        # viz_bbx_list = bbox_list
        # fig, ax = plt.subplots(4, 1, figsize=(5,11))
        
        # ######## viz-0: original feature, original bbx
        # canvas_ori = viz_on_canvas(feature, bbox_list)
        # plt.sca(ax[0])
        # # plt.axis("off")
        # plt.imshow(canvas_ori.canvas)
        # ##########

        C, H, W = feature.size()
        num_cav = bbox_list.shape[0]
        basic_mat = torch.tensor([[1,0,0],[0,1,0]]).unsqueeze(0).to(torch.float32)
        basic_warp_mat = F.affine_grid(basic_mat, [1, C, H, W], align_corners=align_corners).to(feature)

        '''
        create affine matrix:
        ------------
        1  0  -2*t_y/W
        0  1  -2*t_x/H
        0  0    1 
        ------------
        '''
        flow_clone = flow.detach().clone()

        affine_matrices = torch.eye(3).unsqueeze(0).repeat(flow.shape[0], 1, 1)
        flow_clone = -2 * flow_clone / torch.tensor([feature.shape[2], feature.shape[1]]).to(feature)
        # flow_clone = flow_clone[:, [1, 0]]
        affine_matrices[:, :2, 2] = flow_clone 
        
        cav_t_mat = affine_matrices[:, :2, :]   # n, 2, 3
        # print("cav_t_mat", cav_t_mat)

        cav_warp_mat = F.affine_grid(cav_t_mat,
                            [num_cav, C, H, W],
                            align_corners=align_corners).to(feature) # .to() 统一数据格式 float32
        
        ######### viz-1: original feature, original bbx and flowed bbx
        flowed_bbx_list = bbox_list + flow.unsqueeze(1).repeat(1,4,1)  # n, 4, 2
        # viz_bbx_list = torch.cat((bbox_list, flowed_bbx_list), dim=0)
        # canvas_hidden = viz_on_canvas(feature, viz_bbx_list)
        # plt.sca(ax[1])
        # # plt.axis("off") 
        # plt.imshow(canvas_hidden.canvas)
        ##########

        x_min = torch.min(flowed_bbx_list[:,:,0],dim=1)[0] - 1
        x_max = torch.max(flowed_bbx_list[:,:,0],dim=1)[0] + 1
        y_min = torch.min(flowed_bbx_list[:,:,1],dim=1)[0] - 1
        y_max = torch.max(flowed_bbx_list[:,:,1],dim=1)[0] + 1
        x_min_fid = (x_min + 176).to(torch.int) # TODO: 这里面的176需要重新考虑
        x_max_fid = (x_max + 176).to(torch.int)
        y_min_fid = (y_min + 50).to(torch.int)  # TODO: 这里面的50需要重新考虑
        y_max_fid = (y_max + 50).to(torch.int)

        for cav in range(num_cav):
            basic_warp_mat[0,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]] = cav_warp_mat[cav,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]]

        final_feature = F.grid_sample(feature.unsqueeze(0), basic_warp_mat, align_corners=align_corners)[0]
        
        ####### viz-2: warped feature, flowed box and warped 
        # p_0 = torch.stack((x_min, y_min), dim=1).to(torch.int)
        # p_1 = torch.stack((x_min, y_max), dim=1).to(torch.int)
        # p_2 = torch.stack((x_max, y_max), dim=1).to(torch.int)
        # p_3 = torch.stack((x_max, y_min), dim=1).to(torch.int)
        # warp_area_bbox_list = torch.stack((p_0, p_1, p_2, p_3), dim=1)
        # viz_bbx_list = torch.cat((flowed_bbx_list, warp_area_bbox_list), dim=0)
        # canvas_new = viz_on_canvas(final_feature, viz_bbx_list)
        # plt.sca(ax[2]) 
        # # plt.axis("off") 
        # plt.imshow(canvas_new.canvas)
        ############## 

        ####### viz-3: mask area out of warped bbx
        # partial_feature_one = torch.zeros_like(feature)  # C, H, W
        # for cav in range(num_cav):
        #     partial_feature_one[:,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]] = 1
        # masked_final_feature = partial_feature_one * final_feature
        # canvas_hidden = viz_on_canvas(masked_final_feature, warp_area_bbox_list)
        # plt.sca(ax[3]) 
        # # plt.axis("off") 
        # plt.imshow(canvas_hidden.canvas)
        ##############

        # plt.tight_layout()
        # plt.savefig('result_canvas.jpg', transparent=False, dpi=400)
        # plt.clf()

        # fig, axes = plt.subplots(2, 1, figsize=(4, 4))
        # major_ticks_x = np.linspace(0,350,8)
        # minor_ticks_x = np.linspace(0,350,15)
        # major_ticks_y = np.linspace(0,100,3)
        # minor_ticks_y = np.linspace(0,100,5)
        # for i, ax in enumerate(axes):
        #     plt.sca(ax); #plt.axis("off")
        #     ax.set_xticks(major_ticks_x); ax.set_xticks(minor_ticks_x, minor=True)
        #     ax.set_yticks(major_ticks_y); ax.set_yticks(minor_ticks_y, minor=True)
        #     ax.grid(which='major', color='w', linewidth=0.4)
        #     ax.grid(which='minor', color='w', linewidth=0.2, alpha=0.5)
        #     if i==0:
        #         plt.imshow(torch.max(feature, dim=0)[0].cpu())
        #     else:
        #         plt.imshow(torch.max(final_feature, dim=0)[0].cpu())
        # plt.tight_layout()
        # plt.savefig('result_features.jpg', transparent=False, dpi=400)
        # plt.clf()

        return final_feature

    def get_common_elements(self, A, B):
        common_elements_A = []
        common_elements_B = []
        for i, a in enumerate(A):
            for j, b in enumerate(B):
                if a == b:
                    common_elements_A.append(i)
                    common_elements_B.append(j)
        return common_elements_A, common_elements_B


def get_center_points(corner_points):
    corner_points2d = corner_points[:,:4,:2]

    centers_x = torch.mean(corner_points2d[:,:,0],dim=1,keepdim=True)

    centers_y = torch.mean(corner_points2d[:,:,1],dim=1,keepdim=True)

    return torch.cat((centers_x,centers_y), dim=1)

def build_matcher(args):
    return HungarianMatcher(cost_class=args.set_cost_class, cost_bbox=args.set_cost_bbox, cost_giou=args.set_cost_giou)
opencood/tools/matcher.py
MediaBrain-SJTU/CoBEVFlow
class TrajectoryPrediction():
    def __init__(self):
        super().__init__()
    
    def forward(self, input_dict):
        """
            input_dict: 
            {
                [batch_id] : {
                    'ego' : {
                        'past_k_time_diff' : len=k, float
                        [0] {
                            pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                            pred_box_center_tensor : n, 7
                            scores: (n, )
                        }
                        ... 
                        [k-1]
                    }
                    'cav_id' : { ... }
                }
            }  
            output_dict:
            {
                [batch_id] : {
                    'ego' : {
                        'past_k_time_diff' : 
                        [0] {
                            pred_box_3dcorner_tensor: The prediction bounding box tensor after NMS. n, 8, 3
                            pred_box_center_tensor : n, 7
                            scores: (n, )

                        }
                        ... 
                        [k-1]
                    }
                    [cav_id] : { ... }
                }
            }
        """
        return 0
opencood/tools/matcher_new.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu
# Modified: Sizhe Wei

import argparse
import os
import statistics
import sys
sys.path.append(os.getcwd())
import time

import torch
from torch.utils.data import DataLoader
from tensorboardX import SummaryWriter

import importlib
import opencood.hypes_yaml.yaml_utils as yaml_utils
from opencood.tools import train_utils
from opencood.data_utils.datasets import build_dataset

from tqdm import tqdm
from tqdm.contrib import tenumerate
from tqdm.auto import trange

import subprocess

run_test = True
# from opencood.data_utils.datasets.intermediate_fusion_dataset_opv2v_irregular import illegal_path_list
compensation = True 

def train_parser():
    parser = argparse.ArgumentParser(description="synthetic data generation")
    parser.add_argument("--hypes_yaml", "-y", type=str, required=True,
                        help='data generation yaml file needed ')
    parser.add_argument('--model_dir', default='',
                        help='Continued training path')
    parser.add_argument('--fusion_method', '-f', default="intermediate",
                        help='passed to inference.')
    parser.add_argument('--pretrained_path', default='',
                        help='The path of the model need to be fine tuned.')
    parser.add_argument('--device', '-d', default="cuda", help='cuda or cpu')
    parser.add_argument('--two_stage', help='whether to use two stage training', default=0, type=int)
    opt = parser.parse_args()
    return opt


def main():
    opt = train_parser()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(opt.device)

    hypes = yaml_utils.load_yaml(opt.hypes_yaml, opt)

    finetune_flag = False
    if 'is_finetune' in hypes:
        finetune_flag = hypes['is_finetune']
    if finetune_flag:
        finetune_time_interval = int(hypes['binomial_n'] * hypes['binomial_p'])
        print('### Finetune mode, only tune header. ###')

    print('### Dataset Building ... ###')
    start_time = time.time()
    opencood_train_dataset = build_dataset(hypes, visualize=False, train=True)
    opencood_validate_dataset = build_dataset(hypes,
                                              visualize=False,
                                              train=False)

    train_loader = DataLoader(opencood_train_dataset,
                            batch_size=hypes['train_params']['batch_size'],
                            num_workers=16,
                            collate_fn=opencood_train_dataset.collate_batch_train,
                            shuffle=True,
                            pin_memory=True,
                            drop_last=True)
    val_loader = DataLoader(opencood_validate_dataset,
                            batch_size=hypes['train_params']['batch_size'],
                            num_workers=16,
                            collate_fn=opencood_train_dataset.collate_batch_train,
                            shuffle=True,
                            pin_memory=True,
                            drop_last=True)
    end_time = time.time()
    print("=== Time consumed: %.1f minutes. ===" % ((end_time - start_time)/60))
    start_time = time.time()
    
    print('### Creating Model ... ###')
    model = train_utils.create_model(hypes)
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print("device: ", device)

    # record lowest validation loss checkpoint.
    lowest_val_loss = 1e5
    lowest_val_epoch = -1    
    
    # load pre-train model for single view
    is_pre_trained = False
    if 'is_single_pre_trained' in hypes and hypes['is_single_pre_trained']['pre_train_flag']:
        is_pre_trained = True
    if is_pre_trained:
        pretrain_path = hypes['is_single_pre_trained']['pre_train_path']
        initial_epoch = hypes['is_single_pre_trained']['pre_train_epoch']
        pre_train_model = torch.load(os.path.join(pretrain_path, 'net_epoch%d.pth' % initial_epoch))
        model.load_state_dict(pre_train_model, strict=False)
        print("### Pre-trained point pillar {} loaded successfully! ###".format(os.path.join(pretrain_path, 'net_epoch%d.pth' % initial_epoch)))
        fix = hypes['is_single_pre_trained']['pre_train_fix']
        if fix:
            for name, value in model.named_parameters():
                if name in pre_train_model:
                    value.requires_grad = False

    # we assume gpu is necessary
    if torch.cuda.is_available():
        model.to(device)

    # define the loss 
    criterion = train_utils.create_loss(hypes)

    # optimizer setup
    optimizer = train_utils.setup_optimizer(hypes, model, is_pre_trained)
    
    # for fine tune:
    if opt.pretrained_path:
        saved_path = opt.pretrained_path
        pretrained_model_dict = torch.load(saved_path, map_location='cpu')
        diff_keys = {k:v for k, v in pretrained_model_dict.items() if k not in model.state_dict()}
        if diff_keys:
            print(f"!!! PreTrained model has keys: {diff_keys.keys()}, \
                which are not in the model you have created!!!")
        diff_keys = {k:v for k, v in model.state_dict().items() if k not in pretrained_model_dict.keys()}
        if diff_keys:
            print(f"!!! Created model has keys: {diff_keys.keys()}, \
                which are not in the model you have trained!!!")
        model.load_state_dict(pretrained_model_dict, strict=False)
        for name, value in model.named_parameters():
            if name in pretrained_model_dict:
                value.requires_grad = False
    
    # lr scheduler setup
    # if we want to train from last checkpoint.
    if opt.model_dir:
        saved_path = opt.model_dir
        init_epoch, model = train_utils.load_saved_model_diff(saved_path, model, finetune_flag)
        if finetune_flag:
            init_epoch = 0 # if finetune, we set the init_epoch to 10
        scheduler = train_utils.setup_lr_schedular(hypes, optimizer, init_epoch=init_epoch)

    else:
        init_epoch = 0
        log_path = '../logs' # Your log file path
        # if we train the model from scratch, we need to create a folder to save the model
        saved_path = train_utils.setup_train(hypes, log_path)
        scheduler = train_utils.setup_lr_schedular(hypes, optimizer)

    end_time = time.time()
    print("=== Time consumed: %.1f minutes. ===" % ((end_time - start_time)/60))

    # record training
    writer = SummaryWriter(saved_path)

    start_time = time.time()
    print('### Training start! ###')
    epoches = hypes['train_params']['epoches']
    supervise_single_flag = False
    if 'supervise_single_flag' in hypes['train_params'].keys():
        supervise_single_flag = hypes['train_params']['supervise_single_flag']
        print(f"=== supervise_single_flag: {supervise_single_flag} ===")

    ############ For DiscoNet ##############
    if "kd_flag" in hypes.keys():
        kd_flag = True
        teacher_model_name = hypes['kd_flag']['teacher_model'] # point_pillar_disconet_teacher
        teacher_model_config = hypes['kd_flag']['teacher_model_config']
        teacher_checkpoint_path = hypes['kd_flag']['teacher_path']

        # import the model
        model_filename = "opencood.models." + teacher_model_name
        model_lib = importlib.import_module(model_filename)
        teacher_model_class = None
        target_model_name = teacher_model_name.replace('_', '')

        for name, cls in model_lib.__dict__.items():
            if name.lower() == target_model_name.lower():
                teacher_model_class = cls
        
        teacher_model = teacher_model_class(teacher_model_config)
        teacher_model.load_state_dict(torch.load(teacher_checkpoint_path), strict=False)
        
        for p in teacher_model.parameters():
            p.requires_grad_(False)

        if torch.cuda.is_available():
                teacher_model.to(device)

        teacher_model.eval()

    else:
        kd_flag = False

    sample_interval_all_epoch = 0
    for epoch in range(init_epoch, max(epoches, init_epoch)):
        for param_group in optimizer.param_groups:
            print('learning rate %f' % param_group["lr"])
        
        sample_interval = 0
        i = 0
        for i, batch_data in enumerate(train_loader): 
            if batch_data is None:
                continue
            model.train()
            model.zero_grad()
            optimizer.zero_grad()
            batch_data = train_utils.to_device(batch_data, device)

            batch_data['ego']['epoch'] = epoch

            # TODO: dataset parameter is only used for training flow module
            if opt.two_stage:
                ouput_dict = model(batch_data['ego'], opencood_train_dataset)
            else: 
                ouput_dict = model(batch_data['ego'])

            if kd_flag:
                teacher_output_dict = teacher_model(batch_data['ego'])
                ouput_dict.update(teacher_output_dict)

            # # only for SyncNet training
            # final_loss = criterion(ouput_dict, batch_data['ego']['label_dict'])
            # if i % 10 == 0:
            #     criterion.logging(epoch, i, len(train_loader), writer)

            # first argument is always your output dictionary,
            # second argument is always your label dictionary.
            
            if 'with_compensation' in hypes['model']['args'] and \
                hypes['model']['args']['with_compensation']:
                if hypes['model']['args']['with_single_supervise']:
                    final_loss = ouput_dict['recon_loss'] 
                    single_det_loss = criterion(ouput_dict, batch_data['ego']['single_object_label'], mode = 'single')
                    detection_loss = criterion(ouput_dict, batch_data['ego']['label_dict'])
                    final_loss = ouput_dict['recon_loss'] + single_det_loss + detection_loss
                else:
                    detection_loss = criterion(ouput_dict, batch_data['ego']['label_dict'])
                    final_loss = ouput_dict['recon_loss'] + detection_loss
            else:
                final_loss = criterion(ouput_dict, batch_data['ego']['label_dict'])
                if i % 10 == 0:
                    criterion.logging(epoch, i, len(train_loader), writer)
                if supervise_single_flag:
                    final_loss += criterion(ouput_dict, batch_data['ego']['single_object_label'], suffix="_single")
                    if i % 10 == 0:
                        criterion.logging(epoch, i, len(train_loader), writer, suffix="_single")

            # back-propagation
            final_loss.backward()
            optimizer.step()
        
        sample_interval /= i
        sample_interval_all_epoch += sample_interval

        if epoch % hypes['train_params']['eval_freq'] == 0:
            valid_ave_loss = []
            end_time = time.time()
            print('### %d th epoch trained, start validation! Time consumed %.2f ###' % (epoch, (end_time - start_time)/60))
            with torch.no_grad():
                for i, batch_data in tenumerate(val_loader):
                    if batch_data is None:
                        continue
                    model.zero_grad()
                    optimizer.zero_grad()
                    model.eval()

                    batch_data = train_utils.to_device(batch_data, device)
                    batch_data['ego']['epoch'] = epoch
                    # TODO: dataset parameter is only used for training flow module
                    if opt.two_stage:
                        ouput_dict = model(batch_data['ego'], opencood_validate_dataset)
                    else:
                        ouput_dict = model(batch_data['ego'])

                    if kd_flag:
                        teacher_output_dict = teacher_model(batch_data['ego'])
                        ouput_dict.update(teacher_output_dict)

                    final_loss = criterion(ouput_dict,
                                           batch_data['ego']['label_dict'])
                    valid_ave_loss.append(final_loss.item())

            valid_ave_loss = statistics.mean(valid_ave_loss)
            print('At epoch %d, the validation loss is %f' % (epoch,
                                                              valid_ave_loss))
            writer.add_scalar('Validate_Loss', valid_ave_loss, epoch)

            # lowest val loss
            if valid_ave_loss < lowest_val_loss:
                lowest_val_loss = valid_ave_loss
                save_checkpoint_prefix = 'net_epoch_bestval_at'
                if finetune_flag:
                    save_checkpoint_prefix = f'finetune_{finetune_time_interval}_' + save_checkpoint_prefix
                torch.save(model.state_dict(),
                       os.path.join(saved_path,
                                    f'{save_checkpoint_prefix}%d.pth' % (epoch + 1)))
                if lowest_val_epoch != -1 and os.path.exists(os.path.join(saved_path,
                                    f'{save_checkpoint_prefix}%d.pth' % (lowest_val_epoch))):
                        os.remove(os.path.join(saved_path,
                                    f'{save_checkpoint_prefix}%d.pth' % (lowest_val_epoch)))
                lowest_val_epoch = epoch + 1

        if epoch % hypes['train_params']['save_freq'] == 0:
            torch.save(model.state_dict(),
                       os.path.join(saved_path,
                                    'net_epoch%d.pth' % (epoch + 1)))
        scheduler.step(epoch)

    end_time = time.time()
    print("Time consumed: %.1f" % ((end_time - start_time)/60))
    sample_interval_all_epoch /= max(epoches, init_epoch) - init_epoch
    print("Avg sample interval of all epochs: %.2f" % sample_interval_all_epoch) 
    print('Training Finished, checkpoints saved to %s' % saved_path)
    torch.cuda.empty_cache()
    
    if run_test:
        fusion_method = opt.fusion_method
        cmd = f"python opencood/tools/inference_multi_sweep.py --model_dir {saved_path} --fusion_method {fusion_method}"
        print(f"Running command: {cmd}")
        os.system(cmd)

if __name__ == '__main__':
    main()

opencood/tools/train.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Hao Xiang <haxiang@g.ucla.edu>,
# License: TDG-Attribution-NonCommercial-NoDistrib


import glob
import importlib
import yaml
import os
import re
from datetime import datetime
import shutil
import torch
import torch.optim as optim
import numpy
from pathlib import Path

def backup_script(full_path, folders_to_save=["models", "data_utils", "utils", "loss", "tools"]):
    target_folder = os.path.join(full_path, 'scripts')
    if not os.path.exists(target_folder):
        os.mkdir(target_folder)
    
    current_path = os.path.dirname(__file__)  # __file__ refer to this file, then the dirname is "?/tools"

    for folder_name in folders_to_save:
        ttarget_folder = os.path.join(target_folder, folder_name)
        source_folder = os.path.join(current_path, f'../{folder_name}')
        shutil.copytree(source_folder, ttarget_folder)

# def load_saved_model(saved_path, model):
#     """
#     Load saved model if exiseted

#     Parameters
#     __________
#     saved_path : str
#        model saved path
#     model : opencood object
#         The model instance.

#     Returns
#     -------
#     model : opencood object
#         The model instance loaded pretrained params.
#     """
#     assert os.path.exists(saved_path), '{} not found'.format(saved_path)

#     def findLastCheckpoint(save_dir):
#         file_list = glob.glob(os.path.join(save_dir, '*epoch*.pth'))
#         if file_list:
#             epochs_exist = []
#             for file_ in file_list:
#                 result = re.findall(".*epoch(.*).pth.*", file_)
#                 epochs_exist.append(int(result[0]))
#             initial_epoch_ = max(epochs_exist)
#         else:
#             initial_epoch_ = 0
#         return initial_epoch_

#     if os.path.exists(os.path.join(saved_path, 'net_latest.pth')):
#         model.load_state_dict(torch.load(
#             os.path.join(saved_path,
#                          'net_latest.pth')))
#         return 100, model

#     initial_epoch = findLastCheckpoint(saved_path)
#     if initial_epoch > 0:
#         print('resuming by loading epoch %d' % initial_epoch)
#         model.load_state_dict(torch.load(
#             os.path.join(saved_path,
#                          'net_epoch%d.pth' % initial_epoch)), strict=False)

#     return initial_epoch, model

def load_saved_model(saved_path, model):
    """
    Load saved model if exiseted

    Parameters
    __________
    saved_path : str
        model saved path
    model : opencood object
        The model instance.

    Returns
    -------
    model : opencood object
        The model instance loaded pretrained params.
    """
    assert os.path.exists(saved_path), '{} not found'.format(saved_path)

    def findLastCheckpoint(save_dir):
        file_list = glob.glob(os.path.join(save_dir, '*epoch*.pth'))
        if file_list:
            epochs_exist = []
            for file_ in file_list:
                result = re.findall(".*epoch(.*).pth.*", file_)
                epochs_exist.append(int(result[0]))
            initial_epoch_ = max(epochs_exist)
        else:
            initial_epoch_ = 0
        return initial_epoch_

    file_list = glob.glob(os.path.join(saved_path, 'net_epoch_bestval_at*.pth'))
    if file_list:
        assert len(file_list) == 1
        print("resuming best validation model at epoch %d" % \
                eval(file_list[0].split("/")[-1].rstrip(".pth").lstrip("net_epoch_bestval_at")))
        model.load_state_dict(torch.load(file_list[0] , map_location='cpu'), strict=False)
        return eval(file_list[0].split("/")[-1].rstrip(".pth").lstrip("net_epoch_bestval_at")), model

    initial_epoch = findLastCheckpoint(saved_path)
    if initial_epoch > 0:
        print('resuming by loading epoch %d' % initial_epoch)
        model.load_state_dict(torch.load(
            os.path.join(saved_path,
                'net_epoch%d.pth' % initial_epoch), map_location='cpu'), strict=False)

    # ckpt = torch.load(os.path.join(saved_path, 'net_epoch%d.pth' % initial_epoch), map_location='cpu')
    # keys = [x for x in ckpt if 'resnet' in x]
    # ckpt_new = {}
    # for k in keys:
    #     old_key = k.split('layer')
    #     start = old_key[0]
    #     end = old_key[1][1:]
    #     mid = int(old_key[1][0]) - 1
    #     new_key = start + 'layer' + str(mid) + end
    #     ckpt_new[new_key] = ckpt[k]
    # for k in ckpt:
    #     if k not in keys:
    #         ckpt_new[k] = ckpt[k]
    # model.load_state_dict(ckpt_new, strict=False)
    return initial_epoch, model

def load_saved_model_diff(saved_path, model, finetune_flag=False):
    """
    Load saved model, model and checkpoint may not be totally same.

    Parameters
    __________
    saved_path : str
        model saved path
    model : opencood object
        The model instance.

    Returns
    -------
    model : opencood object
        The model instance loaded pretrained params.
    """
    assert os.path.exists(saved_path), '{} not found'.format(saved_path)

    def findLastCheckpoint(save_dir):
        file_list = glob.glob(os.path.join(save_dir, '*epoch*.pth'))
        if file_list:
            epochs_exist = []
            for file_ in file_list:
                result = re.findall(".*epoch(.*).pth.*", file_)
                epochs_exist.append(int(result[0]))
            initial_epoch_ = max(epochs_exist)
        else:
            initial_epoch_ = 0
        return initial_epoch_

    file_list = glob.glob(os.path.join(saved_path, 'net_epoch_bestval_at*.pth'))
    if file_list:
        assert len(file_list) == 1
        print("resuming best validation model at epoch %d" % \
                eval(file_list[0].split("/")[-1].rstrip(".pth").lstrip("net_epoch_bestval_at")))
        trained_model_dict = torch.load(file_list[0] , map_location='cpu')

        # TODO: uncomment lines below to introduce v2xvit v2-based pretrained model
        # old_trained_model_dict = torch.load(file_list[0] , map_location='cpu')
        # trained_model_dict = {}
        # for key, value in old_trained_model_dict.items():
        #     if 'fusion_net.fusion_net.' in key:
        #         new_key = key.replace('fusion_net.fusion_net.', 'fusion_net.')
        #     else:
        #         new_key = key
        #     trained_model_dict[new_key] = value

        # TODO: uncomment lines below to introduce disconet v2-based pretrained model
        # old_trained_model_dict = torch.load(file_list[0] , map_location='cpu')
        # trained_model_dict = {}
        # for key, value in old_trained_model_dict.items():
        #     if 'fusion_net.pixel_weight_layer.' in key:
        #         new_key = key.replace('fusion_net.pixel_weight_layer.', 'pixel_weight_layer.')
        #     else:
        #         new_key = key
        #     trained_model_dict[new_key] = value

        ######## finetune header
        # if finetune_flag:
        # finetune_split_name = ['cls_head.weight', 'cls_head.bias', 'reg_head.weight', 'reg_head.bias']
        # for k in finetune_split_name:
        #     trained_model_dict.update({
        #         'fused_'+k: trained_model_dict[k]})
    #############################

        diff_keys = {k:v for k, v in trained_model_dict.items() if k not in model.state_dict()}
        if diff_keys:
            print(f"!!! Trained model has keys: {diff_keys.keys()}, \
                which are not in the model you have created!!!")
        diff_keys = {k:v for k, v in model.state_dict().items() if k not in trained_model_dict.keys()}
        if diff_keys:
            print(f"!!! Created model has keys: {diff_keys.keys()}, \
                which are not in the model you have trained!!!")
        model.load_state_dict(trained_model_dict, strict=False)
        return eval(file_list[0].split("/")[-1].rstrip(".pth").lstrip("net_epoch_bestval_at")), model

    initial_epoch = findLastCheckpoint(saved_path)
    if initial_epoch > 0:
        print('resuming by loading epoch %d' % initial_epoch)
        trained_model_dict = torch.load(\
            os.path.join(saved_path, 'net_epoch%d.pth' % initial_epoch), map_location='cpu')
        diff_keys = {k:v for k, v in trained_model_dict.items() if k not in model.state_dict()}
        if diff_keys:
            print(f"!!! Trained model has keys: {diff_keys.keys()}, \
                which are not in the model you have created!!!")
        model.load_state_dict(trained_model_dict, strict=False)

    return initial_epoch, model


def load_two_parts_model(saved_path, model):
    """
    This function is used for intermediate+flow method, we should load two parts of pretrained model:
    1. single detection model
    2. fused detection head: self.cls_head_fused & self.reg_head_fused
    """
    file_list = glob.glob(os.path.join(saved_path, 'net_epoch_bestval_at*.pth'))
    assert len(file_list) == 1
    print("loading pretrained single detection model at epoch %d" % \
            eval(file_list[0].split("/")[-1].rstrip(".pth").lstrip("net_epoch_bestval_at")))
    trained_single_model_dict = torch.load(file_list[0], map_location='cpu')

    fused_model_path = '/DB/data/sizhewei/logs/where2comm_irr_no_shift_pretrain/net_epoch_bestval_at18.pth'
    trained_fused_model_dict = torch.load(fused_model_path, map_location='cpu')
    
    final_model_stat_dict = model.state_dict()
    
    state_dict = {k:v for k,v in trained_single_model_dict.items() if k in final_model_stat_dict.keys()}    
    state_dict.update({
        'cls_head_fused.weight': trained_fused_model_dict['cls_head.weight'], 
        'cls_head_fused.bias': trained_fused_model_dict['cls_head.bias'], 
        'reg_head_fused.weight': trained_fused_model_dict['reg_head.weight'], 
        'reg_head_fused.bias': trained_fused_model_dict['reg_head.bias']
    })

    final_model_stat_dict.update(state_dict)

    model.load_state_dict(final_model_stat_dict, strict=False)
    return eval(file_list[0].split("/")[-1].rstrip(".pth").lstrip("net_epoch_bestval_at")), model


def setup_train(hypes, log_path="/DB/data/sizhewei"):
    """
    Create folder for saved model based on current timestep and model name

    Parameters
    ----------
    hypes: dict
        Config yaml dictionary for training:
    """
    model_name = hypes['name']
    current_time = datetime.now()

    folder_name = model_name
    # folder_name 需要包含：
    # model_name, time_delay, num_sweep, batch_size, 
    # if 'time_delay' in hypes:
    #     folder_name = folder_name + '_d_' + str(hypes['time_delay'])
    # else:
    #     folder_name = folder_name + '_d_' + str(0)
    # if 'num_sweep_frames' in hypes:
    #     folder_name = folder_name + '_swps_' + str(hypes['num_sweep_frames'])
    # else:
    #     folder_name = folder_name + '_swps_' + str(1)
    # folder_name = folder_name + '_bs_' + str(hypes['train_params']['batch_size'])
    # model_hypes = hypes['model']['args']
    # if 'base_bev_backbone' in model_hypes and model_hypes['base_bev_backbone']['resnet']:
    #     folder_name = folder_name + '_w_resnet'
    # else:
    #     folder_name = folder_name + '_wo_resnet'
    # if 'rain_model' in model_hypes and  model_hypes['rain_model']['multi_scale']:
    #     folder_name = folder_name + '_w_multiscale'
    # else:
    #     folder_name = folder_name + '_wo_multiscale'
    # if 'exp_name' in hypes and len(str(hypes['exp_name']))>0:
    #     folder_name = folder_name + '_' + str(hypes['exp_name'])
    curr_time = current_time.strftime("_%Y_%m_%d_%H_%M_%S")
    folder_name = folder_name + curr_time
    # current_path = os.path.dirname(__file__)
    current_path = Path(log_path)
    current_path = os.path.join(current_path, 'logs/')

    full_path = os.path.join(current_path, folder_name)
    print("full path is: ", full_path)

    if not os.path.exists(full_path):
        if not os.path.exists(full_path):
            try:
                os.makedirs(full_path)
            except FileExistsError:
                pass
        # save the yaml file
        save_name = os.path.join(full_path, 'config.yaml')
        with open(save_name, 'w') as outfile:
            yaml.dump(hypes, outfile)

    # if not os.path.exists(full_path):
    #     os.makedirs(full_path)
    #     # save the yaml file
    #     save_name = os.path.join(full_path, 'config.yaml')
    #     with open(save_name, 'w') as outfile:
    #         yaml.dump(hypes, outfile)

        backup_script(full_path)

    return full_path


def create_model(hypes):
    """
    Import the module "models/[model_name].py

    Parameters
    __________
    hypes : dict
        Dictionary containing parameters.

    Returns
    -------
    model : opencood,object
        Model object.
    """
    backbone_name = hypes['model']['core_method']
    backbone_config = hypes['model']['args']

    model_filename = "opencood.models." + backbone_name
    model_lib = importlib.import_module(model_filename)
    model = None
    target_model_name = backbone_name.replace('_', '')

    for name, cls in model_lib.__dict__.items():
        if name.lower() == target_model_name.lower():
            model = cls

    if model is None:
        print('backbone not found in models folder. Please make sure you '
              'have a python file named %s and has a class '
              'called %s ignoring upper/lower case' % (model_filename,
                                                       target_model_name))
        exit(0)
    instance = model(backbone_config)
    return instance


def create_loss(hypes):
    """
    Create the loss function based on the given loss name.

    Parameters
    ----------
    hypes : dict
        Configuration params for training.
    Returns
    -------
    criterion : opencood.object
        The loss function.
    """
    loss_func_name = hypes['loss']['core_method']
    loss_func_config = hypes['loss']['args']

    loss_filename = "opencood.loss." + loss_func_name
    loss_lib = importlib.import_module(loss_filename)
    loss_func = None
    target_loss_name = loss_func_name.replace('_', '')

    for name, lfunc in loss_lib.__dict__.items():
        if name.lower() == target_loss_name.lower():
            loss_func = lfunc

    if loss_func is None:
        print('loss function not found in loss folder. Please make sure you '
              'have a python file named %s and has a class '
              'called %s ignoring upper/lower case' % (loss_filename,
                                                       target_loss_name))
        exit(0)

    criterion = loss_func(loss_func_config)
    return criterion


def setup_optimizer(hypes, model, is_pre_trained=False):
    """
    Create optimizer corresponding to the yaml file

    Parameters
    ----------
    hypes : dict
        The training configurations.
    model : opencood model
        The pytorch model
    """
    method_dict = hypes['optimizer']
    optimizer_method = getattr(optim, method_dict['core_method'], None)
    if not optimizer_method:
        raise ValueError('{} is not supported'.format(method_dict['name']))
    #####################################################################
    if is_pre_trained:
        params = filter(lambda p: p.requires_grad, model.parameters())
    else:
        params = model.parameters()
    ######################################################################
    if 'args' in method_dict:
        return optimizer_method(params,
                                lr=method_dict['lr'],
                                **method_dict['args'])
    else:
        return optimizer_method(params,
                                lr=method_dict['lr'])


def setup_lr_schedular(hypes, optimizer, init_epoch=None):
    """
    Set up the learning rate schedular.

    Parameters
    ----------
    hypes : dict
        The training configurations.

    optimizer : torch.optimizer
    """
    lr_schedule_config = hypes['lr_scheduler']
    last_epoch = init_epoch if init_epoch is not None else 0
    

    if lr_schedule_config['core_method'] == 'step':
        from torch.optim.lr_scheduler import StepLR
        step_size = lr_schedule_config['step_size']
        gamma = lr_schedule_config['gamma']
        scheduler = StepLR(optimizer, step_size=step_size, gamma=gamma)

    elif lr_schedule_config['core_method'] == 'multistep':
        from torch.optim.lr_scheduler import MultiStepLR
        milestones = lr_schedule_config['step_size']
        gamma = lr_schedule_config['gamma']
        scheduler = MultiStepLR(optimizer,
                                milestones=milestones,
                                gamma=gamma)

    else:
        from torch.optim.lr_scheduler import ExponentialLR
        gamma = lr_schedule_config['gamma']
        scheduler = ExponentialLR(optimizer, gamma)

    for _ in range(last_epoch):
        scheduler.step()

    return scheduler


def to_device(inputs, device):
    if isinstance(inputs, list):
        return [to_device(x, device) for x in inputs]
    elif isinstance(inputs, dict):        
        return {k: to_device(v, device) for k, v in inputs.items()}
    else:
        if isinstance(inputs, int) or isinstance(inputs, float) \
                or isinstance(inputs, str) or isinstance(inputs, numpy.int64):  
                # AttributeError: 'numpy.int64' object has no attribute 'to'
                # sizhewei @ 2022/10/04
                # 添加类型 numpy.int64 数据集中的 sample_idx 是此类型, numpy.ndarry 数据集中 time_interval 是此类型
            return inputs
        return inputs.to(device)

opencood/tools/train_utils.py
MediaBrain-SJTU/CoBEVFlow
import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
import math
from torch.autograd import Variable

def clones(module, N):
    "Produce N identical layers."
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.embedding = nn.Sequential(nn.Linear(vocab,d_model),nn.ReLU(),nn.Linear(d_model,d_model))

    def forward(self, x):
        return self.embedding(x)

class PositionalEncoding_irregular(nn.Module):

    def __init__(self, d_model, dropout, max_len=5000 ):
        super(PositionalEncoding_irregular, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        self.max_len = max_len
        self.d_model = d_model
        self.linear = nn.Linear(2*d_model,d_model)
         
    def forward(self, x, time):
        ### x:[batch, agent, k, d] time:[batch, agent, k]
        if len(x.shape) > 3:
            last_shape = x.shape
            x = x.reshape(-1,x.shape[-2],x.shape[-1])
            time = time.reshape(-1,time.shape[-1])

        pe = torch.zeros(x.shape[0],time.shape[1], self.d_model).cuda()
        position = time.unsqueeze(2) # 增加维度
        div_term = torch.exp(torch.arange(0., self.d_model, 2) *
                             -(math.log(10000.0) / self.d_model)).cuda()#相对位置公式
         
        pe[:,:, 0::2] = torch.sin(position * div_term)   #取奇数列
        pe[:,:, 1::2] = torch.cos(position * div_term)   #取偶数列

        # x = x + Variable(pe[:, :x.size(1)], requires_grad=False) # embedding 与positional相加
        
        x = self.linear(torch.cat((x,Variable(pe[:, :x.size(1)], requires_grad=False)),dim=-1))
        x = x.reshape(last_shape)
        return self.dropout(x)

class Encoder(nn.Module):
    "Core encoder is a stack of N layers"

    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.N = N

        self.norm = LayerNorm(layer.size) #归一化

    def forward(self, x, mask):
        "Pass the input (and mask) through each layer in turn."
        for i in range(self.N):
            # print('x',x.shape)
            x = self.layers[i](x, mask) 

        return self.norm(x)

class EncoderLayer(nn.Module):
    "Encoder is made up of self-attn and feed forward (defined below)"
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn #sublayer 1 
        self.feed_forward = feed_forward #sublayer 2 
        self.sublayer = clones(SublayerConnection(size, dropout), 2) #拷贝两个SublayerConnection，一个为了attention，一个是独自作为简单的神经网络
        self.size = size

    def forward(self, x, mask):
        "Follow Figure 1 (left) for connections."
        # print('x',x.shape,'attn',self.self_attn(x,x,x,mask).shape)
        
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) #attention层 对应层1
        return self.sublayer[1](x, self.feed_forward) # 对应 层2
        return x


class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
 
    def forward(self, x):
        "Norm"
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2

class SublayerConnection(nn.Module):
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)
 
    def forward(self, x, sublayer):
        # print('sub',x.shape,self.dropout(sublayer(self.norm(x))).shape)
        return x + self.dropout(sublayer(self.norm(x))) #残差连接

class PositionwiseFeedForward(nn.Module):
    "Implements FFN equation."
    def __init__(self, d_model,d_out, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_out)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

class DecoderLayer(nn.Module):
    "Decoder is made of self-attn, src-attn, and feed forward (defined below)"
    def __init__(self, size, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.src_attn = src_attn #解码的attention
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)

    def forward(self, x, src_key,memory, tgt_mask):
        m = memory
        # x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask)) #self-attention
        x = self.sublayer[0](x, lambda x: self.src_attn(x, src_key, m, tgt_mask)) #解码
        return self.sublayer[1](x, self.feed_forward)


class Decoder(nn.Module):
    "Generic N layer decoder with masking."
    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        # print('layersize',layer.size)
        self.norm = LayerNorm(layer.size)
        
    def forward(self, x, src_key ,memory, tgt_mask):
        for layer in self.layers:
            x = layer(x, src_key, memory, tgt_mask) #添加编码的后的结果
        return self.norm(x)

def attention(query, key, value, mask=None, dropout=None):
    "Compute 'Scaled Dot Product Attention'"
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1)) \
             / math.sqrt(d_k)
    
    # print('score',scores.shape,'mask',mask.shape)
    if mask is not None:
        mask = mask.transpose(1,2).cuda()
        scores = scores.masked_fill(mask == 0, -1e9) #mask必须是一个ByteTensor 而且shape必须和 a一样 并且元素只能是 0或者1 ，是将 mask中为1的 元素所在的索引，在a中相同的的索引处替换为 value  ,mask value必须同为tensor
    
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, q_dim, k_dim,v_dim, d_model, dropout=0.1):
        "Take in model size and number of heads."
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0
        # We assume d_v always equals d_k
        self.d_k = d_model // h  # d_v=d_k=d_model/h 
        self.h = h # heads 的数目文中为8
        self.q_dim = q_dim
        self.k_dim = k_dim
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.linear_q = nn.Linear(q_dim,d_model)
        self.linear_k = nn.Linear(k_dim,d_model)
        self.linear_v = nn.Linear(v_dim,d_model)
        self.linear_out = nn.Linear(d_model,q_dim)
        self.attn = None  
        self.dropout = nn.Dropout(p=dropout)


    def forward(self, query, key, value, mask=None):
        "Implements Figure 2"
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)

        
        dim_to_keep = query.size(1)

        # 1) Do all the linear projections in batch from d_model => h x d_k 


        query = self.linear_q(query).view(nbatches,dim_to_keep,-1,self.h,self.d_k).transpose(2,3)
        key = self.linear_k(key).view(nbatches,dim_to_keep,-1,self.h,self.d_k).transpose(2,3)
        value = self.linear_v(value).view(nbatches,dim_to_keep,-1,self.h,self.d_k).transpose(2,3)

        # 2) Apply attention on all the projected vectors in batch. 
        x, self.attn = attention(query, key, value, mask=mask,# 进行attention
                                 dropout=self.dropout)

        # 3) "Concat" using a view and apply a final linear. 
        x = x.transpose(2, 3).contiguous() \
            .view(nbatches, dim_to_keep ,-1, self.h * self.d_k) # 还原序列[batch_size,len,d_model]

        return self.linear_out(x)

class Generator(nn.Module):
    def __init__(self, d_model, out_dim):
        
        super(Generator, self).__init__()
        
        self.proj = nn.Linear(d_model, out_dim)

    def forward(self, x):
        
        return self.proj(x)


class Motion_prediction(nn.Module):
    def __init__(self, encoder, decoder, embedding_src, embedding_tgt, position_src, position_tgt, generator, embed_dim):
        super(Motion_prediction, self).__init__()
        
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = embedding_src
        self.tgt_embed = embedding_tgt
        self.src_pe = position_src
        self.tgt_pe = position_tgt
        self.generator = generator

        

    def forward(self, input, query, time, future_time, mask=None):
        ################################ input: [batch, agent, past_frames, dim]  mask: [batch, agent, past_frames, past_frames]
        input = self.src_embed(input)
        input = self.src_pe(input, time)

        query = self.tgt_embed(query)
        query = self.tgt_pe(query, future_time)

        input_features = self.encoder(input, mask)
        
        prediction_features = self.decoder(query, input_features, input_features,mask)
        predictions = self.generator(prediction_features)
        return predictions



class Motion_interaction(nn.Module):
    def __init__(self, encoder, embedding, generator, neighbor_threshold):
        super(Motion_interaction, self).__init__()
        self.encoder = encoder
        self.src_embed = embedding
        self.generator = generator
        self.neighbor_threshold = neighbor_threshold


    def forward(self, input):
        ####input:[batch,agent,2]
        mask = torch.cdist(input,input)

        mask = torch.where(mask<self.neighbor_threshold,1,0)
        input = input.unsqueeze(1)
        mask = mask.unsqueeze(1)

        input = self.src_embed(input)

        input_features = self.encoder(input, mask)
        predictions = self.generator(input_features)

        return predictions


def make_model(input_dim, output_dim, num_layers=2, d_model=64, d_ff=128, num_heads=2, dropout=0.1,neighbor_shreshold=10):
    c = copy.deepcopy
    attn = MultiHeadedAttention(num_heads, d_model,d_model,d_model,d_model)
    attn_decoder = MultiHeadedAttention(num_heads,d_model,d_model,d_model,d_model)
    ff = PositionwiseFeedForward(d_model, d_model, d_ff,dropout)
    position = PositionalEncoding_irregular(d_model, dropout)

    model_prediction = Motion_prediction(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), num_layers),
        Decoder(DecoderLayer(d_model, c(attn), c(ff), dropout), num_layers),
        Embeddings(d_model, input_dim),
        Embeddings(d_model, output_dim), 
        c(position),
        c(position),
        Generator(d_model, output_dim),
        d_model)

    model_interaction = Motion_interaction(
        Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), num_layers),
        Embeddings(d_model, input_dim), 
        Generator(d_model, output_dim),
        neighbor_shreshold)
    
    for p in model_prediction.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

    for p in model_interaction.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

    return model_prediction, model_interaction




# m1, m2 = make_model(2,2)

# input = torch.ones((7,4,3,2))
# prediction_query = torch.zeros((7,4,1,2))
# time = torch.ones((7,4,3))
# future_time = torch.ones((7,4,1))
# x = m1(input,prediction_query,time,future_time)
# print(x.shape)
# x = x.squeeze(-2)
# print(x.shape)
# x = m2(x)
# print(x.shape)
opencood/tools/traj_prediction.py
MediaBrain-SJTU/CoBEVFlow
"""
Pytorch implementation of the logarithm of the modified Bessel function 
of the 1st kind I(nu, z).
Based mainly on scipy and extends the definition domain compared to 
scipy methods (which can easily provide infinite values). The extension
is done using numerical approximations involving the ratio of Bessel 
functions:   
-	https://arxiv.org/pdf/1606.02008.pdf (Theorems 5,6)
-	https://arxiv.org/pdf/1902.02603.pdf (Appendix)
Jean-Remy Conti
2021
"""

import torch
from scipy import special, pi


def logbessel_I_scipy(nu, z, check = True):
	'''
	Pytorch version of scipy computation of modified Bessel functions 
	of the 1st kind I(nu,z).
	Parameters
	----------
	nu: positive int, float
		Order of modified Bessel function of 1st kind.
	z: int/float or tensor, shape (N,) 
		Argument of Bessel function.
	check: bool
		If True, check if argument of log is non zero.
	
	Return
	------
	result: tensor, shape (N,)
	'''
	if not isinstance(z, torch.Tensor):
		z = torch.tensor(z)
	z = z.reshape(-1)

	result = special.ive(nu, z)
	if check:
		assert len(result[ result == 0]) == 0, ('Bessel functions take ' +
												'value 0 for z = {}'.format(
												z[ result == 0]))
	result = torch.log(result) + z
	return result



def logbessel_I_asymptotic(nu, z):
	'''
	Asymptotic branches of the modified Bessel function of the 1st kind.
	https://arxiv.org/pdf/1902.02603.pdf
	Parameters
	----------
	nu: positive int, float
		Order of modified Bessel function of 1st kind.
	z: tensor, shape (N,) 
		Argument of Bessel function.
	
	Return
	------
	result: tensor, shape (N,)
	'''
	z = z.double()
	eta = (nu + 0.5)/(2* (nu+1) )
	result = torch.zeros(z.shape[0]).double()

	result[ z <= nu ] = (
						nu*torch.log(z[ z <= nu ]) + eta*z[ z <= nu ] 
						- (nu+eta)*torch.log(torch.tensor(2.))
						- torch.log( torch.tensor(special.gamma(nu+1)) )
						)
	
	result[ z > nu ] = (
						z[ z > nu ] - 0.5*torch.log(z[ z > nu ]) 
						- 0.5*torch.log(torch.tensor(2*pi)) 
						)
	return result

def B(alpha, nu, z):
	'''
	https://arxiv.org/pdf/1606.02008.pdf
	'''
	nu = nu.reshape(1,-1)
	z = z.reshape(-1,1)
	lamda = nu + float(alpha-1)/2. 
	delta = nu-0.5 + lamda / (2*torch.sqrt(lamda**2 + z**2))
	return z / (delta + torch.sqrt(delta**2 + z**2) )
def B_tilde(alpha, nu, z):
	'''
	https://arxiv.org/pdf/1606.02008.pdf
	'''
	nu = nu.reshape(1,-1)
	z = z.reshape(-1,1)
	sigma = nu + float(alpha+1)/2.
	delta_p = nu + 0.5 + sigma/(2*torch.sqrt(sigma**2 + z**2))
	delta_m = nu - 0.5 - sigma/(2*torch.sqrt(sigma**2 + z**2))
	return z/( delta_m + torch.sqrt(delta_p**2 + z**2))
def lb_Ak(nu, z):
	'''
	Lower-bound for the ratio of modified Bessel functions of 1st kind. 
	https://arxiv.org/pdf/1606.02008.pdf (Theorems 5 and 6).
	'''
	assert torch.all(nu >= 0)
	nu = nu.reshape(1,-1)
	z = z.reshape(-1,1)
	return B_tilde(0, nu, z)
def ub_Ak(nu, z):
	'''
	Upper-bound for the ratio of modified Bessel functions of 1st kind.
	https://arxiv.org/pdf/1606.02008.pdf (Theorems 5 and 6).
	Return
	------
	ub: tensor, shape (z.shape[0], nu.shape[0])
		Upper-bound for Ak(nu, z).
	'''
	assert torch.all(nu >= 0)
	nu = nu.reshape(1,-1)
	z = z.reshape(-1,1)

	ub = torch.zeros(z.shape[0], nu.shape[1])
	ub[:, nu.reshape(-1) >= 0.5] = torch.min(B(0, nu[ nu >= 0.5 ], z), 
											B_tilde(2, nu[ nu >= 0.5 ], z))
	ub[:, nu.reshape(-1) < 0.5] = B_tilde(2, nu[ nu < 0.5 ], z) 
	return ub

def Ak_approx(nu, z):
	'''
	Approximation of ratio of modified Bessel functions of 1st kind.
	https://arxiv.org/pdf/1902.02603.pdf
	Parameters
	----------
	nu: tensor, shape (N0,)
		Order of modified Bessel functions of 1st kind.
	z: tensor, shape (N1,) 
		Argument of Bessel function. Positive values only.
	
	Return
	------
	tensor, shape (N1, N0)
	'''
	return 0.5*(lb_Ak(nu, z) + ub_Ak(nu, z)) 


def logbessel_I_approx(nu, z):
	'''
	Approximation of the logarithm of the modified Bessel function of 
	1st kind I(nu, z) using their ratio.
	https://arxiv.org/pdf/1902.02603.pdf
	Parameters
	----------
	nu: positive int, float
		Order of modified Bessel function of 1st kind.
	z: tensor, shape (N,) 
		Argument of Bessel function. Positive values only.
	Return
	------
	approx: tensor, shape (N,)
	'''
	assert nu >= 0
	approx = logbessel_I_scipy(nu-int(nu), z, check= True) 
	nu_v = nu - torch.arange(0, int(nu))
	A = Ak_approx(nu_v, z) 
	approx += torch.log(A).sum(axis= 1)
	return approx


def logbessel_I(nu, z, fast = False, check = True):
	'''
	Pytorch implementation of the logarithm of the modified Bessel 
	function of the 1st kind I(nu, z).
	Based mainly on scipy and extends the definition domain compared to 
	scipy methods (which can easily provide infinite values). The 
	extension is done using numerical approximations involving the ratio
	of Bessel functions:   
	-	https://arxiv.org/pdf/1606.02008.pdf (Theorems 5,6)
	-	https://arxiv.org/pdf/1902.02603.pdf (Appendix)
	Parameters
	----------
	nu: positive int, float
		Order of modified Bessel function of 1st kind.
	z: int/float or tensor, shape (N,) 
		Argument of Bessel function.
	fast: bool
		If True, use asymptotic behavior as approximation when main 
		scipy method is not tractable. If False, use tight bounds for 
		the ratio of Bessel functions:
		https://arxiv.org/pdf/1902.02603.pdf
	check: bool
		If True, check if argument of log is non zero and not NaN.
	
	Return
	------
	result: tensor, shape (N,)
	'''
	if not isinstance(z, torch.Tensor):
		z = torch.tensor(z)
	z = z.reshape(-1)

	result = special.ive(nu, z)
	# Indices for which scipy.special.ive is wrong
	bad_idx = torch.arange(result.shape[0])[ result == 0]
	result = torch.log(result) + z
	if fast:
		result[ bad_idx ] = logbessel_I_asymptotic(nu, z[ bad_idx ])
	else: 	 
		result[ bad_idx ] = logbessel_I_approx(nu, z[ bad_idx ])

	if check:
		# If problem with assertion, use a better defined init in 
		# logbessel_approx
		assert len(result[ torch.isnan(result) ]) == 0, ('Bessel functions take ' +
												'NaN value for z = {}'.format(
		 										z[ torch.isnan(result) ]))
		assert len(result[ torch.isinf(result) ]) == 0, ('Bessel functions take ' +
		 										'inf value for z = {}'.format(
		 										z[ torch.isinf(result) ]))
	return result



if __name__ == "__main__":

	import time
	import matplotlib.pyplot as plt

	# ------- Single test ------- #
	print(logbessel_I(nu= 10000, z= 10000))


	# ------- Vectorized tests ------- #
	nu = 1000
	z = torch.arange(1, 200001)

	# Scipy
	print('------\nScipy adaptation:')
	start = time.time()
	scipy_adapt = logbessel_I_scipy(nu, z, check= False)
	end = time.time()
	print('Computation time: ', "%.4f"%(end-start), 's')

	# Asymptotic
	print('------\nAsymptotic computation:')
	start = time.time()
	asympt = logbessel_I_asymptotic(nu, z)
	end = time.time()
	print('Computation time: ', "%.4f"%(end-start), 's')

	# Approximation using ratios of Bessel functions
	print('------\nApproximation via ratios of Bessel functions:')
	start = time.time()
	approx = logbessel_I_approx(nu, z)
	end = time.time()
	print('Computation time: ', "%.4f"%(end-start), 's')

	# Fast extension of scipy
	print('------\nFast method (proposed):')
	start = time.time()
	logbessels_fast = logbessel_I(nu, z, fast= True, check= False)
	end = time.time()
	print('Computation time: ', "%.4f"%(end-start), 's')

	# Precise extension of scipy
	print('------\nPrecise method (proposed):')
	start = time.time()
	logbessels = logbessel_I(nu, z, fast= False, check= False)
	end = time.time()
	print('Computation time: ', "%.4f"%(end-start), 's')


	# ------- Plots ------- #
	linewidth = 4

	# Plot different methods
	plt.plot(z, scipy_adapt, '-', label='scipy', linewidth = linewidth)
	plt.plot(z, asympt, '--', label='asymptotic', linewidth = linewidth)
	plt.plot(z, logbessels, '--', label='ours', linewidth = linewidth)
	plt.xlabel(r'$z$')	
	plt.title(r'$\log[I(\nu = {}, z)]$'.format(nu))
	plt.legend()

	# Plot relative error of approximation
	plt.figure()
	plt.plot(z, torch.abs((approx - scipy_adapt)/scipy_adapt), '-', label=None, linewidth = linewidth)
	plt.xlabel(r'$z$')
	plt.title(r'Relative error of approximation for $\log[I(\nu = {}, z)]$'.format(nu))

	plt.show()
opencood/utils/bessel_utils.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Hao Xiang <haxiang@g.ucla.edu>,
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
Bounding box related utility functions
"""
import sys

import numpy as np

import torch
import torch.nn.functional as F
import opencood.utils.common_utils as common_utils
from opencood.utils.transformation_utils import x1_to_x2, x_to_world
from pyquaternion import Quaternion
import copy


def corner_to_center_torch(corner3d, order='lwh'):
    corner3d_ = corner3d.cpu().numpy()
    return torch.from_numpy(corner_to_center(corner3d_, order)).to(corner3d.device)

def corner_to_center(corner3d, order='lwh'):
    """
    Convert 8 corners to x, y, z, dx, dy, dz, yaw.
    yaw in radians

    Parameters
    ----------
    corner3d : np.ndarray
        (N, 8, 3)

    order : str, for output.
        'lwh' or 'hwl'

    Returns
    -------
    box3d : np.ndarray
        (N, 7)
    """
    assert corner3d.ndim == 3
    batch_size = corner3d.shape[0]

    xyz = np.mean(corner3d[:, [0, 3, 5, 6], :], axis=1)
    h = abs(np.mean(corner3d[:, 4:, 2] - corner3d[:, :4, 2], axis=1,
                    keepdims=True))
    l = (np.sqrt(np.sum((corner3d[:, 0, [0, 1]] - corner3d[:, 3, [0, 1]]) ** 2,
                        axis=1, keepdims=True)) +
         np.sqrt(np.sum((corner3d[:, 2, [0, 1]] - corner3d[:, 1, [0, 1]]) ** 2,
                        axis=1, keepdims=True)) +
         np.sqrt(np.sum((corner3d[:, 4, [0, 1]] - corner3d[:, 7, [0, 1]]) ** 2,
                        axis=1, keepdims=True)) +
         np.sqrt(np.sum((corner3d[:, 5, [0, 1]] - corner3d[:, 6, [0, 1]]) ** 2,
                        axis=1, keepdims=True))) / 4

    w = (np.sqrt(
        np.sum((corner3d[:, 0, [0, 1]] - corner3d[:, 1, [0, 1]]) ** 2, axis=1,
               keepdims=True)) +
         np.sqrt(np.sum((corner3d[:, 2, [0, 1]] - corner3d[:, 3, [0, 1]]) ** 2,
                        axis=1, keepdims=True)) +
         np.sqrt(np.sum((corner3d[:, 4, [0, 1]] - corner3d[:, 5, [0, 1]]) ** 2,
                        axis=1, keepdims=True)) +
         np.sqrt(np.sum((corner3d[:, 6, [0, 1]] - corner3d[:, 7, [0, 1]]) ** 2,
                        axis=1, keepdims=True))) / 4

    theta = (np.arctan2(corner3d[:, 1, 1] - corner3d[:, 2, 1],
                        corner3d[:, 1, 0] - corner3d[:, 2, 0]) +
             np.arctan2(corner3d[:, 0, 1] - corner3d[:, 3, 1],
                        corner3d[:, 0, 0] - corner3d[:, 3, 0]) +
             np.arctan2(corner3d[:, 5, 1] - corner3d[:, 6, 1],
                        corner3d[:, 5, 0] - corner3d[:, 6, 0]) +
             np.arctan2(corner3d[:, 4, 1] - corner3d[:, 7, 1],
                        corner3d[:, 4, 0] - corner3d[:, 7, 0]))[:,
            np.newaxis] / 4

    if order == 'lwh':
        return np.concatenate([xyz, l, w, h, theta], axis=1).reshape(
            batch_size, 7)
    elif order == 'hwl':
        return np.concatenate([xyz, h, w, l, theta], axis=1).reshape(
            batch_size, 7)
    else:
        sys.exit('Unknown order')


def boxes_to_corners2d(boxes3d, order):
    """
      0 -------- 1
      |          |
      |          |
      |          |
      3 -------- 2
    Parameters
    __________
    boxes3d: np.ndarray or torch.Tensor
        (N, 7) [x, y, z, dx, dy, dz, heading], (x, y, z) is the box center.

    order : str
        'lwh' or 'hwl'

    Returns:
        corners2d: np.ndarray or torch.Tensor
        (N, 4, 3), the 4 corners of the bounding box.

    """
    corners3d = boxes_to_corners_3d(boxes3d, order)
    corners2d = corners3d[:, :4, :]
    return corners2d


def boxes2d_to_corners2d(boxes2d, order="lwh"):
    """
      0 -------- 1
      |          |
      |          |
      |          |
      3 -------- 2
    Parameters
    __________
    boxes2d: np.ndarray or torch.Tensor
        (..., 5) [x, y, dx, dy, heading], (x, y) is the box center.

    order : str
        'lwh' or 'hwl'

    Returns:
        corners2d: np.ndarray or torch.Tensor
        (..., 4, 2), the 4 corners of the bounding box.

    """
    assert order == "lwh", \
        "boxes2d_to_corners_2d only supports lwh order for now."
    boxes2d, is_numpy = common_utils.check_numpy_to_torch(boxes2d)
    template = boxes2d.new_tensor((
        [1, -1], [1, 1], [-1, 1], [-1, -1]
    )) / 2
    input_shape = boxes2d.shape
    boxes2d = boxes2d.view(-1, 5)
    corners2d = boxes2d[:, None, 2:4].repeat(1, 4, 1) * template[None, :, :]
    corners2d = common_utils.rotate_points_along_z_2d(corners2d.view(-1, 2),
                                                      boxes2d[:,
                                                      4].repeat_interleave(
                                                          4)).view(-1, 4,
                                                                   2)
    corners2d += boxes2d[:, None, 0:2]
    corners2d = corners2d.view(*(input_shape[:-1]), 4, 2)
    return corners2d


def boxes_to_corners_3d(boxes3d, order):
    """
        4 -------- 5
       /|         /|
      7 -------- 6 .
      | |        | |
      . 0 -------- 1
      |/         |/
      3 -------- 2
    Parameters
    __________
    boxes3d: np.ndarray or torch.Tensor
        (N, 7) [x, y, z, l, w, h, heading], or [x, y, z, h, w, l, heading]

               (x, y, z) is the box center.

    order : str
        'lwh' or 'hwl'

    Returns:
        corners3d: np.ndarray or torch.Tensor
        (N, 8, 3), the 8 corners of the bounding box.


    opv2v's left hand coord 
    
    ^ z
    |
    |
    | . x
    |/
    +-------> y

    """

    boxes3d, is_numpy = common_utils.check_numpy_to_torch(boxes3d)
    boxes3d_ = boxes3d

    if order == 'hwl':
        boxes3d_ = boxes3d[:, [0, 1, 2, 5, 4, 3, 6]]

    template = boxes3d_.new_tensor((
        [1, -1, -1], [1, 1, -1], [-1, 1, -1], [-1, -1, -1],
        [1, -1, 1], [1, 1, 1], [-1, 1, 1], [-1, -1, 1],
    )) / 2

    corners3d = boxes3d_[:, None, 3:6].repeat(1, 8, 1) * template[None, :, :]
    corners3d = common_utils.rotate_points_along_z(corners3d.view(-1, 8, 3),
                                                   boxes3d_[:, 6]).view(-1, 8,
                                                                        3)
    corners3d += boxes3d_[:, None, 0:3]

    return corners3d.numpy() if is_numpy else corners3d


def box3d_to_2d(box3d):
    """
    Convert 3D bounding box to 2D.

    Parameters
    ----------
    box3d : np.ndarray
        (n, 8, 3)

    Returns
    -------
    box2d : np.ndarray
        (n, 4, 2), project 3d to 2d.
    """
    box2d = box3d[:, :4, :2]
    return box2d


def corner2d_to_standup_box(box2d):
    """
    Find the minmaxx, minmaxy for each 2d box. (N, 4, 2) -> (N, 4)
    x1, y1, x2, y2

    Parameters
    ----------
    box2d : np.ndarray
        (n, 4, 2), four corners of the 2d bounding box.

    Returns
    -------
    standup_box2d : np.ndarray
        (n, 4)
    """
    N = box2d.shape[0]
    standup_boxes2d = np.zeros((N, 4))

    standup_boxes2d[:, 0] = np.min(box2d[:, :, 0], axis=1)
    standup_boxes2d[:, 1] = np.min(box2d[:, :, 1], axis=1)
    standup_boxes2d[:, 2] = np.max(box2d[:, :, 0], axis=1)
    standup_boxes2d[:, 3] = np.max(box2d[:, :, 1], axis=1)

    return standup_boxes2d


def corner_to_standup_box_torch(box_corner):
    """
    Find the minmax x and y for each bounding box.

    Parameters
    ----------
    box_corner : torch.Tensor
        Shape: (N, 8, 3) or (N, 4)

    Returns
    -------
    standup_box2d : torch.Tensor
        (n, 4)
    """
    N = box_corner.shape[0]
    standup_boxes2d = torch.zeros((N, 4))

    standup_boxes2d = standup_boxes2d.to(box_corner.device)

    standup_boxes2d[:, 0] = torch.min(box_corner[:, :, 0], dim=1).values
    standup_boxes2d[:, 1] = torch.min(box_corner[:, :, 1], dim=1).values
    standup_boxes2d[:, 2] = torch.max(box_corner[:, :, 0], dim=1).values
    standup_boxes2d[:, 3] = torch.max(box_corner[:, :, 1], dim=1).values

    return standup_boxes2d


def project_box3d(box3d, transformation_matrix):
    """
    Project the 3d bounding box to another coordinate system based on the
    transfomration matrix.

    Parameters
    ----------
    box3d : torch.Tensor or np.ndarray
        3D bounding box, (N, 8, 3)

    transformation_matrix : torch.Tensor or np.ndarray
        Transformation matrix, (4, 4)

    Returns
    -------
    projected_box3d : torch.Tensor
        The projected bounding box, (N, 8, 3)
    """
    assert transformation_matrix.shape == (4, 4)
    box3d, is_numpy = \
        common_utils.check_numpy_to_torch(box3d)
    transformation_matrix, _ = \
        common_utils.check_numpy_to_torch(transformation_matrix)

    # (N, 3, 8)
    box3d_corner = box3d.transpose(1, 2)
    # (N, 1, 8)
    torch_ones = torch.ones((box3d_corner.shape[0], 1, 8))
    torch_ones = torch_ones.to(box3d_corner.device)
    # (N, 4, 8)
    box3d_corner = torch.cat((box3d_corner, torch_ones),
                             dim=1)
    # (N, 4, 8)
    projected_box3d = torch.matmul(transformation_matrix,
                                   box3d_corner)
    # (N, 8, 3)
    projected_box3d = projected_box3d[:, :3, :].transpose(1, 2)

    return projected_box3d if not is_numpy else projected_box3d.numpy()


def project_points_by_matrix_torch(points, transformation_matrix):
    """
    Project the points to another coordinate system based on the
    transfomration matrix. 
    
    IT NOT USED. LATTER ONE WITH THE SAME NAME WILL BE USED.

    Parameters
    ----------
    points : torch.Tensor
        3D points, (N, 3)

    transformation_matrix : torch.Tensor
        Transformation matrix, (4, 4)

    Returns
    -------
    projected_points : torch.Tensor
        The projected points, (N, 3)
    """
    # convert to homogeneous  coordinates via padding 1 at the last dimension.
    # (N, 4)
    points_homogeneous = F.pad(points, (0, 1), mode="constant", value=1)
    # (N, 4)
    projected_points = torch.einsum("ik, jk->ij", points_homogeneous,
                                    transformation_matrix)
    return projected_points[:, :3]


def get_mask_for_boxes_within_range_torch(boxes, gt_range):
    """
    Generate mask to remove the bounding boxes
    outside the range.

    Parameters
    ----------
    boxes : torch.Tensor
        Groundtruth bbx, shape: N,8,3 or N,4,2

    gt_range: list
        [xmin, ymin, zmin, xmax, ymax, zmax]
    Returns
    -------
    mask: torch.Tensor
        The mask for bounding box -- True means the
        bbx is within the range and False means the
        bbx is outside the range.
    """

    # mask out the gt bounding box out fixed range (-140, -40, -3, 140, 40 1)
    device = boxes.device
    boundary_lower_range = \
        torch.Tensor(gt_range[:2]).reshape(1, 1, -1).to(device)
    boundary_higher_range = \
        torch.Tensor(gt_range[3:5]).reshape(1, 1, -1).to(device)

    mask = torch.all(
        torch.all(boxes[:, :, :2] >= boundary_lower_range,
                  dim=-1) & \
        torch.all(boxes[:, :, :2] <= boundary_higher_range,
                  dim=-1), dim=-1)

    return mask


def mask_boxes_outside_range_numpy(boxes, limit_range, order,
                                   min_num_corners=8, return_mask=False):
    """
    Parameters
    ----------
    boxes: np.ndarray
        (N, 7) [x, y, z, dx, dy, dz, heading], (x, y, z) is the box center

    limit_range: list
        [minx, miny, minz, maxx, maxy, maxz]

    min_num_corners: int
        The required minimum number of corners to be considered as in range.

    order : str
        'lwh' or 'hwl'

    return_mask : bool
        Whether return the mask.

    Returns
    -------
    boxes: np.ndarray
        The filtered boxes.
    """
    assert boxes.shape[1] == 8 or boxes.shape[1] == 7

    new_boxes = boxes.copy()
    if boxes.shape[1] == 7:
        new_boxes = boxes_to_corners_3d(new_boxes, order)

    mask = ((new_boxes >= limit_range[0:3]) &
            (new_boxes <= limit_range[3:6])).all(axis=2)
    mask = mask.sum(axis=1) >= min_num_corners  # (N)

    if return_mask:
        return boxes[mask], mask
    return boxes[mask]


def create_bbx(extent):
    """
    Create bounding box with 8 corners under obstacle vehicle reference.

    Parameters
    ----------
    extent : list
        half length, width and height

    Returns
    -------
    bbx : np.array
        The bounding box with 8 corners, shape: (8, 3)
    """

    bbx = np.array([[extent[0], -extent[1], -extent[2]],
                    [extent[0], extent[1], -extent[2]],
                    [-extent[0], extent[1], -extent[2]],
                    [-extent[0], -extent[1], -extent[2]],
                    [extent[0], -extent[1], extent[2]],
                    [extent[0], extent[1], extent[2]],
                    [-extent[0], extent[1], extent[2]],
                    [-extent[0], -extent[1], extent[2]]])

    return bbx


def project_world_objects(object_dict,
                          output_dict,
                          lidar_pose,
                          lidar_range,
                          order):
    """
    Project the objects under world coordinates into another coordinate
    based on the provided extrinsic.

    Parameters
    ----------
    object_dict : dict
        The dictionary contains all objects surrounding a certain cav.

    output_dict : dict
        key: object id, value: object bbx (xyzlwhyaw).

    lidar_pose : list
        (6, ), lidar pose under world coordinate, [x, y, z, roll, yaw, pitch].

    lidar_range : list
         [minx, miny, minz, maxx, maxy, maxz]

    order : str
        'lwh' or 'hwl'
    """
    for object_id, object_content in object_dict.items():
        location = object_content['location']
        rotation = object_content['angle']
        center = object_content['center']
        extent = object_content['extent']

        object_pose = [location[0] + center[0],
                       location[1] + center[1],
                       location[2] + center[2],
                       rotation[0], rotation[1], rotation[2]]


        object2lidar = x1_to_x2(object_pose, lidar_pose)

        # shape (3, 8)
        bbx = create_bbx(extent).T
        # bounding box under ego coordinate shape (4, 8)
        bbx = np.r_[bbx, [np.ones(bbx.shape[1])]]

        # project the 8 corners to world coordinate
        bbx_lidar = np.dot(object2lidar, bbx).T
        bbx_lidar = np.expand_dims(bbx_lidar[:, :3], 0)
        bbx_lidar = corner_to_center(bbx_lidar, order=order)
        bbx_lidar = mask_boxes_outside_range_numpy(bbx_lidar,
                                                   lidar_range,
                                                   order)

        if bbx_lidar.shape[0] > 0:
            output_dict.update({object_id: bbx_lidar})


def project_world_objects_v2x(object_dict,
                          output_dict,
                          reference_lidar_pose,
                          lidar_range,
                          order,
                          lidar_np):
    """
    Project the objects under world coordinates into another coordinate
    based on the provided extrinsic.

    Parameters
    ----------
    object_dict : 
        gt boxes: numpy.ndarray (N,10)
            [x,y,z,dx,dy,dz,w,a,b,c], dxdydz=lwh
        object_ids: numpy.ndarray (N,)

    output_dict : dict
        key: object id, value: object bbx (xyzlwhyaw).

    reference_lidar_pose : list
        (6, ), lidar pose under world coordinate, [x, y, z, roll, yaw, pitch].

    lidar_range : list
         [minx, miny, minz, maxx, maxy, maxz]

    order : str
        'lwh' or 'hwl'

    lidar_np: np.ndarray
        point cloud in ego coord. Used to determine if any lidar point hits the box


    output_dict: [x,y,z, lwh or hwl, yaw] 
    """
    from icecream import ic
    gt_boxes = object_dict['gt_boxes']
    object_ids = object_dict['object_ids']
    for i, object_content in enumerate(gt_boxes):
        x,y,z,dx,dy,dz,w,a,b,c = object_content

        q = Quaternion([w,a,b,c])
        T_world_object = q.transformation_matrix
        T_world_object[:3,3] = object_content[:3]

        T_world_lidar = x_to_world(reference_lidar_pose)

        object2lidar = np.linalg.solve(T_world_lidar, T_world_object) # T_lidar_object


        # shape (3, 8). 
        # or we can use the create_bbx funcion.
        x_corners = dx / 2 * np.array([ 1,  1, -1, -1,  1,  1, -1, -1]) # (8,)
        y_corners = dy / 2 * np.array([-1,  1,  1, -1, -1,  1,  1, -1])
        z_corners = dz / 2 * np.array([-1, -1, -1, -1,  1,  1,  1,  1])

        bbx = np.vstack((x_corners, y_corners, z_corners)) # (3, 8)

        # bounding box under ego coordinate shape (4, 8)
        bbx = np.r_[bbx, [np.ones(bbx.shape[1])]]

        # project the 8 corners to world coordinate
        bbx_lidar = np.dot(object2lidar, bbx).T # (8, 4)
        bbx_lidar = np.expand_dims(bbx_lidar[:, :3], 0) # (1, 8, 3)
        bbx_lidar = corner_to_center(bbx_lidar, order=order)
        
        lidar_range_z_larger = copy.deepcopy(lidar_range)
        lidar_range_z_larger[2] -= 1
        lidar_range_z_larger[5] += 1
        
        bbx_lidar = mask_boxes_outside_range_numpy(bbx_lidar,
                                                   lidar_range_z_larger,
                                                   order)

        
        if bbx_lidar.shape[0] > 0:
            output_dict.update({object_ids[i]: bbx_lidar})


def get_points_in_rotated_box(p, box_corner):
    """
    Get points within a rotated bounding box (2D version).

    Parameters
    ----------
    p : numpy.array
        Points to be tested with shape (N, 2).
    box_corner : numpy.array
        Corners of bounding box with shape (4, 2).

    Returns
    -------
    p_in_box : numpy.array
        Points within the box.

    """
    edge1 = box_corner[1, :] - box_corner[0, :]
    edge2 = box_corner[3, :] - box_corner[0, :]
    p_rel = p - box_corner[0, :].reshape(1, -1)

    l1 = get_projection_length_for_vector_projection(p_rel, edge1)
    l2 = get_projection_length_for_vector_projection(p_rel, edge2)
    # A point is within the box, if and only after projecting the
    # point onto the two edges s.t. p_rel = [edge1, edge2] @ [l1, l2]^T,
    # we have 0<=l1<=1 and 0<=l2<=1.
    mask = np.logical_and(l1 >= 0, l1 <= 1)
    mask = np.logical_and(mask, l2 >= 0)
    mask = np.logical_and(mask, l2 <= 1)
    p_in_box = p[mask, :]
    return p_in_box


def get_points_in_rotated_box_3d(p, box_corner):
    """
    Get points within a rotated bounding box (3D version).

    Parameters
    ----------
    p : numpy.array
        Points to be tested with shape (N, 3).
    box_corner : numpy.array
        Corners of bounding box with shape (8, 3).

    Returns
    -------
    p_in_box : numpy.array
        Points within the box.

    """
    edge1 = box_corner[1, :] - box_corner[0, :]
    edge2 = box_corner[3, :] - box_corner[0, :]
    edge3 = box_corner[4, :] - box_corner[0, :]

    p_rel = p - box_corner[0, :].reshape(1, -1)

    l1 = get_projection_length_for_vector_projection(p_rel, edge1)
    l2 = get_projection_length_for_vector_projection(p_rel, edge2)
    l3 = get_projection_length_for_vector_projection(p_rel, edge3)
    # A point is within the box, if and only after projecting the
    # point onto the two edges s.t. p_rel = [edge1, edge2] @ [l1, l2]^T,
    # we have 0<=l1<=1 and 0<=l2<=1.
    mask1 = np.logical_and(l1 >= 0, l1 <= 1)
    mask2 = np.logical_and(l2 >= 0, l2 <= 1)
    mask3 = np.logical_and(l3 >= 0, l3 <= 1)

    mask = np.logical_and(mask1, mask2)
    mask = np.logical_and(mask, mask3)
    p_in_box = p[mask, :]

    return p_in_box


def get_projection_length_for_vector_projection(a, b):
    """
    Get projection length for the Vector projection of a onto b s.t.
    a_projected = length * b. (2D version) See
    https://en.wikipedia.org/wiki/Vector_projection#Vector_projection_2
    for more details.

    Parameters
    ----------
    a : numpy.array
        The vectors to be projected with shape (N, 2).

    b : numpy.array
        The vector that is projected onto with shape (2).

    Returns
    -------
    length : numpy.array
        The length of projected a with respect to b.
    """
    assert np.sum(b ** 2, axis=-1) > 1e-6
    length = a.dot(b) / np.sum(b ** 2, axis=-1)
    return length


def nms_rotated(boxes, scores, threshold):
    """Performs rorated non-maximum suppression and returns indices of kept
    boxes.

    Parameters
    ----------
    boxes : torch.tensor
        The location preds with shape (N, 4, 2).

    scores : torch.tensor
        The predicted confidence score with shape (N,)

    threshold: float
        IoU threshold to use for filtering.

    Returns
    -------
        An array of index
    """
    if boxes.shape[0] == 0:
        return np.array([], dtype=np.int32)
    boxes = boxes.cpu().detach().numpy()
    scores = scores.cpu().detach().numpy()

    polygons = common_utils.convert_format(boxes)

    top = 1000
    # Get indicies of boxes sorted by scores (highest first)
    ixs = scores.argsort()[::-1][:top]

    pick = []
    while len(ixs) > 0:
        # Pick top box and add its index to the list
        i = ixs[0]
        pick.append(i)
        # Compute IoU of the picked box with the rest
        iou = common_utils.compute_iou(polygons[i], polygons[ixs[1:]])
        # Identify boxes with IoU over the threshold. This
        # returns indices into ixs[1:], so add 1 to get
        # indices into ixs.
        remove_ixs = np.where(iou > threshold)[0] + 1
        # Remove indices of the picked and overlapped boxes.
        ixs = np.delete(ixs, remove_ixs)
        ixs = np.delete(ixs, 0)

    return np.array(pick, dtype=np.int32)


def nms_pytorch(boxes: torch.tensor, thresh_iou: float):
    """
    Apply non-maximum suppression to avoid detecting too many
    overlapping bounding boxes for a given object.

    Parameters
    ----------
    boxes : torch.tensor
        The location preds along with the class predscores,
         Shape: [num_boxes,5].
    thresh_iou : float
        (float) The overlap thresh for suppressing unnecessary boxes.
    Returns
    -------
        A list of index
    """

    # we extract coordinates for every
    # prediction box present in P
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    # we extract the confidence scores as well
    scores = boxes[:, 4]

    # calculate area of every block in P
    areas = (x2 - x1) * (y2 - y1)

    # sort the prediction boxes in P
    # according to their confidence scores
    order = scores.argsort()

    # initialise an empty list for
    # filtered prediction boxes
    keep = []

    while len(order) > 0:

        # extract the index of the
        # prediction with highest score
        # we call this prediction S
        idx = order[-1]

        # push S in filtered predictions list
        keep.append(idx.numpy().item()
                    if not idx.is_cuda else idx.cpu().detach().numpy().item())

        # remove S from P
        order = order[:-1]

        # sanity check
        if len(order) == 0:
            break

        # select coordinates of BBoxes according to
        # the indices in order
        xx1 = torch.index_select(x1, dim=0, index=order)
        xx2 = torch.index_select(x2, dim=0, index=order)
        yy1 = torch.index_select(y1, dim=0, index=order)
        yy2 = torch.index_select(y2, dim=0, index=order)

        # find the coordinates of the intersection boxes
        xx1 = torch.max(xx1, x1[idx])
        yy1 = torch.max(yy1, y1[idx])
        xx2 = torch.min(xx2, x2[idx])
        yy2 = torch.min(yy2, y2[idx])

        # find height and width of the intersection boxes
        w = xx2 - xx1
        h = yy2 - yy1

        # take max with 0.0 to avoid negative w and h
        # due to non-overlapping boxes
        w = torch.clamp(w, min=0.0)
        h = torch.clamp(h, min=0.0)

        # find the intersection area
        inter = w * h

        # find the areas of BBoxes according the indices in order
        rem_areas = torch.index_select(areas, dim=0, index=order)

        # find the union of every prediction T in P
        # with the prediction S
        # Note that areas[idx] represents area of S
        union = (rem_areas - inter) + areas[idx]

        # find the IoU of every prediction in P with S
        IoU = inter / union

        # keep the boxes with IoU less than thresh_iou
        mask = IoU < thresh_iou
        order = order[mask]

    return keep


def remove_large_pred_bbx(bbx_3d):
    """
    Remove large bounding box.

    Parameters
    ----------
    bbx_3d : torch.Tensor
        Predcited 3d bounding box, shape:(N,8,3)

    Returns
    -------
    index : torch.Tensor
        The keep index.
    """
    bbx_x_max = torch.max(bbx_3d[:, :, 0], dim=1)[0]
    bbx_x_min = torch.min(bbx_3d[:, :, 0], dim=1)[0]
    x_len = bbx_x_max - bbx_x_min

    bbx_y_max = torch.max(bbx_3d[:, :, 1], dim=1)[0]
    bbx_y_min = torch.min(bbx_3d[:, :, 1], dim=1)[0]
    y_len = bbx_y_max - bbx_y_min

    bbx_z_max = torch.max(bbx_3d[:, :, 1], dim=1)[0]
    bbx_z_min = torch.min(bbx_3d[:, :, 1], dim=1)[0]
    z_len = bbx_z_max - bbx_z_min

    index = torch.logical_and(x_len <= 6, y_len <= 6)
    index = torch.logical_and(index, z_len)

    return index


def remove_bbx_abnormal_z(bbx_3d):
    """
    Remove bounding box that has negative z axis.

    Parameters
    ----------
    bbx_3d : torch.Tensor
        Predcited 3d bounding box, shape:(N,8,3)

    Returns
    -------
    index : torch.Tensor
        The keep index.
    """
    bbx_z_min = torch.min(bbx_3d[:, :, 2], dim=1)[0]
    bbx_z_max = torch.max(bbx_3d[:, :, 2], dim=1)[0]
    index = torch.logical_and(bbx_z_min >= -3, bbx_z_max <= 1)

    return index


def project_points_by_matrix_torch(points, transformation_matrix):
    """
    Project the points to another coordinate system based on the
    transformation matrix.

    Parameters
    ----------
    points : torch.Tensor
        3D points, (N, 3)
    transformation_matrix : torch.Tensor
        Transformation matrix, (4, 4)
    Returns
    -------
    projected_points : torch.Tensor
        The projected points, (N, 3)
    """
    points, is_numpy = \
        common_utils.check_numpy_to_torch(points)
    transformation_matrix, _ = \
        common_utils.check_numpy_to_torch(transformation_matrix)

    # convert to homogeneous coordinates via padding 1 at the last dimension.
    # (N, 4)
    points_homogeneous = F.pad(points, (0, 1), mode="constant", value=1)
    # (N, 4)
    projected_points = torch.einsum("ik, jk->ij", points_homogeneous,
                                    transformation_matrix)

    return projected_points[:, :3] if not is_numpy \
        else projected_points[:, :3].numpy()


def box_encode(
        boxes,
        anchors,
        encode_angle_to_vector=False,
        encode_angle_with_residual=False,
        smooth_dim=False,
        norm_velo=False
):
    """box encode for VoxelNet
        Args:
            boxes ([N, 7] Tensor): normal boxes: x, y, z, w, l, h, r.
            anchors ([N, 7] Tensor): anchors.
    """

    box_ndim = anchors.shape[-1]

    if box_ndim == 7:
        xa, ya, za, wa, la, ha, ra = torch.split(anchors, 1, dim=-1)
        xg, yg, zg, wg, lg, hg, rg = torch.split(boxes, 1, dim=-1)
    else:
        xa, ya, za, wa, la, ha, vxa, vya, ra = torch.split(anchors, 1, dim=-1)
        xg, yg, zg, wg, lg, hg, vxg, vyg, rg = torch.split(boxes, 1, dim=-1)

    diagonal = torch.sqrt(la ** 2 + wa ** 2)
    xt = (xg - xa) / diagonal
    yt = (yg - ya) / diagonal
    zt = (zg - za) / ha

    if smooth_dim:
        lt = lg / la - 1
        wt = wg / wa - 1
        ht = hg / ha - 1
    else:
        lt = torch.log(lg / la)
        wt = torch.log(wg / wa)
        ht = torch.log(hg / ha)

    ret = [xt, yt, zt, wt, lt, ht]

    if box_ndim > 7:
        if norm_velo:
            vxt = (vxg - vxa) / diagonal
            vyt = (vyg - vya) / diagonal
        else:
            vxt = vxg - vxa
            vyt = vyg - vya
        ret.extend([vxt, vyt])

    if encode_angle_to_vector:
        rgx = torch.cos(rg)
        rgy = torch.sin(rg)
        if encode_angle_with_residual:
            rax = torch.cos(ra)
            ray = torch.sin(ra)
            rtx = rgx - rax
            rty = rgy - ray
            ret.extend([rtx, rty])
        else:
            ret.extend([rgx, rgy])
    else:
        rt = rg - ra
        ret.append(rt)

    return torch.cat(ret, dim=-1)


def box_decode(
        box_encodings,
        anchors,
        encode_angle_to_vector=False,
        encode_angle_with_residual=False,
        bin_loss=False,
        smooth_dim=False,
        norm_velo=False,
):
    """box decode for VoxelNet in lidar
    Args:
        boxes ([N, 7] Tensor): normal boxes: x, y, z, w, l, h, r
        anchors ([N, 7] Tensor): anchors
    """
    box_ndim = anchors.shape[-1]

    if box_ndim == 9:  # False
        xa, ya, za, wa, la, ha, vxa, vya, ra = torch.split(anchors, 1, dim=-1)
        if encode_angle_to_vector:
            xt, yt, zt, wt, lt, ht, vxt, vyt, rtx, rty = torch.split(box_encodings, 1, dim=-1)
        else:
            xt, yt, zt, wt, lt, ht, vxt, vyt, rt = torch.split(box_encodings, 1, dim=-1)

    elif box_ndim == 7:
        xa, ya, za, wa, la, ha, ra = torch.split(anchors, 1, dim=-1)
        if encode_angle_to_vector:  # False
            xt, yt, zt, wt, lt, ht, rtx, rty = torch.split(box_encodings, 1, dim=-1)
        else:
            xt, yt, zt, wt, lt, ht, rt = torch.split(box_encodings, 1, dim=-1)

    diagonal = torch.sqrt(la ** 2 + wa ** 2)
    xg = xt * diagonal + xa
    yg = yt * diagonal + ya
    zg = zt * ha + za

    ret = [xg, yg, zg]

    if smooth_dim:  # False
        lg = (lt + 1) * la
        wg = (wt + 1) * wa
        hg = (ht + 1) * ha
    else:
        lg = torch.exp(lt) * la
        wg = torch.exp(wt) * wa
        hg = torch.exp(ht) * ha
    ret.extend([wg, lg, hg])

    if encode_angle_to_vector:  # False
        if encode_angle_with_residual:
            rax = torch.cos(ra)
            ray = torch.sin(ra)
            rgx = rtx + rax
            rgy = rty + ray
            rg = torch.atan2(rgy, rgx)
        else:
            rg = torch.atan2(rty, rtx)
    else:
        rg = rt + ra

    if box_ndim > 7:  # False
        if norm_velo:
            vxg = vxt * diagonal + vxa
            vyg = vyt * diagonal + vya
        else:
            vxg = vxt + vxa
            vyg = vyt + vya
        ret.extend([vxg, vyg])

    ret.append(rg)

    return torch.cat(ret, dim=-1)


def project_world_objects_dairv2x(object_list,
                          output_dict,
                          lidar_pose,
                          lidar_range,
                          order):
    """
    Project the objects under world coordinates into another coordinate
    based on the provided extrinsic.

    Parameters
    ----------
    object_list : list
        The list contains all objects surrounding a certain cav.

    output_dict : dict
        key: object id, value: object bbx (xyzlwhyaw).

    lidar_pose : list
        (6, ), lidar pose under world coordinate, [x, y, z, roll, yaw, pitch].

    lidar_range : list
         [minx, miny, minz, maxx, maxy, maxz]

    order : str
        'lwh' or 'hwl'
    """
    i = 0

    for object_content in object_list: 
        object_id = i
        i = i + 1
        lidar_to_world = x_to_world(lidar_pose) # T_world_lidar
        world_to_lidar = np.linalg.inv(lidar_to_world)

        corners_world = np.array(object_content['world_8_points']) # [8,3]
        corners_world_homo = np.pad(corners_world, ((0,0), (0,1)), constant_values=1) # [8, 4]
        corners_lidar = (world_to_lidar @ corners_world_homo.T).T 

        lidar_range_z_larger = copy.deepcopy(lidar_range)
        lidar_range_z_larger[2] -= 1
        lidar_range_z_larger[5] += 1

        bbx_lidar = corners_lidar
        bbx_lidar = np.expand_dims(bbx_lidar[:, :3], 0) # [1, 8, 3]
        bbx_lidar = corner_to_center(bbx_lidar, order=order)
        bbx_lidar = mask_boxes_outside_range_numpy(bbx_lidar, lidar_range_z_larger, order)
        if bbx_lidar.shape[0] > 0:
            output_dict.update({object_id: bbx_lidar})


def load_single_objects_dairv2x(object_list,
                          output_dict,
                          lidar_range,
                          order):
    """
    Load the objects under ego coordinates.
    Parameters
    ----------
    object_list : list
        The list contains all objects surrounding a certain cav.
    output_dict : dict
        key: object id, value: object bbx (xyzlwhyaw).
    lidar_range : list
         [minx, miny, minz, maxx, maxy, maxz]
    order : str
        'lwh' or 'hwl'
    """
    i = 0
    for object_content in object_list:        
        object_id = i
        if 'rotation' not in object_content:
            print(object_content)
        x = object_content['3d_location']['x']
        y = object_content['3d_location']['y']
        z = object_content['3d_location']['z']
        l = object_content['3d_dimensions']['l']
        h = object_content['3d_dimensions']['h']
        w = object_content['3d_dimensions']['w']
        rotation = object_content['rotation']

        if l==0 or h ==0 or w==0:
            continue
        i = i + 1

        lidar_range_z_larger = copy.deepcopy(lidar_range)
        lidar_range_z_larger[2] -= 1
        lidar_range_z_larger[5] += 1
        
        bbx_lidar = [x,y,z,h,w,l,rotation] if order=="hwl" else [x,y,z,l,w,h,rotation] # suppose order is in ['hwl', 'lwh']
        bbx_lidar = np.array(bbx_lidar).reshape(1,-1) # [1,7]

        bbx_lidar = mask_boxes_outside_range_numpy(bbx_lidar, lidar_range_z_larger, order)
        if bbx_lidar.shape[0] > 0:
            if object_content['type'] == "Car" or \
               object_content['type'] == "Van" or \
               object_content['type'] == "Truck" or \
               object_content['type'] == "Bus":
                    output_dict.update({object_id: bbx_lidar})
opencood/utils/box_utils.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Hao Xiang <haxiang@g.ucla.edu>,
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
Common utilities
"""

import numpy as np
import torch
from shapely.geometry import Polygon
import json

def read_json(file_path):
    with open(file_path, 'r') as f:
        data = json.load(f)

    return data

def limit_period(val, offset=0.5, period=2*np.pi):
    """
    continous part: 
    [0 - period * offset, period - period * offset)
    """
    # 首先，numpy格式数据转换为torch格式
    val, is_numpy = check_numpy_to_torch(val)
    # 将方位角限制在[-pi, pi]
    ans = val - torch.floor(val / period + offset) * period
    return ans.numpy() if is_numpy else ans


def check_numpy_to_torch(x):
    if isinstance(x, np.ndarray):
        return torch.from_numpy(x).float(), True
    return x, False

def check_torch_to_numpy(x):
    if isinstance(x, torch.tensor):
        return x.cpu().numpy(), True
    return x, False


def check_contain_nan(x):
    if isinstance(x, dict):
        return any(check_contain_nan(v) for k, v in x.items())
    if isinstance(x, list):
        return any(check_contain_nan(itm) for itm in x)
    if isinstance(x, int) or isinstance(x, float):
        return False
    if isinstance(x, np.ndarray):
        return np.any(np.isnan(x))
    return torch.any(x.isnan()).detach().cpu().item()


def rotate_points_along_z(points, angle):
    """
    Args:
        points: (B, N, 3 + C)
        angle: (B), radians, angle along z-axis, angle increases x ==> y
    Returns:

    """
    points, is_numpy = check_numpy_to_torch(points)
    angle, _ = check_numpy_to_torch(angle)

    cosa = torch.cos(angle)
    sina = torch.sin(angle)
    zeros = angle.new_zeros(points.shape[0])
    ones = angle.new_ones(points.shape[0])
    rot_matrix = torch.stack((
        cosa, sina, zeros,
        -sina, cosa, zeros,
        zeros, zeros, ones
    ), dim=1).view(-1, 3, 3).float()
    points_rot = torch.matmul(points[:, :, 0:3].float(), rot_matrix)
    points_rot = torch.cat((points_rot, points[:, :, 3:]), dim=-1)
    return points_rot.numpy() if is_numpy else points_rot


def rotate_points_along_z_2d(points, angle):
    """
    Rorate the points along z-axis.
    Parameters
    ----------
    points : torch.Tensor / np.ndarray
        (N, 2).
    angle : torch.Tensor / np.ndarray
        (N,)

    Returns
    -------
    points_rot : torch.Tensor / np.ndarray
        Rorated points with shape (N, 2)

    """
    points, is_numpy = check_numpy_to_torch(points)
    angle, _ = check_numpy_to_torch(angle)
    cosa = torch.cos(angle)
    sina = torch.sin(angle)
    # (N, 2, 2)
    rot_matrix = torch.stack((cosa, sina, -sina, cosa), dim=1).view(-1, 2,
                                                                    2).float()
    points_rot = torch.einsum("ik, ikj->ij", points.float(), rot_matrix)
    return points_rot.numpy() if is_numpy else points_rot


def remove_ego_from_objects(objects, ego_id):
    """
    Avoid adding ego vehicle to the object dictionary.

    Parameters
    ----------
    objects : dict
        The dictionary contained all objects.

    ego_id : int
        Ego id.
    """
    if ego_id in objects:
        del objects[ego_id]


def retrieve_ego_id(base_data_dict):
    """
    Retrieve the ego vehicle id from sample(origin format).

    Parameters
    ----------
    base_data_dict : dict
        Data sample in origin format.

    Returns
    -------
    ego_id : str
        The id of ego vehicle.
    """
    ego_id = None

    for cav_id, cav_content in base_data_dict.items():
        if cav_content['ego']:
            ego_id = cav_id
            break
    return ego_id


def compute_iou(box, boxes):
    """
    Compute iou between box and boxes list
    Parameters
    ----------
    box : shapely.geometry.Polygon
        Bounding box Polygon.

    boxes : list
        List of shapely.geometry.Polygon.

    Returns
    -------
    iou : np.ndarray
        Array of iou between box and boxes.

    """
    # Calculate intersection areas
    if np.any(np.array([box.union(b).area for b in boxes])==0):
        print('debug')
    iou = [box.intersection(b).area / box.union(b).area for b in boxes]

    return np.array(iou, dtype=np.float32)


def convert_format(boxes_array):
    """
    Convert boxes array to shapely.geometry.Polygon format.
    Parameters
    ----------
    boxes_array : np.ndarray
        (N, 4, 2) or (N, 8, 3).

    Returns
    -------
        list of converted shapely.geometry.Polygon object.

    """
    polygons = [Polygon([(box[i, 0], box[i, 1]) for i in range(4)]) for box in
                boxes_array]
    return np.array(polygons)


def torch_tensor_to_numpy(torch_tensor):
    """
    Convert a torch tensor to numpy.

    Parameters
    ----------
    torch_tensor : torch.Tensor

    Returns
    -------
    A numpy array.
    """
    return torch_tensor.numpy() if not torch_tensor.is_cuda else \
        torch_tensor.cpu().detach().numpy()


def get_voxel_centers(voxel_coords,
                      downsample_times,
                      voxel_size,
                      point_cloud_range):
    """
    Args:
        voxel_coords: (N, 3)
        downsample_times:
        voxel_size:
        point_cloud_range:

    Returns:

    """
    assert voxel_coords.shape[1] == 3
    voxel_centers = voxel_coords[:, [2, 1, 0]].float()  # (xyz)
    voxel_size = torch.tensor(voxel_size, device=voxel_centers.device).float() * downsample_times
    pc_range = torch.tensor(point_cloud_range[0:3], device=voxel_centers.device).float()
    voxel_centers = (voxel_centers + 0.5) * voxel_size + pc_range
    return voxel_centers
opencood/utils/common_utils.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import os

import numpy as np
import torch

from opencood.utils import common_utils
from opencood.hypes_yaml import yaml_utils


def voc_ap(rec, prec):
    """
    VOC 2010 Average Precision.
    """
    rec.insert(0, 0.0)
    rec.append(1.0)
    mrec = rec[:]

    prec.insert(0, 0.0)
    prec.append(0.0)
    mpre = prec[:]

    for i in range(len(mpre) - 2, -1, -1):
        mpre[i] = max(mpre[i], mpre[i + 1])

    i_list = []
    for i in range(1, len(mrec)):
        if mrec[i] != mrec[i - 1]:
            i_list.append(i)

    ap = 0.0
    for i in i_list:
        ap += ((mrec[i] - mrec[i - 1]) * mpre[i])
    return ap, mrec, mpre


def caluclate_tp_fp(det_boxes, det_score, gt_boxes, result_stat, iou_thresh):
    """
    Calculate the true positive and false positive numbers of the current
    frames.

    Parameters
    ----------
    det_boxes : torch.Tensor
        The detection bounding box, shape (N, 8, 3) or (N, 4, 2).
    det_score :torch.Tensor
        The confidence score for each preditect bounding box.
    gt_boxes : torch.Tensor
        The groundtruth bounding box.
    result_stat: dict
        A dictionary contains fp, tp and gt number.
    iou_thresh : float
        The iou thresh.
    """
    # fp, tp and gt in the current frame
    fp = []
    tp = []
    gt = gt_boxes.shape[0]
    if det_boxes is not None:
        # convert bounding boxes to numpy array
        det_boxes = common_utils.torch_tensor_to_numpy(det_boxes)
        det_score = common_utils.torch_tensor_to_numpy(det_score)
        gt_boxes = common_utils.torch_tensor_to_numpy(gt_boxes)

        # sort the prediction bounding box by score
        score_order_descend = np.argsort(-det_score)
        det_score = det_score[score_order_descend]
        det_polygon_list = list(common_utils.convert_format(det_boxes))
        gt_polygon_list = list(common_utils.convert_format(gt_boxes))

        # match prediction and gt bounding box
        for i in range(score_order_descend.shape[0]):
            det_polygon = det_polygon_list[score_order_descend[i]]
            ious = common_utils.compute_iou(det_polygon, gt_polygon_list)

            if len(gt_polygon_list) == 0 or np.max(ious) < iou_thresh:
                fp.append(1)
                tp.append(0)
                continue

            fp.append(0)
            tp.append(1)

            gt_index = np.argmax(ious)
            gt_polygon_list.pop(gt_index)
        result_stat[iou_thresh]['score'] += det_score.tolist()
    result_stat[iou_thresh]['fp'] += fp
    result_stat[iou_thresh]['tp'] += tp
    result_stat[iou_thresh]['gt'] += gt


def calculate_ap(result_stat, iou):
    """
    Calculate the average precision and recall, and save them into a txt.

    Parameters
    ----------
    result_stat : dict
        A dictionary contains fp, tp and gt number.
    iou : float
    """
    iou_5 = result_stat[iou]

    fp = np.array(iou_5['fp'])
    tp = np.array(iou_5['tp'])
    score = np.array(iou_5['score'])
    assert len(fp) == len(tp) and len(tp) == len(score)

    sorted_index = np.argsort(-score)
    fp = fp[sorted_index].tolist()
    tp = tp[sorted_index].tolist()

    gt_total = iou_5['gt']

    cumsum = 0
    for idx, val in enumerate(fp):
        fp[idx] += cumsum
        cumsum += val

    cumsum = 0
    for idx, val in enumerate(tp):
        tp[idx] += cumsum
        cumsum += val

    rec = tp[:]
    for idx, val in enumerate(tp):
        rec[idx] = float(tp[idx]) / gt_total

    prec = tp[:]
    for idx, val in enumerate(tp):
        prec[idx] = float(tp[idx]) / (fp[idx] + tp[idx])

    ap, mrec, mprec = voc_ap(rec[:], prec[:])

    return ap, mrec, mprec


def eval_final_results(result_stat, save_path, noise_level=None, avg_time_delay=-1.0, avg_sample_interval=0.0, avg_time_var=0.0, note='', dataset='o'):
    dump_dict = {}

    ap_30, mrec_30, mpre_30 = calculate_ap(result_stat, 0.30)
    ap_50, mrec_50, mpre_50 = calculate_ap(result_stat, 0.50)
    ap_70, mrec_70, mpre_70 = calculate_ap(result_stat, 0.70)

    if dataset == 'o':
        dump_dict.update({'ap_30': ap_30,
                            'ap_50': ap_50,
                            'ap_70': ap_70,
                            'avg_time_delay': avg_time_delay,
                            'avg_sample_interval': avg_sample_interval, 
                            'avg_time_var': avg_time_var})
    elif dataset == 'd':
        dump_dict.update({'ap_30': ap_30,
                            'ap_50': ap_50,
                            'ap_70': ap_70,
                            'avg_cp_rate': avg_time_delay,
                            'avg_sample_interval': avg_sample_interval})

    if noise_level is None:
        yaml_utils.save_yaml(dump_dict, os.path.join(save_path, 'eval.yaml'))
    else:
        yaml_utils.save_yaml(dump_dict, os.path.join(save_path, f'eval_{noise_level}_{note}.yaml'))

    print('The Average Precision at IOU 0.3 is %.5f, '
          'The Average Precision at IOU 0.5 is %.5f, '
          'The Average Precision at IOU 0.7 is %.5f' % (ap_30, ap_50, ap_70))

    print('=== Avg Time Delay %.2f ==='
          '=== Avg Sample Interval %.2f ==='
          '=== Avg Time Var %.2f ==='% (avg_time_delay, avg_sample_interval, avg_time_var) )

    return ap_30, ap_50, ap_70
opencood/utils/eval_utils.py
MediaBrain-SJTU/CoBEVFlow

# nuScenes dev-kit.
# Code written by Oscar Beijbom, 2018.

import os
import cv2
import copy
import numpy as np
from numba import njit
from typing import Tuple, List
from matplotlib.axes import Axes
from pyquaternion import Quaternion
from collections import OrderedDict
from matplotlib import pyplot as plt
import cv2

import torch

import torch.nn.functional as F 
from opencood.utils import box_utils

def generate_flow_map_and_mask(flow, bbox_list, scale=1.25, shape_list=None, align_corners=False, file_suffix=""):
    """
    Parameters
    -----------
    feature: [C, H, W] at voxel scale
    bbox_list: [num_cav, 4, 3] at cav coodinate system and lidar scale
    flow:[num_cav, 2] at cav coodinate system and lidar scale
        bbox_list & flow : x and y are exactly image coordinate
        ------------> x
        |
        |
        |
        y
    scale: float, scale meters to voxel, feature_length / lidar_range_length = 1.25 or 2.5

    Returns
    -------
    updated_feature: feature after being warped by flow, [C, H, W]
    """
    # flow = torch.tensor([70, 0]).unsqueeze(0).to(feature)

    # only use x and y
    bbox_list = bbox_list[:, :, :2]

    # scale meters to voxel, feature_length / lidar_range_length = 1.25
    flow = flow * scale
    bbox_list = bbox_list * scale

    C, H, W = shape_list
    num_cav = bbox_list.shape[0]
    basic_mat = torch.tensor([[1,0,0],[0,1,0]]).unsqueeze(0).to(torch.float32)
    basic_warp_mat = F.affine_grid(basic_mat, [1, C, H, W], align_corners=align_corners).to(shape_list.device)
    reserved_area = torch.zeros((C, H, W)).to(shape_list.device)  # C, H, W
    if flow.shape[0] == 0 or flow.max() == 0: 
        reserved_area = torch.ones((C, H, W)).to(shape_list.device)  # C, H, W
        return basic_warp_mat,  reserved_area.unsqueeze(0)  # 返回不变的矩阵

    '''
    create affine matrix:
    ------------
    1  0  -2*t_y/W
    0  1  -2*t_x/H
    0  0    1 
    ------------
    '''
    flow_clone = flow.clone()

    affine_matrices = torch.eye(3).unsqueeze(0).repeat(flow.shape[0], 1, 1)
    flow_clone = -2 * flow_clone / torch.tensor([W, H]).to(torch.float32).to(shape_list.device)
    # flow_clone = flow_clone[:, [1, 0]]
    affine_matrices[:, :2, 2] = flow_clone 
    
    cav_t_mat = affine_matrices[:, :2, :]   # n, 2, 3
    # print("cav_t_mat", cav_t_mat)

    cav_warp_mat = F.affine_grid(cav_t_mat,
                        [num_cav, C, H, W],
                        align_corners=align_corners).to(shape_list.device) # .to() 统一数据格式 float32
    
    flowed_bbx_list = bbox_list + flow.unsqueeze(1).repeat(1,4,1)  # n, 4, 2

    x_min = torch.min(flowed_bbx_list[:,:,0],dim=1)[0] - 1
    x_max = torch.max(flowed_bbx_list[:,:,0],dim=1)[0] + 1
    y_min = torch.min(flowed_bbx_list[:,:,1],dim=1)[0] - 1
    y_max = torch.max(flowed_bbx_list[:,:,1],dim=1)[0] + 1
    x_min_fid = (x_min + int(W/2)).to(torch.int)
    x_max_fid = (x_max + int(W/2)).to(torch.int)
    y_min_fid = (y_min + int(H/2)).to(torch.int)
    y_max_fid = (y_max + int(H/2)).to(torch.int)

    for cav in range(num_cav):
        basic_warp_mat[0,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]] = cav_warp_mat[cav,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]]

    # generate mask
    x_min_ori = torch.min(bbox_list[:,:,0],dim=1)[0]
    x_max_ori = torch.max(bbox_list[:,:,0],dim=1)[0]
    y_min_ori = torch.min(bbox_list[:,:,1],dim=1)[0]
    y_max_ori = torch.max(bbox_list[:,:,1],dim=1)[0]
    x_min_fid_ori = (x_min_ori + int(W/2)).to(torch.int)
    x_max_fid_ori = (x_max_ori + int(W/2)).to(torch.int)
    y_min_fid_ori = (y_min_ori + int(H/2)).to(torch.int)
    y_max_fid_ori = (y_max_ori + int(H/2)).to(torch.int)
    # set original location as 0
    for cav in range(num_cav):
        reserved_area[:,y_min_fid_ori[cav]:y_max_fid_ori[cav],x_min_fid_ori[cav]:x_max_fid_ori[cav]] = 0
    # set warped location as 1
    for cav in range(num_cav):
        reserved_area[:,y_min_fid[cav]:y_max_fid[cav],x_min_fid[cav]:x_max_fid[cav]] = 1

    return basic_warp_mat, reserved_area.unsqueeze(0)

def find_id_difference(obj_A, obj_B, id_A, id_B):
    """
    返回两个tensor对应id行的差值
    :param obj_A: 输入的tensor A 大小为 (N_0, 7)
    :param obj_B: 输入的tensor B 大小为 (N_1, 7)
    :param id_A: tensor A中每一行的id 长度为N_0 list类型
    :param id_B: tensor B中每一行的id 长度为N_1 list类型
    :return diff: 两个tensor对应id行的差值 大小为 (N, 3)
    :return A_common: obj_A中对应common_id行的tensor 大小为 (N, 7)
    """
    # 将id_A, id_B转化为字典
    dict_A = dict(zip(id_A, range(len(id_A))))
    dict_B = dict(zip(id_B, range(len(id_B))))

    # 找到A和B的共有id
    common_id = set(dict_A.keys()) & set(dict_B.keys())

    # 生成对应id的tensor索引
    index_A = torch.tensor([dict_A[id] for id in common_id], dtype=torch.long)
    index_B = torch.tensor([dict_B[id] for id in common_id], dtype=torch.long)

    # 取出对应id的tensor行
    A_common = obj_A[index_A, :]
    B_common = obj_B[index_B, :]

    # 计算差值
    try:
        diff = B_common[:, :2] - A_common[:, :2]
    except IndexError:
        B_common = np.expand_dims(B_common, axis=0)
        A_common = np.expand_dims(A_common, axis=0)
        diff = B_common[:, :2] - A_common[:, :2]

    return diff, A_common

def generate_flow_map_szwei(object_stack, object_id_stack, cav_lidar_range, voxel_size, shape_list = torch.tensor([64, 200, 704]), past_k=1):
    '''
    SizheWei@2023/04/20 
    Generate GT flow map for a single agent.
    Args:
        timestamp: 0, 1, 2, the timestamps between previous and current
        object_stack: Dict {timestamp: objects} objects/boxes at each timestamp
        object_id_stack: Dict {timestamp: object_ids} objects/boxes ids at each timestamp
        past_k: the time stamps between previous and current timestamp
    Output:
        bev_flow: np.array [2,H,W]
    '''
    C, H, W = shape_list

    past0_box = object_stack[0]
    curr_box = object_stack[1]
    past0_id = object_id_stack[0]
    curr_id = object_id_stack[1]
    flow, selected_box_3dcenter_past0 = find_id_difference(past0_box, curr_box, past0_id, curr_id)

    selected_box_3dcorner_past0 = box_utils.boxes_to_corners2d(selected_box_3dcenter_past0, order='hwl')

    flow = torch.from_numpy(flow).to(torch.float32)
    selected_box_3dcorner_past0 = torch.from_numpy(selected_box_3dcorner_past0).to(torch.float32)
    flow_map, mask = generate_flow_map_and_mask(flow, selected_box_3dcorner_past0, scale=2.5, shape_list=shape_list)

    # #######################
    # flow_map = flow_map.squeeze(0).permute(2, 0, 1) #(2, H, W)
    # denorm_flow = flow_map.clone()
    # denorm_flow[0, :, :] = (denorm_flow[0, :, :] + 1.0) * (W / 2.0)
    # denorm_flow[1, :, :] = (denorm_flow[1, :, :] + 1.0) * (H / 2.0)
    # x_coord = torch.arange(W).float()   # [0, ..., W]
    # y_coord = torch.arange(H).float()   # [0, ..., H]
    # y, x = torch.meshgrid(y_coord, x_coord)  # [H, W], [H, W]
    # grid = torch.cat([x.unsqueeze(0), y.unsqueeze(0)], dim=0) # [2, H, W]
    # final_flow = denorm_flow - grid
    # final_flow = final_flow.numpy()
    # #######################

    return flow_map, mask

def generate_flow_map(object_stack, object_id_stack, cav_lidar_range, voxel_size, past_k=1):
    '''
    Args:
        timestamp: 0, 1, 2, the timestamps between previous and current
        object_stack: Dict {timestamp: objects} objects/boxes at each timestamp
        object_id_stack: Dict {timestamp: object_ids} objects/boxes ids at each timestamp
        past_k: the time stamps between previous and current timestamp
    Output:
        bev_flow: np.array [H,W,2]
    '''
    curr_box = box_list_to_dict(object_stack[0], object_id_stack[0], cav_lidar_range, voxel_size)
    prev_box = box_list_to_dict(object_stack[past_k], object_id_stack[past_k], cav_lidar_range, voxel_size)
    bev_map, bev_coords = init_bev_map(cav_lidar_range, voxel_size)

    obj_ids = [i for i in prev_box if i in curr_box]
    for obj_id in obj_ids:
        # bev_map = disp_in_box(bev_map, bev_coords, prev_box[obj_id], curr_box[obj_id])
        bev_map = disp_in_box(bev_map, bev_coords, prev_box[obj_id], curr_box[obj_id])
    # bev_map = np.flip(bev_map, axis=0)
    bev_map = np.transpose(bev_map, (2,0,1))
    return bev_map
    # bev_map = np.array(bev_map[:,:,0] * 255, dtype=np.uint8)[...,None].repeat(3, axis=-1)
    # for obj_id in obj_ids:
    #     bev_map = cv2.circle(bev_map, (int(prev_box[obj_id][0]), int(prev_box[obj_id][1])), radius=0, color=(0, 0, 255), thickness=-1)
    # return bev_map

def generate_history_2_flow_map(object_stack, object_id_stack, cav_lidar_range, voxel_size, past_k=1, history_num=2):
    '''
    Args:
        timestamp: 0, 1, 2, the timestamps between previous and current
        object_stack: Dict {timestamp: objects} objects/boxes at each timestamp
        object_id_stack: Dict {timestamp: object_ids} objects/boxes ids at each timestamp
        past_k: the time stamps between previous and current timestamp
    Output:
        bev_flow: np.array [H,W,2]
    '''
    history_box = box_list_to_dict(object_stack[history_num], object_id_stack[history_num], cav_lidar_range, voxel_size)
    curr_box = box_list_to_dict(object_stack[0], object_id_stack[0], cav_lidar_range, voxel_size)
    prev_box = box_list_to_dict(object_stack[1], object_id_stack[1], cav_lidar_range, voxel_size)
    bev_map, bev_coords = init_bev_map(cav_lidar_range, voxel_size)

    obj_ids = [i for i in history_box if i in curr_box]
    for obj_id in obj_ids:
        bev_map = disp_in_box(bev_map, bev_coords, history_box[obj_id], curr_box[obj_id])
    # bev_map = np.flip(bev_map, axis=0)
    bev_map_0 = np.transpose(bev_map, (2,0,1))
    bev_map_0 = np.expand_dims(bev_map_0, axis=0)

    obj_ids = [i for i in history_box if i in prev_box]
    for obj_id in obj_ids:
        bev_map = disp_in_box(bev_map, bev_coords, history_box[obj_id], prev_box[obj_id])
    # bev_map = np.flip(bev_map, axis=0)
    bev_map_1 = np.transpose(bev_map, (2, 0, 1))
    bev_map_1 = np.expand_dims(bev_map_1, axis=0)

    bev_map = np.concatenate((bev_map_1, bev_map_0), axis=0)
    return bev_map

def box_list_to_dict(object_stack, object_id_stack, cav_lidar_range, voxel_size):
    object_dict = OrderedDict()
    for i, obj_id in enumerate(object_id_stack):
        object_dict[obj_id] = global_coord_to_bev(object_stack[i], cav_lidar_range, voxel_size)
    return object_dict

def global_coord_to_bev(box, cav_lidar_range, voxel_size, out_size_factor=1):
    coor_x = (box[0] - cav_lidar_range[0]) / voxel_size[0] / out_size_factor
    coor_y = (box[1] - cav_lidar_range[1]) / voxel_size[1] / out_size_factor
    coor_z = (box[2] - cav_lidar_range[2]) / voxel_size[2] / out_size_factor
    h = box[3] / voxel_size[2] / out_size_factor
    w = box[4] / voxel_size[1] / out_size_factor
    l = box[5] / voxel_size[0] / out_size_factor
    rot = box[6]
    return [coor_x, coor_y, coor_z, h, w, l, rot]

def init_bev_map(cav_lidar_range, voxel_size, out_size_factor=1):
    # initialize bev flow map
    grid_size = (np.array(cav_lidar_range[3:6]) -
                     np.array(cav_lidar_range[0:3])) / np.array(voxel_size)
    grid_size = np.round(grid_size).astype(np.int64)
    W, H = grid_size[:2] // out_size_factor

    bev_map = np.zeros((H, W, 2), dtype=np.float32)
    y_coord = np.repeat(np.arange(H).reshape(H,1), W, axis=1).reshape(H,W,1)
    x_coord = np.repeat(np.arange(W).reshape(1,W), H, axis=0).reshape(H,W,1)
    z_coord = np.zeros_like(x_coord)
    xyz = np.concatenate([x_coord,y_coord,z_coord], axis=-1)
    bev_coords = xyz.reshape([-1,3])
    return bev_map, bev_coords

def disp_in_box(bev_map, points, prev_box, curr_box):
    '''
    bev_map: [H,W,2] np.array
    points: [N, 3] np.array
    box: [x, y, z, h, w, l, heading]
    '''
    prev_bounding_box = Box(center=prev_box[:3],
                                size=[prev_box[4],prev_box[5],prev_box[3]],
                                orientation=Quaternion(axis=[0, 0, 1], angle=prev_box[-1])) # wlh
    curr_bounding_box = Box(center=curr_box[:3],
                                size=[curr_box[4],curr_box[5],curr_box[3]],
                                orientation=Quaternion(axis=[0, 0, 1], angle=curr_box[-1])) # wlh
    idx = point_in_hull_fast(points, prev_bounding_box)
    # idx = point_in_hull_fast(points, curr_bounding_box)
    pc_in_box = points[idx]
    disp = calc_displace_vector(pc_in_box, prev_bounding_box, curr_bounding_box)
    # disp = calc_displace_vector(pc_in_box, curr_bounding_box, prev_bounding_box)
    bev_map[pc_in_box[:,1], pc_in_box[:,0]] = disp[:,:2]
    return bev_map

# The specifications for BEV maps
voxel_size = (0.25, 0.25, 0.4)
area_extents = np.array([[-32., 32.], [-32., 32.], [-3., 2.]])
past_frame_skip = 3  # when generating the BEV maps, how many history frames need to be skipped
future_frame_skip = 0  # when generating the BEV maps, how many future frames need to be skipped
num_past_frames_for_bev_seq = 5  # the number of past frames for BEV map sequence

# ---------------------- Convert the raw data into (dense) BEV maps ----------------------
def convert_to_dense_bev(data_dict):
    num_sweeps = data_dict['num_sweeps']
    times = data_dict['times']
    trans_matrices = data_dict['trans_matrices']

    num_past_sweeps = len(np.where(times >= 0)[0])
    num_future_sweeps = len(np.where(times < 0)[0])
    assert num_past_sweeps + num_future_sweeps == num_sweeps, "The number of sweeps is incorrect!"

    # Load point cloud
    pc_list = []

    for i in range(num_sweeps):
        pc = data_dict['pc_' + str(i)]
        pc_list.append(pc.T)

    # Reorder the pc, and skip sample frames if wanted
    # Currently the past frames in pc_list are stored in the following order [current, current + 1, current + 2, ...]
    # Therefore, we would like to reorder the frames
    tmp_pc_list_1 = pc_list[0:num_past_sweeps:(past_frame_skip + 1)]
    tmp_pc_list_1 = tmp_pc_list_1[::-1]
    tmp_pc_list_2 = pc_list[(num_past_sweeps + future_frame_skip)::(future_frame_skip + 1)]
    pc_list = tmp_pc_list_1 + tmp_pc_list_2  # now the order is: [past frames -> current frame -> future frames]

    num_past_pcs = len(tmp_pc_list_1)
    num_future_pcs = len(tmp_pc_list_2)

    # Discretize the input point clouds, and compute the ground-truth displacement vectors
    # The following two variables contain the information for the
    # compact representation of binary voxels, as described in the paper
    voxel_indices_list = list()
    padded_voxel_points_list = list()

    past_pcs_idx = list(range(num_past_pcs))
    past_pcs_idx = past_pcs_idx[-num_past_frames_for_bev_seq:]  # we typically use 5 past frames (including the current one)
    for i in past_pcs_idx:
        res, voxel_indices = voxelize_occupy(pc_list[i], voxel_size=voxel_size, extents=area_extents, return_indices=True)
        voxel_indices_list.append(voxel_indices)
        padded_voxel_points_list.append(res)

    # Compile the batch of voxels, so that they can be fed into the network.
    # Note that, the padded_voxel_points in this script will only be used for sanity check.
    padded_voxel_points = np.stack(padded_voxel_points_list, axis=0).astype(np.bool)

    # Finally, generate the ground-truth displacement field
    # - all_disp_field_gt: the ground-truth displacement vectors for each grid cell
    # - all_valid_pixel_maps: the masking map for valid pixels, used for loss computation
    # - non_empty_map: the mask which represents the non-empty grid cells, used for loss computation
    # - pixel_cat_map: the map specifying the category for each non-empty grid cell
    # - pixel_indices: the indices of non-empty grid cells, used to generate sparse BEV maps
    # - pixel_instance_map: the map specifying the instance id for each grid cell, used for loss computation
    all_disp_field_gt, all_valid_pixel_maps, non_empty_map, pixel_cat_map, pixel_indices, pixel_instance_map \
        = gen_2d_grid_gt(data_dict, grid_size=voxel_size[0:2], extents=area_extents,
                         frame_skip=future_frame_skip, return_instance_map=True)

    return voxel_indices_list, padded_voxel_points, pixel_indices, pixel_instance_map, all_disp_field_gt,\
        all_valid_pixel_maps, non_empty_map, pixel_cat_map, num_past_frames_for_bev_seq, num_future_pcs, trans_matrices


# ---------------------- Convert the dense BEV data into sparse format ----------------------
# This will significantly reduce the space used for data storage
def convert_to_sparse_bev(dense_bev_data):
    save_voxel_indices_list, save_voxel_points, save_pixel_indices, save_pixel_instance_maps, \
        save_disp_field_gt, save_valid_pixel_maps, save_non_empty_maps, save_pixel_cat_maps, \
        save_num_past_pcs, save_num_future_pcs, save_trans_matrices = dense_bev_data

    save_valid_pixel_maps = save_valid_pixel_maps.astype(np.bool)
    save_voxel_dims = save_voxel_points.shape[1:]
    num_categories = save_pixel_cat_maps.shape[-1]

    sparse_disp_field_gt = save_disp_field_gt[:, save_pixel_indices[:, 0], save_pixel_indices[:, 1], :]
    sparse_valid_pixel_maps = save_valid_pixel_maps[:, save_pixel_indices[:, 0], save_pixel_indices[:, 1]]
    sparse_pixel_cat_maps = save_pixel_cat_maps[save_pixel_indices[:, 0], save_pixel_indices[:, 1]]
    sparse_pixel_instance_maps = save_pixel_instance_maps[save_pixel_indices[:, 0], save_pixel_indices[:, 1]]

    save_data_dict = dict()
    for i in range(len(save_voxel_indices_list)):
        save_data_dict['voxel_indices_' + str(i)] = save_voxel_indices_list[i].astype(np.int32)

    save_data_dict['disp_field'] = sparse_disp_field_gt
    save_data_dict['valid_pixel_map'] = sparse_valid_pixel_maps
    save_data_dict['pixel_cat_map'] = sparse_pixel_cat_maps
    save_data_dict['num_past_pcs'] = save_num_past_pcs
    save_data_dict['num_future_pcs'] = save_num_future_pcs
    save_data_dict['trans_matrices'] = save_trans_matrices
    save_data_dict['3d_dimension'] = save_voxel_dims
    save_data_dict['pixel_indices'] = save_pixel_indices
    save_data_dict['pixel_instance_ids'] = sparse_pixel_instance_maps

    # -------------------------------- Sanity Check --------------------------------
    dims = save_non_empty_maps.shape

    test_disp_field_gt = np.zeros((save_num_future_pcs, dims[0], dims[1], 2), dtype=np.float32)
    test_disp_field_gt[:, save_pixel_indices[:, 0], save_pixel_indices[:, 1], :] = sparse_disp_field_gt[:]
    assert np.all(test_disp_field_gt == save_disp_field_gt), "Error: Mismatch"

    test_valid_pixel_maps = np.zeros((save_num_future_pcs, dims[0], dims[1]), dtype=np.bool)
    test_valid_pixel_maps[:, save_pixel_indices[:, 0], save_pixel_indices[:, 1]] = sparse_valid_pixel_maps[:]
    assert np.all(test_valid_pixel_maps == save_valid_pixel_maps), "Error: Mismatch"

    test_pixel_cat_maps = np.zeros((dims[0], dims[1], num_categories), dtype=np.float32)
    test_pixel_cat_maps[save_pixel_indices[:, 0], save_pixel_indices[:, 1], :] = sparse_pixel_cat_maps[:]
    assert np.all(test_pixel_cat_maps == save_pixel_cat_maps), "Error: Mismatch"

    test_non_empty_map = np.zeros((dims[0], dims[1]), dtype=np.float32)
    test_non_empty_map[save_pixel_indices[:, 0], save_pixel_indices[:, 1]] = 1.0
    assert np.all(test_non_empty_map == save_non_empty_maps), "Error: Mismatch"

    test_pixel_instance_map = np.zeros((dims[0], dims[1]), dtype=np.uint8)
    test_pixel_instance_map[save_pixel_indices[:, 0], save_pixel_indices[:, 1]] = sparse_pixel_instance_maps[:]
    assert np.all(test_pixel_instance_map == save_pixel_instance_maps), "Error: Mismatch"

    for i in range(len(save_voxel_indices_list)):
        indices = save_data_dict['voxel_indices_' + str(i)]
        curr_voxels = np.zeros(save_voxel_dims, dtype=np.bool)
        curr_voxels[indices[:, 0], indices[:, 1], indices[:, 2]] = 1
        assert np.all(curr_voxels == save_voxel_points[i]), "Error: Mismatch"

    return save_data_dict

def point_in_hull_slow(point, hull, tolerance=1e-12):
    """
    Check if a point lies in a convex hull. This implementation is slow.
    :param point: nd.array (1 x d); d: point dimension
    :param hull: The scipy ConvexHull object
    :param tolerance: Used to compare to a small positive constant because of issues of numerical precision
    (otherwise, you may find that a vertex of the convex hull is not in the convex hull)
    """
    return all((np.dot(eq[:-1], point) + eq[-1] <= tolerance) for eq in hull.equations)



class Box:
    """ Simple data class representing a 3d box including, label, score and velocity. """

    def __init__(self,
                 center: List[float],
                 size: List[float],
                 orientation: Quaternion,
                 label: int = np.nan,
                 score: float = np.nan,
                 velocity: Tuple = (np.nan, np.nan, np.nan),
                 name: str = None,
                 token: str = None):
        """
        :param center: Center of box given as x, y, z.
        :param size: Size of box in width, length, height.
        :param orientation: Box orientation.
        :param label: Integer label, optional.
        :param score: Classification score, optional.
        :param velocity: Box velocity in x, y, z direction.
        :param name: Box name, optional. Can be used e.g. for denote category name.
        :param token: Unique string identifier from DB.
        """
        assert not np.any(np.isnan(center))
        assert not np.any(np.isnan(size))
        assert len(center) == 3
        assert len(size) == 3
        assert type(orientation) == Quaternion

        self.center = np.array(center)
        self.wlh = np.array(size)
        self.orientation = orientation
        self.label = int(label) if not np.isnan(label) else label
        self.score = float(score) if not np.isnan(score) else score
        self.velocity = np.array(velocity)
        self.name = name
        self.token = token

    def __eq__(self, other):
        center = np.allclose(self.center, other.center)
        wlh = np.allclose(self.wlh, other.wlh)
        orientation = np.allclose(self.orientation.elements, other.orientation.elements)
        label = (self.label == other.label) or (np.isnan(self.label) and np.isnan(other.label))
        score = (self.score == other.score) or (np.isnan(self.score) and np.isnan(other.score))
        vel = (np.allclose(self.velocity, other.velocity) or
               (np.all(np.isnan(self.velocity)) and np.all(np.isnan(other.velocity))))

        return center and wlh and orientation and label and score and vel

    def __repr__(self):
        repr_str = 'label: {}, score: {:.2f}, xyz: [{:.2f}, {:.2f}, {:.2f}], wlh: [{:.2f}, {:.2f}, {:.2f}], ' \
                   'rot axis: [{:.2f}, {:.2f}, {:.2f}], ang(degrees): {:.2f}, ang(rad): {:.2f}, ' \
                   'vel: {:.2f}, {:.2f}, {:.2f}, name: {}, token: {}'

        return repr_str.format(self.label, self.score, self.center[0], self.center[1], self.center[2], self.wlh[0],
                               self.wlh[1], self.wlh[2], self.orientation.axis[0], self.orientation.axis[1],
                               self.orientation.axis[2], self.orientation.degrees, self.orientation.radians,
                               self.velocity[0], self.velocity[1], self.velocity[2], self.name, self.token)

    @property
    def rotation_matrix(self) -> np.ndarray:
        """
        Return a rotation matrix.
        :return: <np.float: 3, 3>. The box's rotation matrix.
        """
        return self.orientation.rotation_matrix

    def translate(self, x: np.ndarray) -> None:
        """
        Applies a translation.
        :param x: <np.float: 3, 1>. Translation in x, y, z direction.
        """
        self.center += x

    def rotate(self, quaternion: Quaternion) -> None:
        """
        Rotates box.
        :param quaternion: Rotation to apply.
        """
        self.center = np.dot(quaternion.rotation_matrix, self.center)
        self.orientation = quaternion * self.orientation
        self.velocity = np.dot(quaternion.rotation_matrix, self.velocity)

    def corners(self, wlh_factor: float = 1.0) -> np.ndarray:
        """
        Returns the bounding box corners.
        :param wlh_factor: Multiply w, l, h by a factor to scale the box.
        :return: <np.float: 3, 8>. First four corners are the ones facing forward.
            The last four are the ones facing backwards.
        """
        w, l, h = self.wlh * wlh_factor

        # 3D bounding box corners. (Convention: x points forward, y to the left, z up.)
        x_corners = l / 2 * np.array([1,  1,  1,  1, -1, -1, -1, -1])
        y_corners = w / 2 * np.array([1, -1, -1,  1,  1, -1, -1,  1])
        z_corners = h / 2 * np.array([1,  1, -1, -1,  1,  1, -1, -1])
        corners = np.vstack((x_corners, y_corners, z_corners))

        # Rotate
        corners = np.dot(self.orientation.rotation_matrix, corners)

        # Translate
        x, y, z = self.center
        corners[0, :] = corners[0, :] + x
        corners[1, :] = corners[1, :] + y
        corners[2, :] = corners[2, :] + z

        return corners

    def bottom_corners(self) -> np.ndarray:
        """
        Returns the four bottom corners.
        :return: <np.float: 3, 4>. Bottom corners. First two face forward, last two face backwards.
        """
        return self.corners()[:, [2, 3, 7, 6]]

    def render(self,
               axis: Axes,
               view: np.ndarray = np.eye(3),
               normalize: bool = False,
               colors: Tuple = ('b', 'r', 'k'),
               linewidth: float = 2) -> None:
        """
        Renders the box in the provided Matplotlib axis.
        :param axis: Axis onto which the box should be drawn.
        :param view: <np.array: 3, 3>. Define a projection in needed (e.g. for drawing projection in an image).
        :param normalize: Whether to normalize the remaining coordinate.
        :param colors: (<Matplotlib.colors>: 3). Valid Matplotlib colors (<str> or normalized RGB tuple) for front,
            back and sides.
        :param linewidth: Width in pixel of the box sides.
        """
        corners = view_points(self.corners(), view, normalize=normalize)[:2, :]

        def draw_rect(selected_corners, color):
            prev = selected_corners[-1]
            for corner in selected_corners:
                axis.plot([prev[0], corner[0]], [prev[1], corner[1]], color=color, linewidth=linewidth)
                prev = corner

        # Draw the sides
        for i in range(4):
            axis.plot([corners.T[i][0], corners.T[i + 4][0]],
                      [corners.T[i][1], corners.T[i + 4][1]],
                      color=colors[2], linewidth=linewidth)

        # Draw front (first 4 corners) and rear (last 4 corners) rectangles(3d)/lines(2d)
        draw_rect(corners.T[:4], colors[0])
        draw_rect(corners.T[4:], colors[1])

        # Draw line indicating the front
        center_bottom_forward = np.mean(corners.T[2:4], axis=0)
        center_bottom = np.mean(corners.T[[2, 3, 7, 6]], axis=0)
        axis.plot([center_bottom[0], center_bottom_forward[0]],
                  [center_bottom[1], center_bottom_forward[1]],
                  color=colors[0], linewidth=linewidth)

    def render_cv2(self,
                   im: np.ndarray,
                   view: np.ndarray = np.eye(3),
                   normalize: bool = False,
                   colors: Tuple = ((0, 0, 255), (255, 0, 0), (155, 155, 155)),
                   linewidth: int = 2) -> None:
        """
        Renders box using OpenCV2.
        :param im: <np.array: width, height, 3>. Image array. Channels are in BGR order.
        :param view: <np.array: 3, 3>. Define a projection if needed (e.g. for drawing projection in an image).
        :param normalize: Whether to normalize the remaining coordinate.
        :param colors: ((R, G, B), (R, G, B), (R, G, B)). Colors for front, side & rear.
        :param linewidth: Linewidth for plot.
        """
        corners = view_points(self.corners(), view, normalize=normalize)[:2, :]

        def draw_rect(selected_corners, color):
            prev = selected_corners[-1]
            for corner in selected_corners:
                cv2.line(im,
                         (int(prev[0]), int(prev[1])),
                         (int(corner[0]), int(corner[1])),
                         color, linewidth)
                prev = corner

        # Draw the sides
        for i in range(4):
            cv2.line(im,
                     (int(corners.T[i][0]), int(corners.T[i][1])),
                     (int(corners.T[i + 4][0]), int(corners.T[i + 4][1])),
                     colors[2][::-1], linewidth)

        # Draw front (first 4 corners) and rear (last 4 corners) rectangles(3d)/lines(2d)
        draw_rect(corners.T[:4], colors[0][::-1])
        draw_rect(corners.T[4:], colors[1][::-1])

        # Draw line indicating the front
        center_bottom_forward = np.mean(corners.T[2:4], axis=0)
        center_bottom = np.mean(corners.T[[2, 3, 7, 6]], axis=0)
        cv2.line(im,
                 (int(center_bottom[0]), int(center_bottom[1])),
                 (int(center_bottom_forward[0]), int(center_bottom_forward[1])),
                 colors[0][::-1], linewidth)

    def copy(self) -> 'Box':
        """
        Create a copy of self.
        :return: A copy.
        """
        return copy.deepcopy(self)

def point_in_hull_fast(points: np.array, bounding_box: Box):
    """
    Check if a point lies in a bounding box. We first rotate the bounding box to align with axis. Meanwhile, we
    also rotate the whole point cloud. Finally, we just check the membership with the aid of aligned axis.
    This implementation is fast.
    :param points: nd.array (N x d); N: the number of points, d: point dimension
    :param bounding_box: the Box object
    return: The membership of points within the bounding box
    """
    # Make sure it is a unit quaternion
    bounding_box.orientation = bounding_box.orientation.normalised

    # Rotate the point clouds
    pc = bounding_box.orientation.inverse.rotation_matrix @ points.T
    pc = pc.T

    orientation_backup = Quaternion(bounding_box.orientation)  # Deep clone it
    bounding_box.rotate(bounding_box.orientation.inverse)

    corners = bounding_box.corners()

    # Test if the points are in the bounding box
    # idx = np.where((corners[0, 7] <= pc[:, 0]) & (pc[:, 0] <= corners[0, 0]) &
    #                (corners[1, 1] <= pc[:, 1]) & (pc[:, 1] <= corners[1, 0]) &
    #                (corners[2, 2] <= pc[:, 2]) & (pc[:, 2] <= corners[2, 0]))[0]
    idx = np.where((corners[0, 7] <= pc[:, 0]) & (pc[:, 0] <= corners[0, 0]) &
                   (corners[1, 1] <= pc[:, 1]) & (pc[:, 1] <= corners[1, 0]) )[0]

    # recover
    bounding_box.rotate(orientation_backup)

    return idx


def calc_displace_vector(points: np.array, curr_box: Box, next_box: Box):
    """
    Calculate the displacement vectors for the input points.
    This is achieved by comparing the current and next bounding boxes. Specifically, we first rotate
    the input points according to the delta rotation angle, and then translate them. Finally we compute the
    displacement between the transformed points and the input points.
    :param points: The input points, (N x d). Note that these points should be inside the current bounding box.
    :param curr_box: Current bounding box.
    :param next_box: The future next bounding box in the temporal sequence.
    :return: Displacement vectors for the points.
    """
    assert points.shape[1] == 3, "The input points should have dimension 3."

    # Make sure the quaternions are normalized
    curr_box.orientation = curr_box.orientation.normalised
    next_box.orientation = next_box.orientation.normalised

    delta_rotation = curr_box.orientation.inverse * next_box.orientation
    rotated_pc = (delta_rotation.rotation_matrix @ points.T).T
    rotated_curr_center = np.dot(delta_rotation.rotation_matrix, curr_box.center)
    delta_center = next_box.center - rotated_curr_center

    rotated_tranlated_pc = rotated_pc + delta_center

    pc_displace_vectors = rotated_tranlated_pc - points

    return pc_displace_vectors


def get_static_and_moving_cells(batch_disp_field_gt, upper_thresh=0.1, frame_skip=3):
    """
    Get the indices/masks of static and moving cells. Ths speed of static cells is bounded by upper_thresh.
    In particular, for a given cell, if its displacement over the past 1 second (about 20 sample data) is in the
    range [0, upper_thresh], we consider it as static cell, otherwise as moving cell.
    :param batch_disp_field_gt: Batch of ground-truth displacement fields. numpy array, shape (seq len, h, w)
    :param upper_thresh: The speed upper bound
    :param frame_skip: The number of skipped frame in the sweep sequence. This is used for computing the upper bound
        for defining static objects
    """
    upper_bound = (frame_skip + 1) / 20 * upper_thresh
    disp_norm = np.linalg.norm(batch_disp_field_gt, ord=2, axis=-1)

    static_cell_mask = disp_norm <= upper_bound
    static_cell_mask = np.all(static_cell_mask, axis=0)  # along the sequence axis

    moving_cell_mask = np.logical_not(static_cell_mask)

    return static_cell_mask, moving_cell_mask


def voxelize(pts, voxel_size, extents=None, num_T=35, seed: float = None):
    """
    Voxelize the input point cloud. Code modified from https://github.com/Yc174/voxelnet
    Voxels are 3D grids that represent occupancy info.

    The input for the voxelization is expected to be a PointCloud
    with N points in 4 dimension (x,y,z,i). Voxel size is the quantization size for the voxel grid.

    voxel_size: I.e. if voxel size is 1 m, the voxel space will be
    divided up within 1m x 1m x 1m space. This space will be -1 if free/occluded and 1 otherwise.
    min_voxel_coord: coordinates of the minimum on each axis for the voxel grid
    max_voxel_coord: coordinates of the maximum on each axis for the voxel grid
    num_divisions: number of grids in each axis
    leaf_layout: the voxel grid of size (numDivisions) that contain -1 for free, 0 for occupied

    :param pts: Point cloud as N x [x, y, z, i]
    :param voxel_size: Quantization size for the grid, vd, vh, vw
    :param extents: Optional, specifies the full extents of the point cloud.
                    Used for creating same sized voxel grids. Shape (3, 2)
    :param num_T: Number of points in each voxel after sampling/padding
    :param seed: The random seed for fixing the data generation.
    """
    # Check if points are 3D, otherwise early exit
    if pts.shape[1] < 3 or pts.shape[1] > 4:
        raise ValueError("Points have the wrong shape: {}".format(pts.shape))

    if extents is not None:
        if extents.shape != (3, 2):
            raise ValueError("Extents are the wrong shape {}".format(extents.shape))

        filter_idx = np.where((extents[0, 0] < pts[:, 0]) & (pts[:, 0] < extents[0, 1]) &
                              (extents[1, 0] < pts[:, 1]) & (pts[:, 1] < extents[1, 1]) &
                              (extents[2, 0] < pts[:, 2]) & (pts[:, 2] < extents[2, 1]))[0]
        pts = pts[filter_idx]

    # Discretize voxel coordinates to given quantization size
    discrete_pts = np.floor(pts[:, :3] / voxel_size).astype(np.int32)

    # Use Lex Sort, sort by x, then y, then z
    x_col = discrete_pts[:, 0]
    y_col = discrete_pts[:, 1]
    z_col = discrete_pts[:, 2]
    sorted_order = np.lexsort((z_col, y_col, x_col))

    # Save original points in sorted order
    points = pts[sorted_order]
    discrete_pts = discrete_pts[sorted_order]

    # Format the array to c-contiguous array for unique function
    contiguous_array = np.ascontiguousarray(discrete_pts).view(
        np.dtype((np.void, discrete_pts.dtype.itemsize * discrete_pts.shape[1])))

    # The new coordinates are the discretized array with its unique indexes
    _, unique_indices = np.unique(contiguous_array, return_index=True)

    # Sort unique indices to preserve order
    unique_indices.sort()

    voxel_coords = discrete_pts[unique_indices]

    # Number of points per voxel, last voxel calculated separately
    num_points_in_voxel = np.diff(unique_indices)
    num_points_in_voxel = np.append(num_points_in_voxel, discrete_pts.shape[0] - unique_indices[-1])

    # Compute the minimum and maximum voxel coordinates
    if extents is not None:
        min_voxel_coord = np.floor(extents.T[0] / voxel_size)
        max_voxel_coord = np.ceil(extents.T[1] / voxel_size) - 1
    else:
        min_voxel_coord = np.amin(voxel_coords, axis=0)
        max_voxel_coord = np.amax(voxel_coords, axis=0)

    # Get the voxel grid dimensions
    num_divisions = ((max_voxel_coord - min_voxel_coord) + 1).astype(np.int32)

    # Bring the min voxel to the origin
    voxel_indices = (voxel_coords - min_voxel_coord).astype(int)

    # Padding the points within each voxel
    padded_voxel_points = np.zeros([unique_indices.shape[0], num_T, pts.shape[1] + 3], dtype=np.float32)
    padded_voxel_points = padding_voxel(padded_voxel_points, unique_indices, num_points_in_voxel, points, num_T, seed)

    return padded_voxel_points, voxel_indices, num_divisions


@njit
def padding_voxel(padded_voxel_points, unique_indices, num_points_in_voxel, points, num_T, seed):
    if seed is not None:
        np.random.seed(seed)
    for i, v in enumerate(zip(unique_indices, num_points_in_voxel)):
        if v[1] < num_T:
            padded_voxel_points[i, :v[1], :4] = points[v[0]:v[0] + v[1], :]
            middle_points_x = np.mean(points[v[0]:v[0] + v[1], 0])
            middle_points_y = np.mean(points[v[0]:v[0] + v[1], 1])
            middle_points_z = np.mean(points[v[0]:v[0] + v[1], 2])
            padded_voxel_points[i, :v[1], 4] = padded_voxel_points[i, :v[1], 0] - middle_points_x
            padded_voxel_points[i, :v[1], 5] = padded_voxel_points[i, :v[1], 1] - middle_points_y
            padded_voxel_points[i, :v[1], 6] = padded_voxel_points[i, :v[1], 2] - middle_points_z
        else:
            inds = np.random.choice(v[1], num_T)
            padded_voxel_points[i, :, :4] = points[v[0] + inds, :]
            middle_points_x = np.mean(points[v[0] + inds, 0])
            middle_points_y = np.mean(points[v[0] + inds, 1])
            middle_points_z = np.mean(points[v[0] + inds, 2])
            padded_voxel_points[i, :, 4] = padded_voxel_points[i, :, 0] - middle_points_x
            padded_voxel_points[i, :, 5] = padded_voxel_points[i, :, 1] - middle_points_y
            padded_voxel_points[i, :, 6] = padded_voxel_points[i, :, 2] - middle_points_z

    return padded_voxel_points


def gen_2d_grid_gt(data_dict: dict, grid_size: np.array, extents: np.array = None,
                   frame_skip: int = 0, reordered: bool = False, proportion_thresh: float = 0.5,
                   category_num: int = 5, one_hot_thresh: float = 0.8, h_flip: bool = False,
                   min_point_num_per_voxel: int = -1,
                   return_past_2d_disp_gt: bool = False,
                   return_instance_map: bool = False):
    """
    Generate the 2d grid ground-truth for the input point cloud.
    The ground-truth is: the displacement vectors of the occupied pixels in BEV image.
    The displacement is computed w.r.t the current time and the future time
    :param data_dict: The dictionary containing the data information
    :param grid_size: The size of each pixel
    :param extents: The extents of the point cloud on the 2D xy plane. Shape (3, 2)
    :param frame_skip: The number of sample frames that need to be skipped
    :param reordered: Whether need to reorder the results, so that the first element corresponds to the oldest past
        record. This option is only effective when return_past_2d_disp_gt = True.
    :param proportion_thresh: Within a given pixel, only when the proportion of foreground points exceeds this threshold
        will we compute the displacement vector for this pixel.
    :param category_num: The number of categories for points.
    :param one_hot_thresh: When the proportion of the majority points within a cell exceeds this threshold, we
        compute the (hard) one-hot category vector for this cell, otherwise compute the soft category vector.
    :param h_flip: Flip the point clouds horizontally
    :param min_point_num_per_voxel: Minimum point number inside each voxel (cell). If smaller than this threshold, we
        do not compute the displacement vector for this voxel (cell), and set the displacement to zero.
    :param return_past_2d_disp_gt: Whether to compute the ground-truth displacement filed for the past sweeps.
    :param return_instance_map: Whether to return the instance id map.
    :return: The ground-truth displacement field. Shape (num_sweeps, image height, image width, 2).
    """
    num_sweeps = data_dict['num_sweeps']
    times = data_dict['times']
    num_past_sweeps = len(np.where(times >= 0)[0])
    num_future_sweeps = len(np.where(times < 0)[0])
    assert num_past_sweeps + num_future_sweeps == num_sweeps, "The number of sweeps is incorrect!"

    pc_list = []

    for i in range(num_sweeps):
        pc = data_dict['pc_' + str(i)]
        if h_flip:
            pc[0, :] = -pc[0, :]  # flip x coordinate
        pc_list.append(pc.T)

    # Retrieve the instance boxes
    num_instances = data_dict['num_instances']
    instance_box_list = list()
    instance_cat_list = list()  # for instance categories

    for i in range(num_instances):
        instance = data_dict['instance_boxes_' + str(i)]
        category = data_dict['category_' + str(i)]
        instance_box_list.append(instance)
        instance_cat_list.append(category)

    # ----------------------------------------------------
    # Filter and sort the reference point cloud
    refer_pc = pc_list[0]
    refer_pc = refer_pc[:, 0:3]

    if extents is not None:
        if extents.shape != (3, 2):
            raise ValueError("Extents are the wrong shape {}".format(extents.shape))

        filter_idx = np.where((extents[0, 0] < refer_pc[:, 0]) & (refer_pc[:, 0] < extents[0, 1]) &
                              (extents[1, 0] < refer_pc[:, 1]) & (refer_pc[:, 1] < extents[1, 1]) &
                              (extents[2, 0] < refer_pc[:, 2]) & (refer_pc[:, 2] < extents[2, 1]))[0]
        refer_pc = refer_pc[filter_idx]

    # -- Discretize pixel coordinates to given quantization size
    discrete_pts = np.floor(refer_pc[:, 0:2] / grid_size).astype(np.int32)

    # -- Use Lex Sort, sort by x, then y
    x_col = discrete_pts[:, 0]
    y_col = discrete_pts[:, 1]
    sorted_order = np.lexsort((y_col, x_col))

    refer_pc = refer_pc[sorted_order]
    discrete_pts = discrete_pts[sorted_order]

    contiguous_array = np.ascontiguousarray(discrete_pts).view(
        np.dtype((np.void, discrete_pts.dtype.itemsize * discrete_pts.shape[1])))

    # -- The new coordinates are the discretized array with its unique indexes
    _, unique_indices = np.unique(contiguous_array, return_index=True)

    # -- Sort unique indices to preserve order
    unique_indices.sort()
    pixel_coords = discrete_pts[unique_indices]

    # -- Number of points per voxel, last voxel calculated separately
    num_points_in_pixel = np.diff(unique_indices)
    num_points_in_pixel = np.append(num_points_in_pixel, discrete_pts.shape[0] - unique_indices[-1])

    # -- Compute the minimum and maximum voxel coordinates
    if extents is not None:
        min_pixel_coord = np.floor(extents.T[0, 0:2] / grid_size)
        max_pixel_coord = np.ceil(extents.T[1, 0:2] / grid_size) - 1
    else:
        min_pixel_coord = np.amin(pixel_coords, axis=0)
        max_pixel_coord = np.amax(pixel_coords, axis=0)

    # -- Get the voxel grid dimensions
    num_divisions = ((max_pixel_coord - min_pixel_coord) + 1).astype(np.int32)

    # -- Bring the min voxel to the origin
    pixel_indices = (pixel_coords - min_pixel_coord).astype(int)
    # ----------------------------------------------------

    # ----------------------------------------------------
    # Get the point cloud subsets, which are inside different instance bounding boxes
    refer_box_list = list()
    refer_pc_idx_per_bbox = list()
    points_category = np.zeros(refer_pc.shape[0], dtype=np.int)  # store the point categories

    pixel_instance_id = np.zeros(pixel_indices.shape[0], dtype=np.uint8)
    points_instance_id = np.zeros(refer_pc.shape[0], dtype=np.int)

    for i in range(num_instances):
        instance_cat = instance_cat_list[i]
        instance_box = instance_box_list[i]
        instance_box_data = instance_box[0]
        assert not np.isnan(instance_box_data).any(), "In the keyframe, there should not be NaN box annotation!"

        if h_flip:
            tmp_quad = instance_box_data[6:].copy()
            tmp_quad[2] *= -1  # y
            tmp_quad[3] *= -1  # z
            tmp_quad = Quaternion(tmp_quad)

            tmp_center = instance_box_data[0:3].copy()
            tmp_center[0] = -tmp_center[0]
            tmp_box = Box(center=tmp_center, size=instance_box_data[3:6], orientation=Quaternion(tmp_quad))
        else:
            tmp_box = Box(center=instance_box_data[:3], size=instance_box_data[3:6],
                          orientation=Quaternion(instance_box_data[6:]))
        idx = point_in_hull_fast(refer_pc[:, 0:3], tmp_box)
        refer_pc_idx_per_bbox.append(idx)
        refer_box_list.append(tmp_box)

        points_category[idx] = instance_cat
        points_instance_id[idx] = i + 1  # object id starts from 1, background has id 0

    assert np.max(points_instance_id) <= 255, "The instance id exceeds uint8 max."

    if len(refer_pc_idx_per_bbox) > 0:
        refer_pc_idx_inside_box = np.concatenate(refer_pc_idx_per_bbox).tolist()
    else:
        refer_pc_idx_inside_box = []
    refer_pc_idx_outside_box = set(range(refer_pc.shape[0])) - set(refer_pc_idx_inside_box)
    refer_pc_idx_outside_box = list(refer_pc_idx_outside_box)

    # Compute pixel (cell) categories
    pixel_cat = np.zeros([unique_indices.shape[0], category_num], dtype=np.float32)
    most_freq_info = []

    for h, v in enumerate(zip(unique_indices, num_points_in_pixel)):
        pixel_elements_categories = points_category[v[0]:v[0] + v[1]]
        elements_freq = np.bincount(pixel_elements_categories, minlength=category_num)
        assert np.sum(elements_freq) == v[1], "The frequency count is incorrect."

        elements_freq = elements_freq / float(v[1])
        most_freq_cat, most_freq = np.argmax(elements_freq), np.max(elements_freq)
        most_freq_info.append([most_freq_cat, most_freq])

        most_freq_elements_idx = np.where(pixel_elements_categories == most_freq_cat)[0]
        pixel_elements_instance_ids = points_instance_id[v[0]:v[0] + v[1]]
        most_freq_instance_id = pixel_elements_instance_ids[most_freq_elements_idx[0]]

        if most_freq >= one_hot_thresh:
            one_hot_cat = np.zeros(category_num, dtype=np.float32)
            one_hot_cat[most_freq_cat] = 1.0
            pixel_cat[h] = one_hot_cat

            pixel_instance_id[h] = most_freq_instance_id
        else:
            pixel_cat[h] = elements_freq  # use soft category probability vector.

    pixel_cat_map = np.zeros((num_divisions[0], num_divisions[1], category_num), dtype=np.float32)
    pixel_cat_map[pixel_indices[:, 0], pixel_indices[:, 1]] = pixel_cat[:]

    pixel_instance_map = np.zeros((num_divisions[0], num_divisions[1]), dtype=np.uint8)
    pixel_instance_map[pixel_indices[:, 0], pixel_indices[:, 1]] = pixel_instance_id[:]

    # Set the non-zero pixels to 1.0, which will be helpful for loss computation
    # Note that the non-zero pixels correspond to both the foreground and background objects
    non_empty_map = np.zeros((num_divisions[0], num_divisions[1]), dtype=np.float32)
    non_empty_map[pixel_indices[:, 0], pixel_indices[:, 1]] = 1.0

    # Ignore the voxel/pillar which contains number of points less than min_point_num_per_voxel; only for fg points
    cell_pts_num = np.zeros((num_divisions[0], num_divisions[1]), dtype=np.float32)
    cell_pts_num[pixel_indices[:, 0], pixel_indices[:, 1]] = num_points_in_pixel[:]
    tmp_pixel_cat_map = np.argmax(pixel_cat_map, axis=2)
    ignore_mask = np.logical_and(cell_pts_num <= min_point_num_per_voxel, tmp_pixel_cat_map != 0)
    ignore_mask = np.logical_not(ignore_mask)
    ignore_mask = np.expand_dims(ignore_mask, axis=2)

    # Compute the displacement vectors w.r.t. the other sweeps
    all_disp_field_gt_list = list()
    all_valid_pixel_maps_list = list()  # valid pixel map will be used for masking the computation of loss

    # -- Skip some frames if necessary
    past_part = list(range(0, num_past_sweeps, frame_skip + 1))
    future_part = list(range(num_past_sweeps + frame_skip, num_sweeps, frame_skip + 1))
    if return_past_2d_disp_gt:
        zero_disp_field = np.zeros((num_divisions[0], num_divisions[1], 2), dtype=np.float32)
        all_disp_field_gt_list.append(zero_disp_field)  # append once, which corresponds to the current frame
        all_valid_pixel_maps_list.append(non_empty_map)

        frame_considered = np.asarray(past_part + future_part)
        frame_considered = frame_considered[1:]
    else:
        frame_considered = np.asarray(future_part)

    for i in frame_considered:
        curr_disp_vectors = np.zeros_like(refer_pc, dtype=np.float32)
        curr_disp_vectors.fill(np.nan)
        curr_disp_vectors[refer_pc_idx_outside_box,] = 0.0

        # First, for each instance, compute the corresponding points displacement.
        for j in range(num_instances):
            instance_box = instance_box_list[j]
            instance_box_data = instance_box[i]  # This is for the i-th sweep

            if np.isnan(instance_box_data).any():  # It is possible that in this sweep there is no annotation
                continue

            if h_flip:
                tmp_quad = instance_box_data[6:].copy()
                tmp_quad[2] *= -1  # y
                tmp_quad[3] *= -1  # z
                tmp_quad = Quaternion(tmp_quad)

                tmp_center = instance_box_data[0:3].copy()
                tmp_center[0] = -tmp_center[0]
                tmp_box = Box(center=tmp_center, size=instance_box_data[3:6], orientation=Quaternion(tmp_quad))
            else:
                tmp_box = Box(center=instance_box_data[:3], size=instance_box_data[3:6],
                              orientation=Quaternion(instance_box_data[6:]))
            pc_in_bbox_idx = refer_pc_idx_per_bbox[j]
            disp_vectors = calc_displace_vector(refer_pc[pc_in_bbox_idx], refer_box_list[j], tmp_box)

            curr_disp_vectors[pc_in_bbox_idx] = disp_vectors[:]

        # Second, compute the mean displacement vector and category for each non-empty pixel
        disp_field = np.zeros([unique_indices.shape[0], 2], dtype=np.float32)  # we only consider the 2D field

        # We only compute loss for valid pixels where there are corresponding box annotations between two frames
        valid_pixels = np.zeros(unique_indices.shape[0], dtype=np.bool)

        for h, v in enumerate(zip(unique_indices, num_points_in_pixel)):

            # Only when the number of majority points exceeds predefined proportion, we compute
            # the displacement vector for this pixel. Otherwise, We consider it is background (possibly ground plane)
            # and has zero displacement.
            pixel_elements_categories = points_category[v[0]:v[0] + v[1]]
            most_freq_cat, most_freq = most_freq_info[h]

            if most_freq >= proportion_thresh:
                most_freq_cat_idx = np.where(pixel_elements_categories == most_freq_cat)[0]
                most_freq_cat_disp_vectors = curr_disp_vectors[v[0]:v[0] + v[1], :3]
                most_freq_cat_disp_vectors = most_freq_cat_disp_vectors[most_freq_cat_idx]

                if np.isnan(most_freq_cat_disp_vectors).any():  # contains invalid disp vectors
                    valid_pixels[h] = 0.0
                else:
                    mean_disp_vector = np.mean(most_freq_cat_disp_vectors, axis=0)
                    disp_field[h] = mean_disp_vector[0:2]  # ignore the z direction

                    valid_pixels[h] = 1.0

        # Finally, assemble to a 2D image
        disp_field_sparse = np.zeros((num_divisions[0], num_divisions[1], 2), dtype=np.float32)
        disp_field_sparse[pixel_indices[:, 0], pixel_indices[:, 1]] = disp_field[:]
        disp_field_sparse = disp_field_sparse * ignore_mask

        valid_pixel_map = np.zeros((num_divisions[0], num_divisions[1]), dtype=np.float32)
        valid_pixel_map[pixel_indices[:, 0], pixel_indices[:, 1]] = valid_pixels[:]

        all_disp_field_gt_list.append(disp_field_sparse)
        all_valid_pixel_maps_list.append(valid_pixel_map)

    all_disp_field_gt_list = np.stack(all_disp_field_gt_list, axis=0)
    all_valid_pixel_maps_list = np.stack(all_valid_pixel_maps_list, axis=0)

    if reordered and return_past_2d_disp_gt:
        num_past = len(past_part)
        all_disp_field_gt_list[0:num_past] = all_disp_field_gt_list[(num_past - 1)::-1]
        all_valid_pixel_maps_list[0:num_past] = all_valid_pixel_maps_list[(num_past - 1)::-1]

    if return_instance_map:
        return all_disp_field_gt_list, all_valid_pixel_maps_list, non_empty_map, pixel_cat_map, pixel_indices, pixel_instance_map
    else:
        return all_disp_field_gt_list, all_valid_pixel_maps_list, non_empty_map, pixel_cat_map, pixel_indices


def voxelize_occupy(pts, voxel_size, extents=None, return_indices=False):
    """
    Voxelize the input point cloud. We only record if a given voxel is occupied or not, which is just binary indicator.

    The input for the voxelization is expected to be a PointCloud
    with N points in 4 dimension (x,y,z,i). Voxel size is the quantization size for the voxel grid.

    voxel_size: I.e. if voxel size is 1 m, the voxel space will be
    divided up within 1m x 1m x 1m space. This space will be 0 if free/occluded and 1 otherwise.
    min_voxel_coord: coordinates of the minimum on each axis for the voxel grid
    max_voxel_coord: coordinates of the maximum on each axis for the voxel grid
    num_divisions: number of grids in each axis
    leaf_layout: the voxel grid of size (numDivisions) that contain -1 for free, 0 for occupied

    :param pts: Point cloud as N x [x, y, z, i]
    :param voxel_size: Quantization size for the grid, vd, vh, vw
    :param extents: Optional, specifies the full extents of the point cloud.
                    Used for creating same sized voxel grids. Shape (3, 2)
    :param return_indices: Whether to return the non-empty voxel indices.
    """
    # Function Constants
    VOXEL_EMPTY = 0
    VOXEL_FILLED = 1

    # Check if points are 3D, otherwise early exit
    if pts.shape[1] < 3 or pts.shape[1] > 4:
        raise ValueError("Points have the wrong shape: {}".format(pts.shape))

    if extents is not None:
        if extents.shape != (3, 2):
            raise ValueError("Extents are the wrong shape {}".format(extents.shape))

        filter_idx = np.where((extents[0, 0] < pts[:, 0]) & (pts[:, 0] < extents[0, 1]) &
                              (extents[1, 0] < pts[:, 1]) & (pts[:, 1] < extents[1, 1]) &
                              (extents[2, 0] < pts[:, 2]) & (pts[:, 2] < extents[2, 1]))[0]
        pts = pts[filter_idx]

    # Discretize voxel coordinates to given quantization size
    discrete_pts = np.floor(pts[:, :3] / voxel_size).astype(np.int32)

    # Use Lex Sort, sort by x, then y, then z
    x_col = discrete_pts[:, 0]
    y_col = discrete_pts[:, 1]
    z_col = discrete_pts[:, 2]
    sorted_order = np.lexsort((z_col, y_col, x_col))

    # Save original points in sorted order
    discrete_pts = discrete_pts[sorted_order]

    # Format the array to c-contiguous array for unique function
    contiguous_array = np.ascontiguousarray(discrete_pts).view(
        np.dtype((np.void, discrete_pts.dtype.itemsize * discrete_pts.shape[1])))

    # The new coordinates are the discretized array with its unique indexes
    _, unique_indices = np.unique(contiguous_array, return_index=True)

    # Sort unique indices to preserve order
    unique_indices.sort()

    voxel_coords = discrete_pts[unique_indices]

    # Compute the minimum and maximum voxel coordinates
    if extents is not None:
        min_voxel_coord = np.floor(extents.T[0] / voxel_size)
        max_voxel_coord = np.ceil(extents.T[1] / voxel_size) - 1
    else:
        min_voxel_coord = np.amin(voxel_coords, axis=0)
        max_voxel_coord = np.amax(voxel_coords, axis=0)

    # Get the voxel grid dimensions
    num_divisions = ((max_voxel_coord - min_voxel_coord) + 1).astype(np.int32)

    # Bring the min voxel to the origin
    voxel_indices = (voxel_coords - min_voxel_coord).astype(int)

    # Create Voxel Object with -1 as empty/occluded
    leaf_layout = VOXEL_EMPTY * np.ones(num_divisions.astype(int), dtype=np.float32)

    # Fill out the leaf layout
    leaf_layout[voxel_indices[:, 0],
                voxel_indices[:, 1],
                voxel_indices[:, 2]] = VOXEL_FILLED

    if return_indices:
        return leaf_layout, voxel_indices
    else:
        return leaf_layout


def voxelize_pillar_indices(pts, voxel_size, extents=None):
    """
    Voxelize the input point cloud into pillars. We only return the indices

    The input for the voxelization is expected to be a PointCloud
    with N points in 4 dimension (x,y,z,i). Voxel size is the quantization size for the voxel grid.

    voxel_size: I.e. if voxel size is 1 m, the voxel space will be
    divided up within 1m x 1m x 1m space. This space will be 0 if free/occluded and 1 otherwise.
    min_voxel_coord: coordinates of the minimum on each axis for the voxel grid
    max_voxel_coord: coordinates of the maximum on each axis for the voxel grid
    num_divisions: number of grids in each axis
    leaf_layout: the voxel grid of size (numDivisions) that contain -1 for free, 0 for occupied

    :param pts: Point cloud as N x [x, y, z, i]
    :param voxel_size: Quantization size for the grid, vh, vw
    :param extents: Optional, specifies the full extents of the point cloud.
                    Used for creating same sized voxel grids. Shape (3, 2)
    """
    # Check if points are 3D, otherwise early exit
    if pts.shape[1] < 3 or pts.shape[1] > 4:
        raise ValueError("Points have the wrong shape: {}".format(pts.shape))

    if extents is not None:
        if extents.shape != (3, 2):
            raise ValueError("Extents are the wrong shape {}".format(extents.shape))

        filter_idx = np.where((extents[0, 0] < pts[:, 0]) & (pts[:, 0] < extents[0, 1]) &
                              (extents[1, 0] < pts[:, 1]) & (pts[:, 1] < extents[1, 1]) &
                              (extents[2, 0] < pts[:, 2]) & (pts[:, 2] < extents[2, 1]))[0]
        pts = pts[filter_idx]

    # Discretize voxel coordinates to given quantization size
    discrete_pts = np.floor(pts[:, :2] / voxel_size).astype(np.int32)

    # Use Lex Sort, sort by x, then y
    x_col = discrete_pts[:, 0]
    y_col = discrete_pts[:, 1]
    sorted_order = np.lexsort((y_col, x_col))

    # Save original points in sorted order
    points = pts[sorted_order]
    discrete_pts = discrete_pts[sorted_order]

    # Format the array to c-contiguous array for unique function
    contiguous_array = np.ascontiguousarray(discrete_pts).view(
        np.dtype((np.void, discrete_pts.dtype.itemsize * discrete_pts.shape[1])))

    # The new coordinates are the discretized array with its unique indexes
    _, unique_indices = np.unique(contiguous_array, return_index=True)

    # Sort unique indices to preserve order
    unique_indices.sort()

    voxel_coords = discrete_pts[unique_indices]

    # Number of points per voxel, last voxel calculated separately
    num_points_in_pillar = np.diff(unique_indices)
    num_points_in_pillar = np.append(num_points_in_pillar, discrete_pts.shape[0] - unique_indices[-1])

    # Compute the minimum and maximum voxel coordinates
    if extents is not None:
        min_voxel_coord = np.floor(extents.T[0, 0:2] / voxel_size)
    else:
        min_voxel_coord = np.amin(voxel_coords, axis=0)

    # Bring the min voxel to the origin
    voxel_indices = (voxel_coords - min_voxel_coord).astype(int)

    return points, voxel_indices, num_points_in_pillar


def voxelize_point_pillar(pts, grid_size, extents=None, num_points=100, num_pillars=2500, seed=None,
                          is_padded_pillar=False):
    """
    Discretize the input point cloud into pillars.

    The input for the voxelization is expected to be a PointCloud
    with N points in 4 dimension (x,y,z,i). Grid size is the quantization size for the 2d grid.

    grid_size: I.e. if grid size is 1 m, the grid space will be
    divided up within 1m x 1m space. This space will be -1 if free/occluded and 1 otherwise.
    min_grid_coord: coordinates of the minimum on each axis for the 2d grid
    max_grid_coord: coordinates of the maximum on each axis for the 2d grid
    num_divisions: number of grids in each axis
    leaf_layout: the 2d grid of size (numDivisions) that contain -1 for free, 0 for occupied

    :param pts: Point cloud as N x [x, y, z, i]
    :param grid_size: Quantization size for the grid, (vh, vw)
    :param extents: Optional, specifies the full extents of the point cloud.
                    Used for creating same sized voxel grids. Shape (3, 2)
    :param num_points: Number of points in each pillar after sampling/padding
    :param num_pillars: Number of pillars after sampling/padding
    :param seed: Random seed for fixing data generation.
    :param is_padded_pillar: Whether need to pad/sample the pillar
    """
    if seed is not None:
        np.random.seed(seed)

    # Check if points are 3D, otherwise early exit
    if pts.shape[1] < 3 or pts.shape[1] > 4:
        raise ValueError("Points have the wrong shape: {}".format(pts.shape))

    if extents is not None:
        if extents.shape != (3, 2):
            raise ValueError("Extents are the wrong shape {}".format(extents.shape))

        filter_idx = np.where((extents[0, 0] < pts[:, 0]) & (pts[:, 0] < extents[0, 1]) &
                              (extents[1, 0] < pts[:, 1]) & (pts[:, 1] < extents[1, 1]) &
                              (extents[2, 0] < pts[:, 2]) & (pts[:, 2] < extents[2, 1]))[0]
        pts = pts[filter_idx]

    # Discretize point coordinates to given 2d quantization size
    discrete_pts = np.floor(pts[:, :2] / grid_size).astype(np.int32)

    # Use Lex Sort, sort by x, then y. We do not care about z
    x_col = discrete_pts[:, 0]
    y_col = discrete_pts[:, 1]
    sorted_order = np.lexsort((y_col, x_col))

    # Save original points in sorted order
    points = pts[sorted_order]
    discrete_pts = discrete_pts[sorted_order]

    # Format the array to c-contiguous array for unique function
    contiguous_array = np.ascontiguousarray(discrete_pts).view(
        np.dtype((np.void, discrete_pts.dtype.itemsize * discrete_pts.shape[1])))

    # The new coordinates are the discretized array with its unique indexes
    _, unique_indices = np.unique(contiguous_array, return_index=True)

    # Sort unique indices to preserve order
    unique_indices.sort()

    grid_coords = discrete_pts[unique_indices]

    # Number of points per voxel, last voxel calculated separately
    num_points_in_pillar = np.diff(unique_indices)
    num_points_in_pillar = np.append(num_points_in_pillar, discrete_pts.shape[0] - unique_indices[-1])

    # Compute the minimum and maximum voxel coordinates
    if extents is not None:
        min_grid_coord = np.floor(extents.T[0, 0:2] / grid_size)
        max_grid_coord = np.ceil(extents.T[1, 0:2] / grid_size) - 1
    else:
        min_grid_coord = np.amin(grid_coords, axis=0)
        max_grid_coord = np.amax(grid_coords, axis=0)

    # Get the voxel grid dimensions
    num_divisions = ((max_grid_coord - min_grid_coord) + 1).astype(np.int32)

    # Bring the min voxel to the origin
    pixel_indices = (grid_coords - min_grid_coord).astype(int)

    # Padding the points within each voxel
    x_offset = grid_size[0] / 2.0 + extents[0, 0]
    y_offset = grid_size[1] / 2.0 + extents[1, 0]

    padded_grid_points = np.zeros([unique_indices.shape[0], num_points, pts.shape[1] + 3 + 2], dtype=np.float32)
    padded_pillar = np.zeros([num_pillars, num_points, pts.shape[1] + 3 + 2], dtype=np.float32)
    padded_pixel_indices = np.zeros([num_pillars, pixel_indices.shape[1]], dtype=np.int64)

    padded_grid_points = padding_point_pillar(padded_grid_points, unique_indices, num_points, num_points_in_pillar,
                                              points, pixel_indices, grid_size, x_offset, y_offset, seed)

    if is_padded_pillar:

        # Padding or sampling the pillars  TODO: early sampling to avoid unnecessary computation
        if unique_indices.shape[0] < num_pillars:
            padded_pillar[:unique_indices.shape[0], :, :] = padded_grid_points[:]
            padded_pixel_indices[:unique_indices.shape[0], :] = pixel_indices[:]
        else:
            pillar_inds = np.random.choice(unique_indices.shape[0], num_pillars)
            padded_pillar[:, :, :] = padded_grid_points[pillar_inds, :, :]
            padded_pixel_indices[:, :] = pixel_indices[pillar_inds, :]
    else:
        padded_pillar = padded_grid_points
        padded_pixel_indices = pixel_indices

    return padded_pillar, padded_pixel_indices, num_divisions


@njit
def padding_point_pillar(padded_grid_points, unique_indices, num_points, num_points_in_pillar,
                         points, pixel_indices, grid_size, x_offset, y_offset, seed):
    if seed is not None:
        np.random.seed(seed)
    for i, v in enumerate(zip(unique_indices, num_points_in_pillar)):
        if v[1] < num_points:
            padded_grid_points[i, :v[1], :4] = points[v[0]:v[0] + v[1], :]
            middle_points_x = np.mean(points[v[0]:v[0] + v[1], 0])
            middle_points_y = np.mean(points[v[0]:v[0] + v[1], 1])
            middle_points_z = np.mean(points[v[0]:v[0] + v[1], 2])

            padded_grid_points[i, :v[1], 4] = padded_grid_points[i, :v[1], 0] - middle_points_x
            padded_grid_points[i, :v[1], 5] = padded_grid_points[i, :v[1], 1] - middle_points_y
            padded_grid_points[i, :v[1], 6] = padded_grid_points[i, :v[1], 2] - middle_points_z

            center_offsets = np.zeros((v[1], 2), dtype=np.float32)
            center_offsets[:, 0] = padded_grid_points[i, :v[1], 0] - (pixel_indices[i, 0] * grid_size[0] + x_offset)
            center_offsets[:, 1] = padded_grid_points[i, :v[1], 1] - (pixel_indices[i, 1] * grid_size[1] + y_offset)
            padded_grid_points[i, :v[1], 7:] = center_offsets[:]
        else:
            inds = np.random.choice(v[1], num_points)
            padded_grid_points[i, :, :4] = points[v[0] + inds, :]
            middle_points_x = np.mean(points[v[0] + inds, 0])
            middle_points_y = np.mean(points[v[0] + inds, 1])
            middle_points_z = np.mean(points[v[0] + inds, 2])

            padded_grid_points[i, :, 4] = padded_grid_points[i, :, 0] - middle_points_x
            padded_grid_points[i, :, 5] = padded_grid_points[i, :, 1] - middle_points_y
            padded_grid_points[i, :, 6] = padded_grid_points[i, :, 2] - middle_points_z

            center_offsets = np.zeros((num_points, 2), dtype=np.float32)
            center_offsets[:, 0] = padded_grid_points[i, :, 0] - (pixel_indices[i, 0] * grid_size[0] + x_offset)
            center_offsets[:, 1] = padded_grid_points[i, :, 1] - (pixel_indices[i, 1] * grid_size[1] + y_offset)
            padded_grid_points[i, :, 7:] = center_offsets[:]

    return padded_grid_points


def compute_ratio_cat_and_motion(dataset_root=None, frame_skip=3, voxel_size=(0.4, 0.4, 0.4), split='train',
                                 area_extents=np.array([[-30., 30.], [-30., 30.], [-2., 2.]]), num_obj_cat=5,
                                 num_motion_cat=3):
    """
    Compute the ratios between foreground and background (and static and moving) cells. The ratios will be used for
    non-uniform weighting to mitigate the class imbalance during training.
    :param dataset_root: The path to the dataset
    :param frame_skip: The number of frame skipped in a sample sequence
    :param voxel_size: Voxel size, which determines the "image" resolution
    :param split: The data split
    :param area_extents: The area of interest for point cloud
    :param num_obj_cat: The number of object categories.
    :param num_motion_cat: The number of motion categories. Currently it is 2 (ie, static and moving).
    """
    if dataset_root is None:
        dataset_root = '/homes/pwu/_drives/cv0/data/homes/pwu/preprocessed_pc'

    scene_dirs = [d for d in os.listdir(dataset_root) if os.path.isdir(os.path.join(dataset_root, d))]

    if split == 'train':
        scene_dirs = scene_dirs[:len(scene_dirs) // 2]
    else:
        scene_dirs = scene_dirs[len(scene_dirs) // 2:]

    sample_seq_files = []
    for s_dir in scene_dirs:
        sample_dir = os.path.join(dataset_root, s_dir)
        sample_files = [os.path.join(sample_dir, f) for f in os.listdir(sample_dir)
                        if os.path.isfile(os.path.join(sample_dir, f))]

        sample_seq_files += sample_files

    num_sample_seqs = len(sample_seq_files)

    obj_cat_cnt = np.zeros(num_obj_cat, dtype=np.int64)
    motion_cat_cnt = np.zeros(num_motion_cat, dtype=np.int64)

    for idx in range(num_sample_seqs):
        sample_file = sample_seq_files[idx]

        all_disp_field_gt, all_valid_pixel_maps, non_empty_map, pixel_cat_map_gt \
            = gen_2d_grid_gt(sample_file, grid_size=voxel_size[0:2], reordered=True,
                             extents=area_extents, frame_skip=frame_skip)

        # -- Compute speed level ground truth
        motion_status_gt = compute_speed_level(all_disp_field_gt, frame_skip=frame_skip)

        # -- Count the object category number
        max_prob = np.amax(pixel_cat_map_gt, axis=-1)
        filter_mask = max_prob == 1.0
        pixel_cat_map = np.argmax(pixel_cat_map_gt, axis=-1) + 1  # category starts from 1 (background), etc
        pixel_cat_map = (pixel_cat_map * non_empty_map * filter_mask).astype(np.int)

        for i in range(num_obj_cat):
            curr_cat_mask = pixel_cat_map == (i + 1)
            curr_cat_num = np.sum(curr_cat_mask)
            obj_cat_cnt[i] += curr_cat_num

        # -- Count the motion category number
        motion_cat_map = np.argmax(motion_status_gt, axis=-1) + 1  # category starts from 1 (static), etc
        motion_cat_map = (motion_cat_map * non_empty_map * filter_mask).astype(np.int)

        for i in range(num_motion_cat):
            curr_motion_mask = motion_cat_map == (i + 1)
            curr_motion_num = np.sum(curr_motion_mask)
            motion_cat_cnt[i] += curr_motion_num

        print("Finish {}".format(idx))

    print("The category numbers: \n{}".format(obj_cat_cnt))
    print("The motion numbers: \n{}\n".format(motion_cat_cnt))

    # Convert to ratio
    obj_cat_ratio = obj_cat_cnt / np.sum(obj_cat_cnt)
    motion_cat_ratio = motion_cat_cnt / np.sum(motion_cat_cnt)

    print("The category ratios: \n{}".format(obj_cat_ratio))
    print("The motion ratios: \n{}\n".format(motion_cat_ratio))

    # Convert to inverse ratio
    obj_cat_ratio_inverse = np.exp(-obj_cat_ratio)
    motion_cat_ratio_inverse = np.exp(-motion_cat_ratio)

    print("The category reverse ratios: \n{}".format(obj_cat_ratio_inverse))
    print("The motion reverse ratios: \n{}\n".format(motion_cat_ratio_inverse))


def compute_speed_level(all_disp_field_gt, total_future_sweeps=20, frame_skip=3):
    speed_intervals = np.array([[0, 5.0], [5.0, 20.0], [20.0, np.inf]])  # unit: m/s

    selected_future_sweeps = np.arange(0, total_future_sweeps + 1, frame_skip + 1)
    selected_future_sweeps = selected_future_sweeps[1:]
    last_future_sweep_id = selected_future_sweeps[-1]
    distance_intervals = speed_intervals * (last_future_sweep_id / 20.0)

    speed_level = np.zeros((all_disp_field_gt.shape[1], all_disp_field_gt.shape[2],
                            speed_intervals.shape[0]), dtype=np.float32)
    last_frame_disp_norm = np.linalg.norm(all_disp_field_gt, ord=2, axis=-1)
    last_frame_disp_norm = last_frame_disp_norm[-1, :, :]

    for s, d in enumerate(distance_intervals):
        mask = np.logical_and(d[0] <= last_frame_disp_norm, last_frame_disp_norm < d[1])

        one_hot_vector = np.zeros(speed_intervals.shape[0], dtype=np.float32)
        one_hot_vector[s] = 1.0

        speed_level[mask] = one_hot_vector[:]

    return speed_level


def compute_speed_level_with_static(all_disp_field_gt, total_future_sweeps=20, frame_skip=3):
    speed_intervals = np.array([[0.0, 0.0], [0, 5.0], [5.0, 20.0], [20.0, np.inf]])  # unit: m/s

    # First, compute the static and moving cell masks
    all_disp_field_gt_norm = np.linalg.norm(all_disp_field_gt, ord=2, axis=-1)

    upper_thresh = 0.2
    upper_bound = (frame_skip + 1) / 20 * upper_thresh
    selected_future_sweeps = np.arange(0, total_future_sweeps + 1, frame_skip + 1)
    selected_future_sweeps = selected_future_sweeps[1:]

    future_sweeps_disp_field_gt_norm = all_disp_field_gt_norm[-len(selected_future_sweeps):, ...]
    static_cell_mask = future_sweeps_disp_field_gt_norm <= upper_bound
    static_cell_mask = np.all(static_cell_mask, axis=0)  # along the sequence axis
    moving_cell_mask = np.logical_not(static_cell_mask)

    # Next, compute the speed level mask
    last_future_sweep_id = selected_future_sweeps[-1]
    distance_intervals = speed_intervals * (last_future_sweep_id / 20.0)

    speed_level = np.zeros((all_disp_field_gt.shape[1], all_disp_field_gt.shape[2],
                            speed_intervals.shape[0]), dtype=np.float32)
    last_frame_disp_norm = all_disp_field_gt_norm[-1, :, :]

    for s, d in enumerate(distance_intervals):
        if s == 0:
            mask = static_cell_mask
        else:
            mask = np.logical_and(d[0] <= last_frame_disp_norm, last_frame_disp_norm < d[1])
            mask = np.logical_and(mask, moving_cell_mask)

        one_hot_vector = np.zeros(speed_intervals.shape[0], dtype=np.float32)
        one_hot_vector[s] = 1.0

        speed_level[mask] = one_hot_vector[:]

    return speed_level


def classify_speed_level(all_disp_field_gt, total_future_sweeps=20, future_frame_skip=0):
    """
    Classify each cell into static (possibly background) or moving.
    """
    # First, compute the static and moving cell masks
    all_disp_field_gt_norm = np.linalg.norm(all_disp_field_gt, ord=2, axis=-1)

    # Every future_frame_skip frames, if the movement of grid cells does not exceed this thresh (unit: meters),
    # then they are considered as static. This thresh is set to be the maximum perturbation for 1 second.
    upper_thresh = 0.2
    upper_bound = (future_frame_skip + 1) / 20 * upper_thresh
    selected_future_sweeps = np.arange(0, total_future_sweeps + 1, future_frame_skip + 1)
    selected_future_sweeps = selected_future_sweeps[1:]

    future_sweeps_disp_field_gt_norm = all_disp_field_gt_norm[-len(selected_future_sweeps):, ...]
    static_cell_mask = future_sweeps_disp_field_gt_norm <= upper_bound
    static_cell_mask = np.all(static_cell_mask, axis=0)  # along the temporal axis
    moving_cell_mask = np.logical_not(static_cell_mask)

    # Next, compute corresponding one-hot vectors
    motion_cat = np.zeros((all_disp_field_gt.shape[1], all_disp_field_gt.shape[2], 2), dtype=np.float32)
    bg_one_hot_vector = np.zeros(2, dtype=np.float32)
    bg_one_hot_vector[0] = 1.0
    motion_cat[static_cell_mask] = bg_one_hot_vector[:]

    fg_one_hot_vector = np.zeros(2, dtype=np.float32)
    fg_one_hot_vector[1] = 1.0
    motion_cat[moving_cell_mask] = fg_one_hot_vector[:]

    return motion_cat

def view_points(points: np.ndarray, view: np.ndarray, normalize: bool) -> np.ndarray:
    """
    This is a helper class that maps 3d points to a 2d plane. It can be used to implement both perspective and
    orthographic projections. It first applies the dot product between the points and the view. By convention,
    the view should be such that the data is projected onto the first 2 axis. It then optionally applies a
    normalization along the third dimension.
    For a perspective projection the view should be a 3x3 camera matrix, and normalize=True
    For an orthographic projection with translation the view is a 3x4 matrix and normalize=False
    For an orthographic projection without translation the view is a 3x3 matrix (optionally 3x4 with last columns
     all zeros) and normalize=False
    :param points: <np.float32: 3, n> Matrix of points, where each point (x, y, z) is along each column.
    :param view: <np.float32: n, n>. Defines an arbitrary projection (n <= 4).
        The projection should be such that the corners are projected onto the first 2 axis.
    :param normalize: Whether to normalize the remaining coordinate (along the third axis).
    :return: <np.float32: 3, n>. Mapped point. If normalize=False, the third coordinate is the height.
    """

    assert view.shape[0] <= 4
    assert view.shape[1] <= 4
    assert points.shape[0] == 3

    viewpad = np.eye(4)
    viewpad[:view.shape[0], :view.shape[1]] = view

    nbr_points = points.shape[1]

    # Do operation in homogenous coordinates.
    points = np.concatenate((points, np.ones((1, nbr_points))))
    points = np.dot(viewpad, points)
    points = points[:3, :]

    if normalize:
        points = points / points[2:3, :].repeat(3, 0).reshape(3, nbr_points)

    return points
opencood/utils/flow_utils.py
MediaBrain-SJTU/CoBEVFlow
import open3d as o3d
import numpy as np
import cv2
from matplotlib import pyplot as plt
from opencood.utils.subsampling_utils import get_random_subset
from multiprocessing import Process

vis = False

def mask_points_by_range(points, limit_range, return_mask=False):
    if len(limit_range) == 6:
        mask =  (points[:, 0] > limit_range[0]) & \
                (points[:, 0] < limit_range[3]) & \
                (points[:, 1] > limit_range[1]) & \
                (points[:, 1] < limit_range[4]) & \
                (points[:, 2] > limit_range[2]) & \
                (points[:, 2] < limit_range[5])
    elif len(limit_range) == 4:
        mask =  (points[:, 0] > limit_range[0]) & \
                (points[:, 0] < limit_range[2]) & \
                (points[:, 1] > limit_range[1]) & \
                (points[:, 1] < limit_range[3]) 

    points_mask = points[mask]
    
    if return_mask:
        return points_mask, mask
    else:
        return points_mask

def project_bev(pcd_np, lidar_range, voxel_size):
    """ project pcd to bev
    Args:
        pcd_np: np.ndarray, (N, 3)

        lidar_range: list
            range for bev, [x_min, y_min, z_min, x_max, y_max, z_max]

    Return
        bev: np.array, (H, W), 
            H = (y_max - y_min) / voxel_size
            W = (x_max - x_min) / voxel_size

        pcd_np_with_idx: np.ndarray, (N_, 4)
            last index show it belongs to which grid
    """
    [x_min, y_min, z_min, x_max, y_max, z_max] = lidar_range

    pcd_crop_np, mask = mask_points_by_range(pcd_np, lidar_range, return_mask=True)

    pcd_np_with_idx = np.zeros((pcd_np.shape[0], 4))
    pcd_np_with_idx[:,:3] = pcd_np

    H = round((y_max - y_min) / voxel_size)
    W = round((x_max - x_min) / voxel_size)
    # print(f"BEV map with shape ({H}, {W}).")

    bev = np.zeros((H, W), dtype=np.uint8)
    for i, (x,y,z) in enumerate(pcd_np):
        y_idx = int((y - y_min) / voxel_size)
        x_idx = int((x - x_min) / voxel_size)
        if mask[i]:
            bev[y_idx, x_idx] = 255
        pcd_np_with_idx[i][3] = y_idx * W + x_idx

    if vis:
        plt.imshow(bev)
        plt.show()

    return bev, pcd_np_with_idx

def line_detection(bev_img):
    """
    Should we really need detect line?
    Is edge enough to use?
    """
    edges = cv2.Canny(bev_img, 100, 200)
    if vis:
        plt.imshow(edges)
        plt.show()

    rho = 1  # distance resolution in pixels of the Hough grid
    theta = np.pi / 180  # angular resolution in radians of the Hough grid
    threshold = 25  # minimum number of votes (intersections in Hough grid cell)
    min_line_length = 20  # minimum number of pixels making up a line
    max_line_gap = 20  # maximum gap in pixels between connectable line segments

    line_image = np.copy(bev_img) * 0  # creating a blank to draw lines on

    # Run Hough on edge detected image
    # Output "lines" is an array containing endpoints of detected line segments
    lines = cv2.HoughLinesP(edges, rho, theta, threshold, np.array([]),
                        min_line_length, max_line_gap)

    for line in lines:
        for x1,y1,x2,y2 in line:
            cv2.line(line_image,(x1,y1),(x2,y2),(255),1)

    if vis:
        plt.imshow(line_image)
        plt.show()

    return line_image


def get_point_in_voxels(pcd_np, rows, cols, lidar_range, voxel_size, pcd_with_idx):
    """ use indice in image to filter point cloud, then sample within it.
    Args:
        pcd_np: [N, 3]
        rows: [M,] non zero index -> row
        cols: [M,] non zero index -> col
        pcd_with_idx: [N, 4]
    Returns:
        points_select: [N_, 3]
    """
    [x_min, y_min, z_min, x_max, y_max, z_max] = lidar_range
    H = round((y_max - y_min) / voxel_size)
    W = round((x_max - x_min) / voxel_size)

    M = rows.shape[0]
    points_select = np.zeros((0,4))

    for i in range(M):
        # voxel_range = [x_min + voxel_size * cols[i],
        #                 y_min + voxel_size * rows[i], 
        #                 x_min + voxel_size * (cols[i] + 1),
        #                 y_min + voxel_size * (rows[i] + 1)]
        # points_in_voxel = mask_points_by_range(pcd_np, voxel_range)

        # if not points_in_voxel.any():
        #     continue    

        points_in_voxel = pcd_with_idx[pcd_with_idx[:,3]==(rows[i]*W + cols[i])]
        if not points_in_voxel.any():
            continue
        points_select = np.concatenate((points_select, points_in_voxel), axis=0)

    points_select = points_select[:,:3]
    
    return points_select


def get_keypoints(pcd_all_np, pcd_select_np, n_samples, mode = 'farthest'):
    if pcd_select_np.shape[0] >= n_samples:
        keypoints = get_random_subset(pcd_select_np, n_samples, mode)
    else:
        keypoints = get_random_subset(pcd_all_np, n_samples - pcd_select_np.shape[0], mode)
        keypoints = np.concatenate((keypoints, pcd_select_np), axis=0)

    return keypoints

def bev_sample(pcd_np, lidar_range, n_samples, mode, voxel_size=0.2, all_samples=False):
    """
    Args:
        pcd_np: 
            [N, 3] or [N, 4]
        lidar_range: 
            list len = 4 or len = 6, please use this to remove ground
        all_samples: 
            if True, not use n_samples to subsampling
    Returns:
        keypoints: np.ndarray
            [n_samples, 3]
    """

    pcd_np = pcd_np[:,:3]
    bev_img, pcd_with_idx = project_bev(pcd_np, lidar_range, voxel_size)
    lines = line_detection(bev_img)
    rows, cols = np.nonzero(lines)
    points_select = get_point_in_voxels(pcd_np, rows, cols, lidar_range, voxel_size, pcd_with_idx)

    if all_samples:
        keypoints = points_select
    else:
        keypoints = get_keypoints(pcd_np, points_select, n_samples, mode)

    print(keypoints.shape)

    return keypoints

def seq_generate():
    dirs = ["/GPFS/rhome/yifanlu/workspace/dataset/OPV2V/train",
            "/GPFS/rhome/yifanlu/workspace/dataset/OPV2V/validate"
            "/GPFS/rhome/yifanlu/workspace/dataset/OPV2V/test"]


    kp_store_path = '/GPFS/rhome/yifanlu/workspace/OpenCOOD/keypoints_file/bev_keypoints'
    lidar_range = [-140, -80, -1.5, 140, 80, 1]
    n_samples = 1500

    import os
    import opencood.utils.pcd_utils as pcd_utils

    for root_dir in dirs:
        scenario_folders = sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name = sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        # Structure: {scenario_id : {cav_1 : {timestamp1 : {yaml: path,
        # lidar: path, cameras:list of path}}}}

        # loop over all scenarios
        for (i, scenario_folder) in enumerate(scenario_folders):
            # at least 1 cav should show up
            cav_list = sorted([x for x in os.listdir(scenario_folder)
                               if os.path.isdir(
                    os.path.join(scenario_folder, x))])
            assert len(cav_list) > 0

            # loop over all CAV data
            for (j, cav_id) in enumerate(cav_list):

                # save all yaml files to the dictionary
                cav_path = os.path.join(scenario_folder, cav_id)

                yaml_files = \
                    sorted([os.path.join(cav_path, x)
                            for x in os.listdir(cav_path) if
                            x.endswith('.yaml')])
                timestamps = []

                # extract timestamp
                for file in yaml_files:
                    res = file.split('/')[-1]

                    timestamp = res.replace('.yaml', '')
                    timestamps.append(timestamp)


                for timestamp in timestamps:

                    yaml_file = os.path.join(cav_path,
                                                timestamp + '.yaml')
                    lidar_file = os.path.join(cav_path,
                                                timestamp + '.pcd')

                    # when init the dataset, it read over all pcd files.
                    # it maybe slow, but no need to perform keypoint sampling for each time.\
                    kp_path = f"{kp_store_path}/{scenario_folders_name[i]}/{cav_id}/{timestamp}.npy"
                    kp_dir = kp_path.rsplit('/',1)[0] # before filename

                    if not os.path.exists(kp_dir):
                        os.makedirs(kp_dir)

                    if not os.path.exists(kp_path):
                        pcd_np = pcd_utils.pcd_to_np(lidar_file)
                        kp_file = bev_sample(pcd_np,
                                    lidar_range,
                                    n_samples,
                                    mode='uniform',
                                    all_samples=True)

                        np.save(kp_path, kp_file)


def parallel_generate(scenario_folder, scenario_folder_name):

    kp_store_path = '/GPFS/rhome/yifanlu/workspace/OpenCOOD/keypoints_file/bev_keypoints'
    lidar_range = [-140, -80, -1.5, 140, 80, 1]

    cav_list = sorted([x for x in os.listdir(scenario_folder)
                        if os.path.isdir(
            os.path.join(scenario_folder, x))])
    assert len(cav_list) > 0
    print(cav_list)

    # loop over all CAV data
    for (j, cav_id) in enumerate(cav_list):
        print(cav_id)
        # save all yaml files to the dictionary
        cav_path = os.path.join(scenario_folder, cav_id)

        yaml_files = \
            sorted([os.path.join(cav_path, x)
                    for x in os.listdir(cav_path) if
                    x.endswith('.yaml')])
        timestamps = []

        # extract timestamp
        for file in yaml_files:
            res = file.split('/')[-1]

            timestamp = res.replace('.yaml', '')
            timestamps.append(timestamp)


        for timestamp in timestamps:

            yaml_file = os.path.join(cav_path,
                                        timestamp + '.yaml')
            lidar_file = os.path.join(cav_path,
                                        timestamp + '.pcd')

            # when init the dataset, it read over all pcd files.
            # it maybe slow, but no need to perform keypoint sampling for each time.\
            target = [250,500,750,1000,1250,1500,2000,2500]
            kp_paths = [f"{kp_store_path}/{scenario_folder_name}/{cav_id}/{timestamp}.npy"]
            kp_paths += [f"{kp_store_path}_{n_samples}/{scenario_folder_name}/{cav_id}/{timestamp}.npy" for n_samples in target]
            flag = True
            for kp_path in kp_paths:
                if not os.path.exists(kp_path):
                    flag = False
            if flag:
                continue


            pcd_np = pcd_utils.pcd_to_np(lidar_file)[:,:3]

            all_keypoint = bev_sample(pcd_np,
                                lidar_range,
                                np.inf,
                                mode='uniform',
                                all_samples=True)

            kp_path = f"{kp_store_path}/{scenario_folder_name}/{cav_id}/{timestamp}.npy"
            kp_dir = kp_path.rsplit('/',1)[0] # before filename
            if not os.path.exists(kp_dir):
                os.makedirs(kp_dir)

            if not os.path.exists(kp_path):
                np.save(kp_path, all_keypoint)
                print(f"saving to {kp_path}")


            for n_samples in target:
                kp_path = f"{kp_store_path}_{n_samples}/{scenario_folder_name}/{cav_id}/{timestamp}.npy"
                kp_dir = kp_path.rsplit('/',1)[0] # before filename

                if not os.path.exists(kp_dir):
                    os.makedirs(kp_dir)
                
                select_keypoint = get_keypoints(pcd_np, all_keypoint, n_samples)

                if not os.path.exists(kp_path):
                    np.save(kp_path, select_keypoint)
                    print(f"saving to {kp_path}")




if __name__=="__main__":
    dirs = ["/GPFS/rhome/yifanlu/workspace/dataset/OPV2V/train",
            "/GPFS/rhome/yifanlu/workspace/dataset/OPV2V/validate",
            "/GPFS/rhome/yifanlu/workspace/dataset/OPV2V/test"]

    import os
    import opencood.utils.pcd_utils as pcd_utils

    scenario_folders = []
    scenario_folders_name = []

    for root_dir in dirs:
        scenario_folders += sorted([os.path.join(root_dir, x)
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
        scenario_folders_name += sorted([x
                                   for x in os.listdir(root_dir) if
                                   os.path.isdir(os.path.join(root_dir, x))])
    
    

    scenario_folders = ['/GPFS/rhome/yifanlu/workspace/OpenCOOD/dataset_link/validate/2021_08_21_17_30_41']
    scenario_folders_name = ['2021_08_21_17_30_41']
    num = len(scenario_folders)

    for i in range(num):
        p = Process(target=parallel_generate, args=(scenario_folders[i],scenario_folders_name[i]))
        p.start()
opencood/utils/keypoint_utils.py
MediaBrain-SJTU/CoBEVFlow
import torch
import torch.nn as nn

def load_model_dict(model, pretrained_dict):
    """ load pretrained state dict, keys may not match with model

    Args:
        model: nn.Module

        pretrained_dict: collections.OrderedDict
    
    """
    # 1. filter out unnecessary keys
    model_dict = model.state_dict()
    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
    # 2. overwrite entries in the existing state dict
    model_dict.update(pretrained_dict)
    # 3. load the new state dict
    model.load_state_dict(model_dict)
    return model


def weight_init(m):
    if isinstance(m, nn.Linear):
        nn.init.xavier_normal_(m.weight.data, gain=0.1)
        if hasattr(m.bias, 'data'):
            nn.init.constant_(m.bias.data, 0)

    elif isinstance(m, nn.Conv2d):
        nn.init.xavier_normal_(m.weight, gain=0.1)
        # if hasattr(m, 'bias'):
        #     nn.init.constant_(m.bias, 0)

    # elif isinstance(m, nn.BatchNorm2d):
    #     nn.init.xavier_normal_(m.weight, gain=0.05)
    #     nn.init.constant_(m.bias, 0)


opencood/utils/model_utils.py
MediaBrain-SJTU/CoBEVFlow
# transfer all lidar file from .pcd to .npy 
# sizhewei @ <sizhewei@sjtu.edu.cn>
# 2022/12/20

###################
# $ cd ~/percp/OpenCOOD/
# python opencood/utils/my_tools.py
####################

import os
import sys
sys.path.append(os.getcwd())
import shutil
# import open3d as o3d
import numpy as np
import time
import json
import yaml
from pypcd import pypcd
from tqdm import tqdm
from tqdm.contrib import tenumerate
from tqdm.auto import trange

import matplotlib.pyplot as plt

from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import pcd_to_np

def extract_timestamps(yaml_files):
    """
    Given the list of the yaml files, extract the mocked timestamps.

    Parameters
    ----------
    yaml_files : list
        The full path of all yaml files of ego vehicle

    Returns
    -------
    timestamps : list
        The list containing timestamps only.
    """
    timestamps = []

    for file in yaml_files:
        res = file.split('/')[-1]

        timestamp = res.replace('.yaml', '')
        timestamps.append(timestamp)

    return timestamps

def irv2v_to_npy_and_json(split_name="train", root_path="/remote-home/share/V2XSET"):
    # 处理的数据集
    data_seperate_name = split_name
    # pcd 文件路径
    root_dir = root_path + data_seperate_name
    # yaml 文件路径
    # root_annotation_dir = "/GPFS/public/OPV2V_Irregular_V2/dataset_irregular_v2/" + data_seperate_name
    # 新保存的路径
    save_dir = root_path + "_npy/" + data_seperate_name

    # first load all paths of different scenarios
    scenario_folders = sorted([os.path.join(root_dir, x)
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])
    scenario_folders_name = sorted([x
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])

    print("### Transformation Started! ###")

    # loop over all scenarios
    for (i, scenario_folder) in tenumerate(scenario_folders):
        start_time = time.time()
        # copy timestamps npy file
        timestamps_file = os.path.join(root_annotation_dir, scenario_folders_name[i], 'timestamps.npy')
        save_scene_folder = os.path.join(save_dir, scenario_folders_name[i])
        if not os.path.exists(save_scene_folder):
            os.makedirs(save_scene_folder)
        save_timestamps_file = os.path.join(save_scene_folder, 'timestamps.npy')
        if not os.path.exists(save_timestamps_file):
            shutil.copyfile(timestamps_file, save_timestamps_file)

        # iterate all cav in this scenario
        cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))])
        assert len(cav_list) > 0
        for (j, cav_id) in tenumerate(cav_list):
            cav_path_pcd = os.path.join(scenario_folder, cav_id)
            cav_path_yaml = os.path.join(root_annotation_dir, scenario_folders_name[i], cav_id)
            # use the frame number as key, the full path as the values
            yaml_files = \
                sorted([os.path.join(cav_path_yaml, x)
                        for x in os.listdir(cav_path_yaml) if
                        x.endswith('.yaml')])
            timestamps = extract_timestamps(yaml_files)	

            for k, timestamp in tenumerate(timestamps):
                # if k > 30:
                #     break

                # new save dir
                save_path = os.path.join(save_dir, scenario_folders_name[i], cav_id)
                if not os.path.exists(save_path):
                    os.makedirs(save_path)
                    # print("======Path {} created successfully!======".format(save_path))

                # change pcd to npy 
                lidar_file = os.path.join(cav_path_pcd, timestamp + '.pcd')
                new_lidar_file = os.path.join(save_path, timestamp + '.npy')
                if not os.path.exists(new_lidar_file):
                    lidar = pcd_to_np(lidar_file)
                    np.save(new_lidar_file, lidar)

                # change yaml to json
                yaml_file = os.path.join(cav_path_yaml, timestamp + '.yaml')
                json_file = os.path.join(cav_path_yaml, timestamp + '.json')
                new_json_file = os.path.join(save_path, timestamp + '.json')
                if os.path.exists(json_file): # 原数据集有json文件
                    if not os.path.exists(new_json_file):
                        shutil.copyfile(json_file, new_json_file)
                else:
                    if not os.path.exists(new_json_file):
                        params = load_yaml(yaml_file)
                        json_str = json.dumps(params, indent=4)
                        with open(new_json_file, 'w') as json_file:
                            json_file.write(json_str)
                
                # new_yaml_file = os.path.join(save_path, timestamp + '.yaml')
                # if not os.path.exists(new_yaml_file):
                #     shutil.copyfile(yaml_file, new_yaml_file)

                # print(timestamp)
        end_time = time.time()
        print("=== %d-th scenario (%s) finished, consumed %.3f mins. ===" \
            % (i, scenario_folders_name[i], (end_time-start_time)/60.0))

def v2xset_to_npy_and_json(split_name="train", root_path="/remote-home/share/V2XSET"):
    # 处理的数据集
    data_seperate_name = split_name
    # pcd 文件路径
    root_dir = root_path + data_seperate_name
    # yaml 文件路径
    # root_annotation_dir = "/GPFS/public/OPV2V_Irregular_V2/dataset_irregular_v2/" + data_seperate_name
    # 新保存的路径
    save_dir = root_path + "_npy/" + data_seperate_name

    # first load all paths of different scenarios
    scenario_folders = sorted([os.path.join(root_dir, x)
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])
    scenario_folders_name = sorted([x
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])

    print("### Transformation Started! ###")

    # loop over all scenarios
    for (i, scenario_folder) in tenumerate(scenario_folders):
        start_time = time.time()
        # copy timestamps npy file
        timestamps_file = os.path.join(root_annotation_dir, scenario_folders_name[i], 'timestamps.npy')
        save_scene_folder = os.path.join(save_dir, scenario_folders_name[i])
        if not os.path.exists(save_scene_folder):
            os.makedirs(save_scene_folder)
        save_timestamps_file = os.path.join(save_scene_folder, 'timestamps.npy')
        if not os.path.exists(save_timestamps_file):
            shutil.copyfile(timestamps_file, save_timestamps_file)

        # iterate all cav in this scenario
        cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))])
        assert len(cav_list) > 0
        for (j, cav_id) in tenumerate(cav_list):
            cav_path_pcd = os.path.join(scenario_folder, cav_id)
            cav_path_yaml = os.path.join(root_annotation_dir, scenario_folders_name[i], cav_id)
            # use the frame number as key, the full path as the values
            yaml_files = \
                sorted([os.path.join(cav_path_yaml, x)
                        for x in os.listdir(cav_path_yaml) if
                        x.endswith('.yaml')])
            timestamps = extract_timestamps(yaml_files)	

            for k, timestamp in tenumerate(timestamps):
                # if k > 30:
                #     break

                # new save dir
                save_path = os.path.join(save_dir, scenario_folders_name[i], cav_id)
                if not os.path.exists(save_path):
                    os.makedirs(save_path)
                    # print("======Path {} created successfully!======".format(save_path))

                # change pcd to npy 
                lidar_file = os.path.join(cav_path_pcd, timestamp + '.pcd')
                new_lidar_file = os.path.join(save_path, timestamp + '.npy')
                if not os.path.exists(new_lidar_file):
                    lidar = pcd_to_np(lidar_file)
                    np.save(new_lidar_file, lidar)

                # change yaml to json
                yaml_file = os.path.join(cav_path_yaml, timestamp + '.yaml')
                json_file = os.path.join(cav_path_yaml, timestamp + '.json')
                new_json_file = os.path.join(save_path, timestamp + '.json')
                if os.path.exists(json_file): # 原数据集有json文件
                    if not os.path.exists(new_json_file):
                        shutil.copyfile(json_file, new_json_file)
                else:
                    if not os.path.exists(new_json_file):
                        params = load_yaml(yaml_file)
                        json_str = json.dumps(params, indent=4)
                        with open(new_json_file, 'w') as json_file:
                            json_file.write(json_str)
                
                # new_yaml_file = os.path.join(save_path, timestamp + '.yaml')
                # if not os.path.exists(new_yaml_file):
                #     shutil.copyfile(yaml_file, new_yaml_file)

                # print(timestamp)
        end_time = time.time()
        print("=== %d-th scenario (%s) finished, consumed %.3f mins. ===" \
            % (i, scenario_folders_name[i], (end_time-start_time)/60.0))


def create_small_dataset(split_name='train', size_scene=4, size_cav=3, size_time=2):
    """
    TO BE FINISHED !!!
    Create a small dataset from original large one 
    """
    # split name
    data_seperate_name = split_name
    # 原数据的路径
    root_dir = "/GPFS/public/OPV2V_irregular_npy/" + data_seperate_name
    # 新保存的路径
    save_dir = "/GPFS/public/OPV2V_irregular_npy_small/" + data_seperate_name

    # first load all paths of different scenarios
    scenario_folders = sorted([os.path.join(root_dir, x)
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])
    scenario_folders_name = sorted([x
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])

    print("### Transformation Started! ###")

    # loop over all scenarios
    for (i, scenario_folder) in tenumerate(scenario_folders):
        if i==size_scene:
            break

# 用于单独保存子图的函数
def save_subfig(fig,ax,save_path,fig_name):
    bbox = ax.get_tightbbox(fig.canvas.get_renderer()).expanded(1.02, 1.02)
    extent = bbox.transformed(fig.dpi_scale_trans.inverted())
    fig.savefig(save_path+fig_name, bbox_inches=extent, dpi=600, format='pdf')


def export_legend(ax, filename="legend.pdf"):
    fig2 = plt.figure()
    ax2 = fig2.add_subplot()
    ax2.axis('off')
    legend = ax2.legend(*ax.get_legend_handles_labels(), frameon=False, loc='lower center', ncol=4,fontsize='xx-large')
    fig  = legend.figure
    fig.canvas.draw()
    bbox  = legend.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    fig.savefig(filename, dpi=600, bbox_inches=bbox)

if __name__ == "__main__":
    # to_npy_and_json(split_name='validate')
    # create_small_dataset(split_name='train')

    root_dirs = "/remote-home/share/sizhewei/logs"
    note = 'major_irv2v'
    format = 'jpg'
    save_path = f"./opencood/result_{note}.{format}"
    title = "Average Precision curves of different methods on the IRV2V dataset at different average time intervals."
    plt.style.use('ggplot')
    split_list = ['late', 'v2vnet', 'v2xvit', 'disconet', 'where2comm', 'where2comm_syncnet', 'cobevflow'] #
    colors = ['#999999', '#DEC68B', '#BEB8DC', '#82B0D2', '#8ECFC9', '#FFBE7A', 'red'] #'#FA7F6F'] #'#E7DAD2'
    # colors = ['lightskyblue', 'lightseagreen', 'tomato', 'orange', 'gray', 'purple']
    single_split_name = 'single'

    num_delay = 15

    max_x = 500 # unit is ms
    plt.figure()
    fig, ax = plt.subplots(1,2, sharex='col', sharey=False, figsize=(16,6.5))
    plt.subplots_adjust(wspace = 0.2)
    # fig.suptitle(f'{title}', fontsize='x-large', y=0.99)
    # fig.text(0.5, 0.03, 'Mean delay of the most recent frame(ms).', ha='center', fontsize='x-large')
    # fig.text(0.08, 0.5, 'AP@0.50', va='center', rotation='vertical', fontsize='x-large')
    # ax30 = ax[0]; 
    ax50 = ax[0]; ax70 = ax[1]
    
    # for single fusion
    method_name = 'Single'
    # eval_file = os.path.join(root_dirs, f'eval_{single_split_name}.yaml')
    eval_file = '/remote-home/share/sizhewei/logs/opv2v_late_fusion/eval_no_noise_single_delay_0.00.yaml'
    single_color = 'gray'
    with open(eval_file, "r", encoding="utf-8") as f:
        data = yaml.load(f, Loader=yaml.FullLoader)
    # plt.sca(ax30); plt.plot([0,max_x],[data['ap_30'],data['ap_30']], color=single_color, linestyle='--', label=method_name)
    plt.sca(ax50); plt.plot([0,max_x],[data['ap_50'],data['ap_50']], color=single_color, linestyle='--', linewidth=4, label=method_name)
    plt.sca(ax70); plt.plot([0,max_x],[data['ap_70'],data['ap_70']], color=single_color, linestyle='--', linewidth=4, label=method_name)

    for split_i, split_name in enumerate(split_list):
        ap_list = []
        delays = []
        if split_name == 'late':
            method_name = 'Late Fusion'
            file_name = 'opv2v_late_fusion'
            eval_name = 'late_delay'
        elif split_name == 'v2vnet':
            method_name = 'V2VNet'
            file_name = 'opv2v_v2vnet_32ch'
            eval_name = 'v2vnet_32ch_ep23'
        elif split_name == 'v2xvit':
            method_name = 'V2X-ViT'
            file_name = 'opv2v_point_pillar_v2xvit_v1_slren'
            eval_name = 'v2xvit'
        elif split_name == 'disconet':
            method_name = 'DiscoNet'
            file_name = 'opv2v_disconet_new_ylu'
            eval_name = 'disconet'
        elif split_name == 'where2comm':
            method_name = 'Where2comm'
            file_name = 'irv2v_where2comm_cobevflow_w_dir_finetune' #'irv2v_where2comm_max_multiscale_resnet'
            eval_name = 'where2comm'
        elif split_name == 'where2comm_syncnet':
            method_name = 'Where2comm + SyncNet'
            file_name = 'irv2v_where2comm_syncnet_regular_300_past_3_test' #'irv2v_where2comm_syncnet_wo_tm_test'
            eval_name = 'syncnet_e30' # 'syncnet_ir' # 'syncnet_e8'
        elif split_name == 'cobevflow':
            method_name = 'CoBEVFlow (ours)'
            file_name = 'irv2v_where2comm_cobevflow_w_dir_finetune' #'opv2v_where2comm_cobevflow_w_dir'
            eval_name = 'cobevflow_reverse'
        latest_time_delay = -100.00
        for i in tqdm(np.linspace(0, 0.5, 26)):
            if i < 0.1 and i != 0:
                continue
            # log_file = os.path.join(split_name, f"eval_{note}_%.1f.yaml" % i)
            eval_file = os.path.join(root_dirs, file_name, f"eval_no_noise_{eval_name}_%.2f.yaml" % i)
            # if split_name == 'cobevflow':
            #     eval_file = os.path.join('/remote-home/share/sizhewei/logs/irv2v_where2comm_cobevflow/eval_draw_curve', f"eval_no_noise_{eval_name}_%.2f.yaml" % i)
            
            if not os.path.exists(eval_file):
                print(f'eval file {eval_file} not exist!')
                continue
            with open(eval_file, "r", encoding="utf-8") as f:
                data = yaml.load(f, Loader=yaml.FullLoader)
            
            # try:
            #     unit_time_delay = -data['avg_time_delay']
            # except:
            unit_time_delay = i*1000

            if (unit_time_delay - latest_time_delay) < 30:
                continue
            latest_time_delay = unit_time_delay
            
            tmp_aps = []
            tmp_aps.append(data['ap_30'])
            tmp_aps.append(data['ap_50'])
            tmp_aps.append(data['ap_70'])
            ap_list.append(tmp_aps)

            delays.append(unit_time_delay)

        if len(ap_list) == 0:
            continue

        ap_list_np = np.array(ap_list)
        ap_list_np = np.transpose(ap_list_np)

        ap_30_list = list(ap_list_np[0])
        ap_50_list = list(ap_list_np[1])
        ap_70_list = list(ap_list_np[2])
        
        color = colors[split_i]
        # plt.sca(ax30); plt.plot(delays, ap_30_list, color=color, marker='+', label = method_name)
        plt.sca(ax50); plt.plot(delays, ap_50_list, color=color, linewidth='5', label = method_name)#, marker='+') # , marker='+'
        plt.sca(ax70); plt.plot(delays, ap_70_list, color=color, linewidth='5', label = method_name)#, marker='+') # , marker='+'

    xaxis = np.linspace(0, max_x, 6)
    # ax30.set_title('The Results for AP@0.3'); ax30.grid(True); 
    # ax30.set_xticks(xaxis); \
    #     ax30.legend(loc = 'lower left')
    yaxis = np.linspace(0.3, 0.9, 7)
    fontsize = 24
    # font = "Gill Sans MT Condensed"
    # ax50.set_title('The Results for AP@0.5'); 
    ax50.tick_params(axis='both', which='major', labelsize=22)
    ax50.grid(True); 
    ax50.set_xticks(xaxis); ax50.set_yticks(yaxis);
    ax50.set_xlabel('Expectation of time interval (ms)', fontsize=fontsize); ax50.set_ylabel('AP@0.50', fontsize=fontsize);
    # ax50.legend(loc = 'lower left')
    
    yaxis = np.linspace(0.2, 0.9, 8)
    # ax70.set_title('The Results for AP@0.7'); 
    ax70.grid(True); 
    ax70.tick_params(axis='both', which='major', labelsize=22)
    ax70.set_xticks(xaxis); ax70.set_yticks(yaxis); 
    ax70.set_xlabel('Expectation of time interval (ms)', fontsize=fontsize); ax70.set_ylabel('AP@0.70', fontsize=fontsize);
    # ax70.legendloc = 'lower left')

    save_subfig(fig,ax[0],'./opencood/','ir2v2ap50.pdf')
    save_subfig(fig,ax[1],'./opencood/','irv2vap70.pdf')

    # plt.legend(frameon=False)
    plt.savefig(save_path, dpi=600, format=format)
    print("=== Plt save finished!!! ===")

    export_legend(ax70, '/root/percp/OpenCOOD/opencood/legend.pdf')
    # # then create a new image
    # # adjust the figure size as necessary
    # figsize = (7, 7)
    # fig_leg = plt.figure(figsize=figsize)
    # ax_leg = fig_leg.add_subplot(111)
    # # add the legend from the previous axes
    # # ax70.legend(ncol=7)
    # ax_leg.legend(*ax70.get_legend_handles_labels(), loc='center')
    # # hide the axes frame and the x/y labels
    # ax_leg.axis('off')
    # fig_leg.savefig('/root/percp/OpenCOOD/opencood/legend.jpg')

    # plt.title('标题')
    # plt.show()

opencood/utils/my_tools.py
MediaBrain-SJTU/CoBEVFlow
# transfer all lidar file from .pcd to .npy 
# sizhewei @ <sizhewei@sjtu.edu.cn>
# 2022/12/20

###################
# $ cd ~/percp/OpenCOOD/
# python opencood/utils/my_tools.py
####################

import os
import sys
sys.path.append(os.getcwd())
import shutil
# import open3d as o3d
import numpy as np
import time
import json
import yaml
from pypcd import pypcd
from tqdm import tqdm
from tqdm.contrib import tenumerate
from tqdm.auto import trange

import matplotlib.pyplot as plt

from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import pcd_to_np

def extract_timestamps(yaml_files):
    """
    Given the list of the yaml files, extract the mocked timestamps.

    Parameters
    ----------
    yaml_files : list
        The full path of all yaml files of ego vehicle

    Returns
    -------
    timestamps : list
        The list containing timestamps only.
    """
    timestamps = []

    for file in yaml_files:
        res = file.split('/')[-1]

        timestamp = res.replace('.yaml', '')
        timestamps.append(timestamp)

    return timestamps

def to_npy_and_json(split_name="train"):
    # 处理的数据集
    data_seperate_name = split_name
    # pcd 文件路径
    root_dir = "/GPFS/public/OPV2V_Irregular_V2/" + data_seperate_name
    # yaml 文件路径
    root_annotation_dir = "/GPFS/public/OPV2V_Irregular_V2/dataset_irregular_v2/" + data_seperate_name
    # 新保存的路径
    save_dir = "/GPFS/public/OPV2V_irregular_npy/" + data_seperate_name

    # first load all paths of different scenarios
    scenario_folders = sorted([os.path.join(root_dir, x)
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])
    scenario_folders_name = sorted([x
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])

    print("### Transformation Started! ###")

    # loop over all scenarios
    for (i, scenario_folder) in tenumerate(scenario_folders):
        start_time = time.time()
        # copy timestamps npy file
        timestamps_file = os.path.join(root_annotation_dir, scenario_folders_name[i], 'timestamps.npy')
        save_scene_folder = os.path.join(save_dir, scenario_folders_name[i])
        if not os.path.exists(save_scene_folder):
            os.makedirs(save_scene_folder)
        save_timestamps_file = os.path.join(save_scene_folder, 'timestamps.npy')
        if not os.path.exists(save_timestamps_file):
            shutil.copyfile(timestamps_file, save_timestamps_file)

        # iterate all cav in this scenario
        cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))])
        assert len(cav_list) > 0
        for (j, cav_id) in tenumerate(cav_list):
            cav_path_pcd = os.path.join(scenario_folder, cav_id)
            cav_path_yaml = os.path.join(root_annotation_dir, scenario_folders_name[i], cav_id)
            # use the frame number as key, the full path as the values
            yaml_files = \
                sorted([os.path.join(cav_path_yaml, x)
                        for x in os.listdir(cav_path_yaml) if
                        x.endswith('.yaml')])
            timestamps = extract_timestamps(yaml_files)	

            for k, timestamp in tenumerate(timestamps):
                # if k > 30:
                #     break

                # new save dir
                save_path = os.path.join(save_dir, scenario_folders_name[i], cav_id)
                if not os.path.exists(save_path):
                    os.makedirs(save_path)
                    # print("======Path {} created successfully!======".format(save_path))

                # change pcd to npy 
                lidar_file = os.path.join(cav_path_pcd, timestamp + '.pcd')
                new_lidar_file = os.path.join(save_path, timestamp + '.npy')
                if not os.path.exists(new_lidar_file):
                    lidar = pcd_to_np(lidar_file)
                    np.save(new_lidar_file, lidar)

                # change yaml to json
                yaml_file = os.path.join(cav_path_yaml, timestamp + '.yaml')
                json_file = os.path.join(cav_path_yaml, timestamp + '.json')
                new_json_file = os.path.join(save_path, timestamp + '.json')
                if os.path.exists(json_file): # 原数据集有json文件
                    if not os.path.exists(new_json_file):
                        shutil.copyfile(json_file, new_json_file)
                else:
                    if not os.path.exists(new_json_file):
                        params = load_yaml(yaml_file)
                        json_str = json.dumps(params, indent=4)
                        with open(new_json_file, 'w') as json_file:
                            json_file.write(json_str)
                
                # new_yaml_file = os.path.join(save_path, timestamp + '.yaml')
                # if not os.path.exists(new_yaml_file):
                #     shutil.copyfile(yaml_file, new_yaml_file)

                # print(timestamp)
        end_time = time.time()
        print("=== %d-th scenario (%s) finished, consumed %.3f mins. ===" \
            % (i, scenario_folders_name[i], (end_time-start_time)/60.0))

def create_small_dataset(split_name='train', size_scene=4, size_cav=3, size_time=2):
    """
    TO BE FINISHED !!!
    Create a small dataset from original large one 
    """
    # split name
    data_seperate_name = split_name
    # 原数据的路径
    root_dir = "/GPFS/public/OPV2V_irregular_npy/" + data_seperate_name
    # 新保存的路径
    save_dir = "/GPFS/public/OPV2V_irregular_npy_small/" + data_seperate_name

    # first load all paths of different scenarios
    scenario_folders = sorted([os.path.join(root_dir, x)
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])
    scenario_folders_name = sorted([x
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])

    print("### Transformation Started! ###")

    # loop over all scenarios
    for (i, scenario_folder) in tenumerate(scenario_folders):
        if i==size_scene:
            break

# 用于单独保存子图的函数
def save_subfig(fig,ax,save_path,fig_name):
    bbox = ax.get_tightbbox(fig.canvas.get_renderer()).expanded(1.02, 1.02)
    extent = bbox.transformed(fig.dpi_scale_trans.inverted())
    fig.savefig(save_path+fig_name, bbox_inches=extent, dpi=600, format='pdf')


def export_legend(ax, filename="legend.pdf"):
    fig2 = plt.figure()
    ax2 = fig2.add_subplot()
    ax2.axis('off')
    legend = ax2.legend(*ax.get_legend_handles_labels(), frameon=False, loc='lower center', ncol=8,)
    fig  = legend.figure
    fig.canvas.draw()
    bbox  = legend.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    fig.savefig(filename, dpi=600, bbox_inches=bbox)

if __name__ == "__main__":
    # to_npy_and_json(split_name='validate')
    # create_small_dataset(split_name='train')

    root_dirs = "/remote-home/share/sizhewei/logs"
    note = 'major_irv2v'
    format = 'jpg'
    save_path = f"./opencood/result_{note}.{format}"
    title = "Average Precision curves of different methods on the IRV2V dataset at different average time intervals."
    plt.style.use('ggplot')
    split_list = ['late', 'v2vnet', 'v2xvit', 'disconet', 'where2comm', 'where2comm_syncnet', 'cobevflow'] #
    colors = ['#999999', '#DEC68B', '#BEB8DC', '#82B0D2', '#8ECFC9', '#FFBE7A', 'red'] #'#FA7F6F'] #'#E7DAD2'
    # colors = ['lightskyblue', 'lightseagreen', 'tomato', 'orange', 'gray', 'purple']
    single_split_name = 'single'

    num_delay = 15

    max_x = 500 # unit is ms
    plt.figure()
    fig, ax = plt.subplots(1,4, sharex='col', sharey=False, figsize=(90, 16))
    plt.subplots_adjust(wspace = 0.4)
    # fig.suptitle(f'{title}', fontsize='x-large', y=0.99)
    # fig.text(0.5, 0.03, 'Mean delay of the most recent frame(ms).', ha='center', fontsize='x-large')
    # fig.text(0.08, 0.5, 'AP@0.50', va='center', rotation='vertical', fontsize='x-large')
    # ax30 = ax[0]; 
    ax50 = ax[0]; ax70 = ax[1]; dax50 = ax[2]; dax70 = ax[3]
    linewidth = '16'
    fontsize = 60
    # for single fusion
    method_name = 'Single'
    # eval_file = os.path.join(root_dirs, f'eval_{single_split_name}.yaml')
    eval_file = '/remote-home/share/sizhewei/logs/opv2v_late_fusion/eval_no_noise_single_delay_0.00.yaml'
    single_color = 'gray'
    with open(eval_file, "r", encoding="utf-8") as f:
        data = yaml.load(f, Loader=yaml.FullLoader)
    # plt.sca(ax30); plt.plot([0,max_x],[data['ap_30'],data['ap_30']], color=single_color, linestyle='--', label=method_name)
    plt.sca(ax50); plt.plot([0,max_x],[data['ap_50'],data['ap_50']], color=single_color, linestyle='--', linewidth=linewidth, label=method_name)
    plt.sca(ax70); plt.plot([0,max_x],[data['ap_70'],data['ap_70']], color=single_color, linestyle='--', linewidth=linewidth, label=method_name)

    for split_i, split_name in enumerate(split_list):
        ap_list = []
        delays = []
        if split_name == 'late':
            method_name = 'Late Fusion'
            file_name = 'opv2v_late_fusion'
            eval_name = 'late_delay'
        elif split_name == 'v2vnet':
            method_name = 'V2VNet'
            file_name = 'opv2v_v2vnet_32ch'
            eval_name = 'v2vnet_32ch_ep23'
        elif split_name == 'v2xvit':
            method_name = 'V2X-ViT'
            file_name = 'opv2v_point_pillar_v2xvit_v1_slren'
            eval_name = 'v2xvit'
        elif split_name == 'disconet':
            method_name = 'DiscoNet'
            file_name = 'opv2v_disconet_new_ylu'
            eval_name = 'disconet'
        elif split_name == 'where2comm':
            method_name = 'Where2comm'
            file_name = 'irv2v_where2comm_cobevflow_w_dir_finetune' #'irv2v_where2comm_max_multiscale_resnet'
            eval_name = 'where2comm'
        elif split_name == 'where2comm_syncnet':
            method_name = 'Where2comm + SyncNet'
            file_name = 'irv2v_where2comm_syncnet_regular_300_past_3_test' #'irv2v_where2comm_syncnet_wo_tm_test'
            eval_name = 'syncnet_e30' # 'syncnet_ir' # 'syncnet_e8'
        elif split_name == 'cobevflow':
            method_name = 'CoBEVFlow (ours)'
            file_name = 'irv2v_where2comm_cobevflow_w_dir_finetune' #'opv2v_where2comm_cobevflow_w_dir'
            eval_name = 'cobevflow_reverse'
        latest_time_delay = -100.00
        for i in tqdm(np.linspace(0, 0.5, 26)):
            if i < 0.1 and i != 0:
                continue
            # log_file = os.path.join(split_name, f"eval_{note}_%.1f.yaml" % i)
            eval_file = os.path.join(root_dirs, file_name, f"eval_no_noise_{eval_name}_%.2f.yaml" % i)
            # if split_name == 'cobevflow':
            #     eval_file = os.path.join('/remote-home/share/sizhewei/logs/irv2v_where2comm_cobevflow/eval_draw_curve', f"eval_no_noise_{eval_name}_%.2f.yaml" % i)
            
            if not os.path.exists(eval_file):
                print(f'eval file {eval_file} not exist!')
                continue
            with open(eval_file, "r", encoding="utf-8") as f:
                data = yaml.load(f, Loader=yaml.FullLoader)
            
            # try:
            #     unit_time_delay = -data['avg_time_delay']
            # except:
            unit_time_delay = i*1000

            if (unit_time_delay - latest_time_delay) < 30:
                continue
            latest_time_delay = unit_time_delay
            
            tmp_aps = []
            tmp_aps.append(data['ap_30'])
            tmp_aps.append(data['ap_50'])
            tmp_aps.append(data['ap_70'])
            ap_list.append(tmp_aps)

            delays.append(unit_time_delay)

        if len(ap_list) == 0:
            continue

        ap_list_np = np.array(ap_list)
        ap_list_np = np.transpose(ap_list_np)

        ap_30_list = list(ap_list_np[0])
        ap_50_list = list(ap_list_np[1])
        ap_70_list = list(ap_list_np[2])
        
        color = colors[split_i]
        # plt.sca(ax30); plt.plot(delays, ap_30_list, color=color, marker='+', label = method_name)
        plt.sca(ax50); plt.plot(delays, ap_50_list, color=color, linewidth=linewidth, label = method_name)#, marker='+') # , marker='+'
        plt.sca(ax70); plt.plot(delays, ap_70_list, color=color, linewidth=linewidth, label = method_name)#, marker='+') # , marker='+'

    xaxis = np.linspace(0, max_x, 6)
    # ax30.set_title('The Results for AP@0.3'); ax30.grid(True); 
    # ax30.set_xticks(xaxis); \
    #     ax30.legend(loc = 'lower left')
    yaxis = np.linspace(0.3, 0.9, 7)
    # font = "Gill Sans MT Condensed"
    # ax50.set_title('The Results for AP@0.5'); 
    ax50.tick_params(axis='both', which='major', labelsize=40)
    ax50.grid(True); 
    ax50.set_xticks(xaxis); ax50.set_yticks(yaxis);
    ax50.set_xlabel('Expectation of time interval (ms)', fontsize=fontsize, color='black'); ax50.set_ylabel('AP@0.50', fontsize=fontsize, color='black');
    # ax50.legend(loc = 'lower left')
    
    yaxis = np.linspace(0.2, 0.9, 8)
    # ax70.set_title('The Results for AP@0.7'); 
    ax70.grid(True); 
    ax70.tick_params(axis='both', which='major', labelsize=40)
    ax70.set_xticks(xaxis); ax70.set_yticks(yaxis); 
    ax70.set_xlabel('Expectation of time interval (ms)', fontsize=fontsize, color='black'); ax70.set_ylabel('AP@0.70', fontsize=fontsize, color='black');
    # ax70.legendloc = 'lower left')

    # for single fusion
    method_name = 'Single'
    # eval_file = os.path.join(root_dirs, f'eval_{single_split_name}.yaml')
    eval_file = '/remote-home/share/sizhewei/logs/dairv2x_late_fusion_new_ylu/eval_no_noise_no.yaml'
    single_color = 'gray'
    with open(eval_file, "r", encoding="utf-8") as f:
        data = yaml.load(f, Loader=yaml.FullLoader)
    # plt.sca(ax30); plt.plot([0,max_x],[data['ap_30'],data['ap_30']], color=single_color, linestyle='--', label=method_name)
    plt.sca(dax50); plt.plot([0,max_x],[data['ap_50'],data['ap_50']], color=single_color, linestyle='--', linewidth=linewidth, label=method_name)
    plt.sca(dax70); plt.plot([0,max_x],[data['ap_70'],data['ap_70']], color=single_color, linestyle='--', linewidth=linewidth, label=method_name)

    for split_i, split_name in enumerate(split_list):
        ap_list = []
        delays = []
        if split_name == 'late':
            method_name = 'Late Fusion'
            file_name = 'dairv2x_late_fusion_new_ylu'
            eval_name = 'late'
        elif split_name == 'v2vnet':
            method_name = 'V2VNet'
            file_name = 'dairv2x_v2vnet'
            eval_name = 'v2vnet_range_30'
        elif split_name == 'v2xvit':
            method_name = 'V2X-ViT'
            file_name = 'dairv2x_v2xvit_new_ylu'
            eval_name = 'v2xvit_epoch17'
        elif split_name == 'disconet':
            method_name = 'DiscoNet'
            file_name = 'dairv2x_disconet_sren'
            eval_name = 'disconet'
        elif split_name == 'where2comm':
            method_name = 'Where2comm'
            file_name = 'dairv2x_where2comm_baseline'
            eval_name = 'where2comm'
        elif split_name == 'where2comm_syncnet':
            method_name = 'Where2comm + SyncNet'
            file_name = 'dariv2x_where2comm_syncnet'
            eval_name = 'syncnet_e4'
        elif split_name == 'cobevflow':
            method_name = 'CoBEVFlow (ours)'
            file_name = 'dairv2x_where2comm_cobevflow'
            eval_name = 'cobevflow'
        latest_time_delay = -100.00
        for i in tqdm(np.linspace(0, 0.5, 26)):
            if i < 0.1 and i != 0:
                continue
            # log_file = os.path.join(split_name, f"eval_{note}_%.1f.yaml" % i)
            eval_file = os.path.join(root_dirs, file_name, f"eval_no_noise_{eval_name}_%.2f.yaml" % i)
            # if split_name == 'cobevflow':
            #     eval_file = os.path.join('/remote-home/share/sizhewei/logs/irv2v_where2comm_cobevflow/eval_draw_curve', f"eval_no_noise_{eval_name}_%.2f.yaml" % i)
            
            if not os.path.exists(eval_file):
                print(f'eval file {eval_file} not exist!')
                continue
            with open(eval_file, "r", encoding="utf-8") as f:
                data = yaml.load(f, Loader=yaml.FullLoader)
            
            # try:
            #     unit_time_delay = -data['avg_time_delay']
            # except:
            unit_time_delay = i*1000

            if (unit_time_delay - latest_time_delay) < 30:
                continue
            latest_time_delay = unit_time_delay
            
            tmp_aps = []
            tmp_aps.append(data['ap_30'])
            tmp_aps.append(data['ap_50'])
            tmp_aps.append(data['ap_70'])
            ap_list.append(tmp_aps)

            delays.append(unit_time_delay)

        if len(ap_list) == 0:
            continue

        ap_list_np = np.array(ap_list)
        ap_list_np = np.transpose(ap_list_np)

        ap_30_list = list(ap_list_np[0])
        ap_50_list = list(ap_list_np[1])
        ap_70_list = list(ap_list_np[2])
        
        color = colors[split_i]
        # plt.sca(ax30); plt.plot(delays, ap_30_list, color=color, marker='+', label = method_name)
        plt.sca(dax50); plt.plot(delays, ap_50_list, color=color, linewidth=linewidth, label = method_name)#, marker='+') # 
        plt.sca(dax70); plt.plot(delays, ap_70_list, color=color, linewidth=linewidth, label = method_name)#, marker='+') # , marker='+'

    xaxis = np.linspace(0, max_x, 6)
    # ax30.set_title('The Results for AP@0.3'); ax30.grid(True); 
    # ax30.set_xticks(xaxis); \
    #     ax30.legend(loc = 'lower left')
    yaxis = np.linspace(0.60, 0.80, 5)
    # ax50.set_title('The Results for AP@0.5'); 
    dax50.grid(True); 
    dax50.tick_params(axis='both', which='major', labelsize=40)
    dax50.set_xticks(xaxis); dax50.set_yticks(yaxis); 
    dax50.set_xlabel('Expectation of time interval (ms)', fontsize=fontsize, color='black'); dax50.set_ylabel('AP@0.50', fontsize=fontsize, color='black');
    # ax50.legend(ncol=2, loc = 'best') # ax50.legend(loc = 'lower left')
    
    yaxis = np.linspace(0.45, 0.65, 5)
    # ax70.set_title('The Results for AP@0.7'); 
    dax70.grid(True); 
    dax70.tick_params(axis='both', which='major', labelsize=40)
    dax70.set_xticks(xaxis); dax70.set_yticks(yaxis); 
    dax70.set_xlabel('Expectation of time interval (ms)', fontsize=fontsize, color='black'); dax70.set_ylabel('AP@0.70', fontsize=fontsize, color='black');
    
    save_subfig(fig,ax[0],'./opencood/','ir2v2ap50.pdf')
    save_subfig(fig,ax[1],'./opencood/','irv2vap70.pdf')
    save_subfig(fig,ax[2],'./opencood/','dairv2xap50.pdf')
    save_subfig(fig,ax[3],'./opencood/','dairv2xap70.pdf')

    # plt.legend(frameon=False)
    plt.savefig(save_path, dpi=600, format=format)
    print("=== Plt save finished!!! ===")

    export_legend(ax70, '/root/percp/OpenCOOD/opencood/legend.pdf')
    # # then create a new image
    # # adjust the figure size as necessary
    # figsize = (7, 7)
    # fig_leg = plt.figure(figsize=figsize)
    # ax_leg = fig_leg.add_subplot(111)
    # # add the legend from the previous axes
    # # ax70.legend(ncol=7)
    # ax_leg.legend(*ax70.get_legend_handles_labels(), loc='center')
    # # hide the axes frame and the x/y labels
    # ax_leg.axis('off')
    # fig_leg.savefig('/root/percp/OpenCOOD/opencood/legend.jpg')

    # plt.title('标题')
    # plt.show()

opencood/utils/my_tools_all.py
MediaBrain-SJTU/CoBEVFlow
# transfer all lidar file from .pcd to .npy 
# sizhewei @ <sizhewei@sjtu.edu.cn>
# 2022/12/20

###################
# $ cd ~/percp/OpenCOOD/
# python opencood/utils/my_tools.py
####################

import os
import sys
sys.path.append(os.getcwd())
import shutil
# import open3d as o3d
import numpy as np
import time
import json
import yaml
from pypcd import pypcd
from tqdm import tqdm
from tqdm.contrib import tenumerate
from tqdm.auto import trange

import matplotlib.pyplot as plt

from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import pcd_to_np

def extract_timestamps(yaml_files):
    """
    Given the list of the yaml files, extract the mocked timestamps.

    Parameters
    ----------
    yaml_files : list
        The full path of all yaml files of ego vehicle

    Returns
    -------
    timestamps : list
        The list containing timestamps only.
    """
    timestamps = []

    for file in yaml_files:
        res = file.split('/')[-1]

        timestamp = res.replace('.yaml', '')
        timestamps.append(timestamp)

    return timestamps

def to_npy_and_json(split_name="train"):
    # 处理的数据集
    data_seperate_name = split_name
    # pcd 文件路径
    root_dir = "/GPFS/public/OPV2V_Irregular_V2/" + data_seperate_name
    # yaml 文件路径
    root_annotation_dir = "/GPFS/public/OPV2V_Irregular_V2/dataset_irregular_v2/" + data_seperate_name
    # 新保存的路径
    save_dir = "/GPFS/public/OPV2V_irregular_npy/" + data_seperate_name

    # first load all paths of different scenarios
    scenario_folders = sorted([os.path.join(root_dir, x)
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])
    scenario_folders_name = sorted([x
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])

    print("### Transformation Started! ###")

    # loop over all scenarios
    for (i, scenario_folder) in tenumerate(scenario_folders):
        start_time = time.time()
        # copy timestamps npy file
        timestamps_file = os.path.join(root_annotation_dir, scenario_folders_name[i], 'timestamps.npy')
        save_scene_folder = os.path.join(save_dir, scenario_folders_name[i])
        if not os.path.exists(save_scene_folder):
            os.makedirs(save_scene_folder)
        save_timestamps_file = os.path.join(save_scene_folder, 'timestamps.npy')
        if not os.path.exists(save_timestamps_file):
            shutil.copyfile(timestamps_file, save_timestamps_file)

        # iterate all cav in this scenario
        cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))])
        assert len(cav_list) > 0
        for (j, cav_id) in tenumerate(cav_list):
            cav_path_pcd = os.path.join(scenario_folder, cav_id)
            cav_path_yaml = os.path.join(root_annotation_dir, scenario_folders_name[i], cav_id)
            # use the frame number as key, the full path as the values
            yaml_files = \
                sorted([os.path.join(cav_path_yaml, x)
                        for x in os.listdir(cav_path_yaml) if
                        x.endswith('.yaml')])
            timestamps = extract_timestamps(yaml_files)	

            for k, timestamp in tenumerate(timestamps):
                # if k > 30:
                #     break

                # new save dir
                save_path = os.path.join(save_dir, scenario_folders_name[i], cav_id)
                if not os.path.exists(save_path):
                    os.makedirs(save_path)
                    # print("======Path {} created successfully!======".format(save_path))

                # change pcd to npy 
                lidar_file = os.path.join(cav_path_pcd, timestamp + '.pcd')
                new_lidar_file = os.path.join(save_path, timestamp + '.npy')
                if not os.path.exists(new_lidar_file):
                    lidar = pcd_to_np(lidar_file)
                    np.save(new_lidar_file, lidar)

                # change yaml to json
                yaml_file = os.path.join(cav_path_yaml, timestamp + '.yaml')
                json_file = os.path.join(cav_path_yaml, timestamp + '.json')
                new_json_file = os.path.join(save_path, timestamp + '.json')
                if os.path.exists(json_file): # 原数据集有json文件
                    if not os.path.exists(new_json_file):
                        shutil.copyfile(json_file, new_json_file)
                else:
                    if not os.path.exists(new_json_file):
                        params = load_yaml(yaml_file)
                        json_str = json.dumps(params, indent=4)
                        with open(new_json_file, 'w') as json_file:
                            json_file.write(json_str)
                
                # new_yaml_file = os.path.join(save_path, timestamp + '.yaml')
                # if not os.path.exists(new_yaml_file):
                #     shutil.copyfile(yaml_file, new_yaml_file)

                # print(timestamp)
        end_time = time.time()
        print("=== %d-th scenario (%s) finished, consumed %.3f mins. ===" \
            % (i, scenario_folders_name[i], (end_time-start_time)/60.0))

def create_small_dataset(split_name='train', size_scene=4, size_cav=3, size_time=2):
    """
    TO BE FINISHED !!!
    Create a small dataset from original large one 
    """
    # split name
    data_seperate_name = split_name
    # 原数据的路径
    root_dir = "/GPFS/public/OPV2V_irregular_npy/" + data_seperate_name
    # 新保存的路径
    save_dir = "/GPFS/public/OPV2V_irregular_npy_small/" + data_seperate_name

    # first load all paths of different scenarios
    scenario_folders = sorted([os.path.join(root_dir, x)
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])
    scenario_folders_name = sorted([x
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])

    print("### Transformation Started! ###")

    # loop over all scenarios
    for (i, scenario_folder) in tenumerate(scenario_folders):
        if i==size_scene:
            break

# 用于单独保存子图的函数
def save_subfig(fig,ax,save_path,fig_name):
    bbox = ax.get_tightbbox(fig.canvas.get_renderer()).expanded(1.02, 1.02)
    extent = bbox.transformed(fig.dpi_scale_trans.inverted())
    fig.savefig(save_path+fig_name, bbox_inches=extent, dpi=600, format='pdf')

if __name__ == "__main__":
    # to_npy_and_json(split_name='validate')
    # create_small_dataset(split_name='train')

    root_dirs = "/remote-home/share/sizhewei/logs"
    note = 'major_dairv2x'
    format = 'jpg'
    save_path = f"./opencood/result_{note}.{format}"
    title = "Average Precision curves of different methods on the DAIR-V2X dataset at different average time intervals."
    plt.style.use('ggplot')
    split_list = ['late', 'v2vnet', 'v2xvit', 'disconet', 'where2comm', 'where2comm_syncnet', 'cobevflow'] #
    colors = ['#999999', '#DEC68B', '#BEB8DC', '#82B0D2', '#8ECFC9', '#FFBE7A', 'red'] #'#FA7F6F'] # '#E7DAD2',
    # colors = ['lightskyblue', 'lightseagreen', 'tomato', 'orange', 'gray', 'purple']
    single_split_name = 'single'

    num_delay = 15

    max_x = 500 # unit is ms
    plt.figure()
    fig, ax = plt.subplots(1,2, sharex='col', sharey=False, figsize=(16,6.5))
    plt.subplots_adjust(wspace = 0.5)
    # fig.suptitle(f'{title}', fontsize='x-large', y=0.99)
    # fig.text(0.5, 0.03, 'Expectation of time intervals (ms).', ha='center', fontsize='x-large')
    # fig.text(0.08, 0.5, 'AP', va='center', rotation='vertical', fontsize='x-large')
    # ax30 = ax[0]; 
    ax50 = ax[0]; ax70 = ax[1]

    # for single fusion
    method_name = 'Single'
    # eval_file = os.path.join(root_dirs, f'eval_{single_split_name}.yaml')
    eval_file = '/remote-home/share/sizhewei/logs/dairv2x_late_fusion_new_ylu/eval_no_noise_no.yaml'
    single_color = 'gray'
    with open(eval_file, "r", encoding="utf-8") as f:
        data = yaml.load(f, Loader=yaml.FullLoader)
    # plt.sca(ax30); plt.plot([0,max_x],[data['ap_30'],data['ap_30']], color=single_color, linestyle='--', label=method_name)
    plt.sca(ax50); plt.plot([0,max_x],[data['ap_50'],data['ap_50']], color=single_color, linestyle='--', linewidth=4, label=method_name)
    plt.sca(ax70); plt.plot([0,max_x],[data['ap_70'],data['ap_70']], color=single_color, linestyle='--', linewidth=4, label=method_name)

    for split_i, split_name in enumerate(split_list):
        ap_list = []
        delays = []
        if split_name == 'late':
            method_name = 'Late Fusion'
            file_name = 'dairv2x_late_fusion_new_ylu'
            eval_name = 'late'
        elif split_name == 'v2vnet':
            method_name = 'V2VNet'
            file_name = 'dairv2x_v2vnet'
            eval_name = 'v2vnet_range_30'
        elif split_name == 'v2xvit':
            method_name = 'V2X-ViT'
            file_name = 'dairv2x_v2xvit_new_ylu'
            eval_name = 'v2xvit_epoch17'
        elif split_name == 'disconet':
            method_name = 'DiscoNet'
            file_name = 'dairv2x_disconet_sren'
            eval_name = 'disconet'
        elif split_name == 'where2comm':
            method_name = 'Where2comm'
            file_name = 'dairv2x_where2comm_baseline'
            eval_name = 'where2comm'
        elif split_name == 'where2comm_syncnet':
            method_name = 'Where2comm + SyncNet'
            file_name = 'dariv2x_where2comm_syncnet'
            eval_name = 'syncnet_e4'
        elif split_name == 'cobevflow':
            method_name = 'CoBEVFlow (ours)'
            file_name = 'dairv2x_where2comm_cobevflow'
            eval_name = 'cobevflow'
        latest_time_delay = -100.00
        for i in tqdm(np.linspace(0, 0.5, 26)):
            if i < 0.1 and i != 0:
                continue
            # log_file = os.path.join(split_name, f"eval_{note}_%.1f.yaml" % i)
            eval_file = os.path.join(root_dirs, file_name, f"eval_no_noise_{eval_name}_%.2f.yaml" % i)
            # if split_name == 'cobevflow':
            #     eval_file = os.path.join('/remote-home/share/sizhewei/logs/irv2v_where2comm_cobevflow/eval_draw_curve', f"eval_no_noise_{eval_name}_%.2f.yaml" % i)
            
            if not os.path.exists(eval_file):
                print(f'eval file {eval_file} not exist!')
                continue
            with open(eval_file, "r", encoding="utf-8") as f:
                data = yaml.load(f, Loader=yaml.FullLoader)
            
            # try:
            #     unit_time_delay = -data['avg_time_delay']
            # except:
            unit_time_delay = i*1000

            if (unit_time_delay - latest_time_delay) < 30:
                continue
            latest_time_delay = unit_time_delay
            
            tmp_aps = []
            tmp_aps.append(data['ap_30'])
            tmp_aps.append(data['ap_50'])
            tmp_aps.append(data['ap_70'])
            ap_list.append(tmp_aps)

            delays.append(unit_time_delay)

        if len(ap_list) == 0:
            continue

        ap_list_np = np.array(ap_list)
        ap_list_np = np.transpose(ap_list_np)

        ap_30_list = list(ap_list_np[0])
        ap_50_list = list(ap_list_np[1])
        ap_70_list = list(ap_list_np[2])
        
        color = colors[split_i]
        # plt.sca(ax30); plt.plot(delays, ap_30_list, color=color, marker='+', label = method_name)
        plt.sca(ax50); plt.plot(delays, ap_50_list, color=color, linewidth='5', label = method_name)#, marker='+') # 
        plt.sca(ax70); plt.plot(delays, ap_70_list, color=color, linewidth='5', label = method_name)#, marker='+') # , marker='+'


    xaxis = np.linspace(0, max_x, 6)
    # ax30.set_title('The Results for AP@0.3'); ax30.grid(True); 
    # ax30.set_xticks(xaxis); \
    #     ax30.legend(loc = 'lower left')
    yaxis = np.linspace(0.60, 0.80, 5)
    fontsize = 24
    # ax50.set_title('The Results for AP@0.5'); 
    ax50.grid(True); 
    ax50.tick_params(axis='both', which='major', labelsize=22)
    ax50.set_xticks(xaxis); ax50.set_yticks(yaxis); 
    ax50.set_xlabel('Expectation of time interval (ms)', fontsize=fontsize); ax50.set_ylabel('AP@0.50', fontsize=fontsize);
    # ax50.legend(ncol=2, loc = 'best') # ax50.legend(loc = 'lower left')
    
    yaxis = np.linspace(0.45, 0.65, 5)
    # ax70.set_title('The Results for AP@0.7'); 
    ax70.grid(True); 
    ax70.tick_params(axis='both', which='major', labelsize=22)
    ax70.set_xticks(xaxis); ax70.set_yticks(yaxis); 
    ax70.set_xlabel('Expectation of time interval (ms)', fontsize=fontsize); ax70.set_ylabel('AP@0.70', fontsize=fontsize);
    # ax70.legend(ncol=2, loc = 'best') # ax70.legend(loc = 'lower left')

    plt.savefig(save_path, dpi=600, format=format)

    save_subfig(fig,ax[0],'./opencood/','dairv2xap50.pdf')
    save_subfig(fig,ax[1],'./opencood/','dairv2xap70.pdf')

    print("=== Plt save finished!!! ===")
    # plt.title('标题')
    # plt.show()

opencood/utils/my_tools_dairv2x.py
MediaBrain-SJTU/CoBEVFlow
# transfer all lidar file from .pcd to .npy 
# sizhewei @ <sizhewei@sjtu.edu.cn>
# 2022/12/20

###################
# $ cd ~/percp/OpenCOOD/
# python opencood/utils/my_tools.py
####################

import os
import sys
sys.path.append(os.getcwd())
import shutil
# import open3d as o3d
import numpy as np
import time
import json
import yaml
from pypcd import pypcd
from tqdm import tqdm
from tqdm.contrib import tenumerate
from tqdm.auto import trange

import matplotlib.pyplot as plt

from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import pcd_to_np

# 用于单独保存子图的函数
def save_subfig(fig,ax,save_path,fig_name):
    bbox = ax.get_tightbbox(fig.canvas.get_renderer()).expanded(1.02, 1.02)
    extent = bbox.transformed(fig.dpi_scale_trans.inverted())
    fig.savefig(save_path+fig_name, bbox_inches=extent, dpi=600, format='pdf')

if __name__ == "__main__":
    # to_npy_and_json(split_name='validate')
    # create_small_dataset(split_name='train')

    root_dirs = "/remote-home/share/sizhewei/logs"
    note = 'major_irv2v_irregular_f500'
    format = 'jpg'
    save_path = f"./opencood/result_{note}.{format}"
    title = "Average Precision curves of different methods on the IRV2V dataset at different average time intervals."
    
    split_list = ['where2comm_syncnet', 'cobevflow'] #['late', 'v2vnet', 'v2xvit', 'disconet', 'where2comm', 
    colors = ['#FFBE7A', '#FA7F6F'] #'#E7DAD2'['#999999', '#DEC68B', '#BEB8DC', '#82B0D2', '#8ECFC9', 
    # colors = ['lightskyblue', 'lightseagreen', 'tomato', 'orange', 'gray', 'purple']
    single_split_name = 'single'

    num_delay = 15

    max_x = 500 # unit is ms
    plt.figure()
    fig, ax = plt.subplots(1,2, sharex='col', sharey=False, figsize=(18,6))
    fig.suptitle(f'{title}', fontsize='x-large', y=0.99)
    # fig.text(0.5, 0.03, 'Mean delay of the most recent frame(ms).', ha='center', fontsize='x-large')
    # fig.text(0.08, 0.5, 'AP@0.50', va='center', rotation='vertical', fontsize='x-large')
    # ax30 = ax[0]; 
    ax50 = ax[0]; ax70 = ax[1]
    
    # for single fusion
    method_name = 'Single'
    # eval_file = os.path.join(root_dirs, f'eval_{single_split_name}.yaml')
    eval_file = '/remote-home/share/sizhewei/logs/opv2v_late_fusion/eval_no_noise_single_delay_0.00.yaml'
    single_color = 'red'
    with open(eval_file, "r", encoding="utf-8") as f:
        data = yaml.load(f, Loader=yaml.FullLoader)
    # plt.sca(ax30); plt.plot([0,max_x],[data['ap_30'],data['ap_30']], color=single_color, linestyle='--', label=method_name)
    plt.sca(ax50); plt.plot([100,max_x],[data['ap_50'],data['ap_50']], color=single_color, linestyle='--', label=method_name)
    plt.sca(ax70); plt.plot([100,max_x],[data['ap_70'],data['ap_70']], color=single_color, linestyle='--', label=method_name)

    for split_i, split_name in enumerate(split_list):
        ap_list = []
        delays = []
        if split_name == 'late':
            method_name = 'Late Fusion'
            file_name = 'opv2v_late_fusion'
            eval_name = 'late_delay'
        elif split_name == 'v2vnet':
            method_name = 'V2VNet'
            file_name = 'opv2v_v2vnet_32ch'
            eval_name = 'v2vnet_32ch_ep23'
        elif split_name == 'v2xvit':
            method_name = 'V2X-ViT'
            file_name = 'opv2v_point_pillar_v2xvit_v1_slren'
            eval_name = 'v2xvit'
        elif split_name == 'disconet':
            method_name = 'DiscoNet'
            file_name = 'opv2v_disconet_sren'
            eval_name = 'disconet'
        elif split_name == 'where2comm':
            method_name = 'Where2comm'
            file_name = 'irv2v_where2comm_cobevflow_w_dir_finetune' #'irv2v_where2comm_max_multiscale_resnet'
            eval_name = 'where2comm'
        elif split_name == 'where2comm_syncnet':
            method_name = 'Where2comm + SyncNet'
            file_name = 'irv2v_where2comm_syncnet_regular_100_past_2_test'
            eval_name =  'syncnet_ir_exp_f500' # 'syncnet_e8'
        elif split_name == 'cobevflow':
            method_name = 'CoBEVFlow (ours)'
            file_name = 'irv2v_where2comm_cobevflow_w_dir_finetune' #'opv2v_where2comm_cobevflow_w_dir'
            eval_name = 'cobevflow_ir_exp_f500'
        latest_time_delay = -500.00
        for i in tqdm(np.linspace(0.1, 0.9, 9)):
            # log_file = os.path.join(split_name, f"eval_{note}_%.1f.yaml" % i)
            eval_file = os.path.join(root_dirs, file_name, f"eval_no_noise_{eval_name}_%.2f.yaml" % i)
            
            if not os.path.exists(eval_file):
                print(f'eval file {eval_file} not exist!')
                continue
            with open(eval_file, "r", encoding="utf-8") as f:
                data = yaml.load(f, Loader=yaml.FullLoader)
            
            unit_time_delay = i*1000 - 500

            if (unit_time_delay - latest_time_delay) < 30:
                continue
            latest_time_delay = unit_time_delay
            
            tmp_aps = []
            tmp_aps.append(data['ap_30'])
            tmp_aps.append(data['ap_50'])
            tmp_aps.append(data['ap_70'])
            ap_list.append(tmp_aps)

            delays.append(unit_time_delay)

        if len(ap_list) == 0:
            continue

        ap_list_np = np.array(ap_list)
        ap_list_np = np.transpose(ap_list_np)

        ap_30_list = list(ap_list_np[0])
        ap_50_list = list(ap_list_np[1])
        ap_70_list = list(ap_list_np[2])
        
        color = colors[split_i]
        # plt.sca(ax30); plt.plot(delays, ap_30_list, color=color, marker='+', label = method_name)
        plt.sca(ax50); plt.plot(delays, ap_50_list, color=color, linewidth='3', label = method_name) # , marker='+'
        plt.sca(ax70); plt.plot(delays, ap_70_list, color=color, linewidth='3', label = method_name) # , marker='+'

    xaxis = np.linspace(-400, 400, 9)
    # ax30.set_title('The Results for AP@0.3'); ax30.grid(True); 
    # ax30.set_xticks(xaxis); \
    #     ax30.legend(loc = 'lower left')
    yaxis = np.linspace(0.5, 0.9, 5)
    # ax50.set_title('The Results for AP@0.5'); 
    ax50.grid(True); 
    ax50.set_xticks(xaxis); ax50.set_yticks(yaxis); 
    ax50.set_xlabel('Expectation time delay of the latest frame (ms)'); ax50.set_ylabel('AP@0.50');
    ax50.legend(loc = 'best', ncol=3)
    
    yaxis = np.linspace(0.4, 0.8, 5)
    # ax70.set_title('The Results for AP@0.7'); 
    ax70.grid(True); 
    ax70.set_xticks(xaxis); ax70.set_yticks(yaxis); 
    ax70.set_xlabel('Expectation time delay of the latest frame (ms)'); ax70.set_ylabel('AP@0.70');
    ax70.legend(loc = 'best', ncol=3)

    save_subfig(fig,ax[0],'./opencood/','ir2v2ap50-irr.pdf')
    save_subfig(fig,ax[1],'./opencood/','irv2vap70-irr.pdf')

    # plt.legend(frameon=False)
    plt.savefig(save_path, dpi=600, format=format)
    print("=== Plt save finished!!! ===")
    # plt.title('标题')
    # plt.show()

opencood/utils/my_tools_irregular.py
MediaBrain-SJTU/CoBEVFlow
# transfer all lidar file from .pcd to .npy 
# sizhewei @ <sizhewei@sjtu.edu.cn>
# 2022/12/20

###################
# $ cd ~/percp/OpenCOOD/
# python opencood/utils/my_tools.py
####################

import os
import sys
sys.path.append(os.getcwd())
import shutil
# import open3d as o3d
import numpy as np
import time
import json
import yaml
from pypcd import pypcd
from tqdm import tqdm
from tqdm.contrib import tenumerate
from tqdm.auto import trange

import matplotlib.pyplot as plt

from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import pcd_to_np

def extract_timestamps(yaml_files):
    """
    Given the list of the yaml files, extract the mocked timestamps.

    Parameters
    ----------
    yaml_files : list
        The full path of all yaml files of ego vehicle

    Returns
    -------
    timestamps : list
        The list containing timestamps only.
    """
    timestamps = []

    for file in yaml_files:
        res = file.split('/')[-1]

        timestamp = res.replace('.yaml', '')
        timestamps.append(timestamp)

    return timestamps

def v2xset_to_npy_and_json(split_name="train", root_path="/remote-home/share/V2XSET"):
    # 处理的数据集
    data_seperate_name = split_name
    # pcd / yaml 文件路径
    root_dir = os.path.join(root_path, data_seperate_name)
    # 新保存的路径
    save_dir = os.path.join(root_path + "_npy/", data_seperate_name)

    # first load all paths of different scenarios
    scenario_folders = sorted([os.path.join(root_dir, x)
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])
    scenario_folders_name = sorted([x
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])

    print("### Transformation Started! ###")

    # loop over all scenarios
    for (i, scenario_folder) in tenumerate(scenario_folders):
        start_time = time.time()
        # copy timestamps npy file
        # timestamps_file = os.path.join(root_annotation_dir, scenario_folders_name[i], 'timestamps.npy')
        save_scene_folder = os.path.join(save_dir, scenario_folders_name[i])
        if not os.path.exists(save_scene_folder):
            os.makedirs(save_scene_folder)
        
        # save_timestamps_file = os.path.join(save_scene_folder, 'timestamps.npy')
        # if not os.path.exists(save_timestamps_file):
        #     shutil.copyfile(timestamps_file, save_timestamps_file)

        # iterate all cav in this scenario
        cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))])
        assert len(cav_list) > 0
        for (j, cav_id) in tenumerate(cav_list):
            cav_path = os.path.join(scenario_folder, cav_id)
            # use the frame number as key, the full path as the values
            yaml_files = \
                sorted([os.path.join(cav_path, x)
                        for x in os.listdir(cav_path) if
                        x.endswith('.yaml')])
            timestamps = extract_timestamps(yaml_files)	

            for k, timestamp in tenumerate(timestamps):
                # new save dir
                save_path = os.path.join(save_dir, scenario_folders_name[i], cav_id)
                if not os.path.exists(save_path):
                    os.makedirs(save_path)
                    # print("======Path {} created successfully!======".format(save_path))

                # change pcd to npy 
                lidar_file = os.path.join(cav_path, timestamp + '.pcd')
                new_lidar_file = os.path.join(save_path, timestamp + '.npy')
                if not os.path.exists(new_lidar_file):
                    lidar = pcd_to_np(lidar_file)
                    np.save(new_lidar_file, lidar)

                # change yaml to json
                yaml_file = os.path.join(cav_path, timestamp + '.yaml')
                json_file = os.path.join(cav_path, timestamp + '.json')
                new_json_file = os.path.join(save_path, timestamp + '.json')
                if os.path.exists(json_file): # 原数据集有json文件
                    if not os.path.exists(new_json_file):
                        shutil.copyfile(json_file, new_json_file)
                else:
                    if not os.path.exists(new_json_file):
                        params = load_yaml(yaml_file)
                        json_str = json.dumps(params, indent=4)
                        with open(new_json_file, 'w') as json_file:
                            json_file.write(json_str)
                
                # new_yaml_file = os.path.join(save_path, timestamp + '.yaml')
                # if not os.path.exists(new_yaml_file):
                #     shutil.copyfile(yaml_file, new_yaml_file)

                # print(timestamp)
        end_time = time.time()
        print("=== %d-th scenario (%s) finished, consumed %.3f mins. ===" \
            % (i, scenario_folders_name[i], (end_time-start_time)/60.0))


if __name__ == "__main__":
    v2xset_to_npy_and_json(split_name='test')

opencood/utils/my_tools_v2xset.py
MediaBrain-SJTU/CoBEVFlow
# transfer all lidar file from .pcd to .npy 
# sizhewei @ <sizhewei@sjtu.edu.cn>
# 2022/12/20

###################
# $ cd ~/percp/OpenCOOD/
# python opencood/utils/my_tools.py
####################

import os
import sys
sys.path.append(os.getcwd())
import shutil
# import open3d as o3d
import numpy as np
import time
import json
import yaml
from pypcd import pypcd
from tqdm import tqdm
from tqdm.contrib import tenumerate
from tqdm.auto import trange
import glob

import matplotlib.pyplot as plt

from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import pcd_to_np

def extract_timestamps(yaml_files):
    """
    Given the list of the yaml files, extract the mocked timestamps.

    Parameters
    ----------
    yaml_files : list
        The full path of all yaml files of ego vehicle

    Returns
    -------
    timestamps : list
        The list containing timestamps only.
    """
    timestamps = []

    for file in yaml_files:
        res = file.split('/')[-1]

        timestamp = res.replace('.yaml', '')
        timestamps.append(timestamp)

    return timestamps

def to_npy_and_json(split_name="train"):
    # 处理的数据集
    data_seperate_name = split_name
    # pcd 文件路径
    root_dir = "/GPFS/public/OPV2V_Irregular_V2/" + data_seperate_name
    # yaml 文件路径
    root_annotation_dir = "/GPFS/public/OPV2V_Irregular_V2/dataset_irregular_v2/" + data_seperate_name
    # 新保存的路径
    save_dir = "/GPFS/public/OPV2V_irregular_npy/" + data_seperate_name

    # first load all paths of different scenarios
    scenario_folders = sorted([os.path.join(root_dir, x)
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])
    scenario_folders_name = sorted([x
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])

    print("### Transformation Started! ###")

    # loop over all scenarios
    for (i, scenario_folder) in tenumerate(scenario_folders):
        start_time = time.time()
        # copy timestamps npy file
        timestamps_file = os.path.join(root_annotation_dir, scenario_folders_name[i], 'timestamps.npy')
        save_scene_folder = os.path.join(save_dir, scenario_folders_name[i])
        if not os.path.exists(save_scene_folder):
            os.makedirs(save_scene_folder)
        save_timestamps_file = os.path.join(save_scene_folder, 'timestamps.npy')
        if not os.path.exists(save_timestamps_file):
            shutil.copyfile(timestamps_file, save_timestamps_file)

        # iterate all cav in this scenario
        cav_list = sorted([x 
                            for x in os.listdir(scenario_folder) if 
                            os.path.isdir(os.path.join(scenario_folder, x))])
        assert len(cav_list) > 0
        for (j, cav_id) in tenumerate(cav_list):
            cav_path_pcd = os.path.join(scenario_folder, cav_id)
            cav_path_yaml = os.path.join(root_annotation_dir, scenario_folders_name[i], cav_id)
            # use the frame number as key, the full path as the values
            yaml_files = \
                sorted([os.path.join(cav_path_yaml, x)
                        for x in os.listdir(cav_path_yaml) if
                        x.endswith('.yaml')])
            timestamps = extract_timestamps(yaml_files)	

            for k, timestamp in tenumerate(timestamps):
                # if k > 30:
                #     break

                # new save dir
                save_path = os.path.join(save_dir, scenario_folders_name[i], cav_id)
                if not os.path.exists(save_path):
                    os.makedirs(save_path)
                    # print("======Path {} created successfully!======".format(save_path))

                # change pcd to npy 
                lidar_file = os.path.join(cav_path_pcd, timestamp + '.pcd')
                new_lidar_file = os.path.join(save_path, timestamp + '.npy')
                if not os.path.exists(new_lidar_file):
                    lidar = pcd_to_np(lidar_file)
                    np.save(new_lidar_file, lidar)

                # change yaml to json
                yaml_file = os.path.join(cav_path_yaml, timestamp + '.yaml')
                json_file = os.path.join(cav_path_yaml, timestamp + '.json')
                new_json_file = os.path.join(save_path, timestamp + '.json')
                if os.path.exists(json_file): # 原数据集有json文件
                    if not os.path.exists(new_json_file):
                        shutil.copyfile(json_file, new_json_file)
                else:
                    if not os.path.exists(new_json_file):
                        params = load_yaml(yaml_file)
                        json_str = json.dumps(params, indent=4)
                        with open(new_json_file, 'w') as json_file:
                            json_file.write(json_str)
                
                # new_yaml_file = os.path.join(save_path, timestamp + '.yaml')
                # if not os.path.exists(new_yaml_file):
                #     shutil.copyfile(yaml_file, new_yaml_file)

                # print(timestamp)
        end_time = time.time()
        print("=== %d-th scenario (%s) finished, consumed %.3f mins. ===" \
            % (i, scenario_folders_name[i], (end_time-start_time)/60.0))

def create_small_dataset(split_name='train', size_scene=4, size_cav=3, size_time=2):
    """
    TO BE FINISHED !!!
    Create a small dataset from original large one 
    """
    # split name
    data_seperate_name = split_name
    # 原数据的路径
    root_dir = "/GPFS/public/OPV2V_irregular_npy/" + data_seperate_name
    # 新保存的路径
    save_dir = "/GPFS/public/OPV2V_irregular_npy_small/" + data_seperate_name

    # first load all paths of different scenarios
    scenario_folders = sorted([os.path.join(root_dir, x)
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])
    scenario_folders_name = sorted([x
                                for x in os.listdir(root_dir) if
                                os.path.isdir(os.path.join(root_dir, x))])

    print("### Transformation Started! ###")

    # loop over all scenarios
    for (i, scenario_folder) in tenumerate(scenario_folders):
        if i==size_scene:
            break

if __name__ == "__main__":
    # to_npy_and_json(split_name='validate')
    # create_small_dataset(split_name='train')

    root_dirs = "/remote-home/share/sizhewei/logs"
    note = 'ablation'
    save_path = f"./opencood/result_{note}.jpg"
    title = "CoBEVFlow V.S. SyncNet on different Irrgularity."
    
    split_list = ['cobevflow', 'where2comm_syncnet']#, 'GT'] 
    colors = ['purple', 'gray', 'orange']
    single_split_name = 'single'

    num_delay = 15

    max_x = 30 # unit is ms
    plt.figure()
    fig, ax = plt.subplots(1,2, sharex='col', sharey=False, figsize=(18,6))
    fig.suptitle(f'{title}', fontsize='x-large', y=0.99)
    fig.text(0.5, 0.03, 'Variance of time intervals (ms ^ 2).', ha='center', fontsize='x-large')
    fig.text(0.08, 0.5, 'AP', va='center', rotation='vertical', fontsize='x-large')
    # ax30 = ax[0]; 
    ax50 = ax[0]; ax70 = ax[1]
    for split_i, split_name in enumerate(split_list):
        ap_list = []
        vars = []
        if split_name == 'where2comm_syncnet':
            method_name = 'SyncNet'
            file_name = 'irv2v_where2comm_syncnet_regular_300_past_3_test'
            eval_name = 'eval_no_noise_syncnet_var_s100_flip_t'
            eval_file_list = glob.glob(os.path.join(root_dirs, file_name, f'{eval_name}*.yaml'))
        elif split_name == 'cobevflow':
            method_name = 'CoBEVFlow (ours)'
            file_name = 'irv2v_where2comm_cobevflow_var'
            eval_name = 'eval_no_noise_cobevflow_var_linear_s100'
            eval_file_list = glob.glob(os.path.join(root_dirs, file_name, f'{eval_name}*.yaml'))
            print(eval_file_list)
        elif split_name == 'GT':
            method_name = 'Ground Truth'
            file_name = 'irv2v_where2comm_cobevflow_var'
            eval_name = 'cobevflow_var_gt_0.50'
        latest_time_delay = -100.00
        for eval_file in tqdm(eval_file_list):
            
            if not os.path.exists(eval_file):
                print(f'eval file {eval_file} not exist!')
                continue
            with open(eval_file, "r", encoding="utf-8") as f:
                data = yaml.load(f, Loader=yaml.FullLoader)
            
            try:
                avg_time_var = data['avg_time_var']
            except:
                continue
            
            tmp_aps = []
            tmp_aps.append(data['ap_30'])
            tmp_aps.append(data['ap_50'])
            tmp_aps.append(data['ap_70'])
            ap_list.append(tmp_aps)

            vars.append(avg_time_var)

        ap_list_np = np.array(ap_list)
        ap_list_np = np.transpose(ap_list_np)

        ap_30_list = list(ap_list_np[0])
        ap_50_list = list(ap_list_np[1])
        ap_70_list = list(ap_list_np[2])
        
        color = colors[split_i]
        # plt.sca(ax30); plt.plot(delays, ap_30_list, color=color, marker='+', label = method_name)
        plt.sca(ax50); plt.scatter(vars, ap_50_list, color=color, marker='+', s=60, label = method_name)
        plt.sca(ax70); plt.scatter(vars, ap_70_list, color=color, marker='+', s=60, label = method_name)
    
    # for single fusion
    method_name = single_split_name.split('_')[0]
    # eval_file = os.path.join(root_dirs, f'eval_{single_split_name}.yaml')
    eval_file = '/remote-home/share/sizhewei/logs/opv2v_late_fusion/eval_no_noise_single_delay_0.00.yaml'
    single_color = 'red'
    with open(eval_file, "r", encoding="utf-8") as f:
        data = yaml.load(f, Loader=yaml.FullLoader)
    # plt.sca(ax30); plt.plot([0,max_x],[data['ap_30'],data['ap_30']], color=single_color, linestyle='--', label=method_name)
    plt.sca(ax50); plt.plot([0,max_x],[data['ap_50'],data['ap_50']], color=single_color, linestyle='--', label=method_name)
    plt.sca(ax70); plt.plot([0,max_x],[data['ap_70'],data['ap_70']], color=single_color, linestyle='--', label=method_name)

    xaxis = np.linspace(0, max_x, 16)
    # ax30.set_title('The Results for AP@0.3'); ax30.grid(True); 
    # ax30.set_xticks(xaxis); \
    #     ax30.legend(loc = 'lower left')
    yaxis = np.linspace(0.6, 0.9, 5)
    ax50.set_title('The Results for AP@0.5'); ax50.grid(True); 
    ax50.set_xticks(xaxis); ax50.set_yticks(yaxis); 
    ax50.legend(loc = 'lower left')
    
    yaxis = np.linspace(0.5, 0.9, 5)
    ax70.set_title('The Results for AP@0.7'); ax70.grid(True); 
    ax70.set_xticks(xaxis); ax70.set_yticks(yaxis); 
    ax70.legend(loc = 'lower left')

    plt.savefig(save_path)
    print("=== Plt save finished!!! ===")
    # plt.title('标题')
    # plt.show()

opencood/utils/my_tools_var.py
MediaBrain-SJTU/CoBEVFlow
# transfer all lidar file from .pcd to .npy 
# sizhewei @ <sizhewei@sjtu.edu.cn>
# 2022/12/20

###################
# $ cd ~/percp/OpenCOOD/
# python opencood/utils/my_tools.py
####################

import os
import sys
sys.path.append(os.getcwd())
import shutil
# import open3d as o3d
import numpy as np
import time
import json
import yaml
from pypcd import pypcd
from tqdm import tqdm
from tqdm.contrib import tenumerate
from tqdm.auto import trange
import math
import matplotlib.pyplot as plt

from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.utils.pcd_utils import pcd_to_np

# from brokenaxes import brokenaxes
from scipy import interpolate

# 用于单独保存子图的函数
def save_subfig(fig,ax,save_path,fig_name):
    bbox = ax.get_tightbbox(fig.canvas.get_renderer()).expanded(1.02, 1.02)
    extent = bbox.transformed(fig.dpi_scale_trans.inverted())
    fig.savefig(save_path+fig_name, bbox_inches=extent, dpi=600, format='png')


def export_legend(ax, filename="legend.pdf"):
    fig2 = plt.figure()
    ax2 = fig2.add_subplot()
    ax2.axis('off')
    legend = ax2.legend(*ax.get_legend_handles_labels(), frameon=False, loc='lower center', ncol=3,)
    fig  = legend.figure
    fig.canvas.draw()
    bbox  = legend.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    fig.savefig(filename, dpi=600, bbox_inches=bbox)

if __name__ == "__main__":
    # to_npy_and_json(split_name='validate')
    # create_small_dataset(split_name='train')

    root_dirs = "/remote-home/share/sizhewei/logs/bandwidth_curve"
    note = 'major_bandwidth'
    format = 'jpg'
    save_path = f"./opencood/result_{note}.{format}"
    title = "Average Precision curves of different methods on the IRV2V dataset at different average time intervals."
    plt.style.use('ggplot')
    split_list = ['single', 'v2vnet', 'v2xvit', 'disconet']
    marker_list = []
    split_list_sparse = ['where2comm', 'syncnet', 'cobevflow']
    colors = ['gray', '#DEC68B', '#BEB8DC', '#82B0D2']
    # colors = ['#999999', '#DEC68B', '#BEB8DC', '#82B0D2', '#8ECFC9']  #'#FA7F6F'] #'#E7DAD2'
    colors_sparse = ['#8ECFC9', '#FFBE7A', 'red']


    max_x = 17.5 # unit is ms
    linewidth= '7'
    fontsize = 22

    plt.figure()
    fig, ax = plt.subplots(1,2, sharex='col', sharey=False, figsize=(16,7))
    plt.subplots_adjust(wspace = 0.5, hspace = 0.1)

    ax50 = ax[0]; ax70 = ax[1]

    for split_i, split_name in enumerate(split_list):
        if split_name == 'single':
            method_name = 'Single'
            eval_name = 'single_delay'
        elif split_name == 'late':
            method_name = 'Late Fusion'
            file_name = 'opv2v_late_fusion'
            eval_name = 'late_delay'
        elif split_name == 'v2vnet':
            method_name = 'V2VNet'
            file_name = 'opv2v_v2vnet_32ch'
            eval_name = 'v2vnet_32ch_ep23'
        elif split_name == 'v2xvit':
            method_name = 'V2X-ViT'
            file_name = 'opv2v_point_pillar_v2xvit_v1_slren'
            eval_name = 'v2xvit'
        elif split_name == 'disconet':
            method_name = 'DiscoNet'
            file_name = 'opv2v_disconet_new_ylu'
            eval_name = 'disconet'
        
        color = colors[split_i]
        eval_file = os.path.join(root_dirs,f"eval_no_noise_{eval_name}_0.30.yaml")
            
        if not os.path.exists(eval_file):
            print(f'eval file {eval_file} not exist!')
            continue
        with open(eval_file, "r", encoding="utf-8") as f:
            data = yaml.load(f, Loader=yaml.FullLoader)
        bandwith = math.log(200*704, 2)
        if split_name == 'single':
            bandwith = 0
        plt.sca(ax50); plt.scatter(bandwith, data['ap_50'], color=color, marker='o', s=400, label = method_name)
        plt.sca(ax70); plt.scatter(bandwith, data['ap_70'], color=color, marker='o', s=400, label = method_name) 


    for split_i, split_name in enumerate(split_list_sparse):
        ap_list = []
        bw = []
        if split_name == 'where2comm':
            method_name = 'Where2Comm'
            file_name = 'opv2v_late_fusion'
            eval_name = 'where2comm'
        elif split_name == 'syncnet':
            method_name = 'Where2comm + SyncNet'
            eval_name = 'syncnet'
        elif split_name == 'cobevflow':
            method_name = 'CoBEVFlow (ours)'
            file_name = 'opv2v_v2vnet_32ch'
            eval_name = 'cobevflow_regular'

        for i in [5, 10, 15, 20, 25, 40, 3520]:
            eval_file = os.path.join(root_dirs,f"eval_no_noise_{eval_name}_0.30_noise_0_0_0_0_roi_{i}.yaml")
            if i==0:
                eval_file = '/remote-home/share/sizhewei/logs/bandwidth_curve/eval_no_noise_where2comm_no_fusion_0.30_noise_0_0_0_0_roi_0.yaml'
            if not os.path.exists(eval_file):
                print(f'eval file {eval_file} not exist!')
                continue
            with open(eval_file, "r", encoding="utf-8") as f:
                data = yaml.load(f, Loader=yaml.FullLoader)
            
           
            if i==0:
                bw.append(0)
            else:
                bw.append(math.log(i*40, 2))

            tmp_aps = []
            tmp_aps.append(data['ap_30'])
            tmp_aps.append(data['ap_50'])
            tmp_aps.append(data['ap_70'])

            ap_list.append(tmp_aps)

        color = colors_sparse[split_i]
        ap_list_np = np.array(ap_list)
        ap_list_np = np.transpose(ap_list_np)

        ap_30_list = list(ap_list_np[0])
        ap_50_list = list(ap_list_np[1])
        ap_70_list = list(ap_list_np[2])

        # bw_new = np.linspace(0, 17.103, 170)
        
        # bspline = interpolate.make_interp_spline(bw, ap_50_list, k=4)
        # ap_50_list_new = bspline(bw_new)

        # bspline = interpolate.make_interp_spline(bw, ap_70_list, )
        # ap_70_list_new = bspline(bw_new)


        # plt.sca(ax50); plt.plot(bw_new, ap_50_list_new, color=color, linewidth='3', label = method_name, marker='+') # , marker='+'
        # plt.sca(ax70); plt.plot(bw_new, ap_70_list_new, color=color, linewidth='3', label = method_name, marker='+')

        plt.sca(ax50); plt.plot(bw, ap_50_list, color=color, linewidth=linewidth, label = method_name)#, marker='+') # , marker='+'
        plt.sca(ax70); plt.plot(bw, ap_70_list, color=color, linewidth=linewidth, label = method_name)#, marker='+')

    xaxis = np.linspace(0, 15, 5)
    # ax30.set_title('The Results for AP@0.3'); ax30.grid(True); 
    # ax30.set_xticks(xaxis); \
    #     ax30.legend(loc = 'lower left')
    yaxis = np.linspace(0.4, 0.9, 6)
    # font = "Gill Sans MT Condensed"
    # ax50.set_title('The Results for AP@0.5'); 
    ax50.tick_params(axis='both', which='major', labelsize=16, color='gray')
    ax50.grid(True); 
    ax50.set_xticks(xaxis); ax50.set_yticks(yaxis);
    ax50.set_xlabel('Communication volume (log2)', fontsize=fontsize, color='black'); ax50.set_ylabel('AP@0.50', fontsize=fontsize, color='black');
    # ax50.legend(loc = 'lower left')
    
    yaxis = np.linspace(0.3, 0.8, 6)
    # ax70.set_title('The Results for AP@0.7'); 
    ax70.grid(True); 
    ax70.tick_params(axis='both', which='major', labelsize=16)
    ax70.set_xticks(xaxis); ax70.set_yticks(yaxis); 
    ax70.set_xlabel('Communication volume (log2)', fontsize=fontsize, color='black'); ax70.set_ylabel('AP@0.70', fontsize=fontsize, color='black');
    # ax70.legendloc = 'lower left')

    save_subfig(fig,ax[0],'./opencood/','bandwidth-ap50.png')
    save_subfig(fig,ax[1],'./opencood/','bandwidth-ap70.png')

    # plt.legend(frameon=False)
    plt.savefig(save_path, dpi=600, format=format)
    print("=== Plt save finished!!! ===")

    export_legend(ax70, '/root/percp/OpenCOOD/opencood/legend-bandwidth.pdf')
    # # then create a new image
    # # adjust the figure size as necessary
    # figsize = (7, 7)
    # fig_leg = plt.figure(figsize=figsize)
    # ax_leg = fig_leg.add_subplot(111)
    # # add the legend from the previous axes
    # # ax70.legend(ncol=7)
    # ax_leg.legend(*ax70.get_legend_handles_labels(), loc='center')
    # # hide the axes frame and the x/y labels
    # ax_leg.axis('off')
    # fig_leg.savefig('/root/percp/OpenCOOD/opencood/legend.jpg')

    # plt.title('标题')
    # plt.show()

opencood/utils/my_tools_vis_bandwidth.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Hao Xiang <haxiang@g.ucla.edu>,
# License: TDG-Attribution-NonCommercial-NoDistrib


"""
Utility functions related to point cloud
"""

import open3d as o3d
import numpy as np
from pypcd import pypcd

def pcd_to_np(pcd_file):
    """
    Read  pcd and return numpy array.

    Parameters
    ----------
    pcd_file : str
        The pcd file that contains the point cloud.

    Returns
    -------
    pcd : o3d.PointCloud
        PointCloud object, used for visualization
    pcd_np : np.ndarray
        The lidar data in numpy format, shape:(n, 4)

    """
    pcd = o3d.io.read_point_cloud(pcd_file)

    xyz = np.asarray(pcd.points)
    # we save the intensity in the first channel
    intensity = np.expand_dims(np.asarray(pcd.colors)[:, 0], -1)
    pcd_np = np.hstack((xyz, intensity))

    return np.asarray(pcd_np, dtype=np.float32)


def mask_points_by_range(points, limit_range):
    """
    Remove the lidar points out of the boundary.

    Parameters
    ----------
    points : np.ndarray
        Lidar points under lidar sensor coordinate system.

    limit_range : list
        [x_min, y_min, z_min, x_max, y_max, z_max]

    Returns
    -------
    points : np.ndarray
        Filtered lidar points.
    """

    mask = (points[:, 0] > limit_range[0]) & (points[:, 0] < limit_range[3])\
           & (points[:, 1] > limit_range[1]) & (
                   points[:, 1] < limit_range[4]) \
           & (points[:, 2] > limit_range[2]) & (
                   points[:, 2] < limit_range[5])

    points = points[mask]

    return points


def mask_ego_points(points):
    """
    Remove the lidar points of the ego vehicle itself.

    Parameters
    ----------
    points : np.ndarray
        Lidar points under lidar sensor coordinate system.

    Returns
    -------
    points : np.ndarray
        Filtered lidar points.
    """
    mask = (points[:, 0] >= -1.95) & (points[:, 0] <= 2.95) \
           & (points[:, 1] >= -1.1) & (points[:, 1] <= 1.1)
    points = points[np.logical_not(mask)]

    return points


def shuffle_points(points):
    shuffle_idx = np.random.permutation(points.shape[0])
    points = points[shuffle_idx]

    return points


def lidar_project(lidar_data, extrinsic):
    """
    Given the extrinsic matrix, project lidar data to another space.

    Parameters
    ----------
    lidar_data : np.ndarray
        Lidar data, shape: (n, 4)

    extrinsic : np.ndarray
        Extrinsic matrix, shape: (4, 4)

    Returns
    -------
    projected_lidar : np.ndarray
        Projected lida data, shape: (n, 4)
    """

    lidar_xyz = lidar_data[:, :3].T
    # (3, n) -> (4, n), homogeneous transformation
    lidar_xyz = np.r_[lidar_xyz, [np.ones(lidar_xyz.shape[1])]]
    lidar_int = lidar_data[:, 3]

    # transform to ego vehicle space, (3, n)
    project_lidar_xyz = np.dot(extrinsic, lidar_xyz)[:3, :]
    # (n, 3)
    project_lidar_xyz = project_lidar_xyz.T
    # concatenate the intensity with xyz, (n, 4)
    projected_lidar = np.hstack((project_lidar_xyz,
                                 np.expand_dims(lidar_int, -1)))

    return projected_lidar


def projected_lidar_stack(projected_lidar_list):
    """
    Stack all projected lidar together.

    Parameters
    ----------
    projected_lidar_list : list
        The list containing all projected lidar.

    Returns
    -------
    stack_lidar : np.ndarray
        Stack all projected lidar data together.
    """
    stack_lidar = []
    for lidar_data in projected_lidar_list:
        stack_lidar.append(lidar_data)

    return np.vstack(stack_lidar)


def downsample_lidar(pcd_np, num):
    """
    Downsample the lidar points to a certain number.

    Parameters
    ----------
    pcd_np : np.ndarray
        The lidar points, (n, 4).

    num : int
        The downsample target number.

    Returns
    -------
    pcd_np : np.ndarray
        The downsampled lidar points.
    """
    assert pcd_np.shape[0] >= num

    selected_index = np.random.choice((pcd_np.shape[0]),
                                      num,
                                      replace=False)
    pcd_np = pcd_np[selected_index]

    return pcd_np


def downsample_lidar_minimum(pcd_np_list):
    """
    Given a list of pcd, find the minimum number and downsample all
    point clouds to the minimum number.

    Parameters
    ----------
    pcd_np_list : list
        A list of pcd numpy array(n, 4).
    Returns
    -------
    pcd_np_list : list
        Downsampled point clouds.
    """
    minimum = np.Inf

    for i in range(len(pcd_np_list)):
        num = pcd_np_list[i].shape[0]
        minimum = num if minimum > num else minimum

    for (i, pcd_np) in enumerate(pcd_np_list):
        pcd_np_list[i] = downsample_lidar(pcd_np, minimum)

    return pcd_np_list

def  read_pcd(pcd_path):
    pcd = pypcd.PointCloud.from_path(pcd_path)
    time = None
    pcd_np_points = np.zeros((pcd.points, 4), dtype=np.float32)
    pcd_np_points[:, 0] = np.transpose(pcd.pc_data["x"])
    pcd_np_points[:, 1] = np.transpose(pcd.pc_data["y"])
    pcd_np_points[:, 2] = np.transpose(pcd.pc_data["z"])
    pcd_np_points[:, 3] = np.transpose(pcd.pc_data["intensity"]) / 256.0
    del_index = np.where(np.isnan(pcd_np_points))[0]
    pcd_np_points = np.delete(pcd_np_points, del_index, axis=0)
    return pcd_np_points, time
opencood/utils/pcd_utils.py
MediaBrain-SJTU/CoBEVFlow
# Author: Yifan Lu

import numpy as np
import torch
import torch.distributions as dist

def add_noise_data_dict(data_dict, noise_setting):
    """ Update the base data dict. 
        We retrieve lidar_pose and add_noise to it.
        And set a clean pose.
    """
    if noise_setting['add_noise']:
        for cav_id, cav_content in data_dict.items():
            cav_content['params']['lidar_pose_clean'] = cav_content['params']['lidar_pose'] # 6 dof pose

            if "laplace" in noise_setting['args'].keys() and noise_setting['args']['laplace'] is True:
                cav_content['params']['lidar_pose'] = cav_content['params']['lidar_pose'] + \
                                                        generate_noise_laplace( # we just use the same key name
                                                            noise_setting['args']['pos_std'],
                                                            noise_setting['args']['rot_std'],
                                                            noise_setting['args']['pos_mean'],
                                                            noise_setting['args']['rot_mean']
                                                        )
            else:
                cav_content['params']['lidar_pose'] = cav_content['params']['lidar_pose'] + \
                                                            generate_noise(
                                                                noise_setting['args']['pos_std'],
                                                                noise_setting['args']['rot_std'],
                                                                noise_setting['args']['pos_mean'],
                                                                noise_setting['args']['rot_mean']
                                                            )

    else:
        for cav_id, cav_content in data_dict.items():
            cav_content['params']['lidar_pose_clean'] = cav_content['params']['lidar_pose'] # 6 dof pose

            
    return data_dict

def generate_noise(pos_std, rot_std, pos_mean=0, rot_mean=0):
    """ Add localization error to the 6dof pose
        Noise includes position (x,y) and rotation (yaw).
        We use gaussian distribution to generate noise.
    
    Args:

        pos_std : float 
            std of gaussian dist, in meter

        rot_std : float
            std of gaussian dist, in degree

        pos_mean : float
            mean of gaussian dist, in meter

        rot_mean : float
            mean of gaussian dist, in degree
    
    Returns:
        pose_noise: np.ndarray, [6,]
            [x, y, z, roll, yaw, pitch]
    """

    xy = np.random.normal(pos_mean, pos_std, size=(2))
    yaw = np.random.normal(rot_mean, rot_std, size=(1))

    pose_noise = np.array([xy[0], xy[1], 0, 0, yaw[0], 0])

    
    return pose_noise



def generate_noise_laplace(pos_b, rot_b, pos_mu=0, rot_mu=0):
    """ Add localization error to the 6dof pose
        Noise includes position (x,y) and rotation (yaw).
        We use laplace distribution to generate noise.
    
    Args:

        pos_b : float 
            parameter b of laplace dist, in meter

        rot_b : float
            parameter b of laplace dist, in degree

        pos_mu : float
            mean of laplace dist, in meter

        rot_mu : float
            mean of laplace dist, in degree
    
    Returns:
        pose_noise: np.ndarray, [6,]
            [x, y, z, roll, yaw, pitch]
    """

    xy = np.random.laplace(pos_mu, pos_b, size=(2))
    yaw = np.random.laplace(rot_mu, rot_b, size=(1))

    pose_noise = np.array([xy[0], xy[1], 0, 0, yaw[0], 0])
    return pose_noise


def generate_noise_torch(pose, pos_std, rot_std, pos_mean=0, rot_mean=0):
    """ only used for v2vnet robust.
        rotation noise is sampled from von_mises distribution
    
    Args:
        pose : Tensor, [N. 6]
            including [x, y, z, roll, yaw, pitch]

        pos_std : float 
            std of gaussian dist, in meter

        rot_std : float
            std of gaussian dist, in degree

        pos_mean : float
            mean of gaussian dist, in meter

        rot_mean : float
            mean of gaussian dist, in degree
    
    Returns:
        pose_noisy: Tensor, [N, 6]
            noisy pose
    """

    N = pose.shape[0]
    noise = torch.zeros_like(pose, device=pose.device)
    concentration = (180 / (np.pi * rot_std)) ** 2

    noise[:, :2] = torch.normal(pos_mean, pos_std, size=(N, 2), device=pose.device)
    noise[:, 4] = dist.von_mises.VonMises(loc=rot_mean, concentration=concentration).sample((N,)).to(noise.device)


    return noise


def remove_z_axis(T):
    """ remove rotation/translation related to z-axis
    Args:
        T: np.ndarray
            [4, 4]
    Returns:
        T: np.ndarray
            [4, 4]
    """
    T[2,3] = 0 # z-trans
    T[0,2] = 0
    T[1,2] = 0
    T[2,0] = 0
    T[2,1] = 0
    T[2,2] = 1
    
    return T
opencood/utils/pose_utils.py
MediaBrain-SJTU/CoBEVFlow
"""
https://github.com/AmnonDrory/BestBuddiesRegistration/blob/main/code/bb_pc/utils/subsampling.py
"""

import numpy as np
import open3d as o3d
import pandas as pd
from copy import deepcopy

num_features = 3

def calc_bin_inds(PC, n_bins, axis, mode):
    N = PC.shape[0]
    if "adaptive" in mode:
        inds = np.round(np.linspace(0, N, n_bins + 1)).astype(int)
        s = np.sort(PC[:, axis])
        thresh = s[inds[1:]-1]
    else: # "equally_spaced"
        thresh = np.linspace(np.min(PC[:,axis]), np.max(PC[:,axis]),  n_bins + 1)
        thresh = thresh[1:]

    bin_ind = np.zeros(N) + np.nan
    for i in range(n_bins):
        is_cur = (PC[:, axis] <= thresh[i]) & np.isnan(bin_ind)
        bin_ind[is_cur] = i

    assert np.sum(np.isnan(bin_ind)) == 0, "Error: not all samples were assigned to a bin"

    return bin_ind

def voxelGrid_filter_inner(PC, num_samples, mode):

    if "equal_nbins_per_axis" in mode:
        n_bins = int(np.ceil(num_samples ** (1. / 3)))
        n_bins_x = n_bins
        n_bins_y = n_bins
        n_bins_z = n_bins
    else:
        span = []
        for axis in range(3):
            span.append( np.max(PC[:,axis])-np.min(PC[:,axis]) )
        normalized_num_samples = num_samples * (span[0]**2 / (span[1]*span[2]))
        n_bins_x = int(np.ceil(normalized_num_samples ** (1. / 3)))
        n_bins_y = int(np.ceil(n_bins_x * (span[1]/span[0])))
        n_bins_z = int(np.ceil(n_bins_x * (span[2] / span[0])))
        assert (n_bins_x * n_bins_y * n_bins_z) >= num_samples, "Error"
    x_bin_inds = calc_bin_inds(PC, n_bins_x, 0, mode)
    y_bin_inds = calc_bin_inds(PC, n_bins_y, 1, mode)
    z_bin_inds = calc_bin_inds(PC, n_bins_z, 2, mode)

    data = np.hstack([x_bin_inds.reshape([-1,1]),
                      y_bin_inds.reshape([-1,1]),
                      z_bin_inds.reshape([-1,1]),
                      PC])

    df = pd.DataFrame(data, columns=['x_ind', 'y_ind', 'z_ind', 'x', 'y', 'z'])
    newPC = np.array(df.groupby(['x_ind', 'y_ind', 'z_ind']).mean())

    return newPC

def voxelGrid_filter(PC, num_requested_samples, mode):
    """
    Sub-sample a point cloud by defining a grid of voxels, and returning the average point in each one.

    :param PC: Nx3 array, point cloud, each row is a sample
    :param num_samples: numbver of requested samples
    :param mode: list of strings, can contain any of the following:
                 "exact_number" - return exactly num_requested_samples, otherwise may return more than requested number (but never less)
                 "equal_nbins_per_axis" - same number of bins for each axis (x,y,z). Otherwise the bins are cube shaped, and usually a different number of bins fits in each of the dimensions.
                 "adaptive" - smaller bins where there is more data. Otherwise, all bins are the same size.
    :return: newPC - a point cloud with approximately num_requested_samples
    """
    num_samples = num_requested_samples
    N = PC.shape[0]
    done = False
    MAX_ATTEMPTS = 40
    ACCELERATION_FACTOR = 2
    MAX_DIVERGENCE_TIME = 4
    TOLERANCE = 0.05
    rel_history = []
    newPC_history = []
    while not done:
        newPC = voxelGrid_filter_inner(PC, num_samples, mode)
        new_N = newPC.shape[0]
        newPC_history.append(newPC)
        relative_error_in_size = (new_N/float(num_requested_samples)) -1
        rel_history.append(relative_error_in_size)
        if (relative_error_in_size < 0) or (relative_error_in_size > TOLERANCE):
            best_ind = np.argmin(np.abs(rel_history))
            if (len(rel_history) - best_ind > MAX_DIVERGENCE_TIME) and (np.max(rel_history) > 0):
                    done = True
            else:
                num_samples = int(np.ceil(num_samples*float(num_requested_samples)/new_N))
                if (np.max(rel_history) < 0):
                    num_samples = int(ACCELERATION_FACTOR*num_samples)

        else:
            done = True

        if len(rel_history) >= MAX_ATTEMPTS:
            done = True

    if len(rel_history) >= MAX_ATTEMPTS:
        assert False, "voxelGrid_filter could not supply required number of samples"
        print("Error: voxelGrid_filter could not supply required number of samples, recovering")
        best_ind = np.argmax(rel_history)
        return newPC_history[best_ind]

    rel_history_above_only = np.array(rel_history)
    rel_history_above_only[rel_history_above_only<0] = np.inf
    best_ind_above = np.argmin(rel_history_above_only)

    newPC = newPC_history[best_ind_above]
    if 'exact_number' in mode:
        p = np.random.permutation(newPC.shape[0])
        inds = p[:num_requested_samples]
        newPC = newPC[inds,:]

    return newPC

def voxel_filter(pcd, N):
    # pcd is of open3d point cloud class
    if "numpy" in str(type(pcd)):
        tmp = o3d.geometry.PointCloud()
        tmp.points = o3d.utility.Vector3dVector(pcd)
        pcd = tmp
    K = np.shape(pcd.points)[0]
    vs = 1e-3
    while K>N:
        pcd = o3d.geometry.voxel_down_sample(pcd, voxel_size=vs)
        vs *= 2
        K = np.shape(pcd.points)[0]
    return pcd

def calc_distances(p0, points):
    return ((p0 - points) ** 2).sum(axis=1)

def fps_from_given_pc(pts, K, given_pc):
    """
    copied from https://github.com/orendv/learning_to_sample/blob/master/reconstruction/src/sample_net_point_net_ae.py
    :param self:
    :param pts:
    :param K:
    :param given_pc:
    :return:
    """
    farthest_pts = np.zeros((K, 3))
    t = given_pc.shape[0]
    farthest_pts[0:t,:] = given_pc

    distances = calc_distances(farthest_pts[0], pts)
    for i in range(1, t):
        distances = np.minimum(distances, calc_distances(farthest_pts[i,:], pts))

    for i in range(t, K):
        farthest_pts[i,:] = pts[np.argmax(distances),:]
        distances = np.minimum(distances, calc_distances(farthest_pts[i,:], pts))
    return farthest_pts

def get_random_subset(PC, num_samples, mode="farthest", submode=None, allow_overask=False):
    """
    Subsample a point cloud, using either of various methods

    :param PC:
    :param num_samples:
    :param mode:
    :param n_bins:
    :param submode: Relevant for the "r_normalized" and "r_squared_normalized" methods.
    :return:
    """
    if num_samples > PC.shape[0]:
        if allow_overask:
            return PC
        else:
            assert False, "Error: requesting more samples than there are"

    if PC.shape[0] == num_samples:
        result = PC
    if mode == "uniform":
        inds = np.random.permutation(PC.shape[0])[:num_samples]
        result = PC[inds, :]
    elif mode == "farthest":
        first_ind = np.random.permutation(PC.shape[0])[0]
        result = fps_from_given_pc(PC, num_samples, PC[first_ind:(first_ind+1), :])
    elif "voxel" in mode:
        if submode is None:
            submode = ["equal_nbins_per_axis"]

        # The voxelGrid subsampling algorithm has no randomality.
        # we force it to have some by rendomly removing a small subset of the points

        keep_fraction = 0.9
        num_keep = int(PC.shape[0]*keep_fraction)
        if num_samples < num_keep:
            PC = get_random_subset(PC, num_keep, mode="uniform")
        result = voxelGrid_filter(PC, num_samples, submode)

    else:
        assert False, "unknown mode"

    return result

def subsample_fraction(PC, fraction):
    N = PC.shape[0]
    subset_size = int(np.round(N * fraction))
    inds = np.random.permutation(N)[:subset_size]
    return PC[inds,:]


def keep_closest(PC, max_dist):
    R = np.sqrt(np.sum(PC ** 2, axis=1))
    return PC[R <= max_dist, :]


def fit_plane(PC):
    xy1 = deepcopy(PC)
    xy1[:, 2] = 1
    z = PC[:, 2]
    abc, _, _, _ = np.linalg.lstsq(xy1, z, rcond=None)
    return abc


def is_on_plane(PC, abc, thickness):
    all_xy1 = deepcopy(PC)
    all_xy1[:, 2] = 1
    predicted_road_z = np.matmul(all_xy1, abc.reshape([-1, 1])).flatten()
    res = np.abs(PC[:, 2] - predicted_road_z) <= thickness
    return res

def remove_road(PC):
    mode = "plane"  # "constant_height"
    local_PC = keep_closest(PC, 10)
    count, bin_edges = np.histogram(local_PC[:, 2], 100)
    bin_centers = 0.5 * (bin_edges[:-1] + bin_edges[1:])
    ind_of_most_frequent = np.argmax(count)
    road_z = bin_centers[ind_of_most_frequent]
    road_thickness = 0.5  # meters
    if mode == "constant_height":
        is_road = np.abs(PC[:, 2] - road_z) <= road_thickness
    elif mode == "plane":
        raw_is_road = np.abs(local_PC[:, 2] - road_z) <= road_thickness
        raw_road_points = local_PC[raw_is_road, :]
        xy1 = deepcopy(raw_road_points)
        xy1[:, 2] = 1
        z = raw_road_points[:, 2]
        abc, _, _, _ = np.linalg.lstsq(xy1, z, rcond=None)
        all_xy1 = deepcopy(PC)
        all_xy1[:, 2] = 1
        predicted_road_z = np.matmul(all_xy1, abc.reshape([-1, 1])).flatten()
        is_road = np.abs(PC[:, 2] - predicted_road_z) <= road_thickness
    else:
        assert False, "unknown mode"

    return PC[~is_road, :]
opencood/utils/subsampling_utils.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Hao Xiang <haxiang@g.ucla.edu>
# Author: Yifan Lu

# Note:
# 这里的Yaw,Pitch,Roll为CARLA坐标系统，为左手坐标系，但与常规定义的左手坐标系的Yaw,Pitch,Roll 有所不同
# CARLA坐标系统下：Yaw为Z轴正半轴往原点看顺时针旋转为正， Pitch为Y轴正半轴往原点看逆时针旋转为正，Roll为X轴正半轴往原点看逆时针旋转为正
# 常规的左手坐标系，均应该是正半轴往原点看顺时针旋转为正。

# In addition，常规的右手坐标系，均应该是正半轴往原点看逆时针旋转为正

# 这导致了CARLA的（Yaw，Pitch，Roll）做transformation matrix做转化的时候，与常规的公式（http://msl.cs.uiuc.edu/planning/node102.html）在pitch和roll上都有正负号的差异，体现在他们的sin和cos上
# 由于变换矩阵是consist的，在本套代码里就follow CARLA坐标系统的设定。转化回yaw，pitch，roll的时候遵循CARLA这一套
# V2X-Sim和DAIRV2X，都是先获得变换矩阵，再通过这一套转换转为[x,y,z,roll,yaw,pitch]的位姿描述


"""
Transformation utils
"""

from re import X
import numpy as np
import torch
from icecream import ic
from pyquaternion import Quaternion
from opencood.utils.common_utils import check_numpy_to_torch

def pose_to_tfm(pose):
    """ Transform batch of pose to tfm
    Args:
        pose: torch.Tensor or np.ndarray
            [N, 3], x, y, yaw, in degree
            [N, 6], x, y, z, roll, yaw, pitch, in degree

            roll and pitch follows carla coordinate
    Returns:
        tfm: torch.Tensor
            [N, 4, 4] 
    """

    pose_tensor, is_np = check_numpy_to_torch(pose)
    pose = pose_tensor


    if pose.shape[1] == 3:
        N = pose.shape[0]
        x = pose[:,0]
        y = pose[:,1]
        yaw = pose[:,2]

        tfm = torch.eye(4, device=pose.device).view(1,4,4).repeat(N,1,1)
        tfm[:,0,0] = torch.cos(torch.deg2rad(yaw))
        tfm[:,0,1] = - torch.sin(torch.deg2rad(yaw))
        tfm[:,1,0] = torch.sin(torch.deg2rad(yaw))
        tfm[:,1,1] = torch.cos(torch.deg2rad(yaw))
        tfm[:,0,3] = x
        tfm[:,1,3] = y

    elif pose.shape[1] == 6:
        N = pose.shape[0]
        x = pose[:,0]
        y = pose[:,1]
        z = pose[:,2]
        roll = pose[:,3]
        yaw = pose[:,4]
        pitch = pose[:,5]

        c_y = torch.cos(torch.deg2rad(yaw))
        s_y = torch.sin(torch.deg2rad(yaw))
        c_r = torch.cos(torch.deg2rad(roll))
        s_r = torch.sin(torch.deg2rad(roll))
        c_p = torch.cos(torch.deg2rad(pitch))
        s_p = torch.sin(torch.deg2rad(pitch))

        tfm = torch.eye(4, device=pose.device).view(1,4,4).repeat(N,1,1)

        # translation matrix
        tfm[:, 0, 3] = x
        tfm[:, 1, 3] = y
        tfm[:, 2, 3] = z

        # rotation matrix
        tfm[:, 0, 0] = c_p * c_y
        tfm[:, 0, 1] = c_y * s_p * s_r - s_y * c_r
        tfm[:, 0, 2] = -c_y * s_p * c_r - s_y * s_r
        tfm[:, 1, 0] = s_y * c_p
        tfm[:, 1, 1] = s_y * s_p * s_r + c_y * c_r
        tfm[:, 1, 2] = -s_y * s_p * c_r + c_y * s_r
        tfm[:, 2, 0] = s_p
        tfm[:, 2, 1] = -c_p * s_r
        tfm[:, 2, 2] = c_p * c_r

    if is_np:
        tfm = tfm.numpy()

    return tfm




def tfm_to_pose(tfm: np.ndarray):
    """
    turn transformation matrix to [x, y, z, roll, yaw, pitch]
    we use radians format.
    tfm is pose in transformation format, and XYZ order, i.e. roll-pitch-yaw
    """
    # There forumlas are designed from x_to_world, but equal to the one below.
    yaw = np.degrees(np.arctan2(tfm[1,0], tfm[0,0])) # clockwise in carla
    roll = np.degrees(np.arctan2(-tfm[2,1], tfm[2,2])) # but counter-clockwise in carla
    pitch = np.degrees(np.arctan2(tfm[2,0], ((tfm[2,1]**2 + tfm[2,2]**2) ** 0.5)) ) # but counter-clockwise in carla


    # These formulas are designed for consistent axis orientation
    # yaw = np.degrees(np.arctan2(tfm[1,0], tfm[0,0])) # clockwise in carla
    # roll = np.degrees(np.arctan2(tfm[2,1], tfm[2,2])) # but counter-clockwise in carla
    # pitch = np.degrees(np.arctan2(-tfm[2,0], ((tfm[2,1]**2 + tfm[2,2]**2) ** 0.5)) ) # but counter-clockwise in carla

    # roll = - roll
    # pitch = - pitch

    x, y, z = tfm[:3,3]
    return([x, y, z, roll, yaw, pitch])

def tfm_to_xycs_torch(tfm: torch.Tensor):
    """
        similar to tfm_to_pose_torch,
        return x/y/cos(yaw)/sin(yaw)
    """
    x = tfm[:,0,3]
    y = tfm[:,1,3]
    
    cos = tfm[:,0,0]
    sin = tfm[:,1,0]

    pose = torch.stack([x,y,cos,sin]).T # (N, 4)

    return pose

def xycs_to_tfm_torch(xycs: torch.Tensor):
    """
        Args: xycs
            [N, 4]
    """
    N = xycs.shape[0]
    tfm = torch.eye(4, device=xycs.device).view(1,4,4).repeat(N,1,1)

    x, y, cos, sin = xycs[:,0], xycs[:,1], xycs[:,2], xycs[:,3]

    tfm[:,0,0] = cos
    tfm[:,0,1] = - sin
    tfm[:,1,0] = sin
    tfm[:,1,1] = cos
    tfm[:,0,3] = x
    tfm[:,1,3] = y

    return tfm

def tfm_to_pose_torch(tfm: torch.Tensor, dof: int):
    """
    turn transformation matrix to [x, y, z, roll, yaw, pitch]
    we use degree format.
    tfm is pose in transformation format, and XYZ order, i.e. roll-pitch-yaw

    Args:
        tfm: [N, 4, 4]
        dof: 3 or 6
    Returns:
        6dof pose: [N, 6]
    """

    # There forumlas are designed from x_to_world, but equal to the one below.
    yaw = torch.rad2deg(torch.atan2(tfm[:,1,0], tfm[:,0,0])) # clockwise in carla
    roll = torch.rad2deg(torch.atan2(-tfm[:,2,1], tfm[:,2,2])) # but counter-clockwise in carla
    pitch = torch.rad2deg(torch.atan2(tfm[:,2,0], (tfm[:,2,1]**2 + tfm[:,2,2]**2) ** 0.5)) # but counter-clockwise in carla

    # These formulas are designed for consistent axis orientation
    # yaw = torch.rad2deg(torch.atan2(tfm[:,1,0], tfm[:,0,0])) # clockwise in carla
    # roll = torch.rad2deg(torch.atan2(tfm[:,2,1], tfm[:,2,2])) # but counter-clockwise in carla
    # pitch = torch.rad2deg(torch.atan2(-tfm[:,2,0], (tfm[:,2,1]**2 + tfm[:,2,2]**2) ** 0.5)) # but counter-clockwise in carla

    # roll = - roll
    # pitch = - pitch

    x = tfm[:,0,3]
    y = tfm[:,1,3]
    z = tfm[:,2,3]
    
    if dof == 6:
        pose = torch.stack([x,y,z,roll,yaw,pitch]).T # (N, 6)
    elif dof == 3:
        pose = torch.stack([x,y,yaw]).T
    else:
        raise("Only support returning 3dof/6dof pose.")

    return pose


def x_to_world(pose):
    """
    The transformation matrix from x-coordinate system to carla world system
    Also is the pose in world coordinate: T_world_x

    Parameters
    ----------
    pose : list
        [x, y, z, roll, yaw, pitch], degree

    Returns
    -------
    matrix : np.ndarray
        The transformation matrix.
    """
    x, y, z, roll, yaw, pitch = pose[:]

    # used for rotation matrix
    c_y = np.cos(np.radians(yaw))
    s_y = np.sin(np.radians(yaw))
    c_r = np.cos(np.radians(roll))
    s_r = np.sin(np.radians(roll))
    c_p = np.cos(np.radians(pitch))
    s_p = np.sin(np.radians(pitch))

    matrix = np.identity(4)

    # translation matrix
    matrix[0, 3] = x
    matrix[1, 3] = y
    matrix[2, 3] = z

    # rotation matrix
    matrix[0, 0] = c_p * c_y
    matrix[0, 1] = c_y * s_p * s_r - s_y * c_r
    matrix[0, 2] = -c_y * s_p * c_r - s_y * s_r
    matrix[1, 0] = s_y * c_p
    matrix[1, 1] = s_y * s_p * s_r + c_y * c_r
    matrix[1, 2] = -s_y * s_p * c_r + c_y * s_r
    matrix[2, 0] = s_p
    matrix[2, 1] = -c_p * s_r
    matrix[2, 2] = c_p * c_r

    return matrix


def x1_to_x2(x1, x2):
    """
    Transformation matrix from x1 to x2. T_x2_x1

    Parameters
    ----------
    x1 : list
        The pose of x1 under world coordinates.
    x2 : list
        The pose of x2 under world coordinates.

        yaw, pitch, roll in degree

    Returns
    -------
    transformation_matrix : np.ndarray
        The transformation matrix.

    """
    x1_to_world = x_to_world(x1) # wP = x1_to_world * 1P, so x1_to_world is Tw1
    x2_to_world = x_to_world(x2) # Tw2
    world_to_x2 = np.linalg.inv(x2_to_world) # T2w

    transformation_matrix = np.dot(world_to_x2, x1_to_world) # T2w * Tw1 = T21
    return transformation_matrix


def dist_to_continuous(p_dist, displacement_dist, res, downsample_rate):
    """
    Convert points discretized format to continuous space for BEV representation.
    Parameters
    ----------
    p_dist : numpy.array
        Points in discretized coorindates.

    displacement_dist : numpy.array
        Discretized coordinates of bottom left origin.

    res : float
        Discretization resolution.

    downsample_rate : int
        Dowmsamping rate.

    Returns
    -------
    p_continuous : numpy.array
        Points in continuous coorindates.

    """
    p_dist = np.copy(p_dist)
    p_dist = p_dist + displacement_dist
    p_continuous = p_dist * res * downsample_rate
    return p_continuous


def get_pairwise_transformation_torch(lidar_poses, max_cav, record_len, dof):
    """
    Get pair-wise transformation matrix accross different agents.
    Designed for batch data

    Parameters
    ----------
    lidar_poses : tensor, [N, 3] or [N, 6]
        3 or 6 dof pose of lidar.

    max_cav : int
        The maximum number of cav, default 5

    record: list
        shape (B)

    dof: int, 3 or 6

    Return
    ------
    pairwise_t_matrix : np.array
        The pairwise transformation matrix across each cav.
        shape: (B, L, L, 4, 4), L is the max cav number in a scene
        pairwise_t_matrix[i, j] is Tji, i_to_j
    """
    def regroup(x, record_len):
        cum_sum_len = torch.cumsum(record_len, dim=0)
        split_x = torch.tensor_split(x, cum_sum_len[:-1].cpu())
        return split_x

    B = len(record_len)
    lidar_poses_list = regroup(lidar_poses, record_len)

    pairwise_t_matrix = torch.eye(4, device=lidar_poses.device).view(1,1,1,4,4).repeat(B, max_cav, max_cav, 1, 1) # (B, L, L, 4, 4)
    # save all transformation matrix in a list in order first.
    for b in range(B):
        lidar_poses = lidar_poses_list[b]  # [N_cav, 3] or [N_cav, 6]. 
        t_list = pose_to_tfm(lidar_poses)  # Twx, [N_cav, 4, 4]

        for i in range(len(t_list)):
            for j in range(len(t_list)):
                # identity matrix to self
                if i != j:
                    # i->j: TiPi=TjPj, Tj^(-1)TiPi = Pj
                    # t_matrix = np.dot(np.linalg.inv(t_list[j]), t_list[i])
                    t_matrix = torch.linalg.solve(t_list[j], t_list[i])  # Tjw*Twi = Tji
                    pairwise_t_matrix[b][i, j] = t_matrix

    return pairwise_t_matrix


def get_relative_transformation(lidar_poses):
    """
    Args:
        lidar_pose:  np.ndarray
            [N, dof], lidar pose in world coordinate
            N is the agent number, dof is 3/6.

            [x, y, z, roll, yaw, pitch], degree
        
    Returns:
        relative transformation, in ego's coordinate
    """
    N = lidar_poses.shape[0]
    dof = lidar_poses.shape[1]

    if dof == 3:
        full_lidar_poses = np.zeros((N, 6))
        full_lidar_poses[:,[0,1,4]] = lidar_poses
        lidar_poses = full_lidar_poses

    relative_t_matrix = np.eye(4).reshape(1,4,4).repeat(N, axis=0)  # [N, 4, 4]
    for i in range(1, N):
        relative_t_matrix[i] = x1_to_x2(lidar_poses[i], lidar_poses[0])
    
    return relative_t_matrix



def muilt_coord(rotationA2B, translationA2B, rotationB2C, translationB2C):
    rotationA2B = np.array(rotationA2B).reshape(3, 3)
    rotationB2C = np.array(rotationB2C).reshape(3, 3)
    rotation = np.dot(rotationB2C, rotationA2B)
    translationA2B = np.array(translationA2B).reshape(3, 1)
    translationB2C = np.array(translationB2C).reshape(3, 1)
    translation = np.dot(rotationB2C, translationA2B) + translationB2C

    return rotation, translation


def veh_side_rot_and_trans_to_trasnformation_matrix(lidar_to_novatel_json_file,novatel_to_world_json_file):
    matrix = np.empty([4,4])
    rotationA2B = lidar_to_novatel_json_file["transform"]["rotation"]
    translationA2B = lidar_to_novatel_json_file["transform"]["translation"]
    rotationB2C = novatel_to_world_json_file["rotation"]
    translationB2C = novatel_to_world_json_file["translation"]
    rotation,translation = muilt_coord(rotationA2B, translationA2B, rotationB2C, translationB2C)
    matrix[0:3, 0:3] = rotation
    matrix[:, 3][0:3] = np.array(translation)[:, 0]
    matrix[3, 0:3] = 0
    matrix[3, 3] = 1
    
    return matrix

def inf_side_rot_and_trans_to_trasnformation_matrix(json_file,system_error_offset):
    matrix = np.empty([4,4])
    matrix[0:3, 0:3] = json_file["rotation"]
    translation = np.array(json_file["translation"])
    translation[0][0] = translation[0][0] + system_error_offset["delta_x"]
    translation[1][0] = translation[1][0] + system_error_offset["delta_y"]  #为啥有[1][0]??? --> translation是(3,1)的
    matrix[:, 3][0:3] = translation[:, 0]
    matrix[3, 0:3] = 0
    matrix[3, 3] = 1

    return matrix



def test():
    random_pose = np.random.randn(6)
    tfm = x_to_world(random_pose)
    pose_result = tfm_to_pose(tfm)
    tfm2 = x_to_world(pose_result)

    print(random_pose)
    print(pose_result)
    print()
    print(tfm)
    print(tfm2)

if __name__ == "__main__":
    test()
opencood/utils/transformation_utils.py
MediaBrain-SJTU/CoBEVFlow
from operator import gt
import numpy as np
import pickle
from pyquaternion import Quaternion
from matplotlib import pyplot as plt
from icecream import ic
from torch import margin_ranking_loss



def visualize(pred_box_tensor, gt_tensor, pcd, show_vis, save_path, dataset=None):
        """
        Visualize the prediction, ground truth with point cloud together.

        Parameters
        ----------
        pred_box_tensor : torch.Tensor
            (N, 8, 3) prediction.

        gt_tensor : torch.Tensor
            (N, 8, 3) groundtruth bbx

        pcd : torch.Tensor
            PointCloud, (N, 4).

        show_vis : bool
            Whether to show visualization.

        save_path : str
            Save the visualization results to given path.

        dataset : BaseDataset
            opencood dataset object.

        """

        pcd_np = pcd.cpu().numpy()
        pred_box_np = pred_box_tensor.cpu().numpy()
        gt_box_np = gt_tensor.cpu().numpy()

        plt.figure(dpi=400)
        # draw point cloud. It's in lidar coordinate
        plt.scatter(pcd_np[:,0], pcd_np[:,1], s=0.5)

        N = gt_tensor.shape[0]
        for i in range(N):
            plt.plot(gt_box_np[i,:,0], gt_box_np[i,:,1], c= "r", marker='.', linewidth=1, markersize=1.5)

        N = pred_box_tensor.shape[0]
        for i in range(N):
            plt.plot(pred_box_np[i,:,0], pred_box_np[i,:,1], c= "g", marker='.', linewidth=1, markersize=1.5)
        

        plt.savefig(save_path)
        plt.clf()
opencood/visualization/my_vis.py
MediaBrain-SJTU/CoBEVFlow
from .canvas_3d import Canvas_3D
from .canvas_bev import Canvas_BEV

opencood/visualization/simple_plot3d/__init__.py
MediaBrain-SJTU/CoBEVFlow
"""
Written by Jinhyung Park

Simple 3D visualization for 3D points & boxes. Intended as a simple, hackable
alternative to mayavi for certain point cloud tasks.
"""

import numpy as np
import cv2
import copy
from functools import partial
import matplotlib

class Canvas_3D(object):
    def __init__(self,
                 canvas_shape=(500, 1000),
                 camera_center_coords=(-15, 0, 10),
                 camera_focus_coords=(-15 + 0.9396926, 0, 10 - 0.44202014),
                #  camera_center_coords=(-25, 0, 20),
                #  camera_focus_coords=(-25 + 0.9396926, 0, 20 - 0.64202014),
                 focal_length=None,
                 canvas_bg_color=(0, 0, 0), 
                 left_hand=True):
        """
        Args:
            canvas_shape (Tuple[Int]): Canvas image size - height & width.
            camera_center_coords (Tuple[Float]): Location of camera center in
                3D space.
            camera_focus_coords (Tuple[Float]): Intuitively, what point in 3D 
                space is the camera pointed at? These are absolute coordinates,
                *not* relative to camera center.
            focal_length (None | Int):
                None: Half of the max of height & width of canvas_shape. This
                    seems to be a decent default.
                Int: Specified directly.
            canvas_bg_color (Tuple[Int]): RGB (0 ~ 255) of canvas background
                color.
            left_hand: bool
                Since this algorithm is designed for right hand coord. We take -y if True
        """
        
        self.canvas_shape = canvas_shape
        self.H, self.W = self.canvas_shape
        self.canvas_bg_color = canvas_bg_color
        self.left_hand = left_hand
        if left_hand:
            camera_center_coords = list(camera_center_coords)
            camera_center_coords[1] = - camera_center_coords[1]
            camera_center_coords = tuple(camera_center_coords)
            
            camera_focus_coords = list(camera_focus_coords)
            camera_focus_coords[1] = - camera_focus_coords[1]
            camera_focus_coords = tuple(camera_focus_coords)

        self.camera_center_coords = camera_center_coords
        self.camera_focus_coords = camera_focus_coords

        if focal_length is None:
            self.focal_length = max(self.H, self.W) // 2
        else:
            self.focal_length = focal_length

        # Setup extrinsics and intrinsics of this virtual camera.
        self.ext_matrix = self.get_extrinsic_matrix(
            self.camera_center_coords, self.camera_focus_coords)
        self.int_matrix = np.array([
            [self.focal_length, 0, self.W // 2, 0],
            [0, self.focal_length, self.H // 2, 0],
            [0, 0, 1, 0],
        ])
        
        self.clear_canvas()
    
    def get_canvas(self):
        return self.canvas

    def clear_canvas(self):
        self.canvas = np.zeros((self.H, self.W, 3), dtype=np.uint8)
        self.canvas[..., :] = self.canvas_bg_color

    def get_canvas_coords(self, 
                          xyz, 
                          depth_min=0.1,
                          return_depth=False):
        """
        Projects XYZ points onto the canvas and returns the projected canvas
        coordinates.

        Args:
            xyz (ndarray): (N, 3+) array of coordinates. Additional columns
                beyond the first three are ignored.
            depth_min (Float): Only points with a projected depth larger
                than this value are "valid".
            return_depth (Boolean): Whether to additionally return depth of
                projected points.
        Returns:
            canvas_xy (ndarray): (N, 2) array of projected canvas coordinates.
                "x" is dim0, "y" is dim1 of canvas.
            valid_mask (ndarray): (N,) boolean mask indicating which of 
                canvas_xy fits into canvas (are visible from virtual camera).
            depth (ndarray): Optionally returned (N,) array of depth values
        """
        if self.left_hand:
            xyz[:,1] = - xyz[:,1]

        xyz = xyz[:,:3]
        xyz_hom = np.concatenate(
            [xyz, np.ones((xyz.shape[0], 1), dtype=np.float32)], axis=1)
        img_pts = (self.int_matrix @ self.ext_matrix @ xyz_hom.T).T

        depth = img_pts[:, 2]
        xy = img_pts[:, :2] / depth[:, None]
        xy_int = xy.round().astype(np.int32)

        # Flip X and Y so "x" is dim0, "y" is dim1 of canvas
        xy_int = xy_int[:, ::-1] 

        valid_mask = ((depth > depth_min) &
                      (xy_int[:, 0] >= 0) & (xy_int[:, 0] < self.H) & 
                      (xy_int[:, 1] >= 0) & (xy_int[:, 1] < self.W))
        
        if return_depth:
            return xy_int, valid_mask, depth
        else:
            return xy_int, valid_mask
                                      

    def draw_canvas_points(self, 
                           canvas_xy,
                           radius=-1,
                           colors=None,
                           colors_operand=None):
        """
        Draws canvas_xy onto self.canvas.

        Args:
            canvas_xy (ndarray): (N, 2) array of *valid* canvas coordinates.
                "x" is dim0, "y" is dim1 of canvas.
            radius (Int): 
                -1: Each point is visualized as a single pixel.
                r: Each point is visualized as a circle with radius r.
            colors: 
                None: colors all points white.
                Tuple: RGB (0 ~ 255), indicating a single color for all points.
                ndarray: (N, 3) array of RGB values for each point.
                String: Such as "Spectral", uses a matplotlib cmap, with the
                    operand (the value cmap is called on for each point) being 
                    colors_operand.
            colors_operand (ndarray): (N,) array of values cooresponding to 
                canvas_xy, to be used only if colors is a cmap. Unlike 
                Canvas_BEV, cannot be None if colors is a String.
        """
        if len(canvas_xy) == 0:
            return 
            
        if colors is None:
            colors = np.full(
                (len(canvas_xy), 3), fill_value=255, dtype=np.uint8)
        elif isinstance(colors, tuple):
            assert len(colors) == 3
            colors_tmp = np.zeros((len(canvas_xy), 3), dtype=np.uint8)
            colors_tmp[..., :len(colors)] = np.array(colors)
            colors = colors_tmp
        elif isinstance(colors, np.ndarray):
            assert len(colors) == len(canvas_xy)
            colors = colors.astype(np.uint8)
        elif isinstance(colors, str):
            assert colors_operand is not None
            colors = matplotlib.cm.get_cmap(colors)
                    
            # Normalize 0 ~ 1 for cmap
            colors_operand = colors_operand - colors_operand.min()
            colors_operand = colors_operand / colors_operand.max()
        
            # Get cmap colors - note that cmap returns (*input_shape, 4), with
            # colors scaled 0 ~ 1
            colors = (colors(colors_operand)[:, :3] * 255).astype(np.uint8)
        else:
            raise Exception(
                "colors type {} was not an expected type".format(type(colors)))

        if radius == -1:
            self.canvas[canvas_xy[:, 0], canvas_xy[:, 1], :] = colors
        else:
            for color, (x, y) in zip(colors.tolist(), canvas_xy.tolist()):
                self.canvas = cv2.circle(self.canvas, (y, x), radius, color, 
                                         -1, lineType=cv2.LINE_AA)
        
    def draw_lines(self,
                   canvas_xy, 
                   start_xyz,
                   end_xyz,
                   colors=(255, 255, 255),
                   thickness=1):
        """
        Draws lines between provided 3D points.
        
        Args:
            # added from original repo
            canvas_xy (ndarray): (N, 2) array of *valid* canvas coordinates.
                    "x" is dim0, "y" is dim1 of canvas.

            start_xyz (ndarray): Shape (N, 3) of 3D points to start from.
            end_xyz (ndarray): Shape (N, 3) of 3D points to end at. Same length
                as start_xyz.
            colors: 
                None: colors all points white.
                Tuple: RGB (0 ~ 255), indicating a single color for all points.
                ndarray: (N, 3) array of RGB values for each point.
            thickness (Int):
                Thickness of drawn cv2 line.            
        """
        if colors is None:
            colors = np.full(
                (len(canvas_xy), 3), fill_value=255, dtype=np.uint8)
        elif isinstance(colors, tuple):
            assert len(colors) == 3
            colors_tmp = np.zeros((len(canvas_xy), 3), dtype=np.uint8)
            colors_tmp[..., :len(colors)] = np.array(colors)
            colors = colors_tmp
        elif isinstance(colors, np.ndarray):
            assert len(colors) == len(canvas_xy)
            colors = colors.astype(np.uint8)
        else:
            raise Exception(
                "colors type {} was not an expected type".format(type(colors)))
        
        start_pts_xy, start_pts_valid_mask, start_pts_d = \
            self.get_canvas_coords(start_xyz, True)
        end_pts_xy, end_pts_valid_mask, end_pts_d = \
            self.get_canvas_coords(end_xyz, True)

        for idx, (color, start_pt_xy, end_pt_xy) in enumerate(
                zip(colors.tolist(), start_pts_xy.tolist(), 
                    end_pts_xy.tolist())):
                
            if start_pts_valid_mask[idx] and end_pts_valid_mask[idx]:
                self.canvas = cv2.line(self.canvas,
                                    tuple(start_pt_xy[::-1]),
                                    tuple(end_pt_xy[::-1]),
                                    color=color,
                                    thickness=thickness, 
                                    lineType=cv2.LINE_AA)

    def draw_boxes(self,
                   boxes,
                   colors=None,
                   texts=None,
                   depth_min=0.1,
                   draw_incomplete_boxes=False,
                   box_line_thickness=2,
                   box_text_size=0.5,
                   text_corner=1):
        """
        Draws 3D boxes.

        Args:
            boxes (ndarray): Shape (N, 8, 3), corners in 3d
                modified from original repo

            colors: 
                None: colors all points white.
                Tuple: RGB (0 ~ 255), indicating a single color for all points.
                ndarray: (N, 3) array of RGB values for each point.

            texts (List[String]): Length N; text to write next to boxes.

            depth_min (Float): Only box corners with a projected depth larger
                than this value are drawn if draw_incomplete_boxes is True.

            draw_incomplete_boxes (Boolean): If any boxes are incomplete,
                meaning it has a corner out of view based on depth_min, decide
                whether to draw them at all.

            thickness (Int):
                Thickness of drawn cv2 box lines. 

            box_line_thickness (int): cv2 line/text thickness
            box_text_size (float): cv2 putText size
            text_corner (int): 0 ~ 7. Which corner of 3D box to write text at.
        """
        # Setup colors
        if colors is None:
            colors = np.full(
                (len(boxes), 3), fill_value=255, dtype=np.uint8)
        elif isinstance(colors, tuple):
            assert len(colors) == 3
            colors_tmp = np.zeros((len(boxes), 3), dtype=np.uint8)
            colors_tmp[..., :len(colors)] = np.array(colors)
            colors = colors_tmp
        elif isinstance(colors, np.ndarray):
            assert len(colors) == len(boxes)
            colors = colors.astype(np.uint8)
        else:
            raise Exception(
                "colors type {} was not an expected type".format(type(colors)))
        
        

        corners = boxes # N x 8 x 3

        # Now we have corners. Need them on the canvas 2D space.
        corners_xy, valid_mask = self.get_canvas_coords(
            corners.reshape(-1, 3), depth_min=depth_min)
        corners_xy = corners_xy.reshape(-1, 8, 2)
        valid_mask = valid_mask.reshape(-1, 8)

        # Now draw them with lines in correct places
        for i, (color, curr_corners_xy, curr_valid_mask) in enumerate(
            zip(colors.tolist(), corners_xy.tolist(), valid_mask.tolist())):

            if not draw_incomplete_boxes and sum(curr_valid_mask) != 8:
                # Some corner is invalid, don't draw the box at all.
                continue

            for start, end in [(0, 1), (1, 2), (2, 3), (3, 0),
                               (0, 4), (1, 5), (2, 6), (3, 7),
                               (4, 5), (5, 6), (6, 7), (7, 4)]:
                if not (curr_valid_mask[start] and curr_valid_mask[end]):
                    continue # start or end is not valid
                    
                self.canvas = cv2.line(
                    self.canvas, 
                    (curr_corners_xy[start][1], curr_corners_xy[start][0]),
                    (curr_corners_xy[end][1], curr_corners_xy[end][0]),
                    color=color,
                    thickness=box_line_thickness, 
                    lineType=cv2.LINE_AA)
            
            # If even a single line was drawn, add text as well.
            if sum(curr_valid_mask) > 0:
                if texts is not None:
                    self.canvas = cv2.putText(self.canvas,
                                            str(texts[i]),
                                            (curr_corners_xy[text_corner][1], 
                                             curr_corners_xy[text_corner][0]),
                                            cv2.FONT_HERSHEY_SIMPLEX,
                                            box_text_size,
                                            color,
                                            thickness=box_line_thickness)


    @staticmethod
    def cart2sph(xyz):
        x, y, z = xyz[:, 0], xyz[:, 1], xyz[:, 2]
        
        depth = np.linalg.norm(xyz, 2, axis=1)
        az = -np.arctan2(y, x)
        el = np.arcsin(z / depth)
        return az, el, depth

    @staticmethod
    def get_extrinsic_matrix(
        camera_center_coords,
        camera_focus_coords,
    ):
        """
        Args:
            camera_center_coords: (x, y, z) of where camera should be located 
                in 3D space.
            camera_focus_coords: (x, y, z) of where camera should look at from 
                camera_center_coords
            
        Thoughts:
            Remember that in camera coordiantes, pos x is right, pos y is up, 
                pos z is forward.
        """
        center_x, center_y, center_z = camera_center_coords
        focus_x, focus_y, focus_z = camera_focus_coords
        az, el, depth = Canvas_3D.cart2sph(np.array([
            [focus_x - center_x, focus_y - center_y, focus_z - center_z]
        ]))
        az = float(az)
        el = float(el)
        depth = float(depth)
        
        ### First, construct extrinsics
        ## Rotation matrix
        
        z_rot = np.array([
            [np.cos(az), -np.sin(az), 0],
            [np.sin(az), np.cos(az), 0],
            [0, 0, 1]
        ])
        
        # el is rotation around y axis. 
        y_rot = np.array([
            [np.cos(-el), 0, -np.sin(-el)],
            [0, 1, 0],
            [np.sin(-el), 0, np.cos(-el)],
        ])
        
        
        ## Now, how the z_rot and y_rot work (spherical coordiantes), is it 
        ## computes rotations starting from the positive x axis and rotates 
        ## positive x axis to the desired direction. The desired direction is 
        ## the "looking direction" of the camera, which should actually be the 
        ## z-axis. So should convert the points so that the x axis is the new z 
        ## axis, and after the transformations.
        ## Why x -> z for points? If we think about rotating the camera, z 
        ## should become x, so reverse when moving points.
        last_rot = np.array([
            [0, -1, 0],
            [0, 0, -1],
            [1, 0, 0] # x -> z
        ])
        
        # Put them together. Order matters. Make it hom.
        rot_matrix = np.eye(4, dtype=np.float32)
        rot_matrix[:3, :3] = last_rot @ y_rot @ z_rot
        
        ## Translation matrix
        trans_matrix = np.array([
            [1, 0, 0, -center_x],
            [0, 1, 0, -center_y],
            [0, 0, 1, -center_z],
            [0, 0, 0, 1],
        ])
        
        ## Finally, extrinsics matrix. Order matters - do trans then rot
        ext_matrix =  rot_matrix @ trans_matrix
        
        return ext_matrix
opencood/visualization/simple_plot3d/canvas_3d.py
MediaBrain-SJTU/CoBEVFlow
"""
Written by Jinhyung Park

Simple BEV visualization for 3D points & boxes.
"""

import numpy as np
import cv2
import copy
from functools import partial
import matplotlib

class Canvas_BEV(object):
    def __init__(self, 
                 canvas_shape=(1000, 1000),
                 canvas_x_range=(-50, 50),
                 canvas_y_range=(-50, 50),
                 canvas_bg_color=(0, 0, 0),
                 left_hand=False):
        """
        Args:
            canvas_shape (Tuple[int]): Shape of BEV Canvas image. First element
                corresponds to X range, the second element to Y range.
            canvas_x_range (Tuple[int]): Range of X-coords to visualize. X is
                vertical: negative ~ positive is top ~ down.
            canvas_y_range (Tuple[int]): Range of Y-coords to visualize. Y is
                horizontal: negative ~ positive is left ~ right.
            canvas_bg_color (Tuple[int]): RGB (0 ~ 255) of Canvas background
                color.
            left_hand: (bool), whether the point cloud is left-hand coordinate,
                V2X-Sim is right hand, and OPV2V is left hand.

            Vehicle is heading down. This code is naturally for right-hand coordinate.
        """
        
        # Sanity check ratios
        if ((canvas_shape[0] / canvas_shape[1]) != 
            ((canvas_x_range[0] - canvas_x_range[1]) / 
             (canvas_y_range[0] - canvas_y_range[1]))):

            print("Not an error, but the x & y ranges are not "\
                  "proportional to canvas height & width.")
        
        self.canvas_shape = canvas_shape
        self.canvas_x_range = canvas_x_range
        self.canvas_y_range = canvas_y_range
        self.canvas_bg_color = canvas_bg_color
        self.left_hand = left_hand
        
        self.clear_canvas()
    
    def get_canvas(self):
        return self.canvas

    def clear_canvas(self):
        self.canvas = np.zeros((*self.canvas_shape, 3), dtype=np.uint8)
        self.canvas[..., :] = self.canvas_bg_color

    def get_canvas_coords(self, xy):
        """
        Args:
            xy (ndarray): (N, 2+) array of coordinates. Additional columns
                beyond the first two are ignored.
        
        Returns:
            canvas_xy (ndarray): (N, 2) array of xy scaled into canvas 
                coordinates. Invalid locations of canvas_xy are clipped into 
                range. "x" is dim0, "y" is dim1 of canvas.
            valid_mask (ndarray): (N,) boolean mask indicating which of 
                canvas_xy fits into canvas.
        """
        xy = np.copy(xy) # prevent in-place modifications

        x = xy[:, 0]
        y = xy[:, 1]

        if self.left_hand:
            y = -y

        # Get valid mask
        valid_mask = ((x > self.canvas_x_range[0]) & 
                      (x < self.canvas_x_range[1]) &
                      (y > self.canvas_y_range[0]) & 
                      (y < self.canvas_y_range[1]))
        
        # Rescale points
        x = ((x - self.canvas_x_range[0]) / 
             (self.canvas_x_range[1] - self.canvas_x_range[0]))
        x = x * self.canvas_shape[0]
        x = np.clip(np.around(x), 0, 
                    self.canvas_shape[0] - 1).astype(np.int32)
                    
        y = ((y - self.canvas_y_range[0]) / 
             (self.canvas_y_range[1] - self.canvas_y_range[0]))
        y = y * self.canvas_shape[1]
        y = np.clip(np.around(y), 0, 
                    self.canvas_shape[1] - 1).astype(np.int32)
        
        # Return
        canvas_xy = np.stack([x, y], axis=1)

        return canvas_xy, valid_mask
                                      

    def draw_canvas_points(self, 
                           canvas_xy,
                           radius=-1,
                           colors=None,
                           colors_operand=None):
        """
        Draws canvas_xy onto self.canvas.

        Args:
            canvas_xy (ndarray): (N, 2) array of *valid* canvas coordinates.
                "x" is dim0, "y" is dim1 of canvas.
            radius (Int): 
                -1: Each point is visualized as a single pixel.
                r: Each point is visualized as a circle with radius r.
            colors: 
                None: colors all points white.
                Tuple: RGB (0 ~ 255), indicating a single color for all points.
                ndarray: (N, 3) array of RGB values for each point.
                String: Such as "Spectral", uses a matplotlib cmap, with the
                    operand (the value cmap is called on for each point) being 
                    colors_operand. If colors_operand is None, uses normalized
                    distance from (0, 0) of XY point coords.
            colors_operand (ndarray | None): (N,) array of values cooresponding
                to canvas_xy, to be used only if colors is a cmap.
        """
        if len(canvas_xy) == 0:
            return 
            
        if colors is None:
            colors = np.full(
                (len(canvas_xy), 3), fill_value=255, dtype=np.uint8)
        elif isinstance(colors, tuple):
            assert len(colors) == 3
            colors_tmp = np.zeros((len(canvas_xy), 3), dtype=np.uint8)
            colors_tmp[..., :] = np.array(colors)
            colors = colors_tmp
        elif isinstance(colors, np.ndarray):
            assert len(colors) == len(canvas_xy)
            colors = colors.astype(np.uint8)
        elif isinstance(colors, str):
            colors = matplotlib.cm.get_cmap(colors)
            if colors_operand is None:
                # Get distances from (0, 0) (albeit potentially clipped)
                origin_center = self.get_canvas_coords(np.zeros((1, 2)))[0][0]
                colors_operand = np.sqrt(
                    ((canvas_xy - origin_center) ** 2).sum(axis=1))
                    
            # Normalize 0 ~ 1 for cmap
            colors_operand = colors_operand - colors_operand.min()
            colors_operand = colors_operand / colors_operand.max()
        
            # Get cmap colors - note that cmap returns (*input_shape, 4), with
            # colors scaled 0 ~ 1
            colors = (colors(colors_operand)[:, :3] * 255).astype(np.uint8)
        else:
            raise Exception(
                "colors type {} was not an expected type".format(type(colors)))

        if radius == -1:
            self.canvas[canvas_xy[:, 0], canvas_xy[:, 1], :] = colors
        else:
            for color, (x, y) in zip(colors.tolist(), canvas_xy.tolist()):
                self.canvas = cv2.circle(self.canvas, (y, x), radius, color, 
                                         -1, lineType=cv2.LINE_AA)

    def draw_boxes(self,
                   boxes,
                   colors=None,
                   texts=None,
                   box_line_thickness=2,
                   box_text_size=0.5,
                   text_corner=0):
        """
        Draws a set of boxes onto the canvas.
        Args:
            boxes (ndarray): [N, 8, 3] corner 3d
                
            colors: 
                None: colors all points white.
                Tuple: RGB (0 ~ 255), indicating a single color for all points.
                ndarray: (N, 3) array of RGB values for each point.
            texts (List[String]): Length N; text to write next to boxes.
            box_line_thickness (int): cv2 line/text thickness
            box_text_size (float): cv2 putText size
            text_corner (int): 0 ~ 3. Which corner of 3D box to write text at.
        """
        # Setup colors
        if colors is None:
            colors = np.full((len(boxes), 3), fill_value=255, dtype=np.uint8)
        elif isinstance(colors, tuple):
            assert len(colors) == 3
            colors_tmp = np.zeros((len(boxes), 3), dtype=np.uint8)
            colors_tmp[..., :len(colors)] = np.array(colors)
            colors = colors_tmp
        elif isinstance(colors, np.ndarray):
            assert len(colors) == len(boxes)
            colors = colors.astype(np.uint8)
        else:
            raise Exception(
                "colors type {} was not an expected type".format(type(colors)))

        boxes = np.copy(boxes) # prevent in-place modifications
        

        # Translate BEV 4 corners , [N, 4, 2]
        #     4 -------- 5
        #    /|         /|
        #   7 -------- 6 .
        #   | |        | |
        #   . 0 -------- 1
        #   |/         |/
        #   3 -------- 2
        bev_corners = boxes[:,:4,:2]

        ## Transform BEV 4 corners to canvas coords
        bev_corners_canvas, valid_mask = \
            self.get_canvas_coords(bev_corners.reshape(-1, 2))
        bev_corners_canvas = bev_corners_canvas.reshape(*bev_corners.shape)
        valid_mask = valid_mask.reshape(*bev_corners.shape[:-1])

        # At least 1 corner in canvas to draw.
        valid_mask = valid_mask.sum(axis=1) > 0
        bev_corners_canvas = bev_corners_canvas[valid_mask]
        if texts is not None:
            texts = np.array(texts)[valid_mask]

        ## Draw onto canvas
        # Draw the outer boundaries
        idx_draw_pairs = [(0, 1), (1, 2), (2, 3), (3, 0)]
        for i, (color, curr_box_corners) in enumerate(
                zip(colors.tolist(), bev_corners_canvas)):
                
            curr_box_corners = curr_box_corners.astype(np.int32)
            for start, end in idx_draw_pairs:
                self.canvas = cv2.line(self.canvas,
                                       tuple(curr_box_corners[start][::-1]\
                                        .tolist()),
                                       tuple(curr_box_corners[end][::-1]\
                                        .tolist()),
                                       color=color,
                                       thickness=box_line_thickness)
            if texts is not None:
                self.canvas = cv2.putText(self.canvas,
                                          str(texts[i]),
                                          tuple(curr_box_corners[text_corner]\
                                            [::-1].tolist()),
                                          cv2.FONT_HERSHEY_SIMPLEX,
                                          box_text_size,
                                          color=color,
                                          thickness=box_line_thickness)

class Canvas_BEV_heading_right(object):
    def __init__(self, 
                 canvas_shape=(800, 2800),
                 canvas_x_range=(-140, 140),
                 canvas_y_range=(-40, 40),
                 canvas_bg_color=(0, 0, 0),
                 left_hand=True):
        """
        Args:
            canvas_shape (Tuple[int]): Shape of BEV Canvas image. First element
                corresponds to Y range, the second element to X range.
            canvas_x_range (Tuple[int]): Range of X-coords to visualize. X is
                horizontal: negative ~ positive is left ~ right.
            canvas_y_range (Tuple[int]): Range of Y-coords to visualize. Y is
                vertcal: negative ~ positive is top ~ down.
            canvas_bg_color (Tuple[int]): RGB (0 ~ 255) of Canvas background
                color.
            left_hand: (bool), whether the point cloud is left-hand coordinate
                V2X-Sim is right hand, and OPV2V is left hand.

            Different from Canvas_BEV, the vehicle is heading right.
            Naturally this code is designed for left hand coordinate

        """
        
        # Sanity check ratios
        if ((canvas_shape[1] / canvas_shape[0]) != 
            ((canvas_x_range[0] - canvas_x_range[1]) / 
             (canvas_y_range[0] - canvas_y_range[1]))):

            print("Not an error, but the x & y ranges are not "\
                  "proportional to canvas height & width.")
        
        self.canvas_shape = canvas_shape
        self.canvas_x_range = canvas_x_range
        self.canvas_y_range = canvas_y_range
        self.canvas_bg_color = canvas_bg_color
        self.left_hand = left_hand
        
        self.clear_canvas()
    
    def get_canvas(self):
        return self.canvas

    def clear_canvas(self):
        self.canvas = np.zeros((*self.canvas_shape, 3), dtype=np.uint8)
        self.canvas[..., :] = self.canvas_bg_color

    def get_canvas_coords(self, xy):
        """
        Args:
            xy (ndarray): (N, 2+) array of coordinates. Additional columns
                beyond the first two are ignored.
        
        Returns:
            canvas_xy (ndarray): (N, 2) array of xy scaled into canvas 
                coordinates. Invalid locations of canvas_xy are clipped into 
                range. "x" is dim0, "y" is dim1 of canvas.
            valid_mask (ndarray): (N,) boolean mask indicating which of 
                canvas_xy fits into canvas.
        """
        xy = np.copy(xy) # prevent in-place modifications

        x = xy[:, 0]
        y = xy[:, 1]

        if not self.left_hand:
            y = -y

        # Get valid mask
        valid_mask = ((x > self.canvas_x_range[0]) & 
                      (x < self.canvas_x_range[1]) &
                      (y > self.canvas_y_range[0]) & 
                      (y < self.canvas_y_range[1]))

        # Rescale points
        # They are exactly lidar point coordinate
        x = ((x - self.canvas_x_range[0]) / 
             (self.canvas_x_range[1] - self.canvas_x_range[0]))
        x = x * self.canvas_shape[1]
        x = np.clip(np.around(x), 0, 
                    self.canvas_shape[1] - 1).astype(np.int32) # [0,2800-1]
                    
        y = ((y - self.canvas_y_range[0]) / 
             (self.canvas_y_range[1] - self.canvas_y_range[0]))
        y = y * self.canvas_shape[0]
        y = np.clip(np.around(y), 0, 
                    self.canvas_shape[0] - 1).astype(np.int32) # [0,800-1]
        
        # x and y are exactly image coordinate
        # ------------> x
        # |
        # |
        # |
        # y

        canvas_xy = np.stack([x, y], axis=1)

        return canvas_xy, valid_mask
                                      

    def draw_canvas_points(self, 
                           canvas_xy,
                           radius=-1,
                           colors=None,
                           colors_operand=None):
        """
        Draws canvas_xy onto self.canvas.

        Args:
            canvas_xy (ndarray): (N, 2) array of *valid* canvas coordinates.
                
            radius (Int): 
                -1: Each point is visualized as a single pixel.
                r: Each point is visualized as a circle with radius r.
            colors: 
                None: colors all points white.
                Tuple: RGB (0 ~ 255), indicating a single color for all points.
                ndarray: (N, 3) array of RGB values for each point.
                String: Such as "Spectral", uses a matplotlib cmap, with the
                    operand (the value cmap is called on for each point) being 
                    colors_operand. If colors_operand is None, uses normalized
                    distance from (0, 0) of XY point coords.
            colors_operand (ndarray | None): (N,) array of values cooresponding
                to canvas_xy, to be used only if colors is a cmap.
        """
        if len(canvas_xy) == 0:
            return 
            
        if colors is None:
            colors = np.full(
                (len(canvas_xy), 3), fill_value=255, dtype=np.uint8)
        elif isinstance(colors, tuple):
            assert len(colors) == 3
            colors_tmp = np.zeros((len(canvas_xy), 3), dtype=np.uint8)
            colors_tmp[..., :] = np.array(colors)
            colors = colors_tmp
        elif isinstance(colors, np.ndarray):
            assert len(colors) == len(canvas_xy)
            colors = colors.astype(np.uint8)
        elif isinstance(colors, str):
            colors = matplotlib.cm.get_cmap(colors)
            if colors_operand is None:
                # Get distances from (0, 0) (albeit potentially clipped)
                origin_center = self.get_canvas_coords(np.zeros((1, 2)))[0][0]
                colors_operand = np.sqrt(
                    ((canvas_xy - origin_center) ** 2).sum(axis=1))
                    
            # Normalize 0 ~ 1 for cmap
            colors_operand = colors_operand - colors_operand.min()
            colors_operand = colors_operand / colors_operand.max()
        
            # Get cmap colors - note that cmap returns (*input_shape, 4), with
            # colors scaled 0 ~ 1
            colors = (colors(colors_operand)[:, :3] * 255).astype(np.uint8)
        else:
            raise Exception(
                "colors type {} was not an expected type".format(type(colors)))

        # Here the order is different from Canvas_BEV
        if radius == -1:
            self.canvas[canvas_xy[:, 1], canvas_xy[:, 0], :] = colors
        else:
            for color, (x, y) in zip(colors.tolist(), canvas_xy.tolist()):
                self.canvas = cv2.circle(self.canvas, (x, y), radius, color, 
                                         -1, lineType=cv2.LINE_AA)

    def draw_boxes(self,
                   boxes,
                   colors=None,
                   texts=None,
                   box_line_thickness=3,
                   box_text_size=0.5,
                   text_corner=0):
        """
        Draws a set of boxes onto the canvas.
        Args:
            boxes (ndarray): [N, 8, 3] corner 3d
                
            colors: 
                None: colors all points white.
                Tuple: RGB (0 ~ 255), indicating a single color for all points.
                ndarray: (N, 3) array of RGB values for each point.
            texts (List[String]): Length N; text to write next to boxes.
            box_line_thickness (int): cv2 line/text thickness
            box_text_size (float): cv2 putText size
            text_corner (int): 0 ~ 3. Which corner of 3D box to write text at.
        """
        # Setup colors
        if colors is None:
            colors = np.full((len(boxes), 3), fill_value=255, dtype=np.uint8)
        elif isinstance(colors, tuple):
            assert len(colors) == 3
            colors_tmp = np.zeros((len(boxes), 3), dtype=np.uint8)
            colors_tmp[..., :len(colors)] = np.array(colors)
            colors = colors_tmp
        elif isinstance(colors, np.ndarray):
            assert len(colors) == len(boxes)
            colors = colors.astype(np.uint8)
        else:
            raise Exception(
                "colors type {} was not an expected type".format(type(colors)))

        boxes = np.copy(boxes) # prevent in-place modifications
        

        # Translate BEV 4 corners , [N, 4, 2]
        #     4 -------- 5
        #    /|         /|
        #   7 -------- 6 .
        #   | |        | |
        #   . 0 -------- 1
        #   |/         |/
        #   3 -------- 2
        bev_corners = boxes[:,:4,:2]

        ## Transform BEV 4 corners to canvas coords
        bev_corners_canvas, valid_mask = \
            self.get_canvas_coords(bev_corners.reshape(-1, 2))  # [N, 2]
        bev_corners_canvas = bev_corners_canvas.reshape(*bev_corners.shape)  # [N, 4, 2]
        valid_mask = valid_mask.reshape(*bev_corners.shape[:-1]) 

        # At least 1 corner in canvas to draw.
        valid_mask = valid_mask.sum(axis=1) > 0
        bev_corners_canvas = bev_corners_canvas[valid_mask]
        if texts is not None:
            texts = np.array(texts)[valid_mask]

        ## Draw onto canvas
        # Draw the outer boundaries
        idx_draw_pairs = [(0, 1), (1, 2), (2, 3), (3, 0)]
        for i, (color, curr_box_corners) in enumerate(
                zip(colors.tolist(), bev_corners_canvas)):
                
            curr_box_corners = curr_box_corners.astype(np.int32)
            for start, end in idx_draw_pairs:
                # Notice Difference Here
                self.canvas = cv2.line(self.canvas,
                                       tuple(curr_box_corners[start]\
                                        .tolist()),
                                       tuple(curr_box_corners[end]\
                                        .tolist()),
                                       color=color,
                                       thickness=box_line_thickness)
            if texts is not None:
                self.canvas = cv2.putText(self.canvas,
                                          str(texts[i]),
                                          tuple(curr_box_corners[text_corner]\
                                            .tolist()),
                                          cv2.FONT_HERSHEY_SIMPLEX,
                                          box_text_size,
                                          color=color,
                                          thickness=box_line_thickness)

opencood/visualization/simple_plot3d/canvas_bev.py
MediaBrain-SJTU/CoBEVFlow
from matplotlib import pyplot as plt
import numpy as np

import opencood.visualization.simple_plot3d.canvas_3d as canvas_3d
import opencood.visualization.simple_plot3d.canvas_bev as canvas_bev

def visualize(pred_box_tensor, gt_tensor, pcd, pc_range, save_path, method='3d', vis_gt_box=True, vis_pred_box=True, left_hand=False, uncertainty=None):
        """
        Visualize the prediction, ground truth with point cloud together.
        They may be flipped in y axis. Since carla is left hand coordinate, while kitti is right hand.

        Parameters
        ----------
        pred_box_tensor : torch.Tensor
            (N, 8, 3) prediction.

        gt_tensor : torch.Tensor
            (N, 8, 3) groundtruth bbx

        pcd : torch.Tensor
            PointCloud, (N, 4).

        pc_range : list
            [xmin, ymin, zmin, xmax, ymax, zmax]

        save_path : str
            Save the visualization results to given path.

        dataset : BaseDataset
            opencood dataset object.

        method: str, 'bev' or '3d'

        """

        pc_range = [int(i) for i in pc_range]
        pcd_np = pcd.cpu().numpy()

        if vis_pred_box:
            pred_box_np = pred_box_tensor.cpu().numpy()
            pred_name = ['pred'] * pred_box_np.shape[0]
            if uncertainty is not None:
                uncertainty_np = uncertainty.cpu().numpy()
                uncertainty_np = np.exp(uncertainty_np)
                d_a_square = 1.6**2 + 3.9**2
                
                if uncertainty_np.shape[1] == 3:
                    uncertainty_np[:,:2] *= d_a_square
                    uncertainty_np = np.sqrt(uncertainty_np) 
                    # yaw angle is in radian, it's the same in g2o SE2's setting.

                    pred_name = [f'x_u:{uncertainty_np[i,0]:.3f} y_u:{uncertainty_np[i,1]:.3f} a_u:{uncertainty_np[i,2]:.3f}' \
                                    for i in range(uncertainty_np.shape[0])]

                elif uncertainty_np.shape[1] == 2:
                    uncertainty_np[:,:2] *= d_a_square
                    uncertainty_np = np.sqrt(uncertainty_np) # yaw angle is in radian

                    pred_name = [f'x_u:{uncertainty_np[i,0]:.3f} y_u:{uncertainty_np[i,1]:3f}' \
                                    for i in range(uncertainty_np.shape[0])]

                elif uncertainty_np.shape[1] == 7:
                    uncertainty_np[:,:2] *= d_a_square
                    uncertainty_np = np.sqrt(uncertainty_np) # yaw angle is in radian

                    pred_name = [f'x_u:{uncertainty_np[i,0]:.3f} y_u:{uncertainty_np[i,1]:3f} a_u:{uncertainty_np[i,6]:3f}' \
                                    for i in range(uncertainty_np.shape[0])]                    

        if vis_gt_box:
            gt_box_np = gt_tensor.cpu().numpy()
            gt_name = ['gt'] * gt_box_np.shape[0]

        if method == 'bev':
            canvas = canvas_bev.Canvas_BEV_heading_right(canvas_shape=((pc_range[4]-pc_range[1])*10, (pc_range[3]-pc_range[0])*10),
                                            canvas_x_range=(pc_range[0], pc_range[3]), 
                                            canvas_y_range=(pc_range[1], pc_range[4]),
                                            left_hand=left_hand
                                          ) 

            canvas_xy, valid_mask = canvas.get_canvas_coords(pcd_np) # Get Canvas Coords
            canvas.draw_canvas_points(canvas_xy[valid_mask]) # Only draw valid points
            if vis_gt_box:
                canvas.draw_boxes(gt_box_np,colors=(0,255,0))
                # canvas.draw_boxes(gt_box_np,colors=(0,255,0), texts=gt_name)
            if vis_pred_box:
                canvas.draw_boxes(pred_box_np, colors=(255,0,0))
                # canvas.draw_boxes(pred_box_np, colors=(255,0,0), texts=pred_name)

        elif method == '3d':
            canvas = canvas_3d.Canvas_3D(left_hand=left_hand)
            canvas_xy, valid_mask = canvas.get_canvas_coords(pcd_np)
            canvas.draw_canvas_points(canvas_xy[valid_mask])
            if vis_gt_box:
                canvas.draw_boxes(gt_box_np,colors=(0,255,0), texts=gt_name)
            if vis_pred_box:
                canvas.draw_boxes(pred_box_np, colors=(255,0,0), texts=pred_name)
            
        else:
            raise(f"Not Completed for f{method} visualization.")

        plt.axis("off")

        plt.imshow(canvas.canvas)

        plt.tight_layout()
        plt.savefig(save_path, transparent=False, dpi=400)
        plt.clf()

opencood/visualization/simple_vis.py
MediaBrain-SJTU/CoBEVFlow
import torch
import numpy as np
import opencood.visualization.simple_plot3d.canvas_bev as canvas_bev
import sys
import os
import glob
from opencood.utils.transformation_utils import x1_to_x2
from opencood.utils.box_utils import project_box3d
from collections import OrderedDict
from icecream import ic
from matplotlib import pyplot as plt

def draw_trajectory(pt_file, save_path):
    '''
    box_dict : {
        'single_detection_bbx': # dict, [0, 1, ... , N-1],  
            [i]: dict:{
                [0]/[1]/[2]: past 3 frames detection results:{ # 都是在[0]的各自坐标系下的检测结果。
                    pred_box_3dcorner_tensor
                    pred_box_center_tensor
                    scores
                }
            }
        }
        'lidar_pose_0': 过去第一帧, 所有车都的pose
        'lidar_pose_current': 当前帧, 所有车都的pose
        'compensated_results_list': # len=N-1, 每个non-ego补偿出来的box, each of which is [N_obj, 4, 3], 注意这里用的是4个点, 最后一个维度上是 xyz, z是多余的
        'matched_idx_list': matched_idx_list, # len=N-1, 表示每个non-ego的过去2/3帧的匹配id, each of which is [N_obj, 2or3], 比如 ['matched_idx_list'][0] shape(22,2) 表示过去第一帧的22个框与第二帧的22个框的索引的匹配情况
        'single_gt_box_tensor': # list, len=N, 表示ego与non-ego每辆车在current时刻的检测框结果, 例如box_dict['single_gt_box_tensor'][1]大小为[N_obj, 8, 3] 表示第二辆车在current时刻的检测框结果。在各自current时刻坐标系下
        'single_lidar': # list, len=N, 表示ego与non-ego每辆车在current时刻的lidar np
        'gt_range': # [-140.8, -40, -3, 140.8, 40, 1], 表示lidar的范围
    }

    TODO
    1. transform detection results (from past 1 frame coordinate) into current frame
    2. visualzie GT box
    3. visualize detection results
    4. draw lidar point cloud
    5. clustering the same object (matched_idx_list)
    '''
    data = torch.load(pt_file)

    agent_num = len(data["single_detection_bbx"])
    left_hand = True
    pc_range = data['gt_range']

    for agent_id in range(1, agent_num):
        # step0: initailize canvas
        canvas = canvas_bev.Canvas_BEV_heading_right(canvas_shape=(int((pc_range[4]-pc_range[1])*10), int((pc_range[3]-pc_range[0])*10)),
                                                canvas_x_range=(pc_range[0], pc_range[3]), 
                                                canvas_y_range=(pc_range[1], pc_range[4]),
                                                left_hand=left_hand) 
        # step 1: transform detection results
        lidar_pose_past = data['lidar_pose_0'][agent_id].cpu().numpy()
        lidar_pose_current = data['lidar_pose_current'][agent_id].cpu().numpy()

        past_to_current = x1_to_x2(lidar_pose_past, lidar_pose_current)
        past_to_current = torch.from_numpy(past_to_current).to('cuda').float()
        boxes_past_in_past_coor_dict = data['single_detection_bbx'][agent_id] 
        boxes_past_in_current_coor_dict = OrderedDict()

        for past_frame, detections in boxes_past_in_past_coor_dict.items():
            if past_frame == 'past_k_time_diff':
                continue
            boxes_past_in_current_coor_dict[past_frame] = OrderedDict()
            boxes_past_in_current_coor_dict[past_frame]['pred_box_3dcorner_tensor'] = \
                project_box3d(detections['pred_box_3dcorner_tensor'].float(), past_to_current)
            boxes_past_in_current_coor_dict[past_frame]['scores'] = detections['scores']


        # step 2: draw GT box
        gt_box_np = data['single_gt_box_tensor'][agent_id].cpu().numpy()
        canvas.draw_boxes(gt_box_np,colors=(0,255,0), texts=['gt']*gt_box_np.shape[0])

        # step 3: visualze detection result, compensation, Establish inverted indices
        matched_idx_list = data['matched_idx_list']
        match_pairs = matched_idx_list[agent_id - 1]
        label_1 = [''] * boxes_past_in_current_coor_dict[1]['pred_box_3dcorner_tensor'].shape[0]
        label_0 = [''] * boxes_past_in_current_coor_dict[0]['pred_box_3dcorner_tensor'].shape[0]
        
        for pair_idx, pair in enumerate(match_pairs):
            idx_in_0 = pair[0]
            idx_in_1 = pair[1]
            label_1[idx_in_1] = pair_idx
            label_0[idx_in_0] = pair_idx

        ## past 2 frame (key is 1)
        pred_box_np_1 = boxes_past_in_current_coor_dict[1]['pred_box_3dcorner_tensor'].cpu().numpy()
        canvas.draw_boxes(pred_box_np_1,colors=(int(255*0.5),0,0), texts=label_1)

        ## past 1 frame (key is 0)
        pred_box_np_0 = boxes_past_in_current_coor_dict[0]['pred_box_3dcorner_tensor'].cpu().numpy()
        canvas.draw_boxes(pred_box_np_0,colors=(int(255*0.8),0,0), texts=label_0)

        ## compensate box
        compensate_box_np = data['compensated_results_list'][agent_id-1].cpu().numpy()

        if compensate_box_np.shape[0] == 0:
            print("Error: no compensate box: ", pt_file, " ", agent_id)
            return
        compensate_box_np = np.concatenate([compensate_box_np, compensate_box_np], axis=1)
        compensate_box_np = project_box3d(compensate_box_np, past_to_current.cpu().numpy())
        canvas.draw_boxes(compensate_box_np,colors=(0, int(255*0.83), int(255*0.83)), texts=['cp']*compensate_box_np.shape[0])




        # step 4: draw lidar point cloud
        pcd_np = data['single_lidar'][agent_id].cpu().numpy()
        canvas_xy, valid_mask = canvas.get_canvas_coords(pcd_np) # Get Canvas Coords
        canvas.draw_canvas_points(canvas_xy[valid_mask]) # Only draw valid points


        plt.axis("off")
        plt.imshow(canvas.canvas)
        plt.tight_layout()
        save_name = pt_file.split("/")[-1].rstrip(".pt")
        plt.savefig(f"{save_path}{save_name}_{agent_id}.png", dpi=400)
        plt.close()


def main():
    data_dir = '/remote-home/share/sizhewei/logs/irv2v_where2comm_cobevflow_w_dir_finetune/vis_cobevflow_reverse_0.50/' #'/remote-home/share/sizhewei/logs/irv2v_where2comm_cobevflow_w_dir_finetune/vis_cobevflow_521_trick_reverse_new_0.50_0'

    files = glob.glob(os.path.join(data_dir, 'bbx_folder', '*.pt'))
    files.sort()

    save_path = os.path.join(data_dir, 'multi_results')
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    for file in files:
        draw_trajectory(file, save_path)

if __name__ == "__main__":
    main()


opencood/visualization/vis_box_matching.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import os
from torch.utils.data import DataLoader, Subset
from opencood.data_utils import datasets
import torch
from opencood.tools import train_utils, inference_utils
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.visualization import vis_utils, simple_vis
from opencood.data_utils.datasets.late_fusion_dataset_v2x import \
    LateFusionDatasetV2X
from opencood.data_utils.datasets.late_fusion_dataset import \
    LateFusionDataset
import numpy as np

if __name__ == '__main__':
    current_path = os.path.dirname(os.path.realpath(__file__))
    params = load_yaml(os.path.join(current_path,
                                    '../hypes_yaml/visualization_v2x.yaml'))
    output_path = "/GPFS/rhome/yifanlu/OpenCOOD/data_vis/v2x_2.0_new/train"

    opencda_dataset = LateFusionDatasetV2X(params, visualize=True,
                                            train=False)
    len = len(opencda_dataset)
    sampled_indices = np.random.permutation(len)[:100]
    subset = Subset(opencda_dataset, sampled_indices)
    
    data_loader = DataLoader(subset, batch_size=1, num_workers=2,
                             collate_fn=opencda_dataset.collate_batch_test,
                             shuffle=False,
                             pin_memory=False)
    vis_gt_box = True
    vis_pred_box = False
    hypes = params

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    for i, batch_data in enumerate(data_loader):
        print(i)
        batch_data = train_utils.to_device(batch_data, device)
        gt_box_tensor = opencda_dataset.post_processor.generate_gt_bbx(batch_data)

        vis_save_path = os.path.join(output_path, '3d_%05d.png' % i)
        simple_vis.visualize(None,
                            gt_box_tensor,
                            batch_data['ego']['origin_lidar'][0],
                            hypes['postprocess']['gt_range'],
                            vis_save_path,
                            method='3d',
                            vis_gt_box = vis_gt_box,
                            vis_pred_box = vis_pred_box,
                            left_hand=False)
            
        vis_save_path = os.path.join(output_path, 'bev_%05d.png' % i)
        simple_vis.visualize(None,
                            gt_box_tensor,
                            batch_data['ego']['origin_lidar'][0],
                            hypes['postprocess']['gt_range'],
                            vis_save_path,
                            method='bev',
                            vis_gt_box = vis_gt_box,
                            vis_pred_box = vis_pred_box,
                            left_hand=False)
opencood/visualization/vis_data_sequence.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib

import os
import sys
from time import time
sys.path.append(os.getcwd())
from torch.utils.data import DataLoader, Subset
from opencood.data_utils import datasets
import torch
from opencood.tools import train_utils, inference_utils
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.visualization import vis_utils, simple_vis
from opencood.data_utils.datasets.intermediate_fusion_dataset_dair_delay import IntermediateFusionDatasetDAIRAsync
import numpy as np
from tqdm import tqdm

if __name__ == '__main__':
    current_path = os.path.dirname(os.path.realpath(__file__))
    params = load_yaml(os.path.join(current_path,
                                    '../hypes_yaml/visualization_dair.yaml'))
    time_delay = params['time_delay']
    output_path = "/DB/rhome/sizhewei/percp/OpenCOOD/opencood/viz_out/debug"
    if not os.path.exists(output_path):
        os.mkdir(output_path)

    opencda_dataset = IntermediateFusionDatasetDAIRAsync(params, visualize=True,
                                            train=False)
    # dataset_len = len(opencda_dataset)
    # # print(dataset_len)
    # sampled_indices = np.random.permutation(dataset_len)[:2]

    sampled_indices = np.array(range(15))
    print(sampled_indices)
    
    subset = Subset(opencda_dataset, sampled_indices)
    
    data_loader = DataLoader(subset, batch_size=1, num_workers=2,
                             collate_fn=opencda_dataset.collate_batch_test,
                             shuffle=False,
                             pin_memory=False)
    vis_gt_box = True
    vis_pred_box = False
    hypes = params

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # for test
    temp = subset[0]

    for i, batch_data in enumerate(data_loader):
        print(i)
        veh_frame_id = batch_data['ego']['veh_frame_id']
        batch_data['ego'].pop('veh_frame_id')
        batch_data = train_utils.to_device(batch_data, device)
        gt_box_tensor = opencda_dataset.post_processor.generate_gt_bbx(batch_data)

        vis_save_path = os.path.join(output_path, '3d_veh_%s_%02d.png' % (veh_frame_id, time_delay))
        simple_vis.visualize(None,
                            gt_box_tensor,
                            batch_data['ego']['origin_lidar'][0],
                            hypes['postprocess']['gt_range'],
                            vis_save_path,
                            method='3d',
                            vis_gt_box = vis_gt_box,
                            vis_pred_box = vis_pred_box,
                            left_hand=False)
            
        vis_save_path = os.path.join(output_path, 'bev_veh_%s_%02d.png' % (veh_frame_id, time_delay))
        simple_vis.visualize(None,
                            gt_box_tensor,
                            batch_data['ego']['origin_lidar'][0],
                            hypes['postprocess']['gt_range'],
                            vis_save_path,
                            method='bev',
                            vis_gt_box = vis_gt_box,
                            vis_pred_box = vis_pred_box,
                            left_hand=False)

        # vis_save_path = os.path.join(output_path, '3d_inf_%05d.png' % i)
        # simple_vis.visualize(None,
        #                     gt_box_tensor,
        #                     batch_data['ego']['origin_lidar'][0],
        #                     hypes['postprocess']['gt_range'],
        #                     vis_save_path,
        #                     method='bev',
        #                     vis_gt_box = vis_gt_box,
        #                     vis_pred_box = vis_pred_box,
        #                     left_hand=True)
        
        # vis_save_path = os.path.join(output_path, 'bev_inf_%05d.png' % i)
        # simple_vis.visualize(None,
        #                     gt_box_tensor,
        #                     batch_data['ego']['origin_lidar'][0],
        #                     hypes['postprocess']['gt_range'],
        #                     vis_save_path,
        #                     method='bev',
        #                     vis_gt_box = vis_gt_box,
        #                     vis_pred_box = vis_pred_box,
        #                     left_hand=True)

opencood/visualization/vis_data_sequence_dairv2x.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


import os
from torch.utils.data import DataLoader, Subset
from opencood.data_utils import datasets
import torch
from opencood.tools import train_utils, inference_utils
from opencood.hypes_yaml.yaml_utils import load_yaml
from opencood.visualization import vis_utils, simple_vis
from opencood.data_utils.datasets.late_fusion_dataset_v2x import \
    LateFusionDatasetV2X
from opencood.data_utils.datasets.late_fusion_dataset import \
    LateFusionDataset
import numpy as np

if __name__ == '__main__':
    current_path = os.path.dirname(os.path.realpath(__file__))
    params = load_yaml(os.path.join(current_path,
                                    '../hypes_yaml/visualization_opv2v.yaml'))
    output_path = "SPECIFY YOUR PATH"
    left_hand = True

    opencda_dataset = LateFusionDataset(params, visualize=True,
                                            train=False)
    len = len(opencda_dataset)
    sampled_indices = range(1290,1310)
    subset = Subset(opencda_dataset, sampled_indices)
    
    data_loader = DataLoader(subset, batch_size=1, num_workers=2,
                             collate_fn=opencda_dataset.collate_batch_test,
                             shuffle=False,
                             pin_memory=False)
    vis_gt_box = True
    vis_pred_box = False
    hypes = params

    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    for i, batch_data in enumerate(data_loader):
        # batch_data = train_utils.to_device(batch_data, device)
        gt_box_tensor = opencda_dataset.post_processor.generate_gt_bbx(batch_data)

        vis_save_path = os.path.join(output_path, '3d_%05d.png' % i)
        simple_vis.visualize(None,
                            gt_box_tensor,
                            batch_data['ego']['origin_lidar'][0],
                            hypes['postprocess']['gt_range'],
                            vis_save_path,
                            method='3d',
                            vis_gt_box = vis_gt_box,
                            vis_pred_box = vis_pred_box,
                            left_hand=left_hand)

        vis_save_path = os.path.join(output_path, 'bev_%05d_%01d.png' % i)
        print(vis_save_path)
        simple_vis.visualize(None,
                            gt_box_tensor,
                            batch_data['ego']['origin_lidar'][0],
                            hypes['postprocess']['gt_range'],
                            vis_save_path,
                            method='bev',
                            vis_gt_box = vis_gt_box,
                            vis_pred_box = vis_pred_box,
                            left_hand=left_hand)
opencood/visualization/vis_data_sequence_opv2v.py
MediaBrain-SJTU/CoBEVFlow
"""
Draw the feature flow with arrow
"""

import torch
import numpy as np
import os
import sys
import glob
# import seaborn as sns
import matplotlib.pyplot as plt
from matplotlib.colors import hsv_to_rgb


# 箭头长度
arrow_len = 0.5
# 箭头缩放比例
scale = 20
# 色轮colormap，用于表示箭头的方向
cmap = plt.cm.hsv


def draw_feature_flow(file_idx, pt_file, save_path):
    """
    box_dict : {
        'single_detection_bbx': # dict, [0, 1, ... , N-1], 
            [i]: dict:{
                [0]/[1]/[2]: past 3 frames detection results:{  past 3 frames detection results:{ # 都是在[0]的各自坐标系下的检测结果。
                    pred_box_3dcorner_tensor
                    pred_box_center_tensor
                    scores
                }
            }
        }
        'lidar_pose_0': 过去第一帧 所有车的pose [N, 6]
        'lidar_pose_current': 当前帧 所有车的pose [N, 6]
        'matched_idx_list': matched_idx_list, # len=N-1, 表示每个non-ego的过去两帧的匹配id, each of which is [N_obj, 2], 比如 ['matched_idx_list'][0] shape(22,2) 表示过去第一帧的22个框与第二帧的22个框的索引的匹配情况
        'compensated_results_list': # len=N-1, 每个non-ego补偿出来的box, each of which is [N_obj, 4, 3], 注意这里用的是4个点, 最后一个维度上是 xyz, z是多余的
        'single_gt_box_tensor': # list, len=N, 表示ego与non-ego每辆车在current时刻的检测框结果, 例如box_dict['single_gt_box_tensor'][1]大小为[N_obj, 8, 3] 表示第二辆车在current时刻的检测框结果
        'single_lidar': # list, len=N, 表示ego与non-ego每辆车在current时刻的lidar np
        'gt_range': # [-140.8, -40, -3, 140.8, 40, 1], 表示lidar的范围
        'single_updated_feature': tensor, [N, C, H, W], 表示ego与non-ego每辆车在past-0时刻的更新后的feature
        'single_original_feature': tensor, [N, C, H, W], 表示ego与non-ego每辆车在past-0时刻的原始feature
        'single_flow_map': tensor, [N, H, W, 2], 表示ego与non-ego每辆车在past-0时刻的flow map
        'single_reserved_mask': tensor, [N, C, H, W], 表示ego与non-ego每辆车在past-0时刻的reserved mask
    }
    """

    data = torch.load(pt_file)

    agent_num = len(data["single_detection_bbx"])
    left_hand = True
    pc_range = data['gt_range']

    single_flow_map = data['single_flow_map']
    single_flow_map_np = data['single_flow_map'].cpu().numpy()

    H, W = single_flow_map.shape[1], single_flow_map.shape[2]
    identity_flow_np = torch.nn.functional.affine_grid(torch.eye(2, 3).unsqueeze(0), torch.Size((1, 1, H, W))).numpy()

    delta_flow_map_np =  identity_flow_np - single_flow_map_np
    delta_flow_map_np[..., 0] *= 140.8 / 2
    delta_flow_map_np[..., 1] *= 40 / 2

    delta_flow_map_np_abs = np.abs(delta_flow_map_np)

    for i in range(1, agent_num):
        delta_flow_map_np_single = delta_flow_map_np[i]
        # get object center
        obj_center = data['single_detection_bbx'][i][0]['pred_box_center_tensor'].cpu().numpy()
        obj_center_x = (obj_center[:, 0] - (-140.8)) / 0.4
        obj_center_x = obj_center_x.astype(int)
        obj_center_y = (obj_center[:, 1] - (-40)) / 0.4
        obj_center_y = obj_center_y.astype(int)

        obj = np.vstack([obj_center_y, obj_center_x]).T
        print(obj.shape)
        # 箭头的长度和颜色深度的最大值
        arrow_len = np.sqrt(np.sum(delta_flow_map_np_single**2, axis=2)).max()

        # 绘图
        fig, ax = plt.subplots()
        ax.imshow(np.ones_like(delta_flow_map_np_single[:, :, 0], dtype=np.uint8) * 255, cmap='gray_r')

        for arr in obj:
            # 计算箭头的长度
            x = delta_flow_map_np_single[arr[0], arr[1], 0]
            y = delta_flow_map_np_single[arr[0], arr[1], 1]
            length = np.sqrt(x ** 2 + y ** 2) / 0.4 # normalize by grid size

            # 计算箭头的方向
            angle = np.arctan2(y, x)

            feat_hsv = np.zeros((1, 3))

            feat_hsv[:, 0] = (angle + np.pi) / (2 * np.pi)
            feat_hsv[:, 1] = min(length / arrow_len , 1.0)# * scale
            print('===',  length / arrow_len,  '===')
            feat_hsv[:, 2] = 1

            feat_rgb = hsv_to_rgb(feat_hsv)
            print('rgb', feat_rgb)
            print('cmap', cmap(angle / (2 * np.pi)))

            # 绘制箭头
            # plt.arrow(arr[1], arr[0], length * np.cos(angle), length * np.sin(angle), color=cmap(angle / (2 * np.pi)), width=1,head_width=3, head_length=0.5)
            plt.arrow(arr[1], arr[0], length * np.cos(angle), length * np.sin(angle), color=feat_rgb[0], width=0.6,head_width=2, head_length=0.7)

        plt.xticks([])
        plt.yticks([])
        plt.savefig(os.path.join(save_path, "{}_flow_map_arror_{}.png".format(file_idx, i)), dpi=500)
        plt.close()

def main():
    data_dir = '/remote-home/share/sizhewei/logs/irv2v_where2comm_cobevflow_w_dir_finetune/vis_cobevflow_reverse_viz_debug_3_3_3_new_0.30/bbx_folder'
    files = glob.glob(os.path.join(data_dir, '*.pt'))
    files.sort()
    save_path = os.path.join(data_dir, 'flow_map_arrow')
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    for i in range(len(files)):
        draw_feature_flow(i, files[i], save_path)

if __name__ == "__main__":
    main()


opencood/visualization/vis_flow_map.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>, Hao Xiang <haxiang@g.ucla.edu>,
# License: TDG-Attribution-NonCommercial-NoDistrib

import time

import cv2
import numpy as np
import open3d as o3d
import matplotlib
import matplotlib.pyplot as plt

from matplotlib import cm

from opencood.utils import box_utils
from opencood.utils import common_utils

VIRIDIS = np.array(cm.get_cmap('plasma').colors)
VID_RANGE = np.linspace(0.0, 1.0, VIRIDIS.shape[0])


def bbx2linset(bbx_corner, order='hwl', color=(0, 1, 0)):
    """
    Convert the torch tensor bounding box to o3d lineset for visualization.

    Parameters
    ----------
    bbx_corner : torch.Tensor
        shape: (n, 8, 3).

    order : str
        The order of the bounding box if shape is (n, 7)

    color : tuple
        The bounding box color.

    Returns
    -------
    line_set : list
        The list containing linsets.
    """
    if not isinstance(bbx_corner, np.ndarray):
        bbx_corner = common_utils.torch_tensor_to_numpy(bbx_corner)

    if len(bbx_corner.shape) == 2:
        bbx_corner = box_utils.boxes_to_corners_3d(bbx_corner,
                                                   order)

    # Our lines span from points 0 to 1, 1 to 2, 2 to 3, etc...
    lines = [[0, 1], [1, 2], [2, 3], [0, 3],
             [4, 5], [5, 6], [6, 7], [4, 7],
             [0, 4], [1, 5], [2, 6], [3, 7]]

    # Use the same color for all lines
    colors = [list(color) for _ in range(len(lines))]
    bbx_linset = []

    for i in range(bbx_corner.shape[0]):
        bbx = bbx_corner[i]
        # o3d use right-hand coordinate
        bbx[:, :1] = - bbx[:, :1]

        line_set = o3d.geometry.LineSet()
        line_set.points = o3d.utility.Vector3dVector(bbx)
        line_set.lines = o3d.utility.Vector2iVector(lines)
        line_set.colors = o3d.utility.Vector3dVector(colors)
        bbx_linset.append(line_set)

    return bbx_linset


def bbx2oabb(bbx_corner, order='hwl', color=(0, 0, 1)):
    """
    Convert the torch tensor bounding box to o3d oabb for visualization.

    Parameters
    ----------
    bbx_corner : torch.Tensor
        shape: (n, 8, 3).

    order : str
        The order of the bounding box if shape is (n, 7)

    color : tuple
        The bounding box color.

    Returns
    -------
    oabbs : list
        The list containing all oriented bounding boxes.
    """
    if not isinstance(bbx_corner, np.ndarray):
        bbx_corner = common_utils.torch_tensor_to_numpy(bbx_corner)

    if len(bbx_corner.shape) == 2:
        bbx_corner = box_utils.boxes_to_corners_3d(bbx_corner,
                                                   order)
    oabbs = []

    for i in range(bbx_corner.shape[0]):
        bbx = bbx_corner[i]
        # o3d use right-hand coordinate
        bbx[:, :1] = - bbx[:, :1]

        tmp_pcd = o3d.geometry.PointCloud()
        tmp_pcd.points = o3d.utility.Vector3dVector(bbx)

        oabb = tmp_pcd.get_oriented_bounding_box()
        oabb.color = color
        oabbs.append(oabb)

    return oabbs


def bbx2aabb(bbx_center, order):
    """
    Convert the torch tensor bounding box to o3d aabb for visualization.

    Parameters
    ----------
    bbx_center : torch.Tensor
        shape: (n, 7).

    order: str
        hwl or lwh.

    Returns
    -------
    aabbs : list
        The list containing all o3d.aabb
    """
    if not isinstance(bbx_center, np.ndarray):
        bbx_center = common_utils.torch_tensor_to_numpy(bbx_center)
    bbx_corner = box_utils.boxes_to_corners_3d(bbx_center, order)

    aabbs = []

    for i in range(bbx_corner.shape[0]):
        bbx = bbx_corner[i]
        # o3d use right-hand coordinate
        bbx[:, :1] = - bbx[:, :1]

        tmp_pcd = o3d.geometry.PointCloud()
        tmp_pcd.points = o3d.utility.Vector3dVector(bbx)

        aabb = tmp_pcd.get_axis_aligned_bounding_box()
        aabb.color = (0, 0, 1)
        aabbs.append(aabb)

    return aabbs


def linset_assign_list(vis,
                       lineset_list1,
                       lineset_list2,
                       update_mode='update'):
    """
    Associate two lists of lineset.

    Parameters
    ----------
    vis : open3d.Visualizer
    lineset_list1 : list
    lineset_list2 : list
    update_mode : str
        Add or update the geometry.
    """
    for j in range(len(lineset_list1)):
        index = j if j < len(lineset_list2) else -1
        lineset_list1[j] = \
            lineset_assign(lineset_list1[j],
                                     lineset_list2[index])
        if update_mode == 'add':
            vis.add_geometry(lineset_list1[j])
        else:
            vis.update_geometry(lineset_list1[j])


def lineset_assign(lineset1, lineset2):
    """
    Assign the attributes of lineset2 to lineset1.

    Parameters
    ----------
    lineset1 : open3d.LineSet
    lineset2 : open3d.LineSet

    Returns
    -------
    The lineset1 object with 2's attributes.
    """

    lineset1.points = lineset2.points
    lineset1.lines = lineset2.lines
    lineset1.colors = lineset2.colors

    return lineset1


def color_encoding(intensity, mode='intensity'):
    """
    Encode the single-channel intensity to 3 channels rgb color.

    Parameters
    ----------
    intensity : np.ndarray
        Lidar intensity, shape (n,)

    mode : str
        The color rendering mode. intensity, z-value and constant are
        supported.

    Returns
    -------
    color : np.ndarray
        Encoded Lidar color, shape (n, 3)
    """
    assert mode in ['intensity', 'z-value', 'constant']

    if mode == 'intensity':
        intensity_col = 1.0 - np.log(intensity) / np.log(np.exp(-0.004 * 100))
        int_color = np.c_[
            np.interp(intensity_col, VID_RANGE, VIRIDIS[:, 0]),
            np.interp(intensity_col, VID_RANGE, VIRIDIS[:, 1]),
            np.interp(intensity_col, VID_RANGE, VIRIDIS[:, 2])]

    elif mode == 'z-value':
        min_value = -1.5
        max_value = 0.5
        norm = matplotlib.colors.Normalize(vmin=min_value, vmax=max_value)
        cmap = cm.jet
        m = cm.ScalarMappable(norm=norm, cmap=cmap)

        colors = m.to_rgba(intensity)
        colors[:, [2, 1, 0, 3]] = colors[:, [0, 1, 2, 3]]
        colors[:, 3] = 0.5
        int_color = colors[:, :3]

    elif mode == 'constant':
        # regard all point cloud the same color
        int_color = np.ones((intensity.shape[0], 3))
        int_color[:, 0] *= 247 / 255
        int_color[:, 1] *= 244 / 255
        int_color[:, 2] *= 237 / 255

    return int_color


def visualize_single_sample_output_gt(pred_tensor,
                                      gt_tensor,
                                      pcd,
                                      show_vis=True,
                                      save_path='',
                                      mode='constant'):
    """
    Visualize the prediction, groundtruth with point cloud together.

    Parameters
    ----------
    pred_tensor : torch.Tensor
        (N, 8, 3) prediction.

    gt_tensor : torch.Tensor
        (N, 8, 3) groundtruth bbx

    pcd : torch.Tensor
        PointCloud, (N, 4).

    show_vis : bool
        Whether to show visualization.

    save_path : str
        Save the visualization results to given path.

    mode : str
        Color rendering mode.
    """

    def custom_draw_geometry(pcd, pred, gt):
        vis = o3d.visualization.Visualizer()
        vis.create_window()

        opt = vis.get_render_option()
        opt.background_color = np.asarray([0, 0, 0])
        opt.point_size = 1.0

        vis.add_geometry(pcd)
        for ele in pred:
            vis.add_geometry(ele)
        for ele in gt:
            vis.add_geometry(ele)

        vis.run()
        vis.destroy_window()

    origin_lidar = pcd
    if not isinstance(pcd, np.ndarray):
        origin_lidar = common_utils.torch_tensor_to_numpy(pcd)

    origin_lidar_intcolor = \
        color_encoding(origin_lidar[:, -1] if mode == 'intensity'
                       else origin_lidar[:, 2], mode=mode)
    # left -> right hand
    origin_lidar[:, :1] = -origin_lidar[:, :1]

    o3d_pcd = o3d.geometry.PointCloud()
    o3d_pcd.points = o3d.utility.Vector3dVector(origin_lidar[:, :3])
    o3d_pcd.colors = o3d.utility.Vector3dVector(origin_lidar_intcolor)

    oabbs_pred = bbx2oabb(pred_tensor, color=(1, 0, 0))
    oabbs_gt = bbx2oabb(gt_tensor, color=(0, 1, 0))

    visualize_elements = [o3d_pcd] + oabbs_pred + oabbs_gt
    if show_vis:
        custom_draw_geometry(o3d_pcd, oabbs_pred, oabbs_gt)
    if save_path:
        save_o3d_visualization(visualize_elements, save_path)


def visualize_single_sample_output_bev(pred_box, gt_box, pcd, dataset,
                                       show_vis=True,
                                       save_path=''):
    """
    Visualize the prediction, groundtruth with point cloud together in
    a bev format.

    Parameters
    ----------
    pred_box : torch.Tensor
        (N, 4, 2) prediction.

    gt_box : torch.Tensor
        (N, 4, 2) groundtruth bbx

    pcd : torch.Tensor
        PointCloud, (N, 4).

    show_vis : bool
        Whether to show visualization.

    save_path : str
        Save the visualization results to given path.
    """

    if not isinstance(pcd, np.ndarray):
        pcd = common_utils.torch_tensor_to_numpy(pcd)
    if pred_box is not None and not isinstance(pred_box, np.ndarray):
        pred_box = common_utils.torch_tensor_to_numpy(pred_box)
    if gt_box is not None and not isinstance(gt_box, np.ndarray):
        gt_box = common_utils.torch_tensor_to_numpy(gt_box)

    ratio = dataset.params["preprocess"]["args"]["res"]
    L1, W1, H1, L2, W2, H2 = dataset.params["preprocess"]["cav_lidar_range"]
    bev_origin = np.array([L1, W1]).reshape(1, -1)
    # (img_row, img_col)
    bev_map = dataset.project_points_to_bev_map(pcd, ratio)
    # (img_row, img_col, 3)
    bev_map = \
        np.repeat(bev_map[:, :, np.newaxis], 3, axis=-1).astype(np.float32)
    bev_map = bev_map * 255

    if pred_box is not None:
        num_bbx = pred_box.shape[0]
        for i in range(num_bbx):
            bbx = pred_box[i]

            bbx = ((bbx - bev_origin) / ratio).astype(int)
            bbx = bbx[:, ::-1]
            cv2.polylines(bev_map, [bbx], True, (0, 0, 255), 1)

    if gt_box is not None and len(gt_box):
        for i in range(gt_box.shape[0]):
            bbx = gt_box[i][:4, :2]
            bbx = (((bbx - bev_origin)) / ratio).astype(int)
            bbx = bbx[:, ::-1]
            cv2.polylines(bev_map, [bbx], True, (255, 0, 0), 1)

    if show_vis:
        plt.axis("off")
        plt.imshow(bev_map)
        plt.show()
    if save_path:
        plt.axis("off")
        plt.imshow(bev_map)
        plt.savefig(save_path)


def visualize_single_sample_dataloader(batch_data,
                                       o3d_pcd,
                                       order,
                                       key='origin_lidar',
                                       visualize=False,
                                       save_path='',
                                       oabb=False,
                                       mode='constant'):
    """
    Visualize a single frame of a single CAV for validation of data pipeline.

    Parameters
    ----------
    o3d_pcd : o3d.PointCloud
        Open3d PointCloud.

    order : str
        The bounding box order.

    key : str
        origin_lidar for late fusion and stacked_lidar for early fusion.

    visualize : bool
        Whether to visualize the sample.

    batch_data : dict
        The dictionary that contains current timestamp's data.

    save_path : str
        If set, save the visualization image to the path.

    oabb : bool
        If oriented bounding box is used.
    """

    origin_lidar = batch_data[key]
    if not isinstance(origin_lidar, np.ndarray):
        origin_lidar = common_utils.torch_tensor_to_numpy(origin_lidar)
    # we only visualize the first cav for single sample
    if len(origin_lidar.shape) > 2:
        origin_lidar = origin_lidar[0]
    origin_lidar_intcolor = \
        color_encoding(origin_lidar[:, -1] if mode == 'intensity'
                       else origin_lidar[:, 2], mode=mode)

    # left -> right hand
    origin_lidar[:, :1] = -origin_lidar[:, :1]

    o3d_pcd.points = o3d.utility.Vector3dVector(origin_lidar[:, :3])
    o3d_pcd.colors = o3d.utility.Vector3dVector(origin_lidar_intcolor)

    object_bbx_center = batch_data['object_bbx_center']
    object_bbx_mask = batch_data['object_bbx_mask']
    object_bbx_center = object_bbx_center[object_bbx_mask == 1]

    aabbs = bbx2linset(object_bbx_center, order) if not oabb else \
        bbx2oabb(object_bbx_center, order)
    visualize_elements = [o3d_pcd] + aabbs
    if visualize:
        o3d.visualization.draw_geometries(visualize_elements)

    if save_path:
        save_o3d_visualization(visualize_elements, save_path)

    return o3d_pcd, aabbs


def visualize_inference_sample_dataloader(pred_box_tensor,
                                          gt_box_tensor,
                                          origin_lidar,
                                          o3d_pcd,
                                          mode='constant'):
    """
    Visualize a frame during inference for video stream.

    Parameters
    ----------
    pred_box_tensor : torch.Tensor
        (N, 8, 3) prediction.

    gt_box_tensor : torch.Tensor
        (N, 8, 3) groundtruth bbx

    origin_lidar : torch.Tensor
        PointCloud, (N, 4).

    o3d_pcd : open3d.PointCloud
        Used to visualize the pcd.

    mode : str
        lidar point rendering mode.
    """

    if not isinstance(origin_lidar, np.ndarray):
        origin_lidar = common_utils.torch_tensor_to_numpy(origin_lidar)
    # we only visualize the first cav for single sample
    if len(origin_lidar.shape) > 2:
        origin_lidar = origin_lidar[0]
    origin_lidar_intcolor = \
        color_encoding(origin_lidar[:, -1] if mode == 'intensity'
                       else origin_lidar[:, 2], mode=mode)

    if not isinstance(pred_box_tensor, np.ndarray):
        pred_box_tensor = common_utils.torch_tensor_to_numpy(pred_box_tensor)
    if not isinstance(gt_box_tensor, np.ndarray):
        gt_box_tensor = common_utils.torch_tensor_to_numpy(gt_box_tensor)

    # left -> right hand
    origin_lidar[:, :1] = -origin_lidar[:, :1]

    o3d_pcd.points = o3d.utility.Vector3dVector(origin_lidar[:, :3])
    o3d_pcd.colors = o3d.utility.Vector3dVector(origin_lidar_intcolor)

    gt_o3d_box = bbx2linset(gt_box_tensor, order='hwl', color=(0, 1, 0))
    pred_o3d_box = bbx2linset(pred_box_tensor, color=(1, 0, 0))

    return o3d_pcd, pred_o3d_box, gt_o3d_box


def visualize_sequence_dataloader(dataloader, order, color_mode='constant'):
    """
    Visualize the batch data in animation.

    Parameters
    ----------
    dataloader : torch.Dataloader
        Pytorch dataloader

    order : str
        Bounding box order(N, 7).

    color_mode : str
        Color rendering mode.
    """
    vis = o3d.visualization.Visualizer()
    vis.create_window()

    vis.get_render_option().background_color = [0.05, 0.05, 0.05]
    vis.get_render_option().point_size = 1.0
    vis.get_render_option().show_coordinate_frame = True

    # used to visualize lidar points
    vis_pcd = o3d.geometry.PointCloud()
    # used to visualize object bounding box, maximum 50
    vis_aabbs = []
    for _ in range(50):
        vis_aabbs.append(o3d.geometry.LineSet())

    while True:
        for i_batch, sample_batched in enumerate(dataloader):
            print(i_batch)
            pcd, aabbs = \
                visualize_single_sample_dataloader(sample_batched['ego'],
                                                   vis_pcd,
                                                   order,
                                                   mode=color_mode)
            if i_batch == 0:
                vis.add_geometry(pcd)
                for i in range(len(vis_aabbs)):
                    index = i if i < len(aabbs) else -1
                    vis_aabbs[i] = lineset_assign(vis_aabbs[i], aabbs[index])
                    vis.add_geometry(vis_aabbs[i])

            for i in range(len(vis_aabbs)):
                index = i if i < len(aabbs) else -1
                vis_aabbs[i] = lineset_assign(vis_aabbs[i], aabbs[index])
                vis.update_geometry(vis_aabbs[i])

            vis.update_geometry(pcd)
            vis.poll_events()
            vis.update_renderer()
            time.sleep(0.001)

    vis.destroy_window()


def save_o3d_visualization(element, save_path):
    """
    Save the open3d drawing to folder.

    Parameters
    ----------
    element : list
        List of o3d.geometry objects.

    save_path : str
        The save path.
    """
    vis = o3d.visualization.Visualizer()
    vis.create_window()
    for i in range(len(element)):
        vis.add_geometry(element[i])
        vis.update_geometry(element[i])

    vis.poll_events()
    vis.update_renderer()

    vis.capture_screen_image(save_path)
    vis.destroy_window()


def visualize_bev(batch_data):
    bev_input = batch_data["processed_lidar"]["bev_input"]
    label_map = batch_data["label_dict"]["label_map"]
    if not isinstance(bev_input, np.ndarray):
        bev_input = common_utils.torch_tensor_to_numpy(bev_input)

    if not isinstance(label_map, np.ndarray):
        label_map = label_map[0].numpy() if not label_map[0].is_cuda else \
            label_map[0].cpu().detach().numpy()

    if len(bev_input.shape) > 3:
        bev_input = bev_input[0, ...]

    plt.matshow(np.sum(bev_input, axis=0))
    plt.axis("off")
    plt.matshow(label_map[0, :, :])
    plt.axis("off")
    plt.show()

def draw_box_plt(boxes_dec, ax, color=None, linewidth_scale=1.0):
    """
    draw boxes in a given plt ax
    :param boxes_dec: (N, 5) or (N, 7) in metric
    :param ax:
    :return: ax with drawn boxes
    """
    if not len(boxes_dec)>0:
        return ax
    boxes_np= boxes_dec
    if not isinstance(boxes_np, np.ndarray):
        boxes_np = boxes_np.cpu().detach().numpy()
    if boxes_np.shape[-1]>5:
        boxes_np = boxes_np[:, [0, 1, 3, 4, 6]]
    x = boxes_np[:, 0]
    y = boxes_np[:, 1]
    dx = boxes_np[:, 2]
    dy = boxes_np[:, 3]

    x1 = x - dx / 2
    y1 = y - dy / 2
    x2 = x + dx / 2
    y2 = y + dy / 2
    theta = boxes_np[:, 4:5]
    # bl, fl, fr, br
    corners = np.array([[x1, y1],[x1,y2], [x2,y2], [x2, y1]]).transpose(2, 0, 1)
    new_x = (corners[:, :, 0] - x[:, None]) * np.cos(theta) + (corners[:, :, 1]
              - y[:, None]) * (-np.sin(theta)) + x[:, None]
    new_y = (corners[:, :, 0] - x[:, None]) * np.sin(theta) + (corners[:, :, 1]
              - y[:, None]) * (np.cos(theta)) + y[:, None]
    corners = np.stack([new_x, new_y], axis=2)
    for corner in corners:
        ax.plot(corner[[0,1,2,3,0], 0], corner[[0,1,2,3,0], 1], color=color, linewidth=0.5*linewidth_scale)
        # draw front line (
        ax.plot(corner[[2, 3], 0], corner[[2, 3], 1], color=color, linewidth=2*linewidth_scale)
    return ax


def draw_points_boxes_plt(pc_range, points=None, boxes_pred=None, boxes_gt=None, save_path=None,
                          points_c='y.', bbox_gt_c='green', bbox_pred_c='red', return_ax=False, ax=None):
    if ax is None:
        ax = plt.figure(figsize=(15, 6)).add_subplot(1, 1, 1)
        ax.set_aspect('equal', 'box')
        ax.set(xlim=(pc_range[0], pc_range[3]),
               ylim=(pc_range[1], pc_range[4]))
    if points is not None:
        ax.plot(points[:, 0], points[:, 1], points_c, markersize=0.1)
    if (boxes_gt is not None) and len(boxes_gt)>0:
        ax = draw_box_plt(boxes_gt, ax, color=bbox_gt_c)
    if (boxes_pred is not None) and len(boxes_pred)>0:
        ax = draw_box_plt(boxes_pred, ax, color=bbox_pred_c)
    plt.xlabel('x')
    plt.ylabel('y')

    plt.savefig(save_path)
    if return_ax:
        return ax
    plt.close() 
opencood/visualization/vis_utils.py
MediaBrain-SJTU/CoBEVFlow
# -*- coding: utf-8 -*-
# Author: Runsheng Xu <rxx3386@ucla.edu>
# License: TDG-Attribution-NonCommercial-NoDistrib


from os.path import dirname, realpath
from setuptools import setup, find_packages, Distribution
from opencood.version import __version__


def _read_requirements_file():
    """Return the elements in requirements.txt."""
    req_file_path = '%s/requirements.txt' % dirname(realpath(__file__))
    with open(req_file_path) as f:
        return [line.strip() for line in f]


setup(
    name='OpenCOOD',
    version=__version__,
    packages=find_packages(),
    url='https://github.com/ucla-mobility/OpenCDA.git',
    license='MIT',
    author='Runsheng Xu, Hao Xiang',
    author_email='rxx3386@ucla.edu',
    description='An opensource pytorch framework for autonomous driving '
                'cooperative detection',
    long_description=open("README.md").read(),
    install_requires=_read_requirements_file(),
)

setup.py
MediaBrain-SJTU/CoBEVFlow