<reponame>xizaoqu/mmdetection3d
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
from numpy.core.numeric import False_
import torch
from torch import nn
from torch.autograd import Function
from torch.nn.modules.utils import _pair
import numpy as np
import numba as nb
import torch_scatter
from .voxel_layer import dynamic_voxelize, hard_voxelize


class _Voxelization(Function):

    @staticmethod
    def forward(ctx,
                points,
                voxel_size,
                coors_range,
                max_points=35,
                max_voxels=20000,
                deterministic=True):
        """convert kitti points(N, >=3) to voxels.

        Args:
            points: [N, ndim] float tensor. points[:, :3] contain xyz points
                and points[:, 3:] contain other information like reflectivity
            voxel_size: [3] list/tuple or array, float. xyz, indicate voxel
                size
            coors_range: [6] list/tuple or array, float. indicate voxel
                range. format: xyzxyz, minmax
            max_points: int. indicate maximum points contained in a voxel. if
                max_points=-1, it means using dynamic_voxelize
            max_voxels: int. indicate maximum voxels this function create.
                for second, 20000 is a good choice. Users should shuffle points
                before call this function because max_voxels may drop points.
            deterministic: bool. whether to invoke the non-deterministic
                version of hard-voxelization implementations. non-deterministic
                version is considerablly fast but is not deterministic. only
                affects hard voxelization. default True. for more information
                of this argument and the implementation insights, please refer
                to the following links:
                https://github.com/open-mmlab/mmdetection3d/issues/894
                https://github.com/open-mmlab/mmdetection3d/pull/904
                it is an experimental feature and we will appreciate it if
                you could share with us the failing cases.

        Returns:
            voxels: [M, max_points, ndim] float tensor. only contain points
                    and returned when max_points != -1.
            coordinates: [M, 3] int32 tensor, always returned.
            num_points_per_voxel: [M] int32 tensor. Only returned when
                max_points != -1.
        """
        if max_points == -1 or max_voxels == -1:
            coors = points.new_zeros(size=(points.size(0), 3), dtype=torch.int)
            dynamic_voxelize(points, coors, voxel_size, coors_range, 3)
            return coors
        else:
            voxels = points.new_zeros(
                size=(max_voxels, max_points, points.size(1)))
            coors = points.new_zeros(size=(max_voxels, 3), dtype=torch.int)
            num_points_per_voxel = points.new_zeros(
                size=(max_voxels, ), dtype=torch.int)
            voxel_num = hard_voxelize(points, voxels, coors,
                                      num_points_per_voxel, voxel_size,
                                      coors_range, max_points, max_voxels, 3,
                                      deterministic)
            # select the valid voxels
            voxels_out = voxels[:voxel_num]
            coors_out = coors[:voxel_num]
            num_points_per_voxel_out = num_points_per_voxel[:voxel_num]#没啥用？
            return voxels_out, coors_out, num_points_per_voxel_out


voxelization = _Voxelization.apply


class Voxelization(nn.Module):

    def __init__(self,
                 voxel_size,
                 point_cloud_range,
                 max_num_points,
                 max_voxels=20000,
                 deterministic=True):
        super(Voxelization, self).__init__()
        """
        Args:
            voxel_size (list): list [x, y, z] size of three dimension
            point_cloud_range (list):
                [x_min, y_min, z_min, x_max, y_max, z_max]
            max_num_points (int): max number of points per voxel
            max_voxels (tuple or int): max number of voxels in
                (training, testing) time
            deterministic: bool. whether to invoke the non-deterministic
                version of hard-voxelization implementations. non-deterministic
                version is considerablly fast but is not deterministic. only
                affects hard voxelization. default True. for more information
                of this argument and the implementation insights, please refer
                to the following links:
                https://github.com/open-mmlab/mmdetection3d/issues/894
                https://github.com/open-mmlab/mmdetection3d/pull/904
                it is an experimental feature and we will appreciate it if
                you could share with us the failing cases.
        """
        self.voxel_size = voxel_size
        self.point_cloud_range = point_cloud_range
        self.max_num_points = max_num_points
        if isinstance(max_voxels, tuple):
            self.max_voxels = max_voxels
        else:
            self.max_voxels = _pair(max_voxels)
        self.deterministic = deterministic

        point_cloud_range = torch.tensor(
            point_cloud_range, dtype=torch.float32)
        # [0, -40, -3, 70.4, 40, 1]
        voxel_size = torch.tensor(voxel_size, dtype=torch.float32)
        grid_size = (point_cloud_range[3:] -
                     point_cloud_range[:3]) / voxel_size
        grid_size = torch.round(grid_size).long()
        input_feat_shape = grid_size[:2]
        self.grid_size = grid_size
        # the origin shape is as [x-len, y-len, z-len]
        # [w, h, d] -> [d, h, w]
        self.pcd_shape = [*input_feat_shape, 1][::-1]

    def forward(self, input):
        """
        Args:
            input: NC points
        """
        if self.training:
            max_voxels = self.max_voxels[0]
        else:
            max_voxels = self.max_voxels[1]

        return voxelization(input, self.voxel_size, self.point_cloud_range,
                            self.max_num_points, max_voxels,
                            self.deterministic)

    def __repr__(self):
        tmpstr = self.__class__.__name__ + '('
        tmpstr += 'voxel_size=' + str(self.voxel_size)
        tmpstr += ', point_cloud_range=' + str(self.point_cloud_range)
        tmpstr += ', max_num_points=' + str(self.max_num_points)
        tmpstr += ', max_voxels=' + str(self.max_voxels)
        tmpstr += ', deterministic=' + str(self.deterministic)
        tmpstr += ')'
        return tmpstr

class CylinderVoxelization(nn.Module):
    def __init__(self,
                max_num_points,
                point_cloud_range,
                voxel_size,
                rotate_aug = False,
                flip_aug = False,
                scale_aug = False,
                transform_aug = False,
                trans_std = [0.1, 0.1, 0.1]
                ):
        super(CylinderVoxelization, self).__init__()#感觉加了也没啥用，但不加会报错object has no attribute '_parameters'
        self.max_num_points=max_num_points
        self.point_cloud_range=np.array(point_cloud_range)
        if point_cloud_range[1]=='-np.pi':
            point_cloud_range[1]=-np.pi
            point_cloud_range[4]=np.pi
        self.grid_size=np.array(voxel_size)
        self.max_volume_space=point_cloud_range[3:]
        self.min_volume_space=point_cloud_range[:3]
        self.ignore_label=0 #权宜之计
        self.return_test = False#权宜之计 有什么用？
        self.add_sig_dim = False
        self.test_mode = False
        self.fixed_volume_space = True#权宜之计
        
        self.rotate_aug = rotate_aug
        self.flip_aug = flip_aug
        self.scale_aug = scale_aug
        self.transform = transform_aug
        self.trans_std = trans_std

    def forward(self,data):
        data_tuple=[]
        if len(data) == 1:
            self.test_mode = True
            for xzy in data[0]:
                data_tuple.append(self.forward_test([xzy]))
            data_tuple = self.collate_test(data_tuple)
        elif len(data) == 2:
            for xzy, labels in zip(data[0],data[1]):
                data_tuple.append(self.forward_train([xzy,labels]))
            data_tuple = self.collate_train(data_tuple)
        return data_tuple

    def forward_train(self,data):
        'Generates one sample of data'
        'N*C'
        # if len(data) == 1:
        #     xyz = data[0]
        #     self.test_mode = True
        if len(data) == 2:
            xyz, labels = data
            self.test_mode = False
        else:
            raise Exception('Return invalid data tuple')
        # elif len(data) == 3:
        #     xyz, labels, sig = data
        #     if len(sig.shape) == 2: sig = np.squeeze(sig)
        #     #sig?
        

        #xyz = xyz[0]
        xyz = xyz.cpu().numpy()
        if(len(xyz[0])==4):
            sig=xyz[:,3]
            xyz=xyz[:,:3]
            self.add_sig_dim = True
        if self.test_mode is True:
            labels = np.zeros((xyz.shape[0],1),dtype=np.uint8)#pesudo
        else:
            labels = labels.cpu().numpy()
            labels = np.expand_dims(labels,axis=1)

        # random data augmentation by rotation
        if self.rotate_aug:
            rotate_rad = np.deg2rad(np.random.random() * 90) - np.pi / 4
            #print(rotate_rad)
            c, s = np.cos(rotate_rad), np.sin(rotate_rad)
            j = np.matrix([[c, s], [-s, c]])
            xyz[:, :2] = np.dot(xyz[:, :2], j)

        # random data augmentation by flip x , y or x+y
        if self.flip_aug:
            flip_type = np.random.choice(4, 1)
            if flip_type == 1:
                xyz[:, 0] = -xyz[:, 0]
            elif flip_type == 2:
                xyz[:, 1] = -xyz[:, 1]
            elif flip_type == 3:
                xyz[:, :2] = -xyz[:, :2]
        if self.scale_aug:
            noise_scale = np.random.uniform(0.95, 1.05)
            xyz[:, 0] = noise_scale * xyz[:, 0]
            xyz[:, 1] = noise_scale * xyz[:, 1]
        # convert coordinate into polar coordinates

        if self.transform:
            noise_translate = np.array([np.random.normal(0, self.trans_std[0], 1),
                                        np.random.normal(0, self.trans_std[1], 1),
                                        np.random.normal(0, self.trans_std[2], 1)]).T

            xyz[:, 0:3] += noise_translate
        

        # xyz = xyz[:1000,:]
        # sig = sig[:1000]
        # labels = labels[:1000]
        xyz_pol = self.cart2polar(xyz)

        max_bound_r = np.percentile(xyz_pol[:, 0], 100, axis=0)
        min_bound_r = np.percentile(xyz_pol[:, 0], 0, axis=0)
        max_bound = np.max(xyz_pol[:, 1:], axis=0)
        min_bound = np.min(xyz_pol[:, 1:], axis=0)
        max_bound = np.concatenate(([max_bound_r], max_bound))
        min_bound = np.concatenate(([min_bound_r], min_bound))

        if self.fixed_volume_space:
            max_bound = np.asarray(self.max_volume_space)
            min_bound = np.asarray(self.min_volume_space)

        crop_range = max_bound - min_bound
        cur_grid_size = self.grid_size
        intervals = crop_range / (cur_grid_size - 1)

        if (intervals == 0).any(): print("Zero interval!")
        grid_ind = (np.floor((np.clip(xyz_pol, min_bound, max_bound) - min_bound) / intervals)).astype(np.int)#包含大量重复的坐标

        voxel_position = np.zeros(self.grid_size, dtype=np.float32)
        dim_array = np.ones(len(self.grid_size) + 1, int)#用于 reshape
        dim_array[0] = -1
        voxel_position = np.indices(self.grid_size) * intervals.reshape(dim_array) + min_bound.reshape(dim_array)
        voxel_position = self.polar2cat(voxel_position)#??
        #最终还是到cat坐标系

        processed_label = np.ones(self.grid_size, dtype=np.uint8) * self.ignore_label
        label_voxel_pair = np.concatenate([grid_ind, labels], axis=1)
        label_voxel_pair = label_voxel_pair[np.lexsort((grid_ind[:, 0], grid_ind[:, 1], grid_ind[:, 2])), :]
        #为每个grid分配唯一的label
        #print()

        processed_label = nb_process_label(np.copy(processed_label), label_voxel_pair)
        #processed_label = torch.from_numpy(processed_label).cuda()
        data_tuple = (voxel_position, processed_label)

        # center data on each voxel for PTnet
        voxel_centers = (grid_ind.astype(np.float32) + 0.5) * intervals + min_bound
        return_xyz = xyz_pol - voxel_centers
        return_xyz = np.concatenate((return_xyz, xyz_pol, xyz[:, :2]), axis=1)#N×3,N×3,N×2

        if self.add_sig_dim is False:
            return_fea = return_xyz
        elif self.add_sig_dim:
            return_fea = np.concatenate((return_xyz, sig[..., np.newaxis]), axis=1)
        
        # return_fea = torch.from_numpy(return_fea).type(torch.FloatTensor).cuda()
        # labels = [torch.from_numpy(labels).type(torch.LongTensor).cuda()]#权宜之计
        # grid_ind = torch.from_numpy(grid_ind).cuda()
        if self.return_test:
            data_tuple += (grid_ind, labels, return_fea, 0)#origin:index

        else:
            data_tuple += (grid_ind, labels, return_fea)

        #return self.collate_fn_BEV(data_tuple)
        return data_tuple

    def forward_test(self,data):
        'Generates one sample of data'
        'N*C'
        # if len(data) == 1:
        #     xyz = data[0]
        #     self.test_mode = True
        if len(data) == 1:
            xyz = data[0]
            self.test_mode = False
        else:
            raise Exception('Return invalid data tuple')
        # elif len(data) == 3:
        #     xyz, labels, sig = data
        #     if len(sig.shape) == 2: sig = np.squeeze(sig)
        #     #sig?
        

        #xyz = xyz[0]
        xyz = xyz.cpu().numpy()
        if(len(xyz[0])==4):
            sig=xyz[:,3]
            xyz=xyz[:,:3]
            self.add_sig_dim = True


        # random data augmentation by rotation
        if self.rotate_aug:
            rotate_rad = np.deg2rad(np.random.random() * 90) - np.pi / 4
            #print(rotate_rad)
            c, s = np.cos(rotate_rad), np.sin(rotate_rad)
            j = np.matrix([[c, s], [-s, c]])
            xyz[:, :2] = np.dot(xyz[:, :2], j)

        # random data augmentation by flip x , y or x+y
        if self.flip_aug:
            flip_type = np.random.choice(4, 1)
            if flip_type == 1:
                xyz[:, 0] = -xyz[:, 0]
            elif flip_type == 2:
                xyz[:, 1] = -xyz[:, 1]
            elif flip_type == 3:
                xyz[:, :2] = -xyz[:, :2]
        if self.scale_aug:
            noise_scale = np.random.uniform(0.95, 1.05)
            xyz[:, 0] = noise_scale * xyz[:, 0]
            xyz[:, 1] = noise_scale * xyz[:, 1]
        # convert coordinate into polar coordinates

        if self.transform:
            noise_translate = np.array([np.random.normal(0, self.trans_std[0], 1),
                                        np.random.normal(0, self.trans_std[1], 1),
                                        np.random.normal(0, self.trans_std[2], 1)]).T

            xyz[:, 0:3] += noise_translate
        

        # xyz = xyz[:1000,:]
        # sig = sig[:1000]
        # labels = labels[:1000]
        xyz_pol = self.cart2polar(xyz)

        max_bound_r = np.percentile(xyz_pol[:, 0], 100, axis=0)
        min_bound_r = np.percentile(xyz_pol[:, 0], 0, axis=0)
        max_bound = np.max(xyz_pol[:, 1:], axis=0)
        min_bound = np.min(xyz_pol[:, 1:], axis=0)
        max_bound = np.concatenate(([max_bound_r], max_bound))
        min_bound = np.concatenate(([min_bound_r], min_bound))

        if self.fixed_volume_space:
            max_bound = np.asarray(self.max_volume_space)
            min_bound = np.asarray(self.min_volume_space)

        crop_range = max_bound - min_bound
        cur_grid_size = self.grid_size
        intervals = crop_range / (cur_grid_size - 1)

        if (intervals == 0).any(): print("Zero interval!")
        grid_ind = (np.floor((np.clip(xyz_pol, min_bound, max_bound) - min_bound) / intervals)).astype(np.int)#包含大量重复的坐标

        voxel_position = np.zeros(self.grid_size, dtype=np.float32)
        dim_array = np.ones(len(self.grid_size) + 1, int)#用于 reshape
        dim_array[0] = -1
        voxel_position = np.indices(self.grid_size) * intervals.reshape(dim_array) + min_bound.reshape(dim_array)
        voxel_position = self.polar2cat(voxel_position)#??
        #最终还是到cat坐标系

        #processed_label = torch.from_numpy(processed_label).cuda()

        # center data on each voxel for PTnet
        voxel_centers = (grid_ind.astype(np.float32) + 0.5) * intervals + min_bound
        return_xyz = xyz_pol - voxel_centers
        return_xyz = np.concatenate((return_xyz, xyz_pol, xyz[:, :2]), axis=1)#N×3,N×3,N×2

        if self.add_sig_dim is False:
            return_fea = return_xyz
        elif self.add_sig_dim:
            return_fea = np.concatenate((return_xyz, sig[..., np.newaxis]), axis=1)
        
        # return_fea = torch.from_numpy(return_fea).type(torch.FloatTensor).cuda()
        # labels = [torch.from_numpy(labels).type(torch.LongTensor).cuda()]#权宜之计
        # grid_ind = torch.from_numpy(grid_ind).cuda()
        if self.return_test:
            data_tuple = (None, grid_ind, return_fea, 0)#origin:index

        else:
            data_tuple = (None, grid_ind, return_fea)

        #return self.collate_fn_BEV(data_tuple)
        return data_tuple

    
    def collate_train(self,data):
        data2stack = np.stack([d[0] for d in data]).astype(np.float32)
        label2stack = np.stack([d[1] for d in data]).astype(np.int)
        grid_ind_stack = [d[2] for d in data]
        point_label = [d[3] for d in data]
        xyz = [d[4] for d in data]
        return torch.from_numpy(data2stack), \
        torch.from_numpy(label2stack).type(torch.LongTensor).cuda(), \
        [torch.from_numpy(i).cuda() for i in grid_ind_stack], \
        point_label, \
        [torch.from_numpy(i).type(torch.FloatTensor).cuda() for i in xyz]

    def collate_test(self,data):
        #data2stack = np.stack([d[0] for d in data]).astype(np.float32)
        grid_ind_stack = [d[1] for d in data]
        xyz = [d[2] for d in data]
        return None, \
        [torch.from_numpy(i).cuda() for i in grid_ind_stack], \
        [torch.from_numpy(i).type(torch.FloatTensor).cuda() for i in xyz]

    def cart2polar(self,input_xyz):
        rho = np.sqrt(input_xyz[:, 0] ** 2 + input_xyz[:, 1] ** 2)
        phi = np.arctan2(input_xyz[:, 1], input_xyz[:, 0])
        return np.stack((rho, phi, input_xyz[:, 2]), axis=1)

    def polar2cat(self,input_xyz_polar):
        # print(input_xyz_polar.shape)
        x = input_xyz_polar[0] * np.cos(input_xyz_polar[1])
        y = input_xyz_polar[0] * np.sin(input_xyz_polar[1])
        return np.stack((x, y, input_xyz_polar[2]), axis=0)

class Spherical_dataset(CylinderVoxelization):
    def __init__(self,
                max_num_points,
                point_cloud_range,
                voxel_size,
                rotate_aug = False,
                flip_aug = False,
                scale_aug = False,
                transform_aug = False,
                trans_std = [0.1, 0.1, 0.1]
                ):
        super(CylinderVoxelization, self).__init__(
                max_num_points=max_num_points,
                point_cloud_range=point_cloud_range,
                voxel_size=voxel_size,
                rotate_aug = rotate_aug,
                flip_aug = flip_aug,
                scale_aug = scale_aug,
                transform_aug = transform_aug,
                trans_std = trans_std
        )

    def forward(self,data):
        data_tuple=[]
        if len(data) == 1:
            self.test_mode = True
            for xzy in data[0]:
                data_tuple.append(self.forward_single([xzy]))
        elif len(data) == 2:
            for xzy, labels in zip(data[0],data[1]):
                data_tuple.append(self.forward_single([xzy,labels]))
        elif len(data) == 6:
            for xzy, labels in zip(data[0],data[1]):
                data_tuple.append(self.forward_single([xzy,labels]))
        elif len(data) == 7:
            for xzy, labels in zip(data[0],data[1]):
                data_tuple.append(self.forward_single([xzy,labels]))

        return self.collate_fn_BEV(data_tuple)

    def forward_single(self,data):
        self.test_mode = False
        self.add_sig_dim = False

        if len(data) == 1:
            xyz = data[0]
            self.test_mode = True
        elif len(data) == 2:
            xyz, labels = data
        elif len(data) == 6:
            xyz,labels,sig,ins_labels,valid,pcd_fname = data
            if len(sig.shape) == 2: sig = np.squeeze(sig)
        elif len(data) == 7:
            xyz,labels,sig,ins_labels,valid,pcd_fname,minicluster = data
            if len(sig.shape) == 2: sig = np.squeeze(sig)
        else: raise Exception('Return invalid data tuple')

        xyz = xyz.cpu().numpy()
        if(len(xyz[0])==4):
            sig=xyz[:,3]
            xyz=xyz[:,:3]
            self.add_sig_dim = True
        if self.test_mode is True:
            labels = np.zeros((xyz.shape[0],1),dtype=np.uint8)#pesudo
        else:
            labels = labels.cpu().numpy()
            labels = np.expand_dims(labels,axis=1)

        # random data augmentation by rotation
        if self.rotate_aug:
            rotate_rad = np.deg2rad(np.random.random()*360)# why 360
            c, s = np.cos(rotate_rad), np.sin(rotate_rad)
            j = np.matrix([[c, s], [-s, c]])
            xyz[:, :2] = np.dot(xyz[:, :2], j)

        # random data augmentation by flip x , y or x+y
        if self.flip_aug:
            flip_type = np.random.choice(4, 1)
            if flip_type == 1:
                xyz[:, 0] = -xyz[:, 0]
            elif flip_type == 2:
                xyz[:, 1] = -xyz[:, 1]
            elif flip_type == 3:
                xyz[:, :2] = -xyz[:, :2]

        if self.scale_aug:
            noise_scale = np.random.uniform(0.95, 1.05)
            xyz[:, 0] = noise_scale * xyz[:, 0]
            xyz[:, 1] = noise_scale * xyz[:, 1]

        if self.transform:
            noise_translate = np.array([np.random.normal(0, self.trans_std[0], 1),
                                        np.random.normal(0, self.trans_std[1], 1),
                                        np.random.normal(0, self.trans_std[2], 1)]).T

            xyz[:, 0:3] += noise_translate
        

        # xyz = xyz[:1000,:]
        # sig = sig[:1000]
        # labels = labels[:1000]
        # convert coordinate into polar coordinates
        xyz_pol = self.cart2polar(xyz)

        max_bound_r = np.percentile(xyz_pol[:, 0], 100, axis=0)
        min_bound_r = np.percentile(xyz_pol[:, 0], 0, axis=0)
        max_bound = np.max(xyz_pol[:, 1:], axis=0)
        min_bound = np.min(xyz_pol[:, 1:], axis=0)
        max_bound = np.concatenate(([max_bound_r], max_bound))
        min_bound = np.concatenate(([min_bound_r], min_bound))

        if self.fixed_volume_space:
            max_bound = np.asarray(self.max_volume_space)
            min_bound = np.asarray(self.min_volume_space)

        # get grid index
        crop_range = max_bound - min_bound
        cur_grid_size = self.grid_size
        intervals = crop_range / (cur_grid_size - 1) # (size-1) could directly get index starting from 0, very convenient

        if (intervals == 0).any(): print("Zero interval!")
        grid_ind = (np.floor((np.clip(xyz_pol, min_bound, max_bound) - min_bound) / intervals)).astype(np.int) # point-wise grid index

        # process voxel position
        voxel_position = np.zeros(self.grid_size, dtype=np.float32)
        dim_array = np.ones(len(self.grid_size) + 1, int) # for reshape
        dim_array[0] = -1
        voxel_position = np.indices(self.grid_size) * intervals.reshape(dim_array) + min_bound.reshape(dim_array)
        voxel_position = self.polar2cat(voxel_position)#??

        # process labels
        processed_label = np.ones(self.grid_size, dtype=np.uint8) * self.ignore_label
        label_voxel_pair = np.concatenate([grid_ind, labels], axis=1)
        label_voxel_pair = label_voxel_pair[np.lexsort((grid_ind[:, 0], grid_ind[:, 1], grid_ind[:, 2])), :]
        
        #allocate label for every grid
        processed_label = nb_process_label(np.copy(processed_label), label_voxel_pair)
        data_tuple = (voxel_position, processed_label)

        # center data on each voxel for PTnet
        voxel_centers = (grid_ind.astype(np.float32) + 0.5) * intervals + min_bound
        return_xyz = xyz_pol - voxel_centers #TODO: calculate relative coordinate using polar system?
        return_xyz = np.concatenate((return_xyz, xyz_pol, xyz[:, :2]), axis=1)#N×3,N×3,N×2

        if self.add_sig_dim is False:
            return_fea = return_xyz
        elif self.add_sig_dim:
            return_fea = np.concatenate((return_xyz, sig[..., np.newaxis]), axis=1)
        

        if self.return_test:
            data_tuple += (grid_ind, labels, return_fea, 0)#TODO origin:index

        else:
            data_tuple += (grid_ind, labels, return_fea) # (grid-wise coor, grid-wise sem label, point-wise grid index, point-wise sem label, [relative polar coor(3), polar coor(3), cat coor(2), ref signal(1)])

        if len(data) == 6:
            offsets = np.zeros([xyz.shape[0], 3], dtype=np.float32)
            offsets = nb_aggregate_pointwise_center_offset(offsets, xyz, ins_labels, self.center_type)
            data_tuple += (ins_labels, offsets, valid, xyz, pcd_fname) # plus (point-wise instance label, point-wise center offset)

        if len(data) == 7:
            offsets = np.zeros([xyz.shape[0], 3], dtype=np.float32)
            offsets = nb_aggregate_pointwise_center_offset(offsets, xyz, ins_labels, self.center_type)
            data_tuple += (ins_labels, offsets, valid, xyz, pcd_fname, minicluster) # plus (point-wise instance label, point-wise center offset)
        
        return data_tuple
@nb.jit('u1[:,:,:](u1[:,:,:],i8[:,:])', nopython=True, cache=True, parallel=False)
def nb_process_label(processed_label, sorted_label_voxel_pair):
    label_size = 20#权宜之计
    counter = np.zeros((label_size,), dtype=np.uint16)
    counter[sorted_label_voxel_pair[0, 3]] = 1
    cur_sear_ind = sorted_label_voxel_pair[0, :3]
    for i in range(1, sorted_label_voxel_pair.shape[0]):
        cur_ind = sorted_label_voxel_pair[i, :3]
        if not np.all(np.equal(cur_ind, cur_sear_ind)):
            processed_label[cur_sear_ind[0], cur_sear_ind[1], cur_sear_ind[2]] = np.argmax(counter)
            counter = np.zeros((label_size,), dtype=np.uint16)
            cur_sear_ind = cur_ind
        counter[sorted_label_voxel_pair[i, 3]] += 1
    processed_label[cur_sear_ind[0], cur_sear_ind[1], cur_sear_ind[2]] = np.argmax(counter)
    return processed_label

class FastCylinderVoxelization(nn.Module):

    def __init__(self,
                 point_cloud_range,
                 grid_size):
        super(FastCylinderVoxelization, self).__init__()
        """
        Args:
            voxel_size (list): list [x, y, z] size of three dimension
            point_cloud_range (list):
                [x_min, y_min, z_min, x_max, y_max, z_max]
            max_num_points (int): max number of points per voxel
            max_voxels (tuple or int): max number of voxels in
                (training, testing) time
            deterministic: bool. whether to invoke the non-deterministic
                version of hard-voxelization implementations. non-deterministic
                version is considerablly fast but is not deterministic. only
                affects hard voxelization. default True. for more information
                of this argument and the implementation insights, please refer
                to the following links:
                https://github.com/open-mmlab/mmdetection3d/issues/894
                https://github.com/open-mmlab/mmdetection3d/pull/904
                it is an experimental feature and we will appreciate it if
                you could share with us the failing cases.
        """
        self.point_cloud_range = point_cloud_range
        self.class_num = 20
        self.list_point_cloud_range = point_cloud_range
        if point_cloud_range[1]=='-np.pi':
            point_cloud_range[1]=-np.pi
            point_cloud_range[4]=np.pi

        self.point_cloud_range = torch.tensor(
            self.point_cloud_range).cuda()
        # [0, -40, -3, 70.4, 40, 1]
        self.list_grid_size = grid_size
        self.grid_size = torch.tensor(grid_size).cuda()
        input_feat_shape = grid_size[:2]
        # the origin shape is as [x-len, y-len, z-len]
        # [w, h, d] -> [d, h, w]
        self.pcd_shape = [*input_feat_shape, 1][::-1]
        self.ignore_label = 0
        self.return_test = False

    def forward(self,data):
        data_tuple=[]
        if len(data) == 1:
            for xyz in data[0]:
                data_tuple.append(self.forward_test([xyz[:,:3],xyz[:,3]]))
            data_tuple = self.collate_test(data_tuple)
        elif len(data) == 2:
            for xyz, labels in zip(data[0],data[1]):
                data_tuple.append(self.forward_train([xyz[:,:3],labels,xyz[:,3]]))
            data_tuple = self.collate_train(data_tuple)
        return data_tuple

    def collate_train(self,data):
        #data2stack = torch.stack([d[0] for d in data]).astype(np.float32)
        label2stack = torch.stack([d[1].type(torch.LongTensor).cuda() for d in data])#loss need long
        grid_ind_stack = [d[2] for d in data]
        point_label = [d[3] for d in data]
        xyz = [d[4] for d in data]
        return None, label2stack, grid_ind_stack, point_label, xyz
    
    def collate_test(self,data):
        #data2stack = torch.stack([d[0] for d in data]).astype(np.float32)
        grid_ind_stack = [d[1].type(torch.LongTensor).cuda() for d in data]
        xyz = [d[2] for d in data]
        return None, grid_ind_stack, xyz

    def forward_train(self, data):

        if len(data) ==3:
            xyz, labels, sig = data
            xyz_pol = self.cart2polar(xyz)
        else:
            raise Exception('Return invalid data tuple')
        
        max_bound = self.point_cloud_range[3:]
        min_bound = self.point_cloud_range[:3]
        crop_range = max_bound - min_bound
        voxel_size = crop_range/(self.grid_size-1)# (size-1) could directly get index starting from 0, very convenient

        #get pixel-wise cylinder index #TODO replace dynamic_voxelize with simple division and compare speed
        grid_ind = xyz_pol.new_zeros(size=(xyz_pol.size(0), 3), dtype=torch.int)        
        dynamic_voxelize(self.torch_clip(xyz_pol,min_bound,max_bound), grid_ind, voxel_size.cpu().numpy().tolist(), self.list_point_cloud_range, 3)# TODO why need tolist?
        grid_ind = self.swap_xz(grid_ind)#dynamic_voxelize will swap x and z
    
        #get return_fea
        voxel_centers = (grid_ind.float() + 0.5) * voxel_size + min_bound
        return_xyz = xyz_pol - voxel_centers
        return_xyz = torch.cat((return_xyz, xyz_pol, xyz[:, :2]), axis=1)#N×3,N×3,N×2
        return_fea = torch.cat((return_xyz, torch.unsqueeze(sig,1)), axis=1)#TODO choice to add sig or not

        #get grid-wise label
        processed_label = self.aggregate_label(grid_ind,labels)

        data_tuple = (None,processed_label)#TODO rm None
        if self.return_test:
            data_tuple += (grid_ind, labels, return_fea, 0)#TODO origin:index
        else:
            data_tuple += (grid_ind, labels, return_fea)
        return data_tuple

    def forward_test(self, data):

        if len(data) ==2:
            xyz, sig = data
            xyz_pol = self.cart2polar(xyz)
            #if len(sig.shape) == 2: sig = torch.squeeze(sig)
            #input = torch.cat((xyz,labels,sig),1)
        else:
            raise Exception('Return invalid data tuple')

        max_bound = self.point_cloud_range[3:]
        min_bound = self.point_cloud_range[:3]
        crop_range = max_bound - min_bound
        voxel_size = crop_range/(self.grid_size-1)

        grid_ind = xyz_pol.new_zeros(size=(xyz_pol.size(0), 3), dtype=torch.int)# 这里不能long,只能int,和原code有差别
        
        
        #dynamic_voxelize(self.torch_clip(xyz_pol,min_bound,max_bound), grid_ind, self.list_grid_size, self.list_point_cloud_range, 3)
        #cliped = torch.tensor([[49.0,0,1.9]])
        dynamic_voxelize(self.torch_clip(xyz_pol,min_bound,max_bound), grid_ind, voxel_size.cpu().numpy().tolist(), self.list_point_cloud_range, 3)# TODO why need tolist?
        grid_ind = self.swap_xz(grid_ind)
        
        cur_grid_size = self.grid_size

        intervals = crop_range / (cur_grid_size - 1) # (size-1) could directly get index starting from 0, very convenient
        voxel_centers = (grid_ind.float() + 0.5) * intervals + min_bound

        return_xyz = xyz_pol - voxel_centers
        return_xyz = torch.cat((return_xyz, xyz_pol, xyz[:, :2]), axis=1)#N×3,N×3,N×2


        # if len(data) == 2:
        #     return_fea = return_xyz
        # elif len(data) == 3:
        return_fea = torch.cat((return_xyz, torch.unsqueeze(sig,1)), axis=1)

        #processed_label = self.aggregate_label(grid_ind,labels)

        #data_tuple = tuple(None)#TODO rm None

        if self.return_test:#TODO rm None
            data_tuple = (None, grid_ind, return_fea, 0)#origin:index
        else:
            data_tuple = (None, grid_ind, return_fea)
        return data_tuple

    def cart2polar(self, input_xyz):
        rho = torch.sqrt(input_xyz[:, 0] ** 2 + input_xyz[:, 1] ** 2)
        phi = torch.atan2(input_xyz[:, 1], input_xyz[:, 0])
        return torch.stack((rho, phi, input_xyz[:, 2]), axis=1)
    
    def torch_clip(self,src,rmin,rmax):
        for i in range(3):
            src[:,i] = torch.clamp(src[:,i],rmin[i]+0.001,rmax[i]-0.001)
        return src

    def aggregate_label(self, grid_ind, labels):
        # cat_pt_ind = []
        # for i_batch in range(len(coor)):
        #     cat_pt_ind.append(F.pad(coor[i_batch], (1, 0), 'constant', value=i_batch))
        # cat_pt_labels = torch.cat(labels, dim=0)
        # cat_pt_ind = torch.cat(cat_pt_ind, dim=0)#把不同batch的点全部混在了一起，但第一维增加了batch，所以仍然可以区分
        pt_num = grid_ind.shape[0]
        
        unq, unq_inv, unq_cnt = torch.unique(grid_ind, return_inverse=True, return_counts=True, dim=0)
        # process feature
        processed_cat_pt_fea = grid_ind.new_zeros((pt_num,self.class_num))

        processed_cat_pt_fea[list(range(0, processed_cat_pt_fea.shape[0])),labels[list(range(0, processed_cat_pt_fea.shape[0]))].cpu().numpy().tolist()] +=1
        pooled_data = torch_scatter.scatter_sum(processed_cat_pt_fea, unq_inv, dim=0)
        max_labels = torch.argmax(pooled_data,dim=1)
        processed_label = grid_ind.new_ones(self.list_grid_size) * self.ignore_label
        for i in range(unq.shape[0]):
            processed_label[unq[i,0],unq[i,1],unq[i,2]] = max_labels[i]
        
        return processed_label

    def swap_xz(self,grid_ind):
        index = [2,1,0]
        return grid_ind[:,index]
