def generate_labels(self, center, dimension, angle, ref_xyz): box_corner1 = compute_box_3d(center, dimension * 0.5, angle) box_corner2 = compute_box_3d(center, dimension, angle) labels = np.zeros(len(ref_xyz)) inside1 = extract_pc_in_box3d(ref_xyz, box_corner1) inside2 = extract_pc_in_box3d(ref_xyz, box_corner2) labels[inside2] = -1 labels[inside1] = 1 if inside1.sum() == 0: dis = np.sqrt(((ref_xyz - center)**2).sum(1)) argmin = np.argmin(dis) labels[argmin] = 1 return labels
def check_box_frustum(box, P, center, dimension, angle): x1, y1, x2, y2 = box box_corner = compute_box_3d(center, dimension, angle, P) # 8, 3 z1 = np.arange(0, 70, 0.1) xyz1 = np.zeros((len(z1), 3)) xyz1[:, 0] = x1 xyz1[:, 1] = y1 xyz1[:, 2] = z1 xyz1_rect = project_image_to_rect(xyz1, P) xyz1[:, 0] = x2 xyz1[:, 1] = y2 xyz1[:, 2] = z1 xyz2_rect = project_image_to_rect(xyz1, P) xyz1[:, 0] = x1 xyz1[:, 1] = y2 xyz1[:, 2] = z1 xyz3_rect = project_image_to_rect(xyz1, P) xyz1[:, 0] = x2 xyz1[:, 1] = y1 xyz1[:, 2] = z1 xyz4_rect = project_image_to_rect(xyz1, P) fig = plt.figure() ax = fig.gca(projection='3d') draw_box3d(box_corner, ax) draw_points(xyz1_rect, ax) draw_points(xyz2_rect, ax) draw_points(xyz3_rect, ax) draw_points(xyz4_rect, ax) plt.show()
def __getitem__(self, index): rotate_to_center = cfg.DATA.RTC with_extra_feat = cfg.DATA.WITH_EXTRA_FEAT point_set = self.input_list[index].copy() pred_box3d = self.pred_box3d_list[index].copy() pred_box3d_center = (pred_box3d[0, :] + pred_box3d[6, :]) / 2 pred_box3d_angle = self.pred_box3d_angle_list[index] pred_box3d_size = self.pred_box3d_size_list[index].copy() cls_type = self.type_list[index] # assert cls_type in KITTICategory.CLASSES, cls_type # size_class = KITTICategory.CLASSES.index(cls_type) assert cls_type in self.category_info.CLASSES, '%s not in category_info' % cls_type size_class = self.category_info.CLASSES.index(cls_type) # Compute one hot vector if self.one_hot: one_hot_vec = np.zeros((len(self.category_info.CLASSES))) one_hot_vec[size_class] = 1 if rotate_to_center: point_set[:, :3] = self.get_center_view_point( point_set[:, :3], pred_box3d_center, pred_box3d_angle) if not with_extra_feat: point_set = point_set[:, :3] # Resample if self.npoints > 0: # choice = np.random.choice(point_set.shape[0], self.npoints, replace=True) choice = np.random.choice(point_set.shape[0], self.npoints, point_set.shape[0] < self.npoints) else: choice = np.random.permutation(len(point_set.shape[0])) point_set = point_set[choice, :] # P = self.calib_list[index]['P2'].reshape(3, 4) if rotate_to_center: pred_box3d_center_rot, pred_box3d_angle_rot = self.get_center_view_box3d( pred_box3d_center, pred_box3d_angle, pred_box3d_center, pred_box3d_angle) else: pred_box3d_center_rot = pred_box3d_center pred_box3d_angle_rot = pred_box3d_angle pred_box3d = compute_box_3d(pred_box3d_center_rot, pred_box3d_size, pred_box3d_angle_rot) ref1, ref2, ref3, ref4 = self.generate_ref(pred_box3d) if self.from_rgb_detection: data_inputs = { 'point_cloud': np.transpose(point_set, (1, 0)).astype(np.float32), 'center_ref1': np.transpose(ref1, (1, 0)).astype(np.float32), 'center_ref2': np.transpose(ref2, (1, 0)).astype(np.float32), 'center_ref3': np.transpose(ref3, (1, 0)).astype(np.float32), 'center_ref4': np.transpose(ref4, (1, 0)).astype(np.float32), 'rgb_prob': np.array([self.prob_list[index]]).astype(np.float32), 'rot_angle': np.array([pred_box3d_angle]).astype(np.float32), 'ref_center': pred_box3d_center.astype(np.float32), } if not rotate_to_center: data_inputs.update({'rot_angle': torch.zeros(1)}) data_inputs.update({'ref_center': torch.zeros(3)}) if self.one_hot: data_inputs.update({'one_hot': torch.FloatTensor(one_hot_vec)}) return data_inputs # mask_label = self.label_list[index].copy() # mask_label = mask_label[choice] box3d = self.box3d_list[index].copy() heading_angle = self.heading_list[index] box3d_size = self.size_list[index].copy() box3d_center = (box3d[0, :] + box3d[6, :]) / 2 if rotate_to_center: box3d_center, heading_angle = self.get_center_view_box3d( box3d_center, heading_angle, pred_box3d_center, pred_box3d_angle) # box3d_rot = compute_box_3d(box3d_center, box3d_size, heading_angle) # self.check(point_set, ref1, box3d_rot, pred_box3d) 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 ref1[:, 0] *= -1 ref2[:, 0] *= -1 ref3[:, 0] *= -1 ref4[:, 0] *= -1 if self.random_shift: s1 = cfg.DATA.STRIDE[0] l, w, h = self.size_list[index] dist = np.sqrt(np.sum(l**2 + w**2)) shift = np.clip(np.random.randn() * dist * 0.1, -s1 * 2, 2 * s1) point_set[:, 2] += shift box3d_center[2] += shift # angle_class, angle_residual = angle2class(heading_angle, NUM_HEADING_BIN) # self.check(box, P, box3d_center, self.size_list[index], heading_angle) labels = self.generate_labels(box3d_center, box3d_size, heading_angle, ref2) data_inputs = { 'point_cloud': np.transpose(point_set, (1, 0)).astype(np.float32), 'cls_label': labels.astype(np.int64), 'box3d_center': box3d_center.astype(np.float32), 'box3d_heading': np.array([heading_angle], dtype=np.float32), 'box3d_size': box3d_size.astype(np.float32), 'size_class': np.array([size_class], dtype=np.int64), 'center_ref1': np.transpose(ref1, (1, 0)).astype(np.float32), 'center_ref2': np.transpose(ref2, (1, 0)).astype(np.float32), 'center_ref3': np.transpose(ref3, (1, 0)).astype(np.float32), 'center_ref4': np.transpose(ref4, (1, 0)).astype(np.float32), 'rot_angle': np.array([pred_box3d_angle]).astype(np.float32), 'ref_center': pred_box3d_center.astype(np.float32), } if not rotate_to_center: data_inputs.update({'rot_angle': torch.zeros(1)}) if self.one_hot: data_inputs.update({'one_hot': torch.FloatTensor(one_hot_vec)}) return data_inputs