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)