liyang
ae28d55f81
Source: RGB-D -> Dyn-HaMR -> L20 retargeting -> FoundationPose -> reference repair -> SPIDER, documented in docs/PIPELINE_LATEST.md and docs/SETUP_AND_WEIGHTS.md. Adds FoundationPose and nvdiffrast upstream snapshots, requirements/pipeline_venv.txt and the FoundationPose weight manifest/downloader. Assets (Git LFS): weights/ (WiLoR detector, HandFlow denoiser, UniDepth-L), FoundationPose checkpoints, HaMeR checkpoint, Dyn-HaMR HMP model and BMC constraints, L20 URDF/meshes, the 20260915_171525 D405 recording and the two box CADs. MANO models are not redistributed (third_party/hamer/_DATA/data/mano/README.txt). Environments, caches and run outputs excluded. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
136 lines
4.2 KiB
Python
136 lines
4.2 KiB
Python
# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
|
|
#
|
|
# NVIDIA CORPORATION and its licensors retain all intellectual property
|
|
# and proprietary rights in and to this software, related documentation
|
|
# and any modifications thereto. Any use, reproduction, disclosure or
|
|
# distribution of this software and related documentation without an express
|
|
# license agreement from NVIDIA CORPORATION is strictly prohibited.
|
|
|
|
|
|
import os,sys
|
|
from dataclasses import dataclass
|
|
from typing import Iterator, List, Optional, Set, Union
|
|
import numpy as np
|
|
import torch
|
|
code_dir = os.path.dirname(os.path.realpath(__file__))
|
|
sys.path.append(f'{code_dir}/../../../../')
|
|
from Utils import *
|
|
|
|
|
|
@dataclass
|
|
class PoseData:
|
|
"""
|
|
rgb: (h, w, 3) uint8
|
|
depth: (bsz, h, w) float32
|
|
bbox: (4, ) int
|
|
K: (3, 3) float32
|
|
"""
|
|
rgb: np.ndarray = None
|
|
bbox: np.ndarray = None
|
|
K: np.ndarray = None
|
|
depth: Optional[np.ndarray] = None
|
|
object_data = None
|
|
mesh_diameter: float = None
|
|
rgbA: np.ndarray = None
|
|
rgbB: np.ndarray = None
|
|
depthA: np.ndarray = None
|
|
depthB: np.ndarray = None
|
|
maskA = None
|
|
maskB = None
|
|
poseA: np.ndarray = None #(4,4)
|
|
target: float = None
|
|
|
|
def __init__(self, rgbA=None, rgbB=None, depthA=None, depthB=None, maskA=None, maskB=None, normalA=None, normalB=None, xyz_mapA=None, xyz_mapB=None, poseA=None, poseB=None, K=None, target=None, mesh_diameter=None, tf_to_crop=None, crop_mask=None, model_pts=None, label=None, model_scale=None):
|
|
self.rgbA = rgbA #(H,W,3) or (H,W*n_view,3) when multiview
|
|
self.rgbB = rgbB
|
|
self.depthA = depthA
|
|
self.depthB = depthB
|
|
self.poseA = poseA
|
|
self.poseB = poseB
|
|
self.maskA = maskA
|
|
self.maskB = maskB
|
|
self.crop_mask = crop_mask
|
|
self.normalA = normalA
|
|
self.normalB = normalB
|
|
self.xyz_mapA = xyz_mapA
|
|
self.xyz_mapB = xyz_mapB
|
|
self.target = target
|
|
self.K = K
|
|
self.mesh_diameter = mesh_diameter
|
|
self.tf_to_crop = tf_to_crop
|
|
self.model_pts = model_pts
|
|
self.label = label
|
|
self.model_scale = model_scale
|
|
|
|
|
|
@dataclass
|
|
class BatchPoseData:
|
|
"""
|
|
rgbs: (bsz, 3, h, w) torch tensor uint8
|
|
depths: (bsz, h, w) float32
|
|
bboxes: (bsz, 4) int
|
|
K: (bsz, 3, 3) float32
|
|
"""
|
|
|
|
rgbs: torch.Tensor = None
|
|
object_datas = None
|
|
bboxes: torch.Tensor = None
|
|
K: torch.Tensor = None
|
|
depths: Optional[torch.Tensor] = None
|
|
rgbAs = None
|
|
rgbBs = None
|
|
depthAs = None
|
|
depthBs = None
|
|
normalAs = None
|
|
normalBs = None
|
|
poseA = None #(B,4,4)
|
|
poseB = None
|
|
targets = None # Score targets, torch tensor (B)
|
|
|
|
def __init__(self, rgbAs=None, rgbBs=None, depthAs=None, depthBs=None, normalAs=None, normalBs=None, maskAs=None, maskBs=None, poseA=None, poseB=None, xyz_mapAs=None, xyz_mapBs=None, tf_to_crops=None, Ks=None, crop_masks=None, model_pts=None, mesh_diameters=None, labels=None):
|
|
self.rgbAs = rgbAs
|
|
self.rgbBs = rgbBs
|
|
self.depthAs = depthAs
|
|
self.depthBs = depthBs
|
|
self.normalAs = normalAs
|
|
self.normalBs = normalBs
|
|
self.poseA = poseA
|
|
self.poseB = poseB
|
|
self.maskAs = maskAs
|
|
self.maskBs = maskBs
|
|
self.xyz_mapAs = xyz_mapAs
|
|
self.xyz_mapBs = xyz_mapBs
|
|
self.tf_to_crops = tf_to_crops
|
|
self.crop_masks = crop_masks
|
|
self.Ks = Ks
|
|
self.model_pts = model_pts
|
|
self.mesh_diameters = mesh_diameters
|
|
self.labels = labels
|
|
|
|
|
|
def pin_memory(self) -> "BatchPoseData":
|
|
for k in self.__dict__:
|
|
if self.__dict__[k] is not None:
|
|
try:
|
|
self.__dict__[k] = self.__dict__[k].pin_memory()
|
|
except Exception as e:
|
|
pass
|
|
return self
|
|
|
|
def cuda(self):
|
|
for k in self.__dict__:
|
|
if self.__dict__[k] is not None:
|
|
try:
|
|
self.__dict__[k] = self.__dict__[k].cuda()
|
|
except:
|
|
pass
|
|
return self
|
|
|
|
def select_by_indices(self, ids):
|
|
out = BatchPoseData()
|
|
for k in self.__dict__:
|
|
if self.__dict__[k] is not None:
|
|
out.__dict__[k] = self.__dict__[k][ids.to(self.__dict__[k].device)]
|
|
return out
|
|
|