|
| 1 | +#!/usr/bin/env python |
| 2 | + |
| 3 | +import glooey |
| 4 | +import imgviz |
| 5 | +import numpy as np |
| 6 | +import pyglet |
| 7 | +import trimesh |
| 8 | +import trimesh.transformations as tf |
| 9 | +import trimesh.viewer |
| 10 | + |
| 11 | +import octomap |
| 12 | + |
| 13 | + |
| 14 | +def pointcloud_from_depth(depth, fx, fy, cx, cy): |
| 15 | + assert depth.dtype.kind == 'f', 'depth must be float and have meter values' |
| 16 | + |
| 17 | + rows, cols = depth.shape |
| 18 | + c, r = np.meshgrid(np.arange(cols), np.arange(rows), sparse=True) |
| 19 | + valid = ~np.isnan(depth) |
| 20 | + z = np.where(valid, depth, np.nan) |
| 21 | + x = np.where(valid, z * (c - cx) / fx, np.nan) |
| 22 | + y = np.where(valid, z * (r - cy) / fy, np.nan) |
| 23 | + pc = np.dstack((x, y, z)) |
| 24 | + |
| 25 | + return pc |
| 26 | + |
| 27 | + |
| 28 | +def labeled_scene_widget(scene, label): |
| 29 | + vbox = glooey.VBox() |
| 30 | + vbox.add(glooey.Label(text=label, color=(255, 255, 255)), size=0) |
| 31 | + vbox.add(trimesh.viewer.SceneWidget(scene)) |
| 32 | + return vbox |
| 33 | + |
| 34 | + |
| 35 | +def visualize( |
| 36 | + occupied, empty, K, width, height, rgb, pcd, mask, resolution, aabb |
| 37 | +): |
| 38 | + window = pyglet.window.Window( |
| 39 | + width=int(640 * 0.9 * 3), height=int(480 * 0.9) |
| 40 | + ) |
| 41 | + |
| 42 | + @window.event |
| 43 | + def on_key_press(symbol, modifiers): |
| 44 | + if modifiers == 0: |
| 45 | + if symbol == pyglet.window.key.Q: |
| 46 | + window.on_close() |
| 47 | + |
| 48 | + gui = glooey.Gui(window) |
| 49 | + hbox = glooey.HBox() |
| 50 | + hbox.set_padding(5) |
| 51 | + |
| 52 | + camera = trimesh.scene.Camera( |
| 53 | + resolution=(width, height), focal=(K[0, 0], K[1, 1]) |
| 54 | + ) |
| 55 | + camera_marker = trimesh.creation.camera_marker(camera, marker_height=0.1) |
| 56 | + |
| 57 | + # initial camera pose |
| 58 | + camera_transform = np.array( |
| 59 | + [ |
| 60 | + [0.73256052, -0.28776419, 0.6168848, 0.66972396], |
| 61 | + [-0.26470017, -0.95534823, -0.13131483, -0.12390466], |
| 62 | + [0.62712751, -0.06709345, -0.77602162, -0.28781298], |
| 63 | + [0.0, 0.0, 0.0, 1.0], |
| 64 | + ], |
| 65 | + ) |
| 66 | + |
| 67 | + aabb_min, aabb_max = aabb |
| 68 | + bbox = trimesh.path.creation.box_outline( |
| 69 | + aabb_max - aabb_min, |
| 70 | + tf.translation_matrix((aabb_min + aabb_max) / 2), |
| 71 | + ) |
| 72 | + |
| 73 | + geom = trimesh.PointCloud(vertices=pcd[mask], colors=rgb[mask]) |
| 74 | + scene = trimesh.Scene(camera=camera, geometry=[bbox, geom, camera_marker]) |
| 75 | + scene.camera_transform = camera_transform |
| 76 | + hbox.add(labeled_scene_widget(scene, label='pointcloud')) |
| 77 | + |
| 78 | + geom = trimesh.voxel.ops.multibox( |
| 79 | + occupied, pitch=resolution, colors=[1.0, 0, 0, 0.5] |
| 80 | + ) |
| 81 | + scene = trimesh.Scene(camera=camera, geometry=[bbox, geom, camera_marker]) |
| 82 | + scene.camera_transform = camera_transform |
| 83 | + hbox.add(labeled_scene_widget(scene, label='occupied')) |
| 84 | + |
| 85 | + geom = trimesh.voxel.ops.multibox( |
| 86 | + empty, pitch=resolution, colors=[0.5, 0.5, 0.5, 0.5] |
| 87 | + ) |
| 88 | + scene = trimesh.Scene(camera=camera, geometry=[bbox, geom, camera_marker]) |
| 89 | + scene.camera_transform = camera_transform |
| 90 | + hbox.add(labeled_scene_widget(scene, label='empty')) |
| 91 | + |
| 92 | + gui.add(hbox) |
| 93 | + pyglet.app.run() |
| 94 | + |
| 95 | + |
| 96 | +def main(): |
| 97 | + data = imgviz.data.arc2017() |
| 98 | + camera_info = data['camera_info'] |
| 99 | + K = np.array(camera_info['K']).reshape(3, 3) |
| 100 | + rgb = data['rgb'] |
| 101 | + pcd = pointcloud_from_depth( |
| 102 | + data['depth'], fx=K[0, 0], fy=K[1, 1], cx=K[0, 2], cy=K[1, 2] |
| 103 | + ) |
| 104 | + |
| 105 | + nonnan = ~np.isnan(pcd).any(axis=2) |
| 106 | + mask = np.less(pcd[:, :, 2], 2) |
| 107 | + |
| 108 | + resolution = 0.01 |
| 109 | + octree = octomap.OcTree(resolution) |
| 110 | + octree.insertPointCloud( |
| 111 | + pointcloud=pcd[nonnan], |
| 112 | + origin=np.array([0, 0, 0], dtype=float), |
| 113 | + maxrange=2, |
| 114 | + ) |
| 115 | + occupied, empty = octree.extractPointCloud() |
| 116 | + |
| 117 | + aabb_min = octree.getMetricMin() |
| 118 | + aabb_max = octree.getMetricMax() |
| 119 | + |
| 120 | + visualize( |
| 121 | + occupied=occupied, |
| 122 | + empty=empty, |
| 123 | + K=K, |
| 124 | + width=camera_info['width'], |
| 125 | + height=camera_info['height'], |
| 126 | + rgb=rgb, |
| 127 | + pcd=pcd, |
| 128 | + mask=mask, |
| 129 | + resolution=resolution, |
| 130 | + aabb=(aabb_min, aabb_max), |
| 131 | + ) |
| 132 | + |
| 133 | + |
| 134 | +if __name__ == '__main__': |
| 135 | + main() |
0 commit comments