import torch
import cv2
import lietorch
import droid_backends
import time
import argparse
import numpy as np
import open3d as o3d
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, device="cuda:0"):
""" DROID visualization frontend """
torch.cuda.set_device(device)
droid_visualization.video = video
droid_visualization.cameras = {}
droid_visualization.points = {}
droid_visualization.warmup = 8
droid_visualization.scale = 1.0
droid_visualization.ix = 0
droid_visualization.filter_thresh = 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
#file dialog based pointcloud export added#
def export_pointcloud(vis):
gui.Application.instance.initialize()
window = gui.Application.instance.create_window("Export", 350, 600)
def _on_filedlg_cancel():
window.close_dialog()
window.close()
gui.Application.instance.quit()
def _on_filedlg_done(path):
pcd_export(path)
window.close_dialog()
gui.Application.instance.quit()
def exec_file_dialog():
filedlg = gui.FileDialog(gui.FileDialog.SAVE, "Select file", window.theme)
filedlg.add_filter(".ply .xyz .pcd", "PointCloud (.xyz .ply .pcd)")
filedlg.add_filter("", "All files")
filedlg.set_on_cancel(_on_filedlg_cancel)
filedlg.set_on_done(_on_filedlg_done)
window.show_dialog(filedlg)
def pcd_export(path):
print("\nExporting pointcloud as", path)
final_pcd = o3d.geometry.PointCloud()
for p in droid_visualization.points.items():
final_pcd += p[1]
o3d.io.write_point_cloud(path, final_pcd, write_ascii=False)
#vis.capture_depth_point_cloud("/home/bertuser/droidslam_export.ply")
exec_file_dialog()
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 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.get_render_option().load_from_json("misc/renderoption.json")
vis.run()
vis.destroy_window()