def dump_results(args, scanrefer, data, config): dump_dir = os.path.join(CONF.PATH.BASE, args.folder, "vis") os.makedirs(dump_dir, exist_ok=True) # from inputs ids = data['scan_idx'].detach().cpu().numpy() point_clouds = data['point_clouds'].cpu().numpy() test_point_clouds = data["test_point_clouds"].cpu().numpy() # one object batch_size = point_clouds.shape[0] pcl_color = data["pcl_color"].detach().cpu().numpy() test_pcl_color = data["test_pcl_color"].detach().cpu().numpy() if args.use_color: pcl_color = (pcl_color * 256 + MEAN_COLOR_RGB).astype(np.int64) test_pcl_color = (test_pcl_color * 256 + MEAN_COLOR_RGB).astype( np.int64) for i in range(batch_size): # basic info idx = ids[i] scene_id = scanrefer[idx]["scene_id"] object_id = scanrefer[idx]["object_id"] object_name = scanrefer[idx]["object_name"] ann_id = scanrefer[idx]["ann_id"] # scene_output scene_dump_dir = os.path.join(dump_dir, scene_id) # if not os.path.exists(scene_dump_dir): os.makedirs(scene_dump_dir, exist_ok=True) # # Dump the original scene point clouds # mesh = align_mesh(scene_id) # mesh.write(os.path.join(scene_dump_dir, 'mesh.ply')) print("pcl_color[i] in pc.ply", pcl_color[i]) point_clouds = point_clouds[i] #[:200] pcl_color = pcl_color[i] #[:200] test_point_clouds = test_point_clouds[i] test_pcl_color = test_pcl_color[i] jitter_idx = random.random() * 0.45 + 0.8 write_ply_rgb(point_clouds, pcl_color, os.path.join(scene_dump_dir, 'pc_full.ply')) write_ply_rgb(test_point_clouds, test_pcl_color, os.path.join(scene_dump_dir, 'pc_full_obj.ply')) write_ply_rgb(test_point_clouds * jitter_idx, test_pcl_color, os.path.join(scene_dump_dir, 'pc_full_obj2.ply')) # write_ply_rgb(point_clouds[i], pcl_color[i], os.path.join(scene_dump_dir, 'pc.ply')) break
def dump_results(args, scanrefer, data, config): dump_dir = os.path.join(CONF.PATH.OUTPUT, args.folder, "vis") os.makedirs(dump_dir, exist_ok=True) # from inputs ids = data['scan_idx'].detach().cpu().numpy() point_clouds = data['point_clouds'].cpu().numpy() batch_size = point_clouds.shape[0] pcl_color = data["pcl_color"].detach().cpu().numpy() if args.use_color: pcl_color = (pcl_color * 256 + MEAN_COLOR_RGB).astype(np.int64) # from network outputs # detection pred_objectness = torch.argmax(data['objectness_scores'], 2).float().detach().cpu().numpy() pred_center = data['center'].detach().cpu().numpy() # (B,K,3) pred_heading_class = torch.argmax(data['heading_scores'], -1) # B,num_proposal pred_heading_residual = torch.gather( data['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(data['size_scores'], -1) # B,num_proposal pred_size_residual = torch.gather( data['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 # reference pred_ref_scores = data["cluster_ref"].detach().cpu().numpy() pred_ref_scores_softmax = F.softmax( data["cluster_ref"] * torch.argmax(data['objectness_scores'], 2).float() * data['pred_mask'], dim=1).detach().cpu().numpy() # post-processing nms_masks = data['pred_mask'].detach().cpu().numpy() # B,num_proposal # ground truth gt_center = data['center_label'].cpu().numpy() # (B,MAX_NUM_OBJ,3) gt_heading_class = data['heading_class_label'].cpu().numpy() # B,K2 gt_heading_residual = data['heading_residual_label'].cpu().numpy() # B,K2 gt_size_class = data['size_class_label'].cpu().numpy() # B,K2 gt_size_residual = data['size_residual_label'].cpu().numpy() # B,K2,3 # reference gt_ref_labels = data["ref_box_label"].detach().cpu().numpy() for i in range(batch_size): # basic info idx = ids[i] scene_id = scanrefer[idx]["scene_id"] object_id = scanrefer[idx]["object_id"] object_name = scanrefer[idx]["object_name"] ann_id = scanrefer[idx]["ann_id"] # scene_output scene_dump_dir = os.path.join(dump_dir, scene_id) if not os.path.exists(scene_dump_dir): os.mkdir(scene_dump_dir) # # Dump the original scene point clouds mesh = align_mesh(scene_id) mesh.write(os.path.join(scene_dump_dir, 'mesh.ply')) write_ply_rgb(point_clouds[i], pcl_color[i], os.path.join(scene_dump_dir, 'pc.ply')) # filter out the valid ground truth reference box assert gt_ref_labels[i].shape[0] == gt_center[i].shape[0] gt_ref_idx = np.argmax(gt_ref_labels[i], 0) # visualize the gt reference box # NOTE: for each object there should be only one gt reference box object_dump_dir = os.path.join( dump_dir, scene_id, "gt_{}_{}.ply".format(object_id, object_name)) gt_obb = config.param2obb(gt_center[i, gt_ref_idx, 0:3], gt_heading_class[i, gt_ref_idx], gt_heading_residual[i, gt_ref_idx], gt_size_class[i, gt_ref_idx], gt_size_residual[i, gt_ref_idx]) gt_bbox = get_3d_box(gt_obb[3:6], gt_obb[6], gt_obb[0:3]) if not os.path.exists(object_dump_dir): write_bbox( gt_obb, 0, os.path.join(scene_dump_dir, 'gt_{}_{}.ply'.format(object_id, object_name))) # find the valid reference prediction pred_masks = nms_masks[i] * pred_objectness[i] == 1 assert pred_ref_scores[i].shape[0] == pred_center[i].shape[0] pred_ref_idx = np.argmax(pred_ref_scores[i] * pred_masks, 0) assigned_gt = torch.gather( data["ref_box_label"], 1, data["object_assignment"]).detach().cpu().numpy() # visualize the predicted reference box pred_obb = config.param2obb(pred_center[i, pred_ref_idx, 0:3], pred_heading_class[i, pred_ref_idx], pred_heading_residual[i, pred_ref_idx], pred_size_class[i, pred_ref_idx], pred_size_residual[i, pred_ref_idx]) pred_bbox = get_3d_box(pred_obb[3:6], pred_obb[6], pred_obb[0:3]) iou = box3d_iou(gt_bbox, pred_bbox) write_bbox( pred_obb, 1, os.path.join( scene_dump_dir, 'pred_{}_{}_{}_{:.5f}_{:.5f}.ply'.format( object_id, object_name, ann_id, pred_ref_scores_softmax[i, pred_ref_idx], iou)))
def write_bbox(bbox, mode, output_file): """ bbox: (cx, cy, cz, lx, ly, lz, r), center and length in three axis, the last is the rotation output_file: string """ def create_cylinder_mesh(radius, p0, p1, stacks=10, slices=10): import math def compute_length_vec3(vec3): return math.sqrt(vec3[0]*vec3[0] + vec3[1]*vec3[1] + vec3[2]*vec3[2]) def rotation(axis, angle): rot = np.eye(4) c = np.cos(-angle) s = np.sin(-angle) t = 1.0 - c axis /= compute_length_vec3(axis) x = axis[0] y = axis[1] z = axis[2] rot[0,0] = 1 + t*(x*x-1) rot[0,1] = z*s+t*x*y rot[0,2] = -y*s+t*x*z rot[1,0] = -z*s+t*x*y rot[1,1] = 1+t*(y*y-1) rot[1,2] = x*s+t*y*z rot[2,0] = y*s+t*x*z rot[2,1] = -x*s+t*y*z rot[2,2] = 1+t*(z*z-1) return rot verts = [] indices = [] diff = (p1 - p0).astype(np.float32) height = compute_length_vec3(diff) for i in range(stacks+1): for i2 in range(slices): theta = i2 * 2.0 * math.pi / slices pos = np.array([radius*math.cos(theta), radius*math.sin(theta), height*i/stacks]) verts.append(pos) for i in range(stacks): for i2 in range(slices): i2p1 = math.fmod(i2 + 1, slices) indices.append( np.array([(i + 1)*slices + i2, i*slices + i2, i*slices + i2p1], dtype=np.uint32) ) indices.append( np.array([(i + 1)*slices + i2, i*slices + i2p1, (i + 1)*slices + i2p1], dtype=np.uint32) ) transform = np.eye(4) va = np.array([0, 0, 1], dtype=np.float32) vb = diff vb /= compute_length_vec3(vb) axis = np.cross(vb, va) angle = np.arccos(np.clip(np.dot(va, vb), -1, 1)) if angle != 0: if compute_length_vec3(axis) == 0: dotx = va[0] if (math.fabs(dotx) != 1.0): axis = np.array([1,0,0]) - dotx * va else: axis = np.array([0,1,0]) - va[1] * va axis /= compute_length_vec3(axis) transform = rotation(axis, -angle) transform[:3,3] += p0 verts = [np.dot(transform, np.array([v[0], v[1], v[2], 1.0])) for v in verts] verts = [np.array([v[0], v[1], v[2]]) / v[3] for v in verts] return verts, indices def get_bbox_edges(bbox_min, bbox_max): def get_bbox_verts(bbox_min, bbox_max): verts = [ np.array([bbox_min[0], bbox_min[1], bbox_min[2]]), np.array([bbox_max[0], bbox_min[1], bbox_min[2]]), np.array([bbox_max[0], bbox_max[1], bbox_min[2]]), np.array([bbox_min[0], bbox_max[1], bbox_min[2]]), np.array([bbox_min[0], bbox_min[1], bbox_max[2]]), np.array([bbox_max[0], bbox_min[1], bbox_max[2]]), np.array([bbox_max[0], bbox_max[1], bbox_max[2]]), np.array([bbox_min[0], bbox_max[1], bbox_max[2]]) ] return verts box_verts = get_bbox_verts(bbox_min, bbox_max) edges = [ (box_verts[0], box_verts[1]), (box_verts[1], box_verts[2]), (box_verts[2], box_verts[3]), (box_verts[3], box_verts[0]), (box_verts[4], box_verts[5]), (box_verts[5], box_verts[6]), (box_verts[6], box_verts[7]), (box_verts[7], box_verts[4]), (box_verts[0], box_verts[4]), (box_verts[1], box_verts[5]), (box_verts[2], box_verts[6]), (box_verts[3], box_verts[7]) ] return edges def get_bbox_corners(bbox): centers, lengths = bbox[:3], bbox[3:6] xmin, xmax = centers[0] - lengths[0] / 2, centers[0] + lengths[0] / 2 ymin, ymax = centers[1] - lengths[1] / 2, centers[1] + lengths[1] / 2 zmin, zmax = centers[2] - lengths[2] / 2, centers[2] + lengths[2] / 2 corners = [] corners.append(np.array([xmax, ymax, zmax]).reshape(1, 3)) corners.append(np.array([xmax, ymax, zmin]).reshape(1, 3)) corners.append(np.array([xmin, ymax, zmin]).reshape(1, 3)) corners.append(np.array([xmin, ymax, zmax]).reshape(1, 3)) corners.append(np.array([xmax, ymin, zmax]).reshape(1, 3)) corners.append(np.array([xmax, ymin, zmin]).reshape(1, 3)) corners.append(np.array([xmin, ymin, zmin]).reshape(1, 3)) corners.append(np.array([xmin, ymin, zmax]).reshape(1, 3)) corners = np.concatenate(corners, axis=0) # 8 x 3 return corners radius = 0.03 offset = [0,0,0] verts = [] indices = [] colors = [] corners = get_bbox_corners(bbox) box_min = np.min(corners, axis=0) box_max = np.max(corners, axis=0) palette = { 0: [0, 255, 0], # gt 1: [0, 0, 255] # pred } chosen_color = palette[mode] edges = get_bbox_edges(box_min, box_max) for k in range(len(edges)): cyl_verts, cyl_ind = create_cylinder_mesh(radius, edges[k][0], edges[k][1]) cur_num_verts = len(verts) cyl_color = [[c / 255 for c in chosen_color] for _ in cyl_verts] cyl_verts = [x + offset for x in cyl_verts] cyl_ind = [x + cur_num_verts for x in cyl_ind] verts.extend(cyl_verts) indices.extend(cyl_ind) colors.extend(cyl_color) # print("in scenes color", colors) write_ply_rgb(np.array(verts), np.array(colors), output_file)