def draw_and_connect_keypoints(self, keypoints): """Draws keypoints of an instance and follows the rules for keypoint connections to draw lines between appropriate keypoints. This follows color heuristics for line color. Args: keypoints (Tensor): a tensor of shape (K, 2), where K is the number of keypoints and the last dimension corresponds to (x, y). Returns: output (VisImage): image object with visualizations. """ if self.keypoint_type is None: return super().draw_and_connect_keypoints(keypoints) if True: # from center to other points # assume keypoints (K, 2), the last row is object center cx, cy = keypoints[-1, :2] self.draw_circle((cx, cy), color=self.keypoint_color) colors = colormap(rgb=True, maximum=1) for idx in range(len(keypoints) - 1): # draw keypoint x, y = keypoints[idx] self.draw_circle((x, y), color=self.keypoint_color) # default red (1,0,0) # connect keypoints line_color = colors[idx % len(colors)] self.draw_line([cx, x], [cy, y], color=line_color) else: # assume keypoints (K, 2), the last row is object center colors = colormap(rgb=True, maximum=1) for idx in range(len(keypoints) - 1): # draw keypoint x, y = keypoints[idx, :2] xe, ye = keypoints[idx + 1, :2] self.draw_circle((x, y), color=self.keypoint_color) # default red (1,0,0) # connect keypoints line_color = colors[idx % len(colors)] # x_data, y_data self.draw_line([x, xe], [y, ye], color=line_color) x, y = keypoints[-1, :2] self.draw_circle((x, y), color=self.keypoint_color) # default red (1,0,0) return self.output
def draw_pix3d_dict(dataset_dict, class_names=None): """ Draw the instance annotations for an image. Args: dataset_dict (dict): a dict in Detectron2 Dataset format. See DATASETS.md class_names (list[str] or None): `class_names[cateogory_id]` is the name for this category. If not provided, the visualization will not contain class names. """ img = dataset_dict.get("image", None) if img is None: img = cv2.imread(dataset_dict["file_name"]) annos = dataset_dict["annotations"] if not len(annos): return img boxes = np.asarray([ BoxMode.convert(k["bbox"], k["bbox_mode"], BoxMode.XYXY_ABS) for k in annos ]) # Display in largest to smallest order to reduce occlusion areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) sorted_inds = np.argsort(-areas) sorted_boxes = copy.deepcopy(boxes[sorted_inds]) img = draw_boxes(img, sorted_boxes) cmap = colormap() for num, i in enumerate(sorted_inds): anno = annos[i] bbox = anno["bbox"] assert anno["bbox_mode"] in [ BoxMode.XYXY_ABS, BoxMode.XYWH_ABS, ], "Relative coordinates not yet supported in visualization." iscrowd = anno.get("iscrowd", 0) clsid = anno["category_id"] text = class_names[clsid] if class_names is not None else str(clsid) if iscrowd: text = text + "_crowd" img = draw_text(img, (bbox[0], bbox[1] - 2), text) segs = anno.get("segmentation", None) if segs is not None and not iscrowd: segs_color = cmap[num % len(cmap)] mask = cv2.imread(segs) img = draw_mask(img, mask, segs_color, draw_contours=False) kpts = anno.get("keypoints", None) if kpts is not None and not iscrowd: kpts = np.asarray(kpts).reshape(-1, 3)[:, :2] img = draw_keypoints(img, kpts) return img
def draw_axis3d_and_center( self, keypoints, up_color=_RED, right_color=_GREEN, front_color=_BLUE, center_color=_GREY, linewidth=None, draw_points=False, draw_center=False, ): """ 2 | 3 ---1 / 0 Args: keypoints (Tensor): a tensor of shape (4, 2), where K is the number of keypoints and the last dimension corresponds to (x, y). Returns: output (VisImage): image object with visualizations. """ assert keypoints.shape[-1] == 2, keypoints.shape if linewidth is None: linewidth = self._default_font_size / 3 linewidth = max(linewidth, 1) colors = colormap(rgb=True, maximum=1) # connect keypoints for idx, edge in enumerate([(3, 0), (3, 1), (3, 2)]): i, j = edge x1, y1 = keypoints[i, :2] x2, y2 = keypoints[j, :2] if j == 0: color = front_color elif j == 1: color = right_color else: color = up_color self.draw_line([x1, x2], [y1, y2], color=color, linewidth=linewidth) if draw_center: # assume keypoints (K, 2), the last row is object center cx, cy = keypoints[-1, :2] # draw center with center color self.draw_circle((cx, cy), color=center_color) if draw_points: for idx in range(len(keypoints) - 1): # draw keypoint x, y = keypoints[idx, :2] point_color = colors[idx % len(colors)] self.draw_circle((x, y), color=point_color) return self.output
def main(args): global bev_im mp.set_start_method("spawn", force=True) args = get_parser().parse_args() logger = setup_logger() logger.info("Arguments: " + str(args)) cfg = setup_cfg(args) view = True predictor = DefaultPredictor(cfg) metadata = MetadataCatalog.get(cfg.DATASETS.TEST) f_rgb_detections = open('nuscenes_rgb_detections.txt', "a") dataset = nuscenes_object.nuscenes_object( '/raid/datasets/extracted_nuscenes', split='val', velo_kind='lidar_top') if not os.path.exists(os.path.join(args.output_dir)): os.mkdir(os.path.join(args.output_dir)) if not os.path.exists(os.path.join(args.output_dir, 'data')): os.mkdir(os.path.join(args.output_dir, 'data')) current_scene = 0 current_time = 0 for idx in range(0, len(dataset)): name = dataset.get_idx_name(idx)[1:] if current_scene == int(name[:4]) and current_time >= int(name[-4:]): continue else: current_scene = int(name[:4]) current_time = int(name[-4:]) print(name) ims = [] for ii in range(0, 6): # print(ii) im = dataset.get_image_by_name(str(ii) + name) im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) predictions = predictor(im)[0] print(predictions) instances = predictions["instances"].to(torch.device("cpu")) #draw the detections over the original images visualizer = Visualizer(im, metadata) classes = instances.pred_classes.cpu().numpy() class_names = visualizer.metadata.get("thing_classes") # print(instances.pred_classes) # print(class_names) # print(classes) labels = [class_names[i] for i in classes] vis_colors = [ colormap(rgb=True, maximum=1)[i] if i < 74 else (0, 0, 0) for i in classes ] if (view): visualizer.overlay_instances( boxes=instances.pred_boxes, # masks=instances.pred_masks, labels=_create_text_labels(instances.pred_classes, \ instances.scores, \ visualizer.metadata.get("thing_classes", None)), # ['Car', 'Pedestrian', 'Cyclist', 'Motorcyclist']), assigned_colors=vis_colors, alpha=0.5, ) for jj in range(len(instances)): bbox = instances.pred_boxes[jj].get_numpy()[0] output_str = os.path.join( dataset.image_dir, '%s.jpg' % (str(ii) + name)) + " %s %f %.2f %.2f %.2f %.2f\n" % ( labels[jj], instances.scores[jj].cpu().numpy(), bbox[0], bbox[1], bbox[2], bbox[3]) # print(output_str) f_rgb_detections.write(output_str) det_filename = os.path.join(args.output_dir, 'data', '%s.txt' % (str(ii) + name)) with open(det_filename, 'a+') as f: bbox = instances.pred_boxes[jj].get_numpy()[0] output_eval = '%s -1 -1 -10 %.3f %.3f %.3f %.3f -1 -1 -1 -1 -1 -1 -1 %.3f\n' %\ (labels[jj], bbox[0], bbox[1], bbox[2], bbox[3], instances.scores[jj].cpu().numpy()) f.write(output_eval) # print(output_eval) if (view): im_view = np.array(visualizer.output.get_image()[:, :, ::-1]) # im_v = cv2.rectangle(im_view, (0,0), (im_view.shape[1], im_view.shape[0]), colors[ii], thickness = 30) if (ii == 0): ims = [] ims.append(im_view) if (view): h1 = cv2.hconcat((ims[1], ims[0], ims[2])) h2 = cv2.hconcat((ims[5], ims[3], ims[4])) v1 = cv2.vconcat((h1, h2)) cv2.namedWindow('6im', cv2.WINDOW_NORMAL) cv2.imshow('6im', v1) if (view): key = cv2.waitKey(0) if key == 115: cv2.imwrite('%s_6im.png' % name, v1) cv2.imwrite('%s_bev_im.png' % name, bev_im) print('SAVING IMAGES') if key == 27: break # esc to quit
def draw_bbox3d_and_center( self, keypoints, top_color=_RED, middle_color=None, bottom_color=(0.5, 0.5, 0.5), linewidth=None, draw_points=False, draw_center=False, ): """1 -------- 0. /| /| 2 -------- 3 . | | | | . 5 -------- 4 |/ |/ 6 -------- 7 8: center Args: keypoints (Tensor): a tensor of shape (K, 2), where K is the number of keypoints and the last dimension corresponds to (x, y). Returns: output (VisImage): image object with visualizations. """ assert keypoints.shape[-1] == 2, keypoints.shape if linewidth is None: linewidth = self._default_font_size / 3 linewidth = max(linewidth, 1) colors = colormap(rgb=True, maximum=1) # connect keypoints for idx, edge in enumerate([(4, 5), (5, 6), (6, 7), (7, 4)]): i, j = edge x1, y1 = keypoints[i, :2] x2, y2 = keypoints[j, :2] if bottom_color is None: _bottom_color = colors[idx % len(colors)] else: _bottom_color = bottom_color self.draw_line([x1, x2], [y1, y2], color=_bottom_color, linewidth=linewidth) for idx, edge in enumerate([(0, 4), (3, 7), (2, 6), (1, 5)]): i, j = edge x1, y1 = keypoints[i, :2] x2, y2 = keypoints[j, :2] if middle_color is None: edge_color = colors[idx % len(colors)] else: edge_color = middle_color self.draw_line([x1, x2], [y1, y2], color=edge_color, linewidth=linewidth) for edge in [(0, 1), (1, 2), (2, 3), (3, 0)]: i, j = edge x1, y1 = keypoints[i, :2] x2, y2 = keypoints[j, :2] self.draw_line([x1, x2], [y1, y2], color=top_color, linewidth=linewidth) if draw_center: # assume keypoints (K, 2), the last row is object center cx, cy = keypoints[-1, :2] # draw center with edge color self.draw_circle((cx, cy), color=top_color) if draw_points: for idx in range(len(keypoints) - 1): # draw keypoint x, y = keypoints[idx, :2] point_color = colors[idx % len(colors)] self.draw_circle((x, y), color=point_color) return self.output
def readMesh(house_folder, region_id): # Load open3d mesh mesh_o3d = o3d.io.read_triangle_mesh( os.path.join(house_folder, region_id + '.ply')) # Load map: instance -> [segment id], store in aggregation with open(os.path.join(house_folder, region_id + '.semseg.json'), 'r') as f: json_string = f.read().replace('\n', '') data = fix_JSON(json_string) aggregation = np.array(data['segGroups']) # Load mesh: pcd, faces plydata = PlyData.read(os.path.join(house_folder, region_id + '.ply')) vertices = plydata['vertex'] points = np.stack([vertices['x'], vertices['y'], vertices['z']], axis=1) faces = np.array(plydata['face']['vertex_indices']) # Load map: vertex -> segment id, store in segmentation data = json.load( open(os.path.join(house_folder, region_id + '.vsegs.json'), 'r')) segmentation = np.array(data['segIndices']).astype(np.int32) # Extract information from aggretation # groupSegments -> [segment ids] # groupLabels -> category for segment group groupSegments = [] groupLabels = [] for segmentIndex, item in enumerate(aggregation): assert item['id'] == segmentIndex groupSegments.append(item['segments']) groupLabels.append(item['label']) uniqueSegments = np.unique(segmentation).tolist() for segments in groupSegments: for segmentIndex in segments: if segmentIndex in uniqueSegments: uniqueSegments.remove(segmentIndex) for segment in uniqueSegments: groupSegments.append([ segment, ]) groupLabels.append('unannotated') numGroups = len(groupSegments) numPoints = segmentation.shape[0] numPlanes = 1000 ## Segment connections for plane merging later ## e.g. segmentEdges = [(1100, 2689), (4, 1991), (1372, 2238), (390, 2780), (1706, 2022), (207, 3073), (2552, 2841), (1802, 2817), (1117, 2798), (3142, 3463), (999, 1894), (891, 1657), (13, 1905), (2596, 2999), (1340, 1550), (2675, 3723)] ## stores tuples of segment id. segmentEdges = [] for face in faces: segment_1 = segmentation[face[0]] segment_2 = segmentation[face[1]] segment_3 = segmentation[face[2]] if segment_1 != segment_2 or segment_1 != segment_3: if segment_1 != segment_2 and segment_1 != -1 and segment_2 != -1: segmentEdges.append( (min(segment_1, segment_2), max(segment_1, segment_2))) if segment_1 != segment_3 and segment_1 != -1 and segment_3 != -1: segmentEdges.append( (min(segment_1, segment_3), max(segment_1, segment_3))) if segment_2 != segment_3 and segment_2 != -1 and segment_3 != -1: segmentEdges.append( (min(segment_2, segment_3), max(segment_2, segment_3))) segmentEdges = list(set(segmentEdges)) numPlanes = 200 numPlanesPerSegment = 2 segmentRatio = 0.1 planeAreaThreshold = 10 numIterations = 100 numIterationsPair = 1000 planeDiffThreshold = 0.05 fittingErrorThreshold = planeDiffThreshold ## Specify the minimum and maximum number of planes for each object labelNumPlanes = { 'wall': [1, 20], 'floor': [1, 1], 'cabinet': [1, 5], 'bed': [1, 5], 'chair': [1, 4], 'sofa': [1, 10], 'table': [1, 5], 'door': [1, 2], 'window': [1, 2], # 'bookshelf': [1, 5], 'picture': [1, 1], 'counter': [1, 10], 'blinds': [0, 0], # 'desk': [1, 10], # 'shelf': [1, 5], # 'shelves': [1, 5], 'curtain': [0, 0], # 'dresser': [1, 5], 'cushion': [0, 0], # 'pillow': [0, 0], 'mirror': [0, 0], # 'entrance': [1, 1], # 'floor mat': [1, 1], 'clothes': [0, 0], 'ceiling': [1, 5], # 'book': [0, 1], # 'books': [0, 1], # 'refridgerator': [1, 5], 'tv_monitor': [1, 1], # 'television': [1, 1], # 'paper': [0, 1], 'towel': [0, 1], # 'shower curtain': [0, 1], # 'box': [1, 5], # 'whiteboard': [1, 5], # 'person': [0, 0], # 'night stand': [1, 5], 'toilet': [0, 5], 'sink': [0, 5], # 'lamp': [0, 1], 'bathtub': [0, 5], # 'bag': [0, 1], 'misc': [0, 5], # mp3d # 'otherprop': [0, 5], # 'otherstructure': [0, 5], 'furniture': [0, 5], # 'otherfurniture': [0, 5], 'appliances': [0, 5], 'unannotated': [0, 5], 'void': [0, 0], 'chest_of_drawers': [0, 5], 'stairs': [0, 20], 'stool': [0, 2], 'shower': [0, 5], 'column': [0, 4], 'fireplace': [0, 5], 'lighting': [0, 2], 'beam': [0, 3], 'railing': [0, 5], 'shelving': [1, 5], 'gym_equipment': [0, 5], 'seating': [1, 2], 'board_panel': [0, 1], 'objects': [0, 5], 'unlabeled': [0, 5], 'plant': [0, 0], } nonPlanarGroupLabels = ['bicycle', 'bottle', 'water bottle', 'plant'] nonPlanarGroupLabels = {label: True for label in nonPlanarGroupLabels} verticalLabels = ['wall', 'door', 'cabinet'] classMap = loadClassMap() allXYZ = points.reshape(-1, 3) segmentNeighbors = defaultdict(list) for segmentEdge in segmentEdges: segmentNeighbors[segmentEdge[0]].append(segmentEdge[1]) segmentNeighbors[segmentEdge[1]].append(segmentEdge[0]) planeGroups = [] debug = False debugIndex = -1 ## A group corresponds to an instance in the ScanNet annotation for groupIndex, group in enumerate(groupSegments): if debugIndex != -1 and groupIndex != debugIndex: continue if groupLabels[groupIndex] in nonPlanarGroupLabels: groupLabel = groupLabels[groupIndex] minNumPlanes, maxNumPlanes = 0, 0 elif groupLabels[groupIndex] == 'unannotated': groupLabel = 'unannotated' minNumPlanes, maxNumPlanes = labelNumPlanes[groupLabel] elif groupLabels[groupIndex] in classMap: groupLabel = classMap[groupLabels[groupIndex]] minNumPlanes, maxNumPlanes = labelNumPlanes[groupLabel] else: print(groupLabels[groupIndex] + ' not considered') minNumPlanes, maxNumPlanes = 0, 0 groupLabel = '' if maxNumPlanes == 0: pointMasks = [] for segmentIndex in group: pointMasks.append(segmentation == segmentIndex) pointIndices = np.any(np.stack(pointMasks, 0), 0).nonzero()[0] groupPlanes = { 'groupedPlanes': [np.zeros(3)], 'groupedPlanePointIndices': [pointIndices], 'groupNeighbors': [], } planeGroups.append(groupPlanes) continue groupPlanes = [] groupPlanePointIndices = [] groupPlaneSegments = [] ## A group contains multiple segments and we run RANSAC for each segment for segmentIndex in group: segmentMask = segmentation == segmentIndex segmentIndices = segmentMask.nonzero()[0] XYZ = allXYZ[segmentMask.reshape(-1)] if len(XYZ) == 0: continue numPoints = XYZ.shape[0] segmentPlanes = [] segmentPlanePointIndices = [] for c in range(2): if c == 0: ## First try to fit one plane to see if the entire segment is one plane try: plane = fitPlane(XYZ) except: continue diff = np.abs( np.matmul(XYZ, plane) - np.ones(XYZ.shape[0])) / np.linalg.norm(plane) if diff.mean() < fittingErrorThreshold: segmentPlanes.append(plane) segmentPlanePointIndices.append(segmentIndices) break else: ## Run ransac for planeIndex in range(numPlanesPerSegment): if len(XYZ) < planeAreaThreshold: continue bestPlaneInfo = [None, 0, None] for iteration in range(min(XYZ.shape[0], numIterations)): sampledPoints = XYZ[np.random.choice( np.arange(XYZ.shape[0]), size=(3), replace=False)] try: plane = fitPlane(sampledPoints) except: continue diff = np.abs( np.matmul(XYZ, plane) - np.ones(XYZ.shape[0])) / np.linalg.norm(plane) inlierMask = diff < planeDiffThreshold numInliers = inlierMask.sum() if numInliers > bestPlaneInfo[1]: bestPlaneInfo = [plane, numInliers, inlierMask] if bestPlaneInfo[1] < planeAreaThreshold: break pointIndices = segmentIndices[bestPlaneInfo[2]] bestPlane = fitPlane(XYZ[bestPlaneInfo[2]]) segmentPlanes.append(bestPlane) segmentPlanePointIndices.append(pointIndices) outlierMask = np.logical_not(bestPlaneInfo[2]) segmentIndices = segmentIndices[outlierMask] XYZ = XYZ[outlierMask] continue continue if sum([len(indices) for indices in segmentPlanePointIndices ]) < numPoints * 0.5: print('not enough fitted points') if len(segmentIndices) >= planeAreaThreshold: groupPlanes.append(np.zeros(3)) groupPlanePointIndices.append(segmentIndices) groupPlaneSegments.append(set([segmentIndex])) else: groupPlanes += segmentPlanes groupPlanePointIndices += segmentPlanePointIndices for _ in range(len(segmentPlanes)): groupPlaneSegments.append(set([segmentIndex])) continue if len(groupPlanes) > 0: ## Merge planes of each instance groupPlanes = mergePlanesNew(points, groupPlanes, groupPlanePointIndices, groupPlaneSegments, segmentNeighbors, numPlanes=(minNumPlanes, maxNumPlanes), planeDiffThreshold=planeDiffThreshold, planeAreaThreshold=planeAreaThreshold, debug=debugIndex != -1) planeGroups.append(groupPlanes) if debug: colorMap = ColorPalette(segmentation.max() + 2).getColorMap() colorMap[-1] = 0 colorMap[-2] = 255 annotationFolder = 'test/' else: numPlanes = sum([len(group['groupedPlanes']) for group in planeGroups]) segmentationColor = (np.arange(numPlanes) + 1) * 100 colorMap = np.stack([ segmentationColor / (256 * 256), segmentationColor / 256 % 256, segmentationColor % 256 ], axis=1) colorMap[-1] = 255 annotationFolder = os.path.join(OUTPUT_FOLDER, scene_id, 'plane_annotation') os.makedirs(annotationFolder, exist_ok=True) if debug: colors = colorMap[segmentation] writePointCloudFace(annotationFolder + '/segments.ply', np.concatenate([points, colors], axis=-1), faces) groupedSegmentation = np.full(segmentation.shape, fill_value=-1) for segmentIndex in range(len(aggregation)): indices = aggregation[segmentIndex]['segments'] for index in indices: groupedSegmentation[segmentation == index] = segmentIndex continue continue groupedSegmentation = groupedSegmentation.astype(np.int32) colors = colorMap[groupedSegmentation] writePointCloudFace(annotationFolder + '/groups.ply', np.concatenate([points, colors], axis=-1), faces) planes = [] planePointIndices = [] for index, group in enumerate(planeGroups): planes.extend(group['groupedPlanes']) planePointIndices.extend(group['groupedPlanePointIndices']) # Filter out non-planar regions, planeSegmentation consists of all plane id, -1 means non-planar. filtered_plane[plane_id] -> plane params filtered_planes = [] planeSegmentation = np.full(segmentation.shape, fill_value=-1, dtype=np.int32) tmp_idx = 0 for planeIndex, (planePoints, plane_param) in enumerate(zip(planePointIndices, planes)): if np.linalg.norm(planes[planeIndex]) < 1e-4: pass else: planeSegmentation[planePoints] = tmp_idx filtered_planes.append(plane_param) tmp_idx += 1 assert (len(filtered_planes) == tmp_idx) planes = filtered_planes if debug: groupSegmentation = np.full(segmentation.shape, fill_value=-1, dtype=np.int32) structureSegmentation = np.full(segmentation.shape, fill_value=-1, dtype=np.int32) typeSegmentation = np.full(segmentation.shape, fill_value=-1, dtype=np.int32) for planeIndex, planePoints in enumerate(planePointIndices): if len(planeInfo[planeIndex]) > 1: structureSegmentation[planePoints] = planeInfo[planeIndex][1][ 0] typeSegmentation[planePoints] = np.maximum( typeSegmentation[planePoints], planeInfo[planeIndex][1][1] - 2) pass groupSegmentation[planePoints] = planeInfo[planeIndex][0][0] continue colors = colorMap[groupSegmentation] writePointCloudFace(annotationFolder + '/group.ply', np.concatenate([points, colors], axis=-1), faces) colors = colorMap[structureSegmentation] writePointCloudFace(annotationFolder + '/structure.ply', np.concatenate([points, colors], axis=-1), faces) colors = colorMap[typeSegmentation] writePointCloudFace(annotationFolder + '/type.ply', np.concatenate([points, colors], axis=-1), faces) pass planes = np.array(planes) planesD = 1.0 / np.maximum(np.linalg.norm(planes, axis=-1, keepdims=True), 1e-4) planes *= pow(planesD, 2) # mp3d to ai habitat planes = mp3d2habitat(planes) ## Remove boundary faces for rendering purpose removeIndices = [] for faceIndex in range(faces.shape[0]): face = faces[faceIndex] segment_1 = planeSegmentation[face[0]] segment_2 = planeSegmentation[face[1]] segment_3 = planeSegmentation[face[2]] if segment_1 != segment_2 or segment_1 != segment_3: removeIndices.append(faceIndex) pass continue faces = np.delete(faces, removeIndices) colors = colorMap[planeSegmentation] cmap = colormap() / 255 new_color_plane = cmap[planeSegmentation % len(cmap)] new_color_plane[planeSegmentation == -2] = [1, 1, 1] new_color_plane[planeSegmentation == -1] = [1, 1, 1] mesh_o3d.vertex_colors = o3d.utility.Vector3dVector(new_color_plane) o3d.io.write_triangle_mesh( os.path.join(annotationFolder, region_id + ".gltf"), mesh_o3d) np.save(os.path.join(annotationFolder, region_id + "_planeid.npy"), planeSegmentation) if debug: exit(1) pass np.save(os.path.join(annotationFolder, region_id + ".npy"), planes)