|
- import os
- import numpy as np
- import pandas as pd
- from mindspore import dtype as mstype
- import pickle
- NUM_HEADING_BIN = 12
- NUM_SIZE_CLUSTER = 8 # one cluster for each type
- NUM_OBJECT_POINT = 512
- g_type2class={'Car':0, 'Van':1, 'Truck':2, 'Pedestrian':3,
- 'Person_sitting':4, 'Cyclist':5, 'Tram':6, 'Misc':7}
- g_class2type = {g_type2class[t]:t for t in g_type2class}
- g_type2onehotclass = {'Car': 0, 'Pedestrian': 1, 'Cyclist': 2}
- g_type_mean_size = {'Car': np.array([3.88311640418,1.62856739989,1.52563191462]),
- 'Van': np.array([5.06763659,1.9007158,2.20532825]),
- 'Truck': np.array([10.13586957,2.58549199,3.2520595]),
- 'Pedestrian': np.array([0.84422524,0.66068622,1.76255119]),
- 'Person_sitting': np.array([0.80057803,0.5983815,1.27450867]),
- 'Cyclist': np.array([1.76282397,0.59706367,1.73698127]),
- 'Tram': np.array([16.17150617,2.53246914,3.53079012]),
- 'Misc': np.array([3.64300781,1.54298177,1.92320313])}
-
- class FrustumDataset_train:
- def __init__(self):
- self.TRAIN_DATASET = FrustumDataset(npoints=1024, split='train',rotate_to_center=True, random_flip=True, random_shift=True, one_hot=True)
- #self.TRAIN_DATASET = FrustumDataset(npoints=1024, split='train',rotate_to_center=True, random_flip=False, random_shift=False, one_hot=True)
- self.data=np.zeros((4,1024))
- self.label=np.zeros((1024,))
- self.center=np.zeros((3,))
- self.hclass=np.zeros((1,))
- self.hres=np.zeros((1,))
- self.sclass=np.zeros((1,))
- self.sres=np.zeros((3,))
- self.rot_angle=np.zeros((1,))
- self.one_hot_vec=np.zeros((3,))
- def __getitem__(self, index):
- data,label,center,hclass,hres,sclass,sres,rotangle,one_hot_vec=self.TRAIN_DATASET[index]
- self.data=np.squeeze(data.T).astype(np.float32)
- self.one_hot_vec=np.squeeze(one_hot_vec).astype(np.float32)
- self.label=np.squeeze(label).astype(np.float32)
- self.center=np.squeeze(center).astype(np.float32)
- self.hclass=np.squeeze(hclass).astype(np.float32)
- self.hres=np.squeeze(hres).astype(np.float32)
- self.sclass=np.squeeze(sclass).astype(np.float32)
- self.sres=np.squeeze(sres).astype(np.float32)
- #print(self.data.shape, self.one_hot_vec.shape, self.label.shape, self.center.shape, self.hclass.shape, self.hres.shape, self.sclass.shape, self.sres.shape)
- return self.data, self.one_hot_vec, self.label, self.center, self.hclass, self.hres, self.sclass, self.sres
-
- def __len__(self):
- return 73508
-
- class FrustumDataset_test:
- def __init__(self):
- self.TEST_DATASET = FrustumDataset(npoints=1024, split='val',rotate_to_center=True, one_hot=True)
- self.data=np.zeros((4,1024))
- self.label=np.zeros((1024,))
- self.center=np.zeros((3,))
- self.hclass=np.zeros((1,))
- self.hres=np.zeros((1,))
- self.sclass=np.zeros((1,))
- self.sres=np.zeros((3,))
- self.rot_angle=np.zeros((1,))
- self.one_hot_vec=np.zeros((3,))
- def __getitem__(self, index):
- data,label,center,hclass,hres,sclass,sres,rotangle,one_hot_vec=self.TEST_DATASET[index]
- self.data=np.squeeze(data.T).astype(np.float32)
- self.one_hot_vec=np.squeeze(one_hot_vec).astype(np.float32)
- self.label=np.squeeze(label).astype(np.float32)
- self.center=np.squeeze(center).astype(np.float32)
- self.hclass=np.squeeze(hclass).astype(np.float32)
- self.hres=np.squeeze(hres).astype(np.float32)
- self.sclass=np.squeeze(sclass).astype(np.float32)
- self.sres=np.squeeze(sres).astype(np.float32)
- return self.data, self.one_hot_vec, self.label, self.center, self.hclass, self.hres, self.sclass, self.sres
-
- def __len__(self):
- return 15589
- def size2class(size, type_name):
- ''' Convert 3D bounding box size to template class and residuals.
- todo (rqi): support multiple size clusters per type.
-
- Input:
- size: numpy array of shape (3,) for (l,w,h)
- type_name: string
- Output:
- size_class: int scalar
- size_residual: numpy array of shape (3,)
- '''
- size_class = g_type2class[type_name.decode()]
- size_residual = size - g_type_mean_size[type_name.decode()]
- return size_class, size_residual
-
- def class2size(pred_cls, residual):
- ''' Inverse function to size2class. '''
- mean_size = g_type_mean_size[g_class2type[pred_cls]]
- return mean_size + residual
- def angle2class(angle, num_class):
- ''' Convert continuous angle to discrete class and residual.
-
- Input:
- angle: rad scalar, from 0-2pi (or -pi~pi), class center at
- 0, 1*(2pi/N), 2*(2pi/N) ... (N-1)*(2pi/N)
- num_class: int scalar, number of classes N
- Output:
- class_id, int, among 0,1,...,N-1
- residual_angle: float, a number such that
- class*(2pi/N) + residual_angle = angle
- '''
- angle = angle%(2*np.pi)
- assert(angle>=0 and angle<=2*np.pi)
- angle_per_class = 2*np.pi/float(num_class)
- shifted_angle = (angle+angle_per_class/2)%(2*np.pi)
- class_id = int(shifted_angle/angle_per_class)
- residual_angle = shifted_angle - \
- (class_id * angle_per_class + angle_per_class/2)
- return class_id, residual_angle
-
- def class2angle(pred_cls, residual, num_class, to_label_format=True):
- ''' Inverse function to angle2class.
- If to_label_format, adjust angle to the range as in labels.
- '''
- angle_per_class = 2*np.pi/float(num_class)
- angle_center = pred_cls * angle_per_class
- angle = angle_center + residual
- if to_label_format and angle>np.pi:
- angle = angle - 2*np.pi
- return angle
- def rotate_pc_along_y(pc, rot_angle):
- '''
- Input:
- pc: numpy array (N,C), first 3 channels are XYZ
- z is facing forward, x is left ward, y is downward
- rot_angle: rad scalar
- Output:
- pc: updated pc with XYZ rotated
- '''
- cosval = np.cos(rot_angle)
- sinval = np.sin(rot_angle)
- rotmat = np.array([[cosval, -sinval],[sinval, cosval]])
- pc[:,[0,2]] = np.dot(pc[:,[0,2]], np.transpose(rotmat))
- return pc
- class FrustumDataset(object):
- ''' Dataset class for Frustum PointNets training/evaluation.
- Load prepared KITTI data from pickled files, return individual data element
- [optional] along with its annotations.
- '''
- def __init__(self, npoints, split,
- random_flip=False, random_shift=False, rotate_to_center=False,
- overwritten_data_path=None, from_rgb_detection=False, one_hot=False):
- '''
- Input:
- npoints: int scalar, number of points for frustum point cloud.
- split: string, train or val
- random_flip: bool, in 50% randomly flip the point cloud
- in left and right (after the frustum rotation if any)
- random_shift: bool, if True randomly shift the point cloud
- back and forth by a random distance
- rotate_to_center: bool, whether to do frustum rotation
- overwritten_data_path: string, specify pickled file path.
- if None, use default path (with the split)
- from_rgb_detection: bool, if True we assume we do not have
- groundtruth, just return data elements.
- one_hot: bool, if True, return one hot vector
- '''
- self.npoints = npoints
- self.random_flip = random_flip
- self.random_shift = random_shift
- self.rotate_to_center = rotate_to_center
- self.one_hot = one_hot
- if overwritten_data_path is None:
- overwritten_data_path = os.path.join("data/",
- 'frustum_carpedcyc_%s.pickle'%(split))
-
- self.from_rgb_detection = from_rgb_detection
- if from_rgb_detection:
- with open(overwritten_data_path,'rb') as fp:
- self.id_list = pickle.load(fp)
- self.box2d_list = pickle.load(fp,encoding='bytes')
- self.input_list = pickle.load(fp,encoding='bytes')
- self.type_list = pickle.load(fp,encoding='bytes')
- # frustum_angle is clockwise angle from positive x-axis
- self.frustum_angle_list = pickle.load(fp,encoding='bytes')
- self.prob_list = pickle.load(fp,encoding='bytes')
- else:
- with open(overwritten_data_path,'rb') as fp:
- self.id_list = pickle.load(fp)
- self.box2d_list = pickle.load(fp,encoding='bytes')
- self.box3d_list = pickle.load(fp,encoding='bytes')
- self.input_list = pickle.load(fp,encoding='bytes')
- self.label_list = pickle.load(fp,encoding='bytes')
- self.type_list = pickle.load(fp,encoding='bytes')
- self.heading_list = pickle.load(fp,encoding='bytes')
- self.size_list = pickle.load(fp,encoding='bytes')
- # frustum_angle is clockwise angle from positive x-axis
- self.frustum_angle_list = pickle.load(fp,encoding='bytes')
-
- def __len__(self):
- return len(self.input_list)
-
- def __getitem__(self, index):
- ''' Get index-th element from the picked file dataset. '''
- # ------------------------------ INPUTS ----------------------------
- rot_angle = self.get_center_view_rot_angle(index)
-
- # Compute one hot vector
- if self.one_hot:
- cls_type = self.type_list[index].decode()
- #print(type(cls_type),type('Car'))
- assert(cls_type in ['Car', 'Pedestrian', 'Cyclist'])
- one_hot_vec = np.zeros((3))
- one_hot_vec[g_type2onehotclass[cls_type]] = 1
-
- # Get point cloud
- if self.rotate_to_center:
- point_set = self.get_center_view_point_set(index)
- else:
- point_set = self.input_list[index]
- # Resample
- choice = np.random.choice(point_set.shape[0], self.npoints, replace=True)
- point_set = point_set[choice, :]
-
- if self.from_rgb_detection:
- if self.one_hot:
- return point_set, rot_angle, self.prob_list[index], one_hot_vec
- else:
- return point_set, rot_angle, self.prob_list[index]
-
- # ------------------------------ LABELS ----------------------------
- seg = self.label_list[index]
- seg = seg[choice]
-
- # Get center point of 3D box
- if self.rotate_to_center:
- box3d_center = self.get_center_view_box3d_center(index)
- else:
- box3d_center = self.get_box3d_center(index)
-
- # Heading
- if self.rotate_to_center:
- heading_angle = self.heading_list[index] - rot_angle
- else:
- heading_angle = self.heading_list[index]
-
- # Size
- size_class, size_residual = size2class(self.size_list[index],
- self.type_list[index])
-
- # Data Augmentation
- if self.random_flip:
- # note: rot_angle won't be correct if we have random_flip
- # so do not use it in case of random flipping.
- if np.random.random()>0.5: # 50% chance flipping
- point_set[:,0] *= -1
- box3d_center[0] *= -1
- heading_angle = np.pi - heading_angle
- if self.random_shift:
- dist = np.sqrt(np.sum(box3d_center[0]**2+box3d_center[1]**2))
- shift = np.clip(np.random.randn()*dist*0.05, dist*0.8, dist*1.2)
- point_set[:,2] += shift
- box3d_center[2] += shift
-
- angle_class, angle_residual = angle2class(heading_angle,
- NUM_HEADING_BIN)
-
- if self.one_hot:
- return point_set, seg, box3d_center, angle_class, angle_residual,\
- size_class, size_residual, rot_angle, one_hot_vec
- else:
- return point_set, seg, box3d_center, angle_class, angle_residual,\
- size_class, size_residual, rot_angle
-
- def get_center_view_rot_angle(self, index):
- ''' Get the frustum rotation angle, it isshifted by pi/2 so that it
- can be directly used to adjust GT heading angle '''
- return np.pi/2.0 + self.frustum_angle_list[index]
-
- def get_box3d_center(self, index):
- ''' Get the center (XYZ) of 3D bounding box. '''
- box3d_center = (self.box3d_list[index][0,:] + \
- self.box3d_list[index][6,:])/2.0
- return box3d_center
-
- def get_center_view_box3d_center(self, index):
- ''' Frustum rotation of 3D bounding box center. '''
- box3d_center = (self.box3d_list[index][0,:] + \
- self.box3d_list[index][6,:])/2.0
- return rotate_pc_along_y(np.expand_dims(box3d_center,0), \
- self.get_center_view_rot_angle(index)).squeeze()
-
- def get_center_view_box3d(self, index):
- ''' Frustum rotation of 3D bounding box corners. '''
- box3d = self.box3d_list[index]
- box3d_center_view = np.copy(box3d)
- return rotate_pc_along_y(box3d_center_view, \
- self.get_center_view_rot_angle(index))
-
- def get_center_view_point_set(self, index):
- ''' Frustum rotation of point clouds.
- NxC points with first 3 channels as XYZ
- z is facing forward, x is left ward, y is downward
- '''
- # Use np.copy to avoid corrupting original data
- point_set = np.copy(self.input_list[index])
- return rotate_pc_along_y(point_set, \
- self.get_center_view_rot_angle(index))
-
-
|