Update to 2026-09-17 pipeline snapshot; add weights, L20 assets and recording via Git LFS
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>
This commit is contained in:
liyang
@@ -0,0 +1,220 @@
|
||||
# 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,h5py,bisect,io,json
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(f'{code_dir}/../../../../')
|
||||
from Utils import *
|
||||
from learning.datasets.pose_dataset import *
|
||||
|
||||
|
||||
|
||||
|
||||
class PairH5Dataset(torch.utils.data.Dataset):
|
||||
def __init__(self, cfg, h5_file, mode='train', max_num_key=None, cache_data=None):
|
||||
self.cfg = cfg
|
||||
self.h5_file = h5_file
|
||||
self.mode = mode
|
||||
|
||||
logging.info(f"self.h5_file:{self.h5_file}")
|
||||
self.n_perturb = None
|
||||
self.H_ori = None
|
||||
self.W_ori = None
|
||||
self.cache_data = cache_data
|
||||
|
||||
if self.mode=='test':
|
||||
pass
|
||||
else:
|
||||
self.object_keys = []
|
||||
key_file = h5_file.replace('.h5','_keys.pkl')
|
||||
if os.path.exists(key_file):
|
||||
with open(key_file, 'rb') as ff:
|
||||
self.object_keys = pickle.load(ff)
|
||||
logging.info(f'object_keys loaded#:{len(self.object_keys)} from {key_file}')
|
||||
if max_num_key is not None:
|
||||
self.object_keys = self.object_keys[:max_num_key]
|
||||
else:
|
||||
with h5py.File(h5_file, 'r', libver='latest') as hf:
|
||||
for k in hf:
|
||||
self.object_keys.append(k)
|
||||
if max_num_key is not None and len(self.object_keys)>=max_num_key:
|
||||
logging.info("break due to max_num_key")
|
||||
break
|
||||
|
||||
logging.info(f'self.object_keys#:{len(self.object_keys)}, max_num_key:{max_num_key}')
|
||||
|
||||
with h5py.File(h5_file, 'r', libver='latest') as hf:
|
||||
group = hf[self.object_keys[0]]
|
||||
cnt = 0
|
||||
for k_perturb in group:
|
||||
if 'i_perturb' in k_perturb:
|
||||
cnt += 1
|
||||
if 'crop_ratio' in group[k_perturb]:
|
||||
self.cfg['crop_ratio'] = float(group[k_perturb]['crop_ratio'][()])
|
||||
if self.H_ori is None:
|
||||
if 'H_ori' in group[k_perturb]:
|
||||
self.H_ori = int(group[k_perturb]['H_ori'][()])
|
||||
self.W_ori = int(group[k_perturb]['W_ori'][()])
|
||||
else:
|
||||
self.H_ori = 540
|
||||
self.W_ori = 720
|
||||
self.n_perturb = cnt
|
||||
logging.info(f'self.n_perturb:{self.n_perturb}')
|
||||
|
||||
|
||||
def __len__(self):
|
||||
if self.mode=='test':
|
||||
return 1
|
||||
return len(self.object_keys)
|
||||
|
||||
|
||||
|
||||
def transform_depth_to_xyzmap(self, batch:BatchPoseData, H_ori, W_ori, bound=1):
|
||||
bs = len(batch.rgbAs)
|
||||
H,W = batch.rgbAs.shape[-2:]
|
||||
mesh_radius = batch.mesh_diameters.cuda()/2
|
||||
tf_to_crops = batch.tf_to_crops.cuda()
|
||||
crop_to_oris = batch.tf_to_crops.inverse().cuda() #(B,3,3)
|
||||
batch.poseA = batch.poseA.cuda()
|
||||
batch.Ks = batch.Ks.cuda()
|
||||
|
||||
if batch.xyz_mapAs is None:
|
||||
depthAs_ori = kornia.geometry.transform.warp_perspective(batch.depthAs.cuda().expand(bs,-1,-1,-1), crop_to_oris, dsize=(H_ori, W_ori), mode='nearest', align_corners=False)
|
||||
batch.xyz_mapAs = depth2xyzmap_batch(depthAs_ori[:,0], batch.Ks, zfar=np.inf).permute(0,3,1,2) #(B,3,H,W)
|
||||
batch.xyz_mapAs = kornia.geometry.transform.warp_perspective(batch.xyz_mapAs, tf_to_crops, dsize=(H,W), mode='nearest', align_corners=False)
|
||||
batch.xyz_mapAs = batch.xyz_mapAs.cuda()
|
||||
if self.cfg['normalize_xyz']:
|
||||
invalid = batch.xyz_mapAs[:,2:3]<0.001
|
||||
batch.xyz_mapAs = batch.xyz_mapAs-batch.poseA[:,:3,3].reshape(bs,3,1,1)
|
||||
if self.cfg['normalize_xyz']:
|
||||
batch.xyz_mapAs *= 1/mesh_radius.reshape(bs,1,1,1)
|
||||
invalid = invalid.expand(bs,3,-1,-1) | (torch.abs(batch.xyz_mapAs)>=2)
|
||||
batch.xyz_mapAs[invalid.expand(bs,3,-1,-1)] = 0
|
||||
|
||||
if batch.xyz_mapBs is None:
|
||||
depthBs_ori = kornia.geometry.transform.warp_perspective(batch.depthBs.cuda().expand(bs,-1,-1,-1), crop_to_oris, dsize=(H_ori, W_ori), mode='nearest', align_corners=False)
|
||||
batch.xyz_mapBs = depth2xyzmap_batch(depthBs_ori[:,0], batch.Ks, zfar=np.inf).permute(0,3,1,2) #(B,3,H,W)
|
||||
batch.xyz_mapBs = kornia.geometry.transform.warp_perspective(batch.xyz_mapBs, tf_to_crops, dsize=(H,W), mode='nearest', align_corners=False)
|
||||
batch.xyz_mapBs = batch.xyz_mapBs.cuda()
|
||||
if self.cfg['normalize_xyz']:
|
||||
invalid = batch.xyz_mapBs[:,2:3]<0.001
|
||||
batch.xyz_mapBs = batch.xyz_mapBs-batch.poseA[:,:3,3].reshape(bs,3,1,1)
|
||||
if self.cfg['normalize_xyz']:
|
||||
batch.xyz_mapBs *= 1/mesh_radius.reshape(bs,1,1,1)
|
||||
invalid = invalid.expand(bs,3,-1,-1) | (torch.abs(batch.xyz_mapBs)>=2)
|
||||
batch.xyz_mapBs[invalid.expand(bs,3,-1,-1)] = 0
|
||||
|
||||
return batch
|
||||
|
||||
|
||||
|
||||
def transform_batch(self, batch:BatchPoseData, H_ori, W_ori, bound=1):
|
||||
'''Transform the batch before feeding to the network
|
||||
!NOTE the H_ori, W_ori could be different at test time from the training data, and needs to be set
|
||||
'''
|
||||
bs = len(batch.rgbAs)
|
||||
batch.rgbAs = batch.rgbAs.cuda().float()/255.0
|
||||
batch.rgbBs = batch.rgbBs.cuda().float()/255.0
|
||||
|
||||
batch = self.transform_depth_to_xyzmap(batch, H_ori, W_ori, bound=bound)
|
||||
return batch
|
||||
|
||||
|
||||
|
||||
|
||||
class TripletH5Dataset(PairH5Dataset):
|
||||
def __init__(self, cfg, h5_file, mode, max_num_key=None, cache_data=None):
|
||||
super().__init__(cfg, h5_file, mode, max_num_key, cache_data=cache_data)
|
||||
|
||||
|
||||
def transform_depth_to_xyzmap(self, batch:BatchPoseData, H_ori, W_ori, bound=1):
|
||||
bs = len(batch.rgbAs)
|
||||
H,W = batch.rgbAs.shape[-2:]
|
||||
mesh_radius = batch.mesh_diameters.cuda()/2
|
||||
tf_to_crops = batch.tf_to_crops.cuda()
|
||||
crop_to_oris = batch.tf_to_crops.inverse().cuda() #(B,3,3)
|
||||
batch.poseA = batch.poseA.cuda()
|
||||
batch.Ks = batch.Ks.cuda()
|
||||
|
||||
if batch.xyz_mapAs is None:
|
||||
depthAs_ori = kornia.geometry.transform.warp_perspective(batch.depthAs.cuda().expand(bs,-1,-1,-1), crop_to_oris, dsize=(H_ori, W_ori), mode='nearest', align_corners=False)
|
||||
batch.xyz_mapAs = depth2xyzmap_batch(depthAs_ori[:,0], batch.Ks, zfar=np.inf).permute(0,3,1,2) #(B,3,H,W)
|
||||
batch.xyz_mapAs = kornia.geometry.transform.warp_perspective(batch.xyz_mapAs, tf_to_crops, dsize=(H,W), mode='nearest', align_corners=False)
|
||||
batch.xyz_mapAs = batch.xyz_mapAs.cuda()
|
||||
invalid = batch.xyz_mapAs[:,2:3]<0.1
|
||||
batch.xyz_mapAs = (batch.xyz_mapAs-batch.poseA[:,:3,3].reshape(bs,3,1,1))
|
||||
if self.cfg['normalize_xyz']:
|
||||
batch.xyz_mapAs *= 1/mesh_radius.reshape(bs,1,1,1)
|
||||
invalid = invalid.expand(bs,3,-1,-1) | (torch.abs(batch.xyz_mapAs)>=2)
|
||||
batch.xyz_mapAs[invalid.expand(bs,3,-1,-1)] = 0
|
||||
|
||||
if batch.xyz_mapBs is None:
|
||||
depthBs_ori = kornia.geometry.transform.warp_perspective(batch.depthBs.cuda().expand(bs,-1,-1,-1), crop_to_oris, dsize=(H_ori, W_ori), mode='nearest', align_corners=False)
|
||||
batch.xyz_mapBs = depth2xyzmap_batch(depthBs_ori[:,0], batch.Ks, zfar=np.inf).permute(0,3,1,2) #(B,3,H,W)
|
||||
batch.xyz_mapBs = kornia.geometry.transform.warp_perspective(batch.xyz_mapBs, tf_to_crops, dsize=(H,W), mode='nearest', align_corners=False)
|
||||
batch.xyz_mapBs = batch.xyz_mapBs.cuda()
|
||||
invalid = batch.xyz_mapBs[:,2:3]<0.1
|
||||
batch.xyz_mapBs = (batch.xyz_mapBs-batch.poseA[:,:3,3].reshape(bs,3,1,1))
|
||||
if self.cfg['normalize_xyz']:
|
||||
batch.xyz_mapBs *= 1/mesh_radius.reshape(bs,1,1,1)
|
||||
invalid = invalid.expand(bs,3,-1,-1) | (torch.abs(batch.xyz_mapBs)>=2)
|
||||
batch.xyz_mapBs[invalid.expand(bs,3,-1,-1)] = 0
|
||||
|
||||
return batch
|
||||
|
||||
|
||||
def transform_batch(self, batch:BatchPoseData, H_ori, W_ori, bound=1):
|
||||
bs = len(batch.rgbAs)
|
||||
batch.rgbAs = batch.rgbAs.cuda().float()/255.0
|
||||
batch.rgbBs = batch.rgbBs.cuda().float()/255.0
|
||||
|
||||
batch = self.transform_depth_to_xyzmap(batch, H_ori, W_ori, bound=bound)
|
||||
return batch
|
||||
|
||||
|
||||
|
||||
class ScoreMultiPairH5Dataset(TripletH5Dataset):
|
||||
def __init__(self, cfg, h5_file, mode, max_num_key=None, cache_data=None):
|
||||
super().__init__(cfg, h5_file, mode, max_num_key, cache_data=cache_data)
|
||||
if mode in ['train', 'val']:
|
||||
self.cfg['train_num_pair'] = self.n_perturb
|
||||
|
||||
|
||||
class PoseRefinePairH5Dataset(PairH5Dataset):
|
||||
def __init__(self, cfg, h5_file, mode='train', max_num_key=None, cache_data=None):
|
||||
super().__init__(cfg=cfg, h5_file=h5_file, mode=mode, max_num_key=max_num_key, cache_data=cache_data)
|
||||
|
||||
if mode!='test':
|
||||
with h5py.File(h5_file, 'r', libver='latest') as hf:
|
||||
group = hf[self.object_keys[0]]
|
||||
for key_perturb in group:
|
||||
depthA = imageio.imread(group[key_perturb]['depthA'][()])
|
||||
depthB = imageio.imread(group[key_perturb]['depthB'][()])
|
||||
self.cfg['n_view'] = min(self.cfg['n_view'], depthA.shape[1]//depthB.shape[1])
|
||||
logging.info(f'n_view:{self.cfg["n_view"]}')
|
||||
self.trans_normalizer = group[key_perturb]['trans_normalizer'][()]
|
||||
if isinstance(self.trans_normalizer, np.ndarray):
|
||||
self.trans_normalizer = self.trans_normalizer.tolist()
|
||||
self.rot_normalizer = group[key_perturb]['rot_normalizer'][()]/180.0*np.pi
|
||||
logging.info(f'self.trans_normalizer:{self.trans_normalizer}, self.rot_normalizer:{self.rot_normalizer}')
|
||||
break
|
||||
|
||||
|
||||
def transform_batch(self, batch:BatchPoseData, H_ori, W_ori, bound=1):
|
||||
'''Transform the batch before feeding to the network
|
||||
!NOTE the H_ori, W_ori could be different at test time from the training data, and needs to be set
|
||||
'''
|
||||
bs = len(batch.rgbAs)
|
||||
batch.rgbAs = batch.rgbAs.cuda().float()/255.0
|
||||
batch.rgbBs = batch.rgbBs.cuda().float()/255.0
|
||||
|
||||
batch = self.transform_depth_to_xyzmap(batch, H_ori, W_ori, bound=bound)
|
||||
return batch
|
||||
|
||||
@@ -0,0 +1,135 @@
|
||||
# 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
|
||||
|
||||
Reference in New Issue
Block a user