Spaces:
Running
Running
| import torch | |
| import cv2 | |
| import lietorch | |
| import droid_backends | |
| import time | |
| import argparse | |
| import numpy as np | |
| # import os | |
| # os.environ['PYOPENGL_PLATFORM'] = 'egl' | |
| #os.environ['PYOPENGL_PLATFORM'] = 'osmesa' | |
| import open3d as o3d | |
| # o3d.visualization.webrtc_server.enable_webrtc() | |
| from lietorch import SE3 | |
| import geom.projective_ops as pops | |
| CAM_POINTS = np.array([ | |
| [ 0, 0, 0], | |
| [-1, -1, 1.5], | |
| [ 1, -1, 1.5], | |
| [ 1, 1, 1.5], | |
| [-1, 1, 1.5], | |
| [-0.5, 1, 1.5], | |
| [ 0.5, 1, 1.5], | |
| [ 0, 1.2, 1.5]]) | |
| CAM_LINES = np.array([ | |
| [1,2], [2,3], [3,4], [4,1], [1,0], [0,2], [3,0], [0,4], [5,7], [7,6]]) | |
| def white_balance(img): | |
| # from https://stackoverflow.com/questions/46390779/automatic-white-balancing-with-grayworld-assumption | |
| result = cv2.cvtColor(img, cv2.COLOR_BGR2LAB) | |
| avg_a = np.average(result[:, :, 1]) | |
| avg_b = np.average(result[:, :, 2]) | |
| result[:, :, 1] = result[:, :, 1] - ((avg_a - 128) * (result[:, :, 0] / 255.0) * 1.1) | |
| result[:, :, 2] = result[:, :, 2] - ((avg_b - 128) * (result[:, :, 0] / 255.0) * 1.1) | |
| result = cv2.cvtColor(result, cv2.COLOR_LAB2BGR) | |
| return result | |
| def create_camera_actor(g, scale=0.05): | |
| """ build open3d camera polydata """ | |
| camera_actor = o3d.geometry.LineSet( | |
| points=o3d.utility.Vector3dVector(scale * CAM_POINTS), | |
| lines=o3d.utility.Vector2iVector(CAM_LINES)) | |
| color = (g * 1.0, 0.5 * (1-g), 0.9 * (1-g)) | |
| camera_actor.paint_uniform_color(color) | |
| return camera_actor | |
| def create_point_actor(points, colors): | |
| """ open3d point cloud from numpy array """ | |
| point_cloud = o3d.geometry.PointCloud() | |
| point_cloud.points = o3d.utility.Vector3dVector(points) | |
| point_cloud.colors = o3d.utility.Vector3dVector(colors) | |
| return point_cloud | |
| def droid_visualization(video, save_path, device="cuda:0"): | |
| """ DROID visualization frontend """ | |
| torch.cuda.set_device(0) | |
| droid_visualization.video = video | |
| droid_visualization.cameras = {} | |
| droid_visualization.points = {} | |
| droid_visualization.warmup = 8 | |
| droid_visualization.scale = 1.0 | |
| droid_visualization.ix = 0 | |
| print("headless droid_visualization") | |
| droid_visualization.filter_thresh = 0.3 #0.005 | |
| def increase_filter(vis): | |
| droid_visualization.filter_thresh *= 2 | |
| with droid_visualization.video.get_lock(): | |
| droid_visualization.video.dirty[:droid_visualization.video.counter.value] = True | |
| def decrease_filter(vis): | |
| droid_visualization.filter_thresh *= 0.5 | |
| with droid_visualization.video.get_lock(): | |
| droid_visualization.video.dirty[:droid_visualization.video.counter.value] = True | |
| def animation_callback(vis): | |
| cam = vis.get_view_control().convert_to_pinhole_camera_parameters() | |
| with torch.no_grad(): | |
| with video.get_lock(): | |
| t = video.counter.value | |
| dirty_index, = torch.where(video.dirty.clone()) | |
| dirty_index = dirty_index | |
| if len(dirty_index) == 0: | |
| return | |
| video.dirty[dirty_index] = False | |
| # convert poses to 4x4 matrix | |
| poses = torch.index_select(video.poses, 0, dirty_index) | |
| disps = torch.index_select(video.disps, 0, dirty_index) | |
| Ps = SE3(poses).inv().matrix().cpu().numpy() | |
| images = torch.index_select(video.images, 0, dirty_index) | |
| images = images.cpu()[:,[2,1,0],3::8,3::8].permute(0,2,3,1) / 255.0 | |
| points = droid_backends.iproj(SE3(poses).inv().data, disps, video.intrinsics[0]).cpu() | |
| thresh = droid_visualization.filter_thresh * torch.ones_like(disps.mean(dim=[1,2])) | |
| count = droid_backends.depth_filter( | |
| video.poses, video.disps, video.intrinsics[0], dirty_index, thresh) | |
| count = count.cpu() | |
| disps = disps.cpu() | |
| masks = ((count >= 2) & (disps > .5*disps.mean(dim=[1,2], keepdim=True))) | |
| for i in range(len(dirty_index)): | |
| pose = Ps[i] | |
| ix = dirty_index[i].item() | |
| if ix in droid_visualization.cameras: | |
| vis.remove_geometry(droid_visualization.cameras[ix]) | |
| del droid_visualization.cameras[ix] | |
| if ix in droid_visualization.points: | |
| vis.remove_geometry(droid_visualization.points[ix]) | |
| del droid_visualization.points[ix] | |
| ### add camera actor ### | |
| cam_actor = create_camera_actor(True) | |
| cam_actor.transform(pose) | |
| vis.add_geometry(cam_actor) | |
| droid_visualization.cameras[ix] = cam_actor | |
| mask = masks[i].reshape(-1) | |
| pts = points[i].reshape(-1, 3)[mask].cpu().numpy() | |
| clr = images[i].reshape(-1, 3)[mask].cpu().numpy() | |
| ## add point actor ### | |
| point_actor = create_point_actor(pts, clr) | |
| vis.add_geometry(point_actor) | |
| droid_visualization.points[ix] = point_actor | |
| ### Hack to save Point Cloud Data and Camnera results ### | |
| # Save points | |
| pcd_points = o3d.geometry.PointCloud() | |
| for p in droid_visualization.points.items(): | |
| pcd_points += p[1] | |
| o3d.io.write_point_cloud(f"{save_path}/points.ply", pcd_points, write_ascii=False) | |
| # Save pose | |
| pcd_camera = create_camera_actor(True) | |
| for c in droid_visualization.cameras.items(): | |
| pcd_camera += c[1] | |
| o3d.io.write_line_set(f"{save_path}/camera.ply", pcd_camera, write_ascii=False) | |
| ### end ### | |
| # hack to allow interacting with vizualization during inference | |
| if len(droid_visualization.cameras) >= droid_visualization.warmup: | |
| cam = vis.get_view_control().convert_from_pinhole_camera_parameters(cam) | |
| droid_visualization.ix += 1 | |
| vis.poll_events() | |
| vis.update_renderer() | |
| ### create Open3D visualization ### | |
| vis = o3d.visualization.VisualizerWithKeyCallback() | |
| vis.register_animation_callback(animation_callback) | |
| vis.register_key_callback(ord("S"), increase_filter) | |
| vis.register_key_callback(ord("A"), decrease_filter) | |
| vis.create_window(height=540, width=960) | |
| # vis.create_window(height=512, width=384) | |
| vis.get_render_option().load_from_json("thirdparty/DROID-SLAM//misc/renderoption.json") | |
| vis.run() | |
| vis.destroy_window() | |