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

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

third_party/Dyn-HaMR/dyn-hamr/util/tensor.py

@@ -0,0 +1,139 @@
+import torch
+import numpy as np
+
+
+def move_to(obj, device):
+    if isinstance(obj, torch.Tensor):
+        return obj.to(device)
+    if isinstance(obj, dict):
+        return {k: move_to(v, device) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [move_to(x, device) for x in obj]
+    return obj  # otherwise do nothing
+
+
+def detach_all(obj):
+    if isinstance(obj, torch.Tensor):
+        return obj.detach()
+    if isinstance(obj, dict):
+        return {k: detach_all(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [detach_all(x) for x in obj]
+    return obj  # otherwise do nothing
+
+
+def to_torch(obj):
+    if isinstance(obj, np.ndarray):
+        return torch.from_numpy(obj).float()
+    if isinstance(obj, dict):
+        return {k: to_torch(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [to_torch(x) for x in obj]
+    return obj
+
+
+def to_np(obj):
+    if isinstance(obj, torch.Tensor):
+        return obj.numpy()
+    if isinstance(obj, dict):
+        return {k: to_np(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [to_np(x) for x in obj]
+    return obj
+
+
+def get_device(i=0):
+    device = f"cuda:{i}" if torch.cuda.is_available() else "cpu"
+    return torch.device(device)
+
+
+def scatter_intervals(tensor, start, end, T=None):
+    """
+    Scatter the tensor contents into intervals from start to end
+    output tensor indexed from 0 to end.max()
+    :param tensor (B, S, *)
+    :param start (B) start indices
+    :param end (B) end indices
+    :param T (int, optional) max length
+    returns (B, T, *) scattered tensor
+    """
+    assert isinstance(tensor, torch.Tensor) and tensor.ndim >= 2
+    if T is not None:
+        end = torch.where(end < 0, T + 1 + end, end)
+        assert torch.all(end <= T)
+    else:
+        T = end.max()
+
+    B, S, *dims = tensor.shape
+    start, end = start.long(), end.long()
+    # get idcs that go past the last time step so we don't have repeat indices in scatter
+    idcs = time_segment_idcs(start, end, clip=False)  # (B, S)
+    # mask out the extra padding
+    mask = idcs >= end[:, None]
+    tensor[mask] = 0
+
+    idcs = idcs.reshape(B, S, *(1,) * len(dims)).repeat(1, 1, *dims)
+    output = torch.zeros(
+        B, idcs.max() + 1, *dims, device=tensor.device, dtype=tensor.dtype
+    )
+    output.scatter_(1, idcs, tensor)
+    # slice out the extra segments
+    return output[:, :T]
+
+
+def get_scatter_mask(start, end, T):
+    """
+    get the mask of selected intervals
+    """
+    B = start.shape[0]
+    start, end = start.long(), end.long()
+    assert torch.all(end <= T)
+    idcs = time_segment_idcs(start, end, clip=True)
+    mask = torch.zeros(B, T, device=start.device, dtype=torch.bool)
+    mask.scatter_(1, idcs, 1)
+    return mask
+
+
+def select_intervals(series, start, end):
+    """
+    Select slices of a tensor from start to end
+    will pad uneven sequences to all the max segment length
+    :param series (B, T, *)
+    :param start (B)
+    :param end (B)
+    returns (B, S, *) selected segments, S = max(end - start)
+    """
+    B, T, *dims = series.shape
+    assert torch.all(end <= T)
+    sel = time_segment_idcs(start, end, clip=True)
+    S = sel.shape[1]
+    sel = sel.reshape(B, S, *(1,) * len(dims)).repeat(1, 1, *dims)
+    return torch.gather(series, 1, sel)
+
+
+def get_select_mask(start, end):
+    """
+    get the mask of unpadded elementes for the selected time segments
+    e.g. sel[mask] are the unpadded elements
+    :param start (B)
+    :param end (B)
+    """
+    idcs = time_segment_idcs(start, end, clip=False)
+    return idcs < end[:, None]  # (B, S)
+
+
+def time_segment_idcs(start, end, clip=True):
+    """
+    :param start (B)
+    :param end (B)
+    returns (B, S) long tensor of indices, where S = max(end - start)
+    """
+    start, end = start.long(), end.long()
+    S = (end - start).max()
+    seg = torch.arange(S, dtype=torch.int64, device=start.device)
+    idcs = start[:, None] + seg[None, :]  # (B, S)
+    if clip:
+        # clip at the lengths of each track
+        imax = end[:, None] - 1
+        idcs = idcs.clamp(max=imax)
+    return idcs