def viz_obb(pc, label, mask, angle_classes, angle_residuals, size_classes, size_residuals, name=''): """ Visualize oriented bounding box ground truth pc: (N,3) label: (K,3) K == MAX_NUM_OBJ mask: (K,) angle_classes: (K,) angle_residuals: (K,) size_classes: (K,) size_residuals: (K,3) """ oriented_boxes = [] K = label.shape[0] for i in range(K): if mask[i] == 0: continue obb = np.zeros(7) obb[0:3] = label[i, 0:3] heading_angle = 0 # hard code to 0 box_size = DC.mean_size_arr[size_classes[i], :] + size_residuals[i, :] obb[3:6] = box_size obb[6] = -1 * heading_angle print(obb) oriented_boxes.append(obb) pc_util.write_oriented_bbox(oriented_boxes, 'gt_obbs{}.ply'.format(name)) pc_util.write_ply(label[mask == 1, :], 'gt_centroids{}.ply'.format(name))
def viz_obb(pc, label, mask, angle_classes, angle_residuals, size_classes, size_residuals): """ Visualize oriented bounding box ground truth pc: (N,3) label: (K,3) K == MAX_NUM_OBJ mask: (K,) angle_classes: (K,) angle_residuals: (K,) size_classes: (K,) size_residuals: (K,3) """ oriented_boxes = [] K = label.shape[0] for i in range(K): if mask[i] == 0: continue obb = np.zeros(7) obb[0:3] = label[i, 0:3] heading_angle = DC.class2angle(angle_classes[i], angle_residuals[i]) box_size = DC.class2size(size_classes[i], size_residuals[i]) obb[3:6] = box_size obb[6] = -1 * heading_angle print(obb) oriented_boxes.append(obb) pc_util.write_oriented_bbox(oriented_boxes, 'gt_obbs.ply') pc_util.write_ply(label[mask == 1, :], 'gt_centroids.ply')
def data_viz(data_dir, dump_dir=os.path.join(BASE_DIR, 'data_viz_dump')): ''' Examine and visualize ycbgrasp dataset. ''' ycb = ycb_object(data_dir) idxs = np.array(range(0, len(ycb))) #np.random.seed(0) #np.random.shuffle(idxs) if not os.path.exists(dump_dir): os.mkdir(dump_dir) for idx in range(len(ycb)): if idx % 500: continue data_idx = idxs[idx] print('data index: ', data_idx) pc = ycb.get_pointcloud(data_idx) objects = ycb.get_label_objects(data_idx) oriented_boxes = [] for obj in objects: object_pc, inds = ycbgrasp_utils.get_object_points( pc, obj.classname) pc_util.write_ply( object_pc, os.path.join(dump_dir, str(idx) + '_' + obj.classname + '_pc.ply')) if len(object_pc) > 300: obb = np.zeros((7)) obb[0:3] = obj.centroid obb[3:6] = np.array([obj.l, obj.w, obj.h]) * 2 obb[6] = obj.heading_angle oriented_boxes.append(obb) if len(oriented_boxes) > 0: oriented_boxes = np.vstack(tuple(oriented_boxes)) pc_util.write_oriented_bbox( oriented_boxes, os.path.join(dump_dir, str(idx) + '_obbs.ply')) pc = pc[:, 0:3] pc_util.write_ply(pc, os.path.join(dump_dir, str(idx) + '_pc.ply')) print('Complete!')
def data_viz(data_dir, dump_dir=os.path.join(BASE_DIR, 'data_viz_dump')): ''' Examine and visualize SUN RGB-D data. ''' sunrgbd = sunrgbd_object(data_dir) idxs = np.array(range(1, len(sunrgbd) + 1)) np.random.seed(0) np.random.shuffle(idxs) for idx in range(len(sunrgbd)): data_idx = idxs[idx] print('-' * 10, 'data index: ', data_idx) pc = sunrgbd.get_depth(data_idx) print('Point cloud shape:', pc.shape) # Project points to image calib = sunrgbd.get_calibration(data_idx) uv, d = calib.project_upright_depth_to_image(pc[:, 0:3]) print('Point UV:', uv) print('Point depth:', d) import matplotlib.pyplot as plt cmap = plt.cm.get_cmap('hsv', 256) cmap = np.array([cmap(i) for i in range(256)])[:, :3] * 255 img = sunrgbd.get_image(data_idx) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for i in range(uv.shape[0]): depth = d[i] color = cmap[int(120.0 / depth), :] cv2.circle(img, (int(np.round(uv[i, 0])), int(np.round(uv[i, 1]))), 2, color=tuple(color), thickness=-1) if not os.path.exists(dump_dir): os.mkdir(dump_dir) Image.fromarray(img).save(os.path.join(dump_dir, 'img_depth.jpg')) # Load box labels objects = sunrgbd.get_label_objects(data_idx) print('Objects:', objects) # Draw 2D boxes on image img = sunrgbd.get_image(data_idx) img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for i, obj in enumerate(objects): cv2.rectangle(img, (int(obj.xmin), int(obj.ymin)), (int(obj.xmax), int(obj.ymax)), (0, 255, 0), 2) cv2.putText(img, '%d %s' % (i, obj.classname), (max(int(obj.xmin), 15), max(int(obj.ymin), 15)), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 0, 0), 2) Image.fromarray(img).save(os.path.join(dump_dir, 'img_box2d.jpg')) # Dump OBJ files for the colored point cloud for num_point in [10000, 20000, 40000, 80000]: sampled_pcrgb = pc_util.random_sampling(pc, num_point) pc_util.write_ply_rgb( sampled_pcrgb[:, 0:3], (sampled_pcrgb[:, 3:] * 256).astype(np.int8), os.path.join(dump_dir, 'pcrgb_%dk.obj' % (num_point // 1000))) # Dump OBJ files for 3D bounding boxes # l,w,h correspond to dx,dy,dz # heading angle is from +X rotating towards -Y # (+X is degree, -Y is 90 degrees) oriented_boxes = [] for obj in objects: obb = np.zeros((7)) obb[0:3] = obj.centroid # Some conversion to map with default setting of w,l,h # and angle in box dumping obb[3:6] = np.array([obj.l, obj.w, obj.h]) * 2 obb[6] = -1 * obj.heading_angle print('Object cls, heading, l, w, h:',\ obj.classname, obj.heading_angle, obj.l, obj.w, obj.h) oriented_boxes.append(obb) if len(oriented_boxes) > 0: oriented_boxes = np.vstack(tuple(oriented_boxes)) pc_util.write_oriented_bbox(oriented_boxes, os.path.join(dump_dir, 'obbs.ply')) else: print('-' * 30) continue # Draw 3D boxes on depth points box3d = [] ori3d = [] for obj in objects: corners_3d_image, corners_3d = sunrgbd_utils.compute_box_3d( obj, calib) ori_3d_image, ori_3d = sunrgbd_utils.compute_orientation_3d( obj, calib) print('Corners 3D: ', corners_3d) box3d.append(corners_3d) ori3d.append(ori_3d) pc_box3d = np.concatenate(box3d, 0) pc_ori3d = np.concatenate(ori3d, 0) print(pc_box3d.shape) print(pc_ori3d.shape) pc_util.write_ply(pc_box3d, os.path.join(dump_dir, 'box3d_corners.ply')) pc_util.write_ply(pc_ori3d, os.path.join(dump_dir, 'box3d_ori.ply')) print('-' * 30) print('Point clouds and bounding boxes saved to PLY files under %s' % (dump_dir)) print('Type anything to continue to the next sample...') input()
def dump_results(end_points, dump_dir, config, inference_switch=False): ''' Dump results. Args: end_points: dict {..., pred_mask} pred_mask is a binary mask array of size (batch_size, num_proposal) computed by running NMS and empty box removal Returns: None ''' if not os.path.exists(dump_dir): os.system('mkdir %s'%(dump_dir)) # INPUT point_clouds = end_points['point_clouds'].cpu().numpy() batch_size = point_clouds.shape[0] # NETWORK OUTPUTS seed_xyz = end_points['seed_xyz'].detach().cpu().numpy() # (B,num_seed,3) if 'vote_xyz' in end_points: aggregated_vote_xyz = end_points['aggregated_vote_xyz'].detach().cpu().numpy() vote_xyz = end_points['vote_xyz'].detach().cpu().numpy() # (B,num_seed,3) aggregated_vote_xyz = end_points['aggregated_vote_xyz'].detach().cpu().numpy() objectness_scores = end_points['objectness_scores'].detach().cpu().numpy() # (B,K,2) pred_center = end_points['center'].detach().cpu().numpy() # (B,K,3) pred_heading_class = torch.argmax(end_points['heading_scores'], -1) # B,num_proposal pred_heading_residual = torch.gather(end_points['heading_residuals'], 2, pred_heading_class.unsqueeze(-1)) # B,num_proposal,1 pred_heading_class = pred_heading_class.detach().cpu().numpy() # B,num_proposal pred_heading_residual = pred_heading_residual.squeeze(2).detach().cpu().numpy() # B,num_proposal pred_size_class = torch.argmax(end_points['size_scores'], -1) # B,num_proposal pred_size_residual = torch.gather(end_points['size_residuals'], 2, pred_size_class.unsqueeze(-1).unsqueeze(-1).repeat(1,1,1,3)) # B,num_proposal,1,3 pred_size_residual = pred_size_residual.squeeze(2).detach().cpu().numpy() # B,num_proposal,3 # OTHERS pred_mask = end_points['pred_mask'] # B,num_proposal idx_beg = 0 for i in range(batch_size): pc = point_clouds[i,:,:] objectness_prob = softmax(objectness_scores[i,:,:])[:,1] # (K,) # Dump various point clouds pc_util.write_ply(pc, os.path.join(dump_dir, '%06d_pc.ply'%(idx_beg+i))) pc_util.write_ply(seed_xyz[i,:,:], os.path.join(dump_dir, '%06d_seed_pc.ply'%(idx_beg+i))) #pc_util.write_ply_rgb(pc[:,0:3], (pc[:,3:]*256).astype(np.int8),os.path.join(dump_dir, '%06d_rgb_pc.ply'%(idx_beg+i))) #jason color if 'vote_xyz' in end_points: pc_util.write_ply(end_points['vote_xyz'][i,:,:], os.path.join(dump_dir, '%06d_vgen_pc.ply'%(idx_beg+i))) pc_util.write_ply(aggregated_vote_xyz[i,:,:], os.path.join(dump_dir, '%06d_aggregated_vote_pc.ply'%(idx_beg+i))) pc_util.write_ply(aggregated_vote_xyz[i,:,:], os.path.join(dump_dir, '%06d_aggregated_vote_pc.ply'%(idx_beg+i))) pc_util.write_ply(pred_center[i,:,0:3], os.path.join(dump_dir, '%06d_proposal_pc.ply'%(idx_beg+i))) if np.sum(objectness_prob>DUMP_CONF_THRESH)>0: pc_util.write_ply(pred_center[i,objectness_prob>DUMP_CONF_THRESH,0:3], os.path.join(dump_dir, '%06d_confident_proposal_pc.ply'%(idx_beg+i))) # Dump predicted bounding boxes if np.sum(objectness_prob>DUMP_CONF_THRESH)>0: num_proposal = pred_center.shape[1] obbs = [] for j in range(num_proposal): obb = config.param2obb(pred_center[i,j,0:3], pred_heading_class[i,j], pred_heading_residual[i,j], pred_size_class[i,j], pred_size_residual[i,j]) obbs.append(obb) if len(obbs)>0: obbs = np.vstack(tuple(obbs)) # (num_proposal, 7) pc_util.write_oriented_bbox(obbs[objectness_prob>DUMP_CONF_THRESH,:], os.path.join(dump_dir, '%06d_pred_confident_bbox.ply'%(idx_beg+i))) pc_util.write_oriented_bbox(obbs[np.logical_and(objectness_prob>DUMP_CONF_THRESH, pred_mask[i,:]==1),:], os.path.join(dump_dir, '%06d_pred_confident_nms_bbox.ply'%(idx_beg+i))) pc_util.write_oriented_bbox(obbs[pred_mask[i,:]==1,:], os.path.join(dump_dir, '%06d_pred_nms_bbox.ply'%(idx_beg+i))) pc_util.write_oriented_bbox(obbs, os.path.join(dump_dir, '%06d_pred_bbox.ply'%(idx_beg+i))) # Return if it is at inference time. No dumping of groundtruths if inference_switch: return # LABELS gt_center = end_points['center_label'].cpu().numpy() # (B,MAX_NUM_OBJ,3) gt_mask = end_points['box_label_mask'].cpu().numpy() # B,K2 gt_heading_class = end_points['heading_class_label'].cpu().numpy() # B,K2 gt_heading_residual = end_points['heading_residual_label'].cpu().numpy() # B,K2 gt_size_class = end_points['size_class_label'].cpu().numpy() # B,K2 gt_size_residual = end_points['size_residual_label'].cpu().numpy() # B,K2,3 objectness_label = end_points['objectness_label'].detach().cpu().numpy() # (B,K,) objectness_mask = end_points['objectness_mask'].detach().cpu().numpy() # (B,K,) for i in range(batch_size): if np.sum(objectness_label[i,:])>0: pc_util.write_ply(pred_center[i,objectness_label[i,:]>0,0:3], os.path.join(dump_dir, '%06d_gt_positive_proposal_pc.ply'%(idx_beg+i))) if np.sum(objectness_mask[i,:])>0: pc_util.write_ply(pred_center[i,objectness_mask[i,:]>0,0:3], os.path.join(dump_dir, '%06d_gt_mask_proposal_pc.ply'%(idx_beg+i))) pc_util.write_ply(gt_center[i,:,0:3], os.path.join(dump_dir, '%06d_gt_centroid_pc.ply'%(idx_beg+i))) pc_util.write_ply_color(pred_center[i,:,0:3], objectness_label[i,:], os.path.join(dump_dir, '%06d_proposal_pc_objectness_label.obj'%(idx_beg+i))) # Dump GT bounding boxes obbs = [] for j in range(gt_center.shape[1]): if gt_mask[i,j] == 0: continue obb = config.param2obb(gt_center[i,j,0:3], gt_heading_class[i,j], gt_heading_residual[i,j], gt_size_class[i,j], gt_size_residual[i,j]) obbs.append(obb) if len(obbs)>0: obbs = np.vstack(tuple(obbs)) # (num_gt_objects, 7) pc_util.write_oriented_bbox(obbs, os.path.join(dump_dir, '%06d_gt_bbox.ply'%(idx_beg+i))) # OPTIONALL, also dump prediction and gt details if 'batch_pred_map_cls' in end_points: for ii in range(batch_size): fout = open(os.path.join(dump_dir, '%06d_pred_map_cls.txt'%(ii)), 'w') for t in end_points['batch_pred_map_cls'][ii]: fout.write(str(t[0])+' ') fout.write(",".join([str(x) for x in list(t[1].flatten())])) fout.write(' '+str(t[2])) fout.write('\n') fout.close() if 'batch_gt_map_cls' in end_points: for ii in range(batch_size): fout = open(os.path.join(dump_dir, '%06d_gt_map_cls.txt'%(ii)), 'w') for t in end_points['batch_gt_map_cls'][ii]: fout.write(str(t[0])+' ') fout.write(",".join([str(x) for x in list(t[1].flatten())])) fout.write('\n') fout.close()
raise NotImplementedError if FLAGS.dump_dir is None: dump_dir = 'dump_'+FLAGS.dataset else: dump_dir = FLAGS.dump_dir dump_dir = os.path.join(ROOT_DIR, dump_dir) if not os.path.exists(dump_dir): os.mkdir(dump_dir) for name in all_scan_names: scene_name = name[3:15] # if scene_name != '000001': # continue obbs = np.load(os.path.join(FLAGS.path,name)) try: print(obbs[:, -1]) except: continue boxes = obbs[:, :-1] sem_cls = obbs[:, -1] - 1 print('(%d, %d)' % (obbs.shape[0], obbs.shape[1])) case_dump_dir = os.path.join(dump_dir, scene_name) if not os.path.exists(case_dump_dir): os.mkdir(case_dump_dir) # for l in np.unique(sem_cls): # mask = (sem_cls == l) # if np.sum(mask)>0: pc_util.write_oriented_bbox(boxes, os.path.join(case_dump_dir, 'pred_confident_nms_bbox.ply'))
all_scan_names = list( set([ os.path.basename(x) for x in os.listdir(FLAGS.path) if x.endswith('gt.mat') ])) dump_dir = FLAGS.dump_dir if not os.path.exists(dump_dir): os.mkdir(dump_dir) for name in all_scan_names: scene_name = name[3:15] if scene_name != 'scene0488_00': continue obbs = sio.loadmat(os.path.join(FLAGS.path, name))['gt'] obbs = np.array(obbs) print(obbs[:, -1]) boxes = obbs[:, :-1] sem_cls = obbs[:, -1] - 1 print('(%d, %d)' % (obbs.shape[0], obbs.shape[1])) case_dump_dir = os.path.join(dump_dir, scene_name) if not os.path.exists(case_dump_dir): os.mkdir(case_dump_dir) for l in np.unique(sem_cls): mask = (sem_cls == l) if np.sum(mask) > 0: pc_util.write_oriented_bbox( boxes[mask, :], os.path.join(case_dump_dir, '%d_gt_bbox.ply' % (l)))