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>
221 lines
9.5 KiB
Python
221 lines
9.5 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,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
|
|
|