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,222 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from collections import OrderedDict
+
+from modules.extractor import BasicEncoder
+from modules.corr import CorrBlock
+from modules.gru import ConvGRU
+from modules.clipping import GradientClip
+
+from lietorch import SE3
+from geom.ba import BA
+
+import geom.projective_ops as pops
+from geom.graph_utils import graph_to_edge_list, keyframe_indicies
+
+from torch_scatter import scatter_mean
+
+
+def cvx_upsample(data, mask):
+    """ upsample pixel-wise transformation field """
+    batch, ht, wd, dim = data.shape
+    data = data.permute(0, 3, 1, 2)
+    mask = mask.view(batch, 1, 9, 8, 8, ht, wd)
+    mask = torch.softmax(mask, dim=2)
+
+    up_data = F.unfold(data, [3,3], padding=1)
+    up_data = up_data.view(batch, dim, 9, 1, 1, ht, wd)
+
+    up_data = torch.sum(mask * up_data, dim=2)
+    up_data = up_data.permute(0, 4, 2, 5, 3, 1)
+    up_data = up_data.reshape(batch, 8*ht, 8*wd, dim)
+
+    return up_data
+
+def upsample_disp(disp, mask):
+    batch, num, ht, wd = disp.shape
+    disp = disp.view(batch*num, ht, wd, 1)
+    mask = mask.view(batch*num, -1, ht, wd)
+    return cvx_upsample(disp, mask).view(batch, num, 8*ht, 8*wd)
+
+
+class GraphAgg(nn.Module):
+    def __init__(self):
+        super(GraphAgg, self).__init__()
+        self.conv1 = nn.Conv2d(128, 128, 3, padding=1)
+        self.conv2 = nn.Conv2d(128, 128, 3, padding=1)
+        self.relu = nn.ReLU(inplace=True)
+
+        self.eta = nn.Sequential(
+            nn.Conv2d(128, 1, 3, padding=1),
+            GradientClip(),
+            nn.Softplus())
+
+        self.upmask = nn.Sequential(
+            nn.Conv2d(128, 8*8*9, 1, padding=0))
+
+    def forward(self, net, ii):
+        batch, num, ch, ht, wd = net.shape
+        net = net.view(batch*num, ch, ht, wd)
+
+        _, ix = torch.unique(ii, return_inverse=True)
+        net = self.relu(self.conv1(net))
+
+        net = net.view(batch, num, 128, ht, wd)
+        net = scatter_mean(net, ix, dim=1)
+        net = net.view(-1, 128, ht, wd)
+
+        net = self.relu(self.conv2(net))
+
+        eta = self.eta(net).view(batch, -1, ht, wd)
+        upmask = self.upmask(net).view(batch, -1, 8*8*9, ht, wd)
+
+        return .01 * eta, upmask
+
+
+class UpdateModule(nn.Module):
+    def __init__(self):
+        super(UpdateModule, self).__init__()
+        cor_planes = 4 * (2*3 + 1)**2
+
+        self.corr_encoder = nn.Sequential(
+            nn.Conv2d(cor_planes, 128, 1, padding=0),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True))
+
+        self.flow_encoder = nn.Sequential(
+            nn.Conv2d(4, 128, 7, padding=3),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 64, 3, padding=1),
+            nn.ReLU(inplace=True))
+
+        self.weight = nn.Sequential(
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 2, 3, padding=1),
+            GradientClip(),
+            nn.Sigmoid())
+
+        self.delta = nn.Sequential(
+            nn.Conv2d(128, 128, 3, padding=1),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(128, 2, 3, padding=1),
+            GradientClip())
+
+        self.gru = ConvGRU(128, 128+128+64)
+        self.agg = GraphAgg()
+
+    def forward(self, net, inp, corr, flow=None, ii=None, jj=None):
+        """ RaftSLAM update operator """
+
+        batch, num, ch, ht, wd = net.shape
+
+        if flow is None:
+            flow = torch.zeros(batch, num, 4, ht, wd, device=net.device)
+
+        output_dim = (batch, num, -1, ht, wd)
+        net = net.view(batch*num, -1, ht, wd)
+        inp = inp.view(batch*num, -1, ht, wd)        
+        corr = corr.view(batch*num, -1, ht, wd)
+        flow = flow.view(batch*num, -1, ht, wd)
+
+        corr = self.corr_encoder(corr)
+        flow = self.flow_encoder(flow)
+        net = self.gru(net, inp, corr, flow)
+
+        ### update variables ###
+        delta = self.delta(net).view(*output_dim)
+        weight = self.weight(net).view(*output_dim)
+
+        delta = delta.permute(0,1,3,4,2)[...,:2].contiguous()
+        weight = weight.permute(0,1,3,4,2)[...,:2].contiguous()
+
+        net = net.view(*output_dim)
+
+        if ii is not None:
+            eta, upmask = self.agg(net, ii.to(net.device))
+            return net, delta, weight, eta, upmask
+
+        else:
+            return net, delta, weight
+
+
+class DroidNet(nn.Module):
+    def __init__(self):
+        super(DroidNet, self).__init__()
+        self.fnet = BasicEncoder(output_dim=128, norm_fn='instance')
+        self.cnet = BasicEncoder(output_dim=256, norm_fn='none')
+        self.update = UpdateModule()
+
+
+    def extract_features(self, images):
+        """ run feeature extraction networks """
+
+        # normalize images
+        images = images[:, :, [2,1,0]] / 255.0
+        mean = torch.as_tensor([0.485, 0.456, 0.406], device=images.device)
+        std = torch.as_tensor([0.229, 0.224, 0.225], device=images.device)
+        images = images.sub_(mean[:, None, None]).div_(std[:, None, None])
+
+        fmaps = self.fnet(images)
+        net = self.cnet(images)
+        
+        net, inp = net.split([128,128], dim=2)
+        net = torch.tanh(net)
+        inp = torch.relu(inp)
+        return fmaps, net, inp
+
+
+    def forward(self, Gs, images, disps, intrinsics, graph=None, num_steps=12, fixedp=2):
+        """ Estimates SE3 or Sim3 between pair of frames """
+
+        u = keyframe_indicies(graph)
+        ii, jj, kk = graph_to_edge_list(graph)
+
+        ii = ii.to(device=images.device, dtype=torch.long)
+        jj = jj.to(device=images.device, dtype=torch.long)
+
+        fmaps, net, inp = self.extract_features(images)
+        net, inp = net[:,ii], inp[:,ii]
+        corr_fn = CorrBlock(fmaps[:,ii], fmaps[:,jj], num_levels=4, radius=3)
+
+        ht, wd = images.shape[-2:]
+        coords0 = pops.coords_grid(ht//8, wd//8, device=images.device)
+        
+        coords1, _ = pops.projective_transform(Gs, disps, intrinsics, ii, jj)
+        target = coords1.clone()
+
+        Gs_list, disp_list, residual_list = [], [], []
+        for step in range(num_steps):
+            Gs = Gs.detach()
+            disps = disps.detach()
+            coords1 = coords1.detach()
+            target = target.detach()
+
+            # extract motion features
+            corr = corr_fn(coords1)
+            resd = target - coords1
+            flow = coords1 - coords0
+
+            motion = torch.cat([flow, resd], dim=-1)
+            motion = motion.permute(0,1,4,2,3).clamp(-64.0, 64.0)
+
+            net, delta, weight, eta, upmask = \
+                self.update(net, inp, corr, motion, ii, jj)
+
+            target = coords1 + delta
+
+            for i in range(2):
+                Gs, disps = BA(target, weight, eta, Gs, disps, intrinsics, ii, jj, fixedp=2)
+
+            coords1, valid_mask = pops.projective_transform(Gs, disps, intrinsics, ii, jj)
+            residual = (target - coords1)
+
+            Gs_list.append(Gs)
+            disp_list.append(upsample_disp(disps, upmask))
+            residual_list.append(valid_mask * residual)
+
+
+        return Gs_list, disp_list, residual_list