Spaces:
Running
on
Zero
Running
on
Zero
| import numpy as np | |
| import matplotlib | |
| def visualize_depth(depth: np.ndarray, | |
| depth_min=None, | |
| depth_max=None, | |
| percentile=5, | |
| ret_minmax=False, | |
| cmap='Spectral'): | |
| if depth_min is None: depth_min = np.percentile(depth, percentile) | |
| if depth_max is None: depth_max = np.percentile(depth, 100 - percentile) | |
| if depth_min == depth_max: | |
| depth_min = depth_min - 1e-6 | |
| depth_max = depth_max + 1e-6 | |
| cm = matplotlib.colormaps[cmap] | |
| depth = ((depth - depth_min) / (depth_max - depth_min)).clip(0, 1) | |
| img_colored_np = cm(depth[None], bytes=False)[:, :, :, 0:3] # value from 0 to 1 | |
| img_colored_np = (img_colored_np[0] * 255.0).astype(np.uint8) | |
| if ret_minmax: | |
| return img_colored_np, depth_min, depth_max | |
| else: | |
| return img_colored_np | |
| def unproject_depth(depth, | |
| ixt, | |
| depth_min=0.01, | |
| depth_max=None, | |
| color=None, | |
| ext=None, | |
| conf=None, | |
| ret_pcd=False, | |
| clip_box=None): | |
| height, width = depth.shape | |
| x = np.arange(0, width) | |
| y = np.arange(0, height) | |
| xx, yy = np.meshgrid(x, y) | |
| xx = xx.reshape(-1) | |
| yy = yy.reshape(-1) | |
| zz = depth.reshape(-1) | |
| mask = np.ones_like(xx, dtype=np.bool_) | |
| if depth_min is not None: | |
| mask &= zz >= depth_min | |
| if depth_max is not None: | |
| mask &= zz <= depth_max | |
| if conf is not None: | |
| mask &= conf.reshape(-1) == 2 | |
| xx = xx[mask] | |
| yy = yy[mask] | |
| zz = zz[mask] | |
| pcd = np.stack([xx, yy, np.ones_like(xx)], axis=1) | |
| pcd = pcd * zz[:, None] | |
| pcd = np.dot(pcd, np.linalg.inv(ixt).T) | |
| if ext is not None: | |
| pcd = np.concatenate([pcd, np.ones((pcd.shape[0], 1))], axis=1) | |
| pcd = np.dot(pcd, np.linalg.inv(ext).T) | |
| new_mask = np.ones_like(pcd[:, 0]).astype(np.bool_) | |
| if clip_box is not None: | |
| assert len(clip_box) == 6 | |
| for i, val in enumerate(clip_box): | |
| if val is None: | |
| continue | |
| if i == 0: new_mask &= (pcd[:, 0] <= val) | |
| elif i == 1: new_mask &= (pcd[:, 1] <= val) | |
| elif i == 2: new_mask &= (pcd[:, 2] <= val) | |
| elif i == 3: new_mask &= (pcd[:, 0] >= val) | |
| elif i == 4: new_mask &= (pcd[:, 1] >= val) | |
| elif i == 5: new_mask &= (pcd[:, 2] >= val) | |
| if color is not None: | |
| if color.dtype == np.uint8: | |
| color = color.astype(np.float32) / 255. | |
| if ret_pcd: | |
| points = pcd | |
| import open3d as o3d | |
| pcd = o3d.geometry.PointCloud() | |
| pcd.points = o3d.utility.Vector3dVector(points[:, :3][new_mask]) | |
| pcd.colors = o3d.utility.Vector3dVector(color.reshape(-1, 3)[mask][new_mask]) | |
| else: | |
| return pcd[:, :3][new_mask], color.reshape(-1, 3)[mask][new_mask] | |
| else: | |
| if ret_pcd: | |
| import open3d as o3d | |
| points = pcd | |
| pcd = o3d.geometry.PointCloud() | |
| pcd.points = o3d.utility.Vector3dVector(pcd[:, :3][new_mask]) | |
| else: | |
| return pcd[:, :3][new_mask] | |
| return pcd | |
| if __name__ == '__main__': | |
| depth = np.random.rand(100, 100) | |
| visualize_depth(depth) |