def __getitem__(self, idx): """ Returns a dict with following keys: point_clouds: (N,3+C) center_label: (MAX_NUM_OBJ,3) for GT box center XYZ sem_cls_label: (MAX_NUM_OBJ,) semantic class index angle_class_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_HEADING_BIN-1 angle_residual_label: (MAX_NUM_OBJ,) size_classe_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_SIZE_CLUSTER size_residual_label: (MAX_NUM_OBJ,3) box_label_mask: (MAX_NUM_OBJ) as 0/1 with 1 indicating a unique box point_votes: (N,3) with votes XYZ point_votes_mask: (N,) with 0/1 with 1 indicating the point is in one of the object's OBB. scan_idx: int scan index in scan_names list pcl_color: unused """ scan_name = self.scan_names[idx] mesh_vertices = np.load( os.path.join(self.data_path, scan_name) + '_vert.npy') instance_labels = np.load( os.path.join(self.data_path, scan_name) + '_ins_label.npy') semantic_labels = np.load( os.path.join(self.data_path, scan_name) + '_sem_label.npy') instance_bboxes = np.load( os.path.join(self.data_path, scan_name) + '_bbox.npy') if not self.use_color: point_cloud = mesh_vertices[:, 0:3] # do not use color for now pcl_color = mesh_vertices[:, 3:6] else: point_cloud = mesh_vertices[:, 0:6] point_cloud[:, 3:] = (point_cloud[:, 3:] - MEAN_COLOR_RGB) / 256.0 if self.use_height: floor_height = np.percentile(point_cloud[:, 2], 0.99) height = point_cloud[:, 2] - floor_height point_cloud = np.concatenate( [point_cloud, np.expand_dims(height, 1)], 1) # ------------------------------- LABELS ------------------------------ target_bboxes = np.zeros((MAX_NUM_OBJ, 6)) target_bboxes_mask = np.zeros((MAX_NUM_OBJ)) angle_classes = np.zeros((MAX_NUM_OBJ, )) angle_residuals = np.zeros((MAX_NUM_OBJ, )) size_classes = np.zeros((MAX_NUM_OBJ, )) size_residuals = np.zeros((MAX_NUM_OBJ, 3)) point_cloud, choices = pc_util.random_sampling(point_cloud, self.num_points, return_choices=True) instance_labels = instance_labels[choices] semantic_labels = semantic_labels[choices] pcl_color = pcl_color[choices] target_bboxes_mask[0:instance_bboxes.shape[0]] = 1 target_bboxes[0:instance_bboxes.shape[0], :] = instance_bboxes[:, 0:6] # ------------------------------- DATA AUGMENTATION ------------------------------ if self.augment: if np.random.random() > 0.5: # Flipping along the YZ plane point_cloud[:, 0] = -1 * point_cloud[:, 0] target_bboxes[:, 0] = -1 * target_bboxes[:, 0] if np.random.random() > 0.5: # Flipping along the XZ plane point_cloud[:, 1] = -1 * point_cloud[:, 1] target_bboxes[:, 1] = -1 * target_bboxes[:, 1] # Rotation along up-axis/Z-axis rot_angle = (np.random.random() * np.pi / 18) - np.pi / 36 # -5 ~ +5 degree rot_mat = pc_util.rotz(rot_angle) point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3], np.transpose(rot_mat)) target_bboxes = rotate_aligned_boxes(target_bboxes, rot_mat) # compute votes *AFTER* augmentation # generate votes # Note: since there's no map between bbox instance labels and # pc instance_labels (it had been filtered # in the data preparation step) we'll compute the instance bbox # from the points sharing the same instance label. point_votes = np.zeros([self.num_points, 3]) point_votes_mask = np.zeros(self.num_points) for i_instance in np.unique(instance_labels): # find all points belong to that instance ind = np.where(instance_labels == i_instance)[0] # find the semantic label if semantic_labels[ind[0]] in DC.nyu40ids: x = point_cloud[ind, :3] center = 0.5 * (x.min(0) + x.max(0)) point_votes[ind, :] = center - x point_votes_mask[ind] = 1.0 point_votes = np.tile(point_votes, (1, 3)) # make 3 votes identical class_ind = [ np.where(DC.nyu40ids == x)[0][0] for x in instance_bboxes[:, -1] ] # NOTE: set size class as semantic class. Consider use size2class. size_classes[0:instance_bboxes.shape[0]] = class_ind size_residuals[0:instance_bboxes.shape[0], :] = \ target_bboxes[0:instance_bboxes.shape[0], 3:6] - DC.mean_size_arr[class_ind,:] # keep the same nums of points for each cloud mesh_vertices, _ = pc_util.random_sampling(mesh_vertices, 50000, return_choices=True) ret_dict = {} ret_dict['mesh_vertices'] = mesh_vertices.astype(np.float32) ret_dict['point_clouds'] = point_cloud.astype(np.float32) ret_dict['center_label'] = target_bboxes.astype(np.float32)[:, 0:3] ret_dict['heading_class_label'] = angle_classes.astype(np.int64) ret_dict['heading_residual_label'] = angle_residuals.astype(np.float32) ret_dict['size_class_label'] = size_classes.astype(np.int64) ret_dict['size_residual_label'] = size_residuals.astype(np.float32) target_bboxes_semcls = np.zeros((MAX_NUM_OBJ)) target_bboxes_semcls[0:instance_bboxes.shape[0]] = \ [DC.nyu40id2class[x] for x in instance_bboxes[:,-1][0:instance_bboxes.shape[0]]] ret_dict['sem_cls_label'] = target_bboxes_semcls.astype(np.int64) ret_dict['box_label_mask'] = target_bboxes_mask.astype(np.float32) ret_dict['vote_label'] = point_votes.astype(np.float32) ret_dict['vote_label_mask'] = point_votes_mask.astype(np.int64) ret_dict['scan_idx'] = np.array(idx).astype(np.int64) ret_dict['pcl_color'] = pcl_color return ret_dict
def __getitem__(self, idx): """ Returns a dict with following keys: point_clouds: (N,3+C) scan_idx: int scan index in scan_names list """ scan_name = self.scan_names[idx] mesh_vertices = np.load( os.path.join(self.data_path, scan_name) + '_pc.npz')['pc'] # Nx6 if not self.use_color: raw_point_cloud = mesh_vertices[:, 0:3] # do not use color for now else: raw_point_cloud = mesh_vertices[:, 0:6] raw_point_cloud[:, 3:] = (raw_point_cloud[:, 3:] - MEAN_COLOR_RGB) / 256.0 if self.use_height: floor_height = np.percentile(raw_point_cloud[:, 2], 0.99) height = raw_point_cloud[:, 2] - floor_height raw_point_cloud = np.concatenate( [raw_point_cloud, np.expand_dims(height, 1)], 1) point_cloud, choices = pc_util.random_sampling(raw_point_cloud, self.num_points, return_choices=True) #ema_point_cloud = pc_util.random_sampling(raw_point_cloud, self.num_points, return_choices=False) ema_point_cloud = point_cloud.copy() # 2021.2.28 # ------------------------------- DATA AUGMENTATION ------------------------------ flip_x_axis = 0 flip_y_axis = 0 flip_x_axis_ema = 0 # 2021.2.28 flip_y_axis_ema = 0 # 2021.2.28 rot_mat = np.identity(3) scale_ratio = np.ones((1, 3)) if self.augment: if np.random.random() > 0.5: # Flipping along the YZ plane flip_x_axis = 1 point_cloud[:, 0] = -1 * point_cloud[:, 0] if np.random.random() > 0.5: # 2021.2.28 # Flipping along the YZ plane flip_x_axis_ema = 1 ema_point_cloud[:, 0] = -1 * ema_point_cloud[:, 0] if np.random.random() > 0.5: # Flipping along the XZ plane flip_y_axis = 1 point_cloud[:, 1] = -1 * point_cloud[:, 1] if np.random.random() > 0.5: # 2021.2.28 # Flipping along the XZ plane flip_y_axis_ema = 1 ema_point_cloud[:, 1] = -1 * ema_point_cloud[:, 1] # Rotation along up-axis/Z-axis rot_angle = (np.random.random() * np.pi / 3) - np.pi / 6 # -30 ~ +30 degree rot_mat = pc_util.rotz(rot_angle) point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3], np.transpose(rot_mat)) # Augment point cloud scale: 0.85x-1.15x scale_ratio = np.random.random() * 0.3 + 0.85 scale_ratio = np.expand_dims(np.tile(scale_ratio, 3), 0) point_cloud[:, 0:3] *= scale_ratio if self.use_height: point_cloud[:, -1] *= scale_ratio[0, 0] ret_dict = {} ret_dict['point_clouds'] = point_cloud.astype(np.float32) ret_dict['scan_idx'] = np.array(idx).astype(np.int64) ret_dict['supervised_mask'] = np.array(0).astype(np.int64) ret_dict['ema_point_clouds'] = ema_point_cloud.astype(np.float32) ret_dict['flip_x_axis'] = np.array(flip_x_axis).astype(np.int64) ret_dict['flip_y_axis'] = np.array(flip_y_axis).astype(np.int64) ret_dict['rot_mat'] = rot_mat.astype(np.float32) ret_dict['scale'] = np.array(scale_ratio).astype(np.float32) ret_dict['flip_x_axis_ema'] = np.array(flip_x_axis_ema).astype( np.int64) # 2021.2.28 ret_dict['flip_y_axis_ema'] = np.array(flip_y_axis_ema).astype( np.int64) # 2021.2.28 return ret_dict
def __getitem__(self, idx): """ Returns a dict with following keys: point_clouds: (N,3+C) center_label: (MAX_NUM_OBJ,3) for GT box center XYZ sem_cls_label: (MAX_NUM_OBJ,) semantic class index angle_class_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_HEADING_BIN-1 angle_residual_label: (MAX_NUM_OBJ,) size_classe_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_SIZE_CLUSTER size_residual_label: (MAX_NUM_OBJ,3) box_label_mask: (MAX_NUM_OBJ) as 0/1 with 1 indicating a unique box point_votes: (N,3) with votes XYZ point_votes_mask: (N,) with 0/1 with 1 indicating the point is in one of the object's OBB. scan_idx: int scan index in scan_names list pcl_color: unused """ scan_name = self.scan_names[idx] mesh_vertices = np.load(os.path.join(self.data_path, scan_name)+'_vert.npy') meta_vertices = np.load(os.path.join(self.data_path, scan_name)+'_all_noangle_40cls.npy') ### Need to change the name here instance_labels = meta_vertices[:,-2] semantic_labels = meta_vertices[:,-1] if not self.use_color: point_cloud = mesh_vertices[:,0:3] # do not use color for now pcl_color = mesh_vertices[:,3:6] else: point_cloud = mesh_vertices[:,0:6] point_cloud[:,3:] = (point_cloud[:,3:]-MEAN_COLOR_RGB)/256.0 pcl_color = (point_cloud[:,3:]-MEAN_COLOR_RGB)/256.0 if self.use_height: floor_height = np.percentile(point_cloud[:,2],0.99) height = point_cloud[:,2] - floor_height point_cloud = np.concatenate([point_cloud, np.expand_dims(height, 1)],1) # ------------------------------- LABELS ------------------------------ target_bboxes = np.zeros((MAX_NUM_OBJ, 6)) target_bboxes_mask = np.zeros((MAX_NUM_OBJ)) angle_classes = np.zeros((MAX_NUM_OBJ,)) angle_label = np.zeros((MAX_NUM_OBJ,)) angle_residuals = np.zeros((MAX_NUM_OBJ,)) size_classes = np.zeros((MAX_NUM_OBJ,)) size_residuals = np.zeros((MAX_NUM_OBJ, 3)) ### For statistics surface_cue = np.zeros((MAX_NUM_OBJ)) line_cue = np.zeros((MAX_NUM_OBJ,)) before_sample = np.unique(instance_labels) while True: orig_point_cloud = np.copy(point_cloud) temp_point_cloud, choices = pc_util.random_sampling(orig_point_cloud, self.num_points, return_choices=True) after_sample = np.unique(instance_labels[choices]) if np.array_equal(before_sample, after_sample): point_cloud = temp_point_cloud break instance_labels = instance_labels[choices] semantic_labels = semantic_labels[choices] meta_vertices = meta_vertices[choices] pcl_color = pcl_color[choices] # ------------------------------- DATA AUGMENTATION ------------------------------ if self.augment: if np.random.random() > 0.5: # Flipping along the YZ plane point_cloud[:,0] = -1 * point_cloud[:,0] # target_bboxes[:,0] = -1 * target_bboxes[:,0] meta_vertices[:, 0] = -1 * meta_vertices[:, 0] meta_vertices[:, 6] = -1 * meta_vertices[:, 6] if np.random.random() > 0.5: # Flipping along the XZ plane point_cloud[:,1] = -1 * point_cloud[:,1] # target_bboxes[:,1] = -1 * target_bboxes[:,1] meta_vertices[:, 1] = -1 * meta_vertices[:, 1] meta_vertices[:, 6] = -1 * meta_vertices[:, 6] # Rotation along up-axis/Z-axis rot_angle = (np.random.random()*np.pi/18) - np.pi/36 # -5 ~ +5 degree rot_mat = pc_util.rotz(rot_angle).astype(np.float32) point_cloud[:,0:3] = np.dot(point_cloud[:,0:3], np.transpose(rot_mat)) meta_vertices[:, :6] = rotate_aligned_boxes(meta_vertices[:, :6], rot_mat) meta_vertices[:, 6] += rot_angle # ------------------------------- Plane and point ------------------------------ # compute votes *AFTER* augmentation # generate votes # Note: since there's no map between bbox instance labels and # pc instance_labels (it had been filtered # in the data preparation step) we'll compute the instance bbox # from the points sharing the same instance label. point_votes = np.zeros([self.num_points, 3]) point_votes_mask = np.zeros(self.num_points) point_boundary_mask_z = np.zeros(self.num_points) point_boundary_mask_xy = np.zeros(self.num_points) point_boundary_offset_z = np.zeros([self.num_points, 3]) point_boundary_offset_xy = np.zeros([self.num_points, 3]) point_boundary_sem_z = np.zeros([self.num_points, 3+2+1]) point_boundary_sem_xy = np.zeros([self.num_points, 3+1+1]) point_line_mask = np.zeros(self.num_points) point_line_offset = np.zeros([self.num_points, 3]) point_line_sem = np.zeros([self.num_points, 3+1]) point_sem_label = np.zeros(self.num_points) selected_instances = [] selected_centers = [] selected_centers_support = [] selected_centers_bsupport = [] obj_meta = [] counter = -1 for i_instance in np.unique(instance_labels): # find all points belong to that instance ind = np.where(instance_labels == i_instance)[0] if semantic_labels[ind[0]] in DC.nyu40ids: counter += 1 idx_instance = counter x = point_cloud[ind,:3] ### Meta information here meta = meta_vertices[ind[0]] obj_meta.append(meta) ### Get the centroid here center = meta[:3] point_votes[ind, :] = center - x point_votes_mask[ind] = 1.0 point_sem_label[ind] = DC.nyu40id2class_sem[meta[-1]] ### Corners corners, xmin, ymin, zmin, xmax, ymax, zmax = params2bbox(center, meta[3], meta[4], meta[5], meta[6]) ## Get lower four lines plane_lower_temp = np.array([0,0,1,-corners[6,-1]]) para_points = np.array([corners[1], corners[3], corners[5], corners[7]]) newd = np.sum(para_points * plane_lower_temp[:3], 1) if check_upright(para_points) and plane_lower_temp[0]+plane_lower_temp[1] < LOWER_THRESH: plane_lower = np.array([0,0,1,plane_lower_temp[-1]]) plane_upper = np.array([0,0,1,-np.mean(newd)]) else: import pdb;pdb.set_trace() print ("error with upright") if check_z(plane_upper, para_points) == False: import pdb;pdb.set_trace() ### Get the boundary points here alldist = np.abs(np.sum(x*plane_lower[:3], 1) + plane_lower[-1]) mind = np.min(alldist) sel = np.abs(alldist - mind) < DIST_THRESH ## Get lower four lines line_sel1, line_sel2, line_sel3, line_sel4 = get_linesel(x[sel], xmin, xmax, ymin, ymax) if np.sum(line_sel1) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel1]] = 1.0 linecenter = np.mean(x[sel][line_sel1], axis=0) linecenter[1] = (ymin+ymax)/2.0 point_line_offset[ind[sel][line_sel1]] = linecenter - x[sel][line_sel1] point_line_sem[ind[sel][line_sel1]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) if np.sum(line_sel2) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel2]] = 1.0 linecenter = np.mean(x[sel][line_sel2], axis=0) linecenter[1] = (ymin+ymax)/2.0 point_line_offset[ind[sel][line_sel2]] = linecenter - x[sel][line_sel2] point_line_sem[ind[sel][line_sel2]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) if np.sum(line_sel3) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel3]] = 1.0 linecenter = np.mean(x[sel][line_sel3], axis=0) linecenter[0] = (xmin+xmax)/2.0 point_line_offset[ind[sel][line_sel3]] = linecenter - x[sel][line_sel3] point_line_sem[ind[sel][line_sel3]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) if np.sum(line_sel4) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel4]] = 1.0 linecenter = np.mean(x[sel][line_sel4], axis=0) linecenter[0] = (xmin+xmax)/2.0 point_line_offset[ind[sel][line_sel4]] = linecenter - x[sel][line_sel4] point_line_sem[ind[sel][line_sel4]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) ### Set the surface labels here if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH: center = np.array([(xmin+xmax)/2.0, (ymin+ymax)/2.0, np.mean(x[sel][:,2])]) sel_global = ind[sel] point_boundary_mask_z[sel_global] = 1.0 point_boundary_sem_z[sel_global] = np.array([center[0], center[1], center[2], xmax - xmin, ymax - ymin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) point_boundary_offset_z[sel_global] = center - x[sel] ### Get the boundary points here alldist = np.abs(np.sum(x*plane_upper[:3], 1) + plane_upper[-1]) mind = np.min(alldist) sel = np.abs(alldist - mind) < DIST_THRESH ## Get upper four lines line_sel1, line_sel2, line_sel3, line_sel4 = get_linesel(x[sel], xmin, xmax, ymin, ymax) if np.sum(line_sel1) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel1]] = 1.0 linecenter = np.mean(x[sel][line_sel1], axis=0) linecenter[1] = (ymin+ymax)/2.0 point_line_offset[ind[sel][line_sel1]] = linecenter - x[sel][line_sel1] point_line_sem[ind[sel][line_sel1]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) if np.sum(line_sel2) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel2]] = 1.0 linecenter = np.mean(x[sel][line_sel2], axis=0) linecenter[1] = (ymin+ymax)/2.0 point_line_offset[ind[sel][line_sel2]] = linecenter - x[sel][line_sel2] point_line_sem[ind[sel][line_sel2]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) if np.sum(line_sel3) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel3]] = 1.0 linecenter = np.mean(x[sel][line_sel3], axis=0) linecenter[0] = (xmin+xmax)/2.0 point_line_offset[ind[sel][line_sel3]] = linecenter - x[sel][line_sel3] point_line_sem[ind[sel][line_sel3]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) if np.sum(line_sel4) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel4]] = 1.0 linecenter = np.mean(x[sel][line_sel4], axis=0) linecenter[0] = (xmin+xmax)/2.0 point_line_offset[ind[sel][line_sel4]] = linecenter - x[sel][line_sel4] point_line_sem[ind[sel][line_sel4]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH: center = np.array([(xmin+xmax)/2.0, (ymin+ymax)/2.0, np.mean(x[sel][:,2])]) sel_global = ind[sel] point_boundary_mask_z[sel_global] = 1.0 point_boundary_sem_z[sel_global] = np.array([center[0], center[1], center[2], xmax - xmin, ymax - ymin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) point_boundary_offset_z[sel_global] = center - x[sel] ## Get left two lines v1 = corners[3] - corners[2] v2 = corners[2] - corners[0] cp = np.cross(v1, v2) d = -np.dot(cp,corners[0]) a,b,c = cp plane_left_temp = np.array([a, b, c, d]) para_points = np.array([corners[4], corners[5], corners[6], corners[7]]) ### Normalize xy here plane_left_temp /= np.linalg.norm(plane_left_temp[:3]) newd = np.sum(para_points * plane_left_temp[:3], 1) if plane_left_temp[2] < LOWER_THRESH: plane_left = plane_left_temp#np.array([cls,res,tempsign,plane_left_temp[-1]]) plane_right = np.array([plane_left_temp[0], plane_left_temp[1], plane_left_temp[2], -np.mean(newd)]) else: import pdb;pdb.set_trace() print ("error with upright") ### Get the boundary points here alldist = np.abs(np.sum(x*plane_left[:3], 1) + plane_left[-1]) mind = np.min(alldist) sel = np.abs(alldist - mind) < DIST_THRESH ## Get upper four lines line_sel1, line_sel2 = get_linesel2(x[sel], ymin, ymax, zmin, zmax, axis=1) if np.sum(line_sel1) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel1]] = 1.0 linecenter = np.mean(x[sel][line_sel1], axis=0) linecenter[2] = (zmin+zmax)/2.0 point_line_offset[ind[sel][line_sel1]] = linecenter - x[sel][line_sel1] point_line_sem[ind[sel][line_sel1]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) if np.sum(line_sel2) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel2]] = 1.0 linecenter = np.mean(x[sel][line_sel2], axis=0) linecenter[2] = (zmin+zmax)/2.0 point_line_offset[ind[sel][line_sel2]] = linecenter - x[sel][line_sel2] point_line_sem[ind[sel][line_sel2]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH: center = np.array([np.mean(x[sel][:,0]), np.mean(x[sel][:,1]), (zmin+zmax)/2.0]) sel_global = ind[sel] point_boundary_mask_xy[sel_global] = 1.0 point_boundary_sem_xy[sel_global] = np.array([center[0], center[1], center[2], zmax - zmin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) point_boundary_offset_xy[sel_global] = center - x[sel] ### Get the boundary points here alldist = np.abs(np.sum(x*plane_right[:3], 1) + plane_right[-1]) mind = np.min(alldist) sel = np.abs(alldist - mind) < DIST_THRESH line_sel1, line_sel2 = get_linesel2(x[sel], ymin, ymax, zmin, zmax, axis=1) if np.sum(line_sel1) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel1]] = 1.0 linecenter = np.mean(x[sel][line_sel1], axis=0) linecenter[2] = (zmin+zmax)/2.0 point_line_offset[ind[sel][line_sel1]] = linecenter - x[sel][line_sel1] point_line_sem[ind[sel][line_sel1]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) if np.sum(line_sel2) > NUM_POINT_LINE: point_line_mask[ind[sel][line_sel2]] = 1.0 linecenter = np.mean(x[sel][line_sel2], axis=0) linecenter[2] = (zmin+zmax)/2.0 point_line_offset[ind[sel][line_sel2]] = linecenter - x[sel][line_sel2] point_line_sem[ind[sel][line_sel2]] = np.array([linecenter[0], linecenter[1], linecenter[2], np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH: center = np.array([np.mean(x[sel][:,0]), np.mean(x[sel][:,1]), (zmin+zmax)/2.0]) sel_global = ind[sel] point_boundary_mask_xy[sel_global] = 1.0 point_boundary_sem_xy[sel_global] = np.array([center[0], center[1], center[2], zmax - zmin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) point_boundary_offset_xy[sel_global] = center - x[sel] ### Get the boundary points here v1 = corners[0] - corners[4] v2 = corners[4] - corners[5] cp = np.cross(v1, v2) d = -np.dot(cp,corners[5]) a,b,c = cp plane_front_temp = np.array([a, b, c, d]) para_points = np.array([corners[2], corners[3], corners[6], corners[7]]) plane_front_temp /= np.linalg.norm(plane_front_temp[:3]) newd = np.sum(para_points * plane_front_temp[:3], 1) if plane_front_temp[2] < LOWER_THRESH: plane_front = plane_front_temp plane_back = np.array([plane_front_temp[0], plane_front_temp[1], plane_front_temp[2], -np.mean(newd)]) else: import pdb;pdb.set_trace() print ("error with upright") ### Get the boundary points here alldist = np.abs(np.sum(x*plane_front[:3], 1) + plane_front[-1]) mind = np.min(alldist) sel = np.abs(alldist - mind) < DIST_THRESH if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH: center = np.array([np.mean(x[sel][:,0]), np.mean(x[sel][:,1]), (zmin+zmax)/2.0]) sel_global = ind[sel] point_boundary_mask_xy[sel_global] = 1.0 point_boundary_sem_xy[sel_global] = np.array([center[0], center[1], center[2], zmax - zmin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) point_boundary_offset_xy[sel_global] = center - x[sel] ### Get the boundary points here alldist = np.abs(np.sum(x*plane_back[:3], 1) + plane_back[-1]) mind = np.min(alldist) sel = np.abs(alldist - mind) < DIST_THRESH if np.sum(sel) > NUM_POINT and np.var(alldist[sel]) < VAR_THRESH: center = np.array([np.mean(x[sel][:,0]), np.mean(x[sel][:,1]), (zmin+zmax)/2.0]) sel_global = ind[sel] point_boundary_mask_xy[sel_global] = 1.0 point_boundary_sem_xy[sel_global] = np.array([center[0], center[1], center[2], zmax - zmin, np.where(DC.nyu40ids == meta_vertices[ind[0],-1])[0][0]]) point_boundary_offset_xy[sel_global] = center - x[sel] num_instance = len(obj_meta) obj_meta = np.array(obj_meta) obj_meta = obj_meta.reshape(-1, 9) target_bboxes_mask[0:num_instance] = 1 target_bboxes[0:num_instance,:6] = obj_meta[:,0:6] class_ind = [np.where(DC.nyu40ids == x)[0][0] for x in obj_meta[:,-1]] # NOTE: set size class as semantic class. Consider use size2class. size_classes[0:num_instance] = class_ind size_residuals[0:num_instance, :] = \ target_bboxes[0:num_instance, 3:6] - DC.mean_size_arr[class_ind,:] point_votes = np.tile(point_votes, (1, 3)) # make 3 votes identical point_sem_label = np.tile(np.expand_dims(point_sem_label, -1), (1, 3)) # make 3 votes identical ret_dict = {} ret_dict['point_clouds'] = point_cloud.astype(np.float32) ret_dict['center_label'] = target_bboxes.astype(np.float32)[:,0:3] ret_dict['size_label'] = target_bboxes.astype(np.float32)[:,3:6] ret_dict['heading_label'] = angle_label.astype(np.float32) ret_dict['heading_class_label'] = angle_classes.astype(np.int64) ret_dict['heading_residual_label'] = angle_residuals.astype(np.float32) ret_dict['size_class_label'] = size_classes.astype(np.int64) ret_dict['size_residual_label'] = size_residuals.astype(np.float32) if self.use_height: ret_dict['floor_height'] = floor_height target_bboxes_semcls = np.zeros((MAX_NUM_OBJ)) target_bboxes_semcls[0:num_instance] = \ [DC.nyu40id2class[x] for x in obj_meta[:,-1][0:obj_meta.shape[0]]] ret_dict['sem_cls_label'] = target_bboxes_semcls.astype(np.int64) ret_dict['point_sem_cls_label'] = point_sem_label.astype(np.int64) ret_dict['box_label_mask'] = target_bboxes_mask.astype(np.float32) ret_dict['point_boundary_mask_z'] = point_boundary_mask_z.astype(np.float32) ret_dict['point_boundary_mask_xy'] = point_boundary_mask_xy.astype(np.float32) ret_dict['point_boundary_offset_z'] = point_boundary_offset_z.astype(np.float32) ret_dict['point_boundary_offset_xy'] = point_boundary_offset_xy.astype(np.float32) ret_dict['point_boundary_sem_z'] = point_boundary_sem_z.astype(np.float32) ret_dict['point_boundary_sem_xy'] = point_boundary_sem_xy.astype(np.float32) ret_dict['point_line_mask'] = point_line_mask.astype(np.float32) ret_dict['point_line_offset'] = point_line_offset.astype(np.float32) ret_dict['point_line_sem'] = point_line_sem.astype(np.float32) ret_dict['vote_label'] = point_votes.astype(np.float32) ret_dict['vote_label_mask'] = point_votes_mask.astype(np.int64) ret_dict['scan_idx'] = np.array(idx).astype(np.int64) ret_dict['pcl_color'] = pcl_color ret_dict['num_instance'] = num_instance return ret_dict
def __getitem__(self, idx): """ Returns a dict with following keys: point_clouds: (N,3+C) center_label: (MAX_NUM_OBJ,3) for GT box center XYZ sem_cls_label: (MAX_NUM_OBJ,) semantic class index angle_class_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_HEADING_BIN-1 angle_residual_label: (MAX_NUM_OBJ,) size_classe_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_SIZE_CLUSTER size_residual_label: (MAX_NUM_OBJ,3) box_label_mask: (MAX_NUM_OBJ) as 0/1 with 1 indicating a unique box point_votes: (N,3) with votes XYZ point_votes_mask: (N,) with 0/1 with 1 indicating the point is in one of the object's OBB. scan_idx: int scan index in scan_names list """ scan_name = self.scan_names[idx] mesh_vertices = np.load( os.path.join(self.data_path, scan_name) + '_vert.npy') instance_labels = np.load( os.path.join(self.data_path, scan_name) + '_ins_label.npy') semantic_labels = np.load( os.path.join(self.data_path, scan_name) + '_sem_label.npy') instance_bboxes = np.load( os.path.join(self.data_path, scan_name) + '_bbox.npy') if not self.use_color: raw_point_cloud = mesh_vertices[:, 0:3] # do not use color for now else: raw_point_cloud = mesh_vertices[:, 0:6] raw_point_cloud[:, 3:] = (raw_point_cloud[:, 3:] - MEAN_COLOR_RGB) / 256.0 if self.use_height: floor_height = np.percentile(raw_point_cloud[:, 2], 0.99) height = raw_point_cloud[:, 2] - floor_height raw_point_cloud = np.concatenate( [raw_point_cloud, np.expand_dims(height, 1)], 1) # ------------------------------- LABELS ------------------------------ target_bboxes = np.zeros((MAX_NUM_OBJ, 6)) target_bboxes_mask = np.zeros((MAX_NUM_OBJ)) angle_classes = np.zeros((MAX_NUM_OBJ, )) angle_residuals = np.zeros((MAX_NUM_OBJ, )) size_classes = np.zeros((MAX_NUM_OBJ, )) size_residuals = np.zeros((MAX_NUM_OBJ, 3)) target_bboxes_semcls = np.zeros((MAX_NUM_OBJ)) point_cloud, choices = pc_util.random_sampling(raw_point_cloud, self.num_points, return_choices=True) #ema_point_cloud = pc_util.random_sampling(raw_point_cloud, self.num_points, return_choices=False) ema_point_cloud = point_cloud.copy() # 2021.2.28 instance_labels = instance_labels[choices] semantic_labels = semantic_labels[choices] target_bboxes_mask[0:instance_bboxes.shape[0]] = 1 target_bboxes[0:instance_bboxes.shape[0], :] = instance_bboxes[:, 0:6] # ------------------------------- DATA AUGMENTATION ------------------------------ flip_x_axis = 0 flip_y_axis = 0 flip_x_axis_ema = 0 flip_y_axis_ema = 0 rot_mat = np.identity(3) scale_ratio = np.ones((1, 3)) if self.augment: if np.random.random() > 0.5: # Flipping along the YZ plane flip_x_axis = 1 point_cloud[:, 0] = -1 * point_cloud[:, 0] target_bboxes[:, 0] = -1 * target_bboxes[:, 0] if np.random.random() > 0.5: # 2021.2.28 # Flipping along the YZ plane for ema flip_x_axis_ema = 1 ema_point_cloud[:, 0] = -1 * ema_point_cloud[:, 0] if np.random.random() > 0.5: # Flipping along the XZ plane flip_y_axis = 1 point_cloud[:, 1] = -1 * point_cloud[:, 1] target_bboxes[:, 1] = -1 * target_bboxes[:, 1] if np.random.random() > 0.5: # 2021.2.28 # Flipping along the XZ plane for ema flip_y_axis_ema = 1 ema_point_cloud[:, 1] = -1 * ema_point_cloud[:, 1] # Rotation along up-axis/Z-axis rot_angle = (np.random.random() * np.pi / 18) - np.pi / 36 # -5 ~ +5 degree rot_mat = pc_util.rotz(rot_angle) point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3], np.transpose(rot_mat)) target_bboxes = rotate_aligned_boxes(target_bboxes, rot_mat) # Augment point cloud scale: 0.85x-1.15x scale_ratio = np.random.random() * 0.3 + 0.85 scale_ratio = np.expand_dims(np.tile(scale_ratio, 3), 0) point_cloud[:, 0:3] *= scale_ratio target_bboxes[:, 0:3] *= scale_ratio target_bboxes[:, 3:6] *= scale_ratio if self.use_height: point_cloud[:, -1] *= scale_ratio[0, 0] # compute votes *AFTER* augmentation # generate votes # Note: since there's no map between bbox instance labels and # pc instance_labels (it had been filtered # in the data preparation step) we'll compute the instance bbox # from the points sharing the same instance label. point_votes = np.zeros([self.num_points, 3]) point_votes_mask = np.zeros(self.num_points) for i_instance in np.unique(instance_labels): # find all points belong to that instance ind = np.where(instance_labels == i_instance)[0] # find the semantic label if semantic_labels[ind[0]] in DC.nyu40ids: x = point_cloud[ind, :3] center = 0.5 * (x.min(0) + x.max(0)) point_votes[ind, :] = center - x point_votes_mask[ind] = 1.0 point_votes = np.tile(point_votes, (1, 3)) # make 3 votes identical class_ind = [ np.where(DC.nyu40ids == x)[0][0] for x in instance_bboxes[:, -1] ] # NOTE: set size class as semantic class. Consider use size2class. size_classes[0:instance_bboxes.shape[0]] = class_ind size_residuals[0:instance_bboxes.shape[0], :] = \ target_bboxes[0:instance_bboxes.shape[0], 3:6] - DC.mean_size_arr[class_ind, :] target_bboxes_semcls[0:instance_bboxes.shape[0]] = class_ind ret_dict = {} ret_dict['point_clouds'] = point_cloud.astype(np.float32) ret_dict['center_label'] = target_bboxes.astype(np.float32)[:, 0:3] ret_dict['heading_class_label'] = angle_classes.astype(np.int64) ret_dict['heading_residual_label'] = angle_residuals.astype(np.float32) ret_dict['size_class_label'] = size_classes.astype(np.int64) ret_dict['size_residual_label'] = size_residuals.astype(np.float32) ret_dict['sem_cls_label'] = target_bboxes_semcls.astype(np.int64) ret_dict['box_label_mask'] = target_bboxes_mask.astype(np.float32) ret_dict['vote_label'] = point_votes.astype(np.float32) ret_dict['vote_label_mask'] = point_votes_mask.astype(np.int64) ret_dict['scan_idx'] = np.array(idx).astype(np.int64) ret_dict['supervised_mask'] = np.array(1).astype(np.int64) scene_label = np.zeros(DC.num_class) unique_class_ind = list(set(class_ind)) for ind in unique_class_ind: scene_label[int(ind)] = 1 ret_dict['scene_label'] = scene_label.astype(np.float32) ret_dict['ema_point_clouds'] = ema_point_cloud.astype(np.float32) ret_dict['flip_x_axis'] = np.array(flip_x_axis).astype(np.int64) ret_dict['flip_y_axis'] = np.array(flip_y_axis).astype(np.int64) ret_dict['rot_mat'] = rot_mat.astype(np.float32) ret_dict['scale'] = np.array(scale_ratio).astype(np.float32) ret_dict['flip_x_axis_ema'] = np.array(flip_x_axis_ema).astype( np.int64) #2021.2.28 ret_dict['flip_y_axis_ema'] = np.array(flip_y_axis_ema).astype( np.int64) #2021.2.28 return ret_dict
def __getitem__(self, index: int): id_scan = self.scene_list[index] assert (id_scan not in self.error_scan) id_scan_path = os.path.join(self.data_path, id_scan) point_cloud = np.load( os.path.join(id_scan_path, '{}.npy'.format('point_cloud'))) ins_vert = np.load( os.path.join(id_scan_path, '{}.npy'.format('ins_vert'))).squeeze(1) ins_bbox = np.load(os.path.join(id_scan_path, '{}.npy'.format('bbox'))) points = point_cloud # (N, 3) center = ins_bbox[:, 0:3] # (B, 10) bbox_length = ins_bbox[:, 3:6] # (B, 3) sem_cls = ins_bbox[:, 6:7] # (B, 1) symmetry = ins_bbox[:, 7:8] # (B, 1) K = center.shape[0] # LABELS if points.shape[0] > self.num_points: points, choices = pc_util.random_sampling(points, self.num_points, return_choices=True) ins_vert = ins_vert[choices] elif points.shape[0] < self.num_points: print('false data') target_bboxes = np.zeros((MAX_NUM_OBJ, 6), dtype=np.float32) target_bboxes_mask = np.zeros((MAX_NUM_OBJ), dtype=np.int64) # target_center = np.zeros((MAX_NUM_OBJ, 3), dtype=np.float32) # target_rot_q = np.zeros((MAX_NUM_OBJ, 4), dtype=np.float32) # target_rot_6d = np.zeros((MAX_NUM_OBJ, 6), dtype=np.float32) # target_scale = np.zeros((MAX_NUM_OBJ, 3), dtype=np.float32) target_sem_cls = np.zeros((MAX_NUM_OBJ, ), dtype=np.int64) target_sym = np.zeros((MAX_NUM_OBJ, ), dtype=np.int64) target_size_classes = np.zeros((MAX_NUM_OBJ, )) target_size_residuals = np.zeros((MAX_NUM_OBJ, 3)) # target_center[:K] = center[:,0:3] # target_rot_q[:K] = alignments[:,3:7] # for k in range(K): # target_rot_6d[k] = from_q_to_6d(alignments[k,3:7]) # target_scale[:K] = alignments[:,7:10] target_sem_cls[:K] = sem_cls.squeeze(1) target_sym[:K] = symmetry.squeeze(1) target_bboxes[:K, 0:3] = center target_bboxes[:K, 3:6] = bbox_length target_bboxes_mask[:K] = 1 # ------------------------------- DATA AUGMENTATION ------------------------------ if self.augment: if np.random.random() > 0.5: # Flipping along the YZ plane points[:, 0] = -1 * points[:, 0] target_bboxes[:, 0] = -1 * target_bboxes[:, 0] if np.random.random() > 0.5: # Flipping along the XZ plane points[:, 1] = -1 * points[:, 1] target_bboxes[:, 1] = -1 * target_bboxes[:, 1] # Rotation along up-axis/Z-axis rot_angle = (np.random.random() * np.pi / 18) - np.pi / 36 # -5 ~ +5 degree rot_mat = pc_util.rotz(rot_angle) points[:, 0:3] = np.dot(points[:, 0:3], np.transpose(rot_mat)) target_bboxes = rotate_aligned_boxes(target_bboxes, rot_mat) target_center = target_bboxes[:, 0:3] # ====== GENERATE VOTES ====== # compute votes *AFTER* augmentation # NOTE: i_ins: (1,B) not (0,B-1) point_votes = np.zeros([self.num_points, 3]) point_votes_mask = np.zeros(self.num_points) for i_ins in np.unique(ins_vert): i_ins -= 1 if target_sem_cls[i_ins] in NOT_CARED_IDS or i_ins < 0: continue ind = np.where(ins_vert == i_ins + 1)[0] x = points[ind, :3] point_votes[ind, :] = x - target_center[i_ins] point_votes_mask[ind] = 1.0 point_votes = np.tile(point_votes, (1, 3)) target_size_classes[:K] = target_sem_cls[:K] target_size_residuals[:K, :3] =\ target_bboxes[:K, 3:6] - DC.mean_size_arr[target_sem_cls[:K], :] # ====== LABELS ====== label = {} label['point_clouds'] = points.astype(np.float32) label['center_label'] = target_center.astype(np.float32) label['heading_class_label'] = np.zeros( (MAX_NUM_OBJ, )).astype(np.int64) label['heading_residual_label'] = np.zeros( (MAX_NUM_OBJ, )).astype(np.float32) label['size_class_label'] = target_size_classes.astype(np.int64) label['size_residual_label'] = target_size_residuals.astype(np.float32) label['sem_cls_label'] = target_sem_cls.astype(np.int64) label['box_label_mask'] = target_bboxes_mask.astype(np.float32) label['vote_label'] = point_votes.astype(np.float32) label['vote_label_mask'] = point_votes_mask.astype(np.int64) label['scan_idx'] = np.array(index).astype(np.int64) return label
def __getitem__(self, idx): """ Returns a dict with following keys: point_clouds: (N,3+C) scan_idx: int scan index in scan_names list """ scan_name = self.scan_names[idx] mesh_vertices = np.load( os.path.join(self.data_path, scan_name) + '_pc.npz')['pc'] # Nx6 if not self.use_color: raw_point_cloud = mesh_vertices[:, 0:3] # do not use color for now else: raw_point_cloud = mesh_vertices[:, 0:6] raw_point_cloud[:, 3:] = (raw_point_cloud[:, 3:] - MEAN_COLOR_RGB) / 256.0 if self.use_height: floor_height = np.percentile(raw_point_cloud[:, 2], 0.99) height = raw_point_cloud[:, 2] - floor_height raw_point_cloud = np.concatenate( [raw_point_cloud, np.expand_dims(height, 1)], 1) ret_dict = {} ema_point_cloud = pc_util.random_sampling(raw_point_cloud, self.num_points, return_choices=False) ret_dict['ema_point_clouds'] = ema_point_cloud.astype(np.float32) bboxes = np.load( os.path.join(self.data_path, scan_name) + '_bbox.npy') # K,8 target_bboxes = np.zeros((MAX_NUM_OBJ, 6)) target_bboxes_mask = np.zeros((MAX_NUM_OBJ)) angle_classes = np.zeros((MAX_NUM_OBJ, )) angle_residuals = np.zeros((MAX_NUM_OBJ, )) size_classes = np.zeros((MAX_NUM_OBJ, )) size_residuals = np.zeros((MAX_NUM_OBJ, 3)) target_bboxes_semcls = np.zeros((MAX_NUM_OBJ)) target_bboxes_mask[0:bboxes.shape[0]] = 1 target_bboxes[0:bboxes.shape[0], :] = bboxes[:, 0:6] for i in range(bboxes.shape[0]): bbox = bboxes[i] semantic_class = bbox[7] angle_class, angle_residual = DC.angle2class(bbox[6]) # NOTE: The mean size stored in size2class is of full length of box edges, # while in sunrgbd_data.py data dumping we dumped *half* length l,w,h.. so have to time it by 2 here box3d_size = bbox[3:6] * 2 size_class, size_residual = DC.size2class( box3d_size, DC.class2type[semantic_class]) angle_classes[i] = angle_class angle_residuals[i] = angle_residual size_classes[i] = size_class size_residuals[i] = size_residual target_bboxes_semcls[i] = semantic_class if self.load_labels: ret_dict['center_label'] = target_bboxes.astype(np.float32)[:, 0:3] ret_dict['heading_class_label'] = angle_classes.astype(np.int64) ret_dict['heading_residual_label'] = angle_residuals.astype( np.float32) ret_dict['size_class_label'] = size_classes.astype(np.int64) ret_dict['size_residual_label'] = size_residuals.astype(np.float32) ret_dict['sem_cls_label'] = target_bboxes_semcls.astype(np.int64) ret_dict['box_label_mask'] = target_bboxes_mask.astype(np.float32) point_cloud, choices = pc_util.random_sampling(raw_point_cloud, self.num_points, return_choices=True) flip_x_axis = 0 flip_y_axis = 0 rot_angle = 0 rot_mat = np.identity(3) scale_ratio = np.ones((1, 3)) if self.augment: if np.random.random() > 0.5: # Flipping along the YZ plane flip_x_axis = 1 point_cloud[:, 0] = -1 * point_cloud[:, 0] # Rotation along up-axis/Z-axis rot_angle = (np.random.random() * np.pi / 3) - np.pi / 6 # -30 ~ +30 degree rot_mat = pc_util.rotz(rot_angle) point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3], np.transpose(rot_mat)) # Augment point cloud scale: 0.85x-1.15x scale_ratio = np.random.random() * 0.3 + 0.85 scale_ratio = np.expand_dims(np.tile(scale_ratio, 3), 0) point_cloud[:, 0:3] *= scale_ratio if self.use_height: point_cloud[:, -1] *= scale_ratio[0, 0] ret_dict['point_clouds'] = point_cloud.astype(np.float32) ret_dict['flip_x_axis'] = np.array(flip_x_axis).astype(np.int64) ret_dict['flip_y_axis'] = np.array(flip_y_axis).astype(np.int64) ret_dict['rot_mat'] = rot_mat.astype(np.float32) ret_dict['rot_angle'] = np.array(rot_angle).astype(np.float32) ret_dict['scale'] = np.array(scale_ratio).astype(np.float32) ret_dict['scan_idx'] = np.array(idx).astype(np.int64) ret_dict['supervised_mask'] = np.array(0).astype(np.int64) return ret_dict
def __getitem__(self, idx): """ Returns a dict with following keys: point_clouds: (N,3+C) center_label: (MAX_NUM_OBJ,3) for GT box center XYZ sem_cls_label: (MAX_NUM_OBJ,) semantic class index angle_class_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_HEADING_BIN-1 angle_residual_label: (MAX_NUM_OBJ,) size_classe_label: (MAX_NUM_OBJ,) with int values in 0,...,NUM_SIZE_CLUSTER size_residual_label: (MAX_NUM_OBJ,3) box_label_mask: (MAX_NUM_OBJ) as 0/1 with 1 indicating a unique box point_votes: (N,3) with votes XYZ point_votes_mask: (N,) with 0/1 with 1 indicating the point is in one of the object's OBB. scan_idx: int scan index in scan_names list pcl_color: unused """ file = self.files[idx] scan_name = file.split('.')[0] h5file = h5py.File(os.path.join(self.data_path, file), 'r') mesh_vertices = np.array(h5file['point_cloud'], dtype=np.float32) instance_labels = np.array(h5file['instance'], dtype=np.int32) semantic_labels = np.array(h5file['semantic'], dtype=np.int32) instance_bboxes = np.array(h5file['bboxes'], dtype=np.float32) instance_bboxes = instance_bboxes[:, :8] h5file.close() # convert PC to z is up. mesh_vertices[:, [0, 1, 2]] = mesh_vertices[:, [0, 2, 1]] # convert annotations to z is up. instance_bboxes[:, [0, 1, 2]] = instance_bboxes[:, [0, 2, 1]] instance_bboxes[:, [3, 4, 5]] = instance_bboxes[:, [3, 5, 4]] if not self.use_color: point_cloud = mesh_vertices[:, 0:3] # do not use color for now else: point_cloud = mesh_vertices[:, 0:6] point_cloud[:, 3:] = point_cloud[:, 3:] - (MEAN_COLOR_RGB) / 256.0 if self.use_height: floor_height = np.percentile(point_cloud[:, 2], 0.99) height = point_cloud[:, 2] - floor_height point_cloud = np.concatenate( [point_cloud, np.expand_dims(height, 1)], 1) # ------------------------------- LABELS ------------------------------ target_bboxes = np.zeros((MAX_NUM_OBJ_FT, 6)) target_bboxes_mask = np.zeros((MAX_NUM_OBJ_FT)) angle_classes = np.zeros((MAX_NUM_OBJ_FT, )) angle_residuals = np.zeros((MAX_NUM_OBJ_FT, )) size_classes = np.zeros((MAX_NUM_OBJ_FT, )) size_residuals = np.zeros((MAX_NUM_OBJ_FT, 3)) target_bboxes_mask[0:instance_bboxes.shape[0]] = 1 target_bboxes[0:instance_bboxes.shape[0], :] = instance_bboxes[:, 0:6] # ------------------------------- DATA AUGMENTATION ------------------------------ if self.augment: if np.random.random() > 0.5: # Flipping along the YZ plane point_cloud[:, 0] = -1 * point_cloud[:, 0] target_bboxes[:, 0] = -1 * target_bboxes[:, 0] if np.random.random() > 0.5: # Flipping along the XZ plane point_cloud[:, 1] = -1 * point_cloud[:, 1] target_bboxes[:, 1] = -1 * target_bboxes[:, 1] # Rotation along up-axis/Z-axis rot_angle = (np.random.random() * np.pi / 18) - np.pi / 36 # -5 ~ +5 degree rot_mat = pc_util.rotz(rot_angle) point_cloud[:, 0:3] = np.dot(point_cloud[:, 0:3], np.transpose(rot_mat)) target_bboxes = rotate_aligned_boxes(target_bboxes, rot_mat) # compute votes *AFTER* augmentation # generate votes # Note: since there's no map between bbox instance labels and # pc instance_labels (it had been filtered # in the data preparation step) we'll compute the instance bbox # from the points sharing the same instance label. num_points = mesh_vertices.shape[0] point_votes = np.zeros([num_points, 3]) point_votes_mask = np.zeros(num_points) #DEBUG: votes_rgb = np.zeros([num_points, 3]) for i_instance in np.unique(instance_labels): # ignore points not associated with a box #if i_instance not in instance_bboxes_instance_labels: continue # find all points belong to that instance ind = np.where(instance_labels == i_instance)[0] # find the semantic label #TODO: change classe labels if not (semantic_labels[ind[0]] == -1): x = point_cloud[ind, :3] center = 0.5 * (x.min(0) + x.max(0)) point_votes[ind, :] = center - x point_votes_mask[ind] = 1.0 #DEBUG debug_color = np.random.uniform(0, 1, 3) votes_rgb[ind, :] = debug_color point_votes = np.tile(point_votes, (1, 3)) # make 3 votes identical # NOTE: set size class as semantic class. Consider use size2class. size_classes[0:instance_bboxes.shape[0]] = instance_bboxes[:, -1] instance_bboxes_sids = instance_bboxes[:, -1] instance_bboxes_sids = instance_bboxes_sids.astype(np.int) size_residuals[0:instance_bboxes.shape[0], :] = \ target_bboxes[0:instance_bboxes.shape[0], 3:6] - DC.mean_size_arr[instance_bboxes_sids,:] #TODO: update angle_classes + residuals angle_residuals[0:instance_bboxes.shape[0]] = instance_bboxes[:, 6] #DEBUG #DEBUG #DEBUG #DEBUG #DEBUG # DEBUG Point cloud # DEBUG: viz PC with colors associated to votes # -- -- debug_pcd = o3d.geometry.PointCloud() # -- -- debug_point_cloud = point_cloud[:,:3].astype(np.float32) # -- -- debug_pcd.points = o3d.utility.Vector3dVector(debug_point_cloud) # -- -- debug_pcd.colors = o3d.utility.Vector3dVector(votes_rgb) # -- -- o3d.io.write_point_cloud('debug_dataloader_pc.ply', # -- -- debug_pcd) # -- -- min_bound = debug_pcd.get_min_bound() # -- -- max_bound = debug_pcd.get_max_bound() # -- -- print(min_bound) # -- -- print(max_bound) # -- -- # Dump GT bounding boxes # -- -- gt_center = target_bboxes.astype(np.float32)[:,0:3] # (B,MAX_NUM_OBJ_FT,3) # -- -- gt_mask = target_bboxes_mask.astype(np.float32) # B,K2 # -- -- gt_heading_class = angle_classes.astype(np.int64) # B,K2 # -- -- gt_heading_residual = angle_residuals.astype(np.float32) # B,K2 # -- -- gt_size_class = size_classes.astype(np.int64) # B,K2 # -- -- gt_size_residual = size_residuals.astype(np.float32) # B,K2,3 # -- -- obbs = [] # -- -- for j in range(gt_center.shape[0]): # -- -- if gt_mask[j] == 0: continue # -- -- obb = DC.param2obb(gt_center[j,0:3], # -- -- gt_heading_class[j], # -- -- gt_heading_residual[j], # -- -- gt_size_class[j], # -- -- gt_size_residual[j]) # -- -- obbs.append(obb) # -- -- print('stored bboxes: ', len(obbs)) # -- -- if len(obbs)>0: # -- -- obbs = np.vstack(tuple(obbs)) # (num_gt_objects, 7) # -- -- pc_util.write_oriented_bbox(obbs, 'debug_dataloader_boxes.ply') # -- -- stop #DEBUG #DEBUG #DEBUG #DEBUG #DEBUG #DEBUG #DEBUG ret_dict = {} ret_dict['point_clouds'] = point_cloud.astype(np.float32) ret_dict['center_label'] = target_bboxes.astype(np.float32)[:, 0:3] ret_dict['heading_class_label'] = angle_classes.astype(np.int64) ret_dict['heading_residual_label'] = angle_residuals.astype(np.float32) ret_dict['size_class_label'] = size_classes.astype(np.int64) ret_dict['size_residual_label'] = size_residuals.astype(np.float32) target_bboxes_semcls = np.zeros((MAX_NUM_OBJ_FT)) target_bboxes_semcls[0:instance_bboxes.shape[0]] = instance_bboxes[:, -1] ret_dict['sem_cls_label'] = target_bboxes_semcls.astype(np.int64) ret_dict['box_label_mask'] = target_bboxes_mask.astype(np.float32) ret_dict['vote_label'] = point_votes.astype(np.float32) ret_dict['vote_label_mask'] = point_votes_mask.astype(np.int64) ret_dict['scan_idx'] = np.array(idx).astype(np.int64) ret_dict['scan_name'] = scan_name if self.cache_crops_filename: files_crops = [x for x in self.files_crops if scan_name in x] files_crops = [ os.path.join(self.root_data_crops_path, self.split_set, 'votenet_inputs', x) for x in files_crops ] ret_dict['files_crops'] = files_crops return ret_dict