Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,154 @@
|
||||
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
|
||||
|
||||
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()
|
||||
Reference in New Issue
Block a user