Initial snapshot: hand motion pipeline (Dyn-HaMR + dex-retargeting + SPIDER)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-09-15 14:13:49 +08:00
commit 7e4ef6f98b
7002 changed files with 1452249 additions and 0 deletions
@@ -0,0 +1,89 @@
+import torch
+import lietorch
+import numpy as np
+
+from droid_net import DroidNet
+from depth_video import DepthVideo
+from motion_filter import MotionFilter
+from droid_frontend import DroidFrontend
+from droid_backend import DroidBackend
+from trajectory_filler import PoseTrajectoryFiller
+
+from collections import OrderedDict
+from torch.multiprocessing import Process
+
+
+class Droid:
+    def __init__(self, args):
+        super(Droid, self).__init__()
+        self.load_weights(args.weights)
+        self.args = args
+        self.disable_vis = args.disable_vis
+
+        # store images, depth, poses, intrinsics (shared between processes)
+        self.video = DepthVideo(args.image_size, args.buffer, stereo=args.stereo)
+
+        # filter incoming frames so that there is enough motion
+        self.filterx = MotionFilter(self.net, self.video, thresh=args.filter_thresh)
+
+        # frontend process
+        self.frontend = DroidFrontend(self.net, self.video, self.args)
+        
+        # backend process
+        self.backend = DroidBackend(self.net, self.video, self.args)
+
+        # visualizer
+        if not self.disable_vis:
+            from visualization import droid_visualization
+            self.visualizer = Process(target=droid_visualization, args=(self.video,))
+            self.visualizer.start()
+
+        # post processor - fill in poses for non-keyframes
+        self.traj_filler = PoseTrajectoryFiller(self.net, self.video)
+
+
+    def load_weights(self, weights):
+        """ load trained model weights """
+
+        print(weights)
+        self.net = DroidNet()
+        state_dict = OrderedDict([
+            (k.replace("module.", ""), v) for (k, v) in torch.load(weights).items()])
+
+        state_dict["update.weight.2.weight"] = state_dict["update.weight.2.weight"][:2]
+        state_dict["update.weight.2.bias"] = state_dict["update.weight.2.bias"][:2]
+        state_dict["update.delta.2.weight"] = state_dict["update.delta.2.weight"][:2]
+        state_dict["update.delta.2.bias"] = state_dict["update.delta.2.bias"][:2]
+
+        self.net.load_state_dict(state_dict)
+        self.net.to("cuda:0").eval()
+
+    def track(self, tstamp, image, depth=None, intrinsics=None):
+        """ main thread - update map """
+
+        with torch.no_grad():
+            # check there is enough motion
+            self.filterx.track(tstamp, image, depth, intrinsics)
+
+            # local bundle adjustment
+            self.frontend()
+
+            # global bundle adjustment
+            # self.backend()
+
+    def terminate(self, stream=None):
+        """ terminate the visualization process, return poses [t, q] """
+
+        del self.frontend
+
+        torch.cuda.empty_cache()
+        print("#" * 32)
+        self.backend(7)
+
+        torch.cuda.empty_cache()
+        print("#" * 32)
+        self.backend(12)
+
+        camera_trajectory = self.traj_filler(stream)
+        return camera_trajectory.inv().data.cpu().numpy()
+