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,296 @@
|
||||
# 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 functools
|
||||
import os,sys,kornia
|
||||
import time
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(f'{code_dir}/../../')
|
||||
import numpy as np
|
||||
import torch
|
||||
from omegaconf import OmegaConf
|
||||
from learning.models.refine_network import RefineNet
|
||||
from learning.datasets.h5_dataset import *
|
||||
from Utils import *
|
||||
from datareader import *
|
||||
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def make_crop_data_batch(render_size, ob_in_cams, mesh, rgb, depth, K, crop_ratio, xyz_map, normal_map=None, mesh_diameter=None, cfg=None, glctx=None, mesh_tensors=None, dataset:PoseRefinePairH5Dataset=None):
|
||||
logging.info("Welcome make_crop_data_batch")
|
||||
H,W = depth.shape[:2]
|
||||
args = []
|
||||
method = 'box_3d'
|
||||
tf_to_crops = compute_crop_window_tf_batch(pts=mesh.vertices, H=H, W=W, poses=ob_in_cams, K=K, crop_ratio=crop_ratio, out_size=(render_size[1], render_size[0]), method=method, mesh_diameter=mesh_diameter)
|
||||
|
||||
logging.info("make tf_to_crops done")
|
||||
|
||||
B = len(ob_in_cams)
|
||||
poseA = torch.as_tensor(ob_in_cams, dtype=torch.float, device='cuda')
|
||||
|
||||
bs = 512
|
||||
rgb_rs = []
|
||||
depth_rs = []
|
||||
normal_rs = []
|
||||
xyz_map_rs = []
|
||||
|
||||
bbox2d_crop = torch.as_tensor(np.array([0, 0, cfg['input_resize'][0]-1, cfg['input_resize'][1]-1]).reshape(2,2), device='cuda', dtype=torch.float)
|
||||
bbox2d_ori = transform_pts(bbox2d_crop, tf_to_crops.inverse()).reshape(-1,4)
|
||||
|
||||
for b in range(0,len(poseA),bs):
|
||||
extra = {}
|
||||
rgb_r, depth_r, normal_r = nvdiffrast_render(K=K, H=H, W=W, ob_in_cams=poseA[b:b+bs], context='cuda', get_normal=cfg['use_normal'], glctx=glctx, mesh_tensors=mesh_tensors, output_size=cfg['input_resize'], bbox2d=bbox2d_ori[b:b+bs], use_light=True, extra=extra)
|
||||
rgb_rs.append(rgb_r)
|
||||
depth_rs.append(depth_r[...,None])
|
||||
normal_rs.append(normal_r)
|
||||
xyz_map_rs.append(extra['xyz_map'])
|
||||
rgb_rs = torch.cat(rgb_rs, dim=0).permute(0,3,1,2) * 255
|
||||
depth_rs = torch.cat(depth_rs, dim=0).permute(0,3,1,2) #(B,1,H,W)
|
||||
xyz_map_rs = torch.cat(xyz_map_rs, dim=0).permute(0,3,1,2) #(B,3,H,W)
|
||||
Ks = torch.as_tensor(K, device='cuda', dtype=torch.float).reshape(1,3,3)
|
||||
if cfg['use_normal']:
|
||||
normal_rs = torch.cat(normal_rs, dim=0).permute(0,3,1,2) #(B,3,H,W)
|
||||
|
||||
logging.info("render done")
|
||||
|
||||
rgbBs = kornia.geometry.transform.warp_perspective(torch.as_tensor(rgb, dtype=torch.float, device='cuda').permute(2,0,1)[None].expand(B,-1,-1,-1), tf_to_crops, dsize=render_size, mode='bilinear', align_corners=False)
|
||||
if rgb_rs.shape[-2:]!=cfg['input_resize']:
|
||||
rgbAs = kornia.geometry.transform.warp_perspective(rgb_rs, tf_to_crops, dsize=render_size, mode='bilinear', align_corners=False)
|
||||
else:
|
||||
rgbAs = rgb_rs
|
||||
if xyz_map_rs.shape[-2:]!=cfg['input_resize']:
|
||||
xyz_mapAs = kornia.geometry.transform.warp_perspective(xyz_map_rs, tf_to_crops, dsize=render_size, mode='nearest', align_corners=False)
|
||||
else:
|
||||
xyz_mapAs = xyz_map_rs
|
||||
xyz_mapBs = kornia.geometry.transform.warp_perspective(torch.as_tensor(xyz_map, device='cuda', dtype=torch.float).permute(2,0,1)[None].expand(B,-1,-1,-1), tf_to_crops, dsize=render_size, mode='nearest', align_corners=False) #(B,3,H,W)
|
||||
|
||||
if cfg['use_normal']:
|
||||
normalAs = kornia.geometry.transform.warp_perspective(normal_rs, tf_to_crops, dsize=render_size, mode='nearest', align_corners=False)
|
||||
normalBs = kornia.geometry.transform.warp_perspective(torch.as_tensor(normal_map, dtype=torch.float, device='cuda').permute(2,0,1)[None].expand(B,-1,-1,-1), tf_to_crops, dsize=render_size, mode='nearest', align_corners=False)
|
||||
else:
|
||||
normalAs = None
|
||||
normalBs = None
|
||||
|
||||
logging.info("warp done")
|
||||
|
||||
mesh_diameters = torch.ones((len(rgbAs)), dtype=torch.float, device='cuda')*mesh_diameter
|
||||
pose_data = BatchPoseData(rgbAs=rgbAs, rgbBs=rgbBs, depthAs=None, depthBs=None, normalAs=normalAs, normalBs=normalBs, poseA=poseA, poseB=None, xyz_mapAs=xyz_mapAs, xyz_mapBs=xyz_mapBs, tf_to_crops=tf_to_crops, Ks=Ks, mesh_diameters=mesh_diameters)
|
||||
pose_data = dataset.transform_batch(batch=pose_data, H_ori=H, W_ori=W, bound=1)
|
||||
|
||||
logging.info("pose batch data done")
|
||||
|
||||
return pose_data
|
||||
|
||||
|
||||
|
||||
class PoseRefinePredictor:
|
||||
def __init__(self,):
|
||||
logging.info("welcome")
|
||||
self.amp = True
|
||||
self.run_name = "2023-10-28-18-33-37"
|
||||
model_name = 'model_best.pth'
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
ckpt_dir = f'{code_dir}/../../weights/{self.run_name}/{model_name}'
|
||||
|
||||
self.cfg = OmegaConf.load(f'{code_dir}/../../weights/{self.run_name}/config.yml')
|
||||
|
||||
self.cfg['ckpt_dir'] = ckpt_dir
|
||||
self.cfg['enable_amp'] = True
|
||||
|
||||
########## Defaults, to be backward compatible
|
||||
if 'use_normal' not in self.cfg:
|
||||
self.cfg['use_normal'] = False
|
||||
if 'use_mask' not in self.cfg:
|
||||
self.cfg['use_mask'] = False
|
||||
if 'use_BN' not in self.cfg:
|
||||
self.cfg['use_BN'] = False
|
||||
if 'c_in' not in self.cfg:
|
||||
self.cfg['c_in'] = 4
|
||||
if 'crop_ratio' not in self.cfg or self.cfg['crop_ratio'] is None:
|
||||
self.cfg['crop_ratio'] = 1.2
|
||||
if 'n_view' not in self.cfg:
|
||||
self.cfg['n_view'] = 1
|
||||
if 'trans_rep' not in self.cfg:
|
||||
self.cfg['trans_rep'] = 'tracknet'
|
||||
if 'rot_rep' not in self.cfg:
|
||||
self.cfg['rot_rep'] = 'axis_angle'
|
||||
if 'zfar' not in self.cfg:
|
||||
self.cfg['zfar'] = 3
|
||||
if 'normalize_xyz' not in self.cfg:
|
||||
self.cfg['normalize_xyz'] = False
|
||||
if isinstance(self.cfg['zfar'], str) and 'inf' in self.cfg['zfar'].lower():
|
||||
self.cfg['zfar'] = np.inf
|
||||
if 'normal_uint8' not in self.cfg:
|
||||
self.cfg['normal_uint8'] = False
|
||||
logging.info(f"self.cfg: \n {OmegaConf.to_yaml(self.cfg)}")
|
||||
|
||||
self.dataset = PoseRefinePairH5Dataset(cfg=self.cfg, h5_file='', mode='test')
|
||||
self.model = RefineNet(cfg=self.cfg, c_in=self.cfg['c_in']).cuda()
|
||||
|
||||
logging.info(f"Using pretrained model from {ckpt_dir}")
|
||||
ckpt = torch.load(ckpt_dir)
|
||||
if 'model' in ckpt:
|
||||
ckpt = ckpt['model']
|
||||
self.model.load_state_dict(ckpt)
|
||||
|
||||
self.model.cuda().eval()
|
||||
logging.info("init done")
|
||||
self.last_trans_update = None
|
||||
self.last_rot_update = None
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def predict(self, rgb, depth, K, ob_in_cams, xyz_map, normal_map=None, get_vis=False, mesh=None, mesh_tensors=None, glctx=None, mesh_diameter=None, iteration=5):
|
||||
'''
|
||||
@rgb: np array (H,W,3)
|
||||
@ob_in_cams: np array (N,4,4)
|
||||
'''
|
||||
torch.set_default_tensor_type('torch.cuda.FloatTensor')
|
||||
logging.info(f'ob_in_cams:{ob_in_cams.shape}')
|
||||
tf_to_center = np.eye(4)
|
||||
ob_centered_in_cams = ob_in_cams
|
||||
mesh_centered = mesh
|
||||
|
||||
logging.info(f'self.cfg.use_normal:{self.cfg.use_normal}')
|
||||
if not self.cfg.use_normal:
|
||||
normal_map = None
|
||||
|
||||
crop_ratio = self.cfg['crop_ratio']
|
||||
logging.info(f"trans_normalizer:{self.cfg['trans_normalizer']}, rot_normalizer:{self.cfg['rot_normalizer']}")
|
||||
bs = 1024
|
||||
|
||||
B_in_cams = torch.as_tensor(ob_centered_in_cams, device='cuda', dtype=torch.float)
|
||||
|
||||
|
||||
if mesh_tensors is None:
|
||||
mesh_tensors = make_mesh_tensors(mesh_centered)
|
||||
|
||||
rgb_tensor = torch.as_tensor(rgb, device='cuda', dtype=torch.float)
|
||||
depth_tensor = torch.as_tensor(depth, device='cuda', dtype=torch.float)
|
||||
xyz_map_tensor = torch.as_tensor(xyz_map, device='cuda', dtype=torch.float)
|
||||
trans_normalizer = self.cfg['trans_normalizer']
|
||||
if not isinstance(trans_normalizer, float):
|
||||
trans_normalizer = torch.as_tensor(list(trans_normalizer), device='cuda', dtype=torch.float).reshape(1,3)
|
||||
|
||||
for _ in range(iteration):
|
||||
logging.info("making cropped data")
|
||||
pose_data = make_crop_data_batch(self.cfg.input_resize, B_in_cams, mesh_centered, rgb_tensor, depth_tensor, K, crop_ratio=crop_ratio, normal_map=normal_map, xyz_map=xyz_map_tensor, cfg=self.cfg, glctx=glctx, mesh_tensors=mesh_tensors, dataset=self.dataset, mesh_diameter=mesh_diameter)
|
||||
B_in_cams = []
|
||||
for b in range(0, pose_data.rgbAs.shape[0], bs):
|
||||
A = torch.cat([pose_data.rgbAs[b:b+bs].cuda(), pose_data.xyz_mapAs[b:b+bs].cuda()], dim=1).float()
|
||||
B = torch.cat([pose_data.rgbBs[b:b+bs].cuda(), pose_data.xyz_mapBs[b:b+bs].cuda()], dim=1).float()
|
||||
logging.info("forward start")
|
||||
with torch.cuda.amp.autocast(enabled=self.amp):
|
||||
output = self.model(A,B)
|
||||
for k in output:
|
||||
output[k] = output[k].float()
|
||||
logging.info("forward done")
|
||||
if self.cfg['trans_rep']=='tracknet':
|
||||
if not self.cfg['normalize_xyz']:
|
||||
trans_delta = torch.tanh(output["trans"])*trans_normalizer
|
||||
else:
|
||||
trans_delta = output["trans"]
|
||||
|
||||
elif self.cfg['trans_rep']=='deepim':
|
||||
def project_and_transform_to_crop(centers):
|
||||
uvs = (pose_data.Ks[b:b+bs]@centers.reshape(-1,3,1)).reshape(-1,3)
|
||||
uvs = uvs/uvs[:,2:3]
|
||||
uvs = (pose_data.tf_to_crops[b:b+bs]@uvs.reshape(-1,3,1)).reshape(-1,3)
|
||||
return uvs[:,:2]
|
||||
|
||||
rot_delta = output["rot"]
|
||||
z_pred = output['trans'][:,2]*pose_data.poseA[b:b+bs][...,2,3]
|
||||
uvA_crop = project_and_transform_to_crop(pose_data.poseA[b:b+bs][...,:3,3])
|
||||
uv_pred_crop = uvA_crop + output['trans'][:,:2]*self.cfg['input_resize'][0]
|
||||
uv_pred = transform_pts(uv_pred_crop, pose_data.tf_to_crops[b:b+bs].inverse().cuda())
|
||||
center_pred = torch.cat([uv_pred, torch.ones((len(rot_delta),1), dtype=torch.float, device='cuda')], dim=-1)
|
||||
center_pred = (pose_data.Ks[b:b+bs].inverse().cuda()@center_pred.reshape(len(rot_delta),3,1)).reshape(len(rot_delta),3) * z_pred.reshape(len(rot_delta),1)
|
||||
trans_delta = center_pred-pose_data.poseA[b:b+bs][...,:3,3]
|
||||
|
||||
else:
|
||||
trans_delta = output["trans"]
|
||||
|
||||
if self.cfg['rot_rep']=='axis_angle':
|
||||
rot_mat_delta = torch.tanh(output["rot"])*self.cfg['rot_normalizer']
|
||||
rot_mat_delta = so3_exp_map(rot_mat_delta).permute(0,2,1)
|
||||
elif self.cfg['rot_rep']=='6d':
|
||||
rot_mat_delta = rotation_6d_to_matrix(output['rot']).permute(0,2,1)
|
||||
else:
|
||||
raise RuntimeError
|
||||
|
||||
if self.cfg['normalize_xyz']:
|
||||
trans_delta *= (mesh_diameter/2)
|
||||
|
||||
B_in_cam = egocentric_delta_pose_to_pose(pose_data.poseA[b:b+bs], trans_delta=trans_delta, rot_mat_delta=rot_mat_delta)
|
||||
B_in_cams.append(B_in_cam)
|
||||
|
||||
B_in_cams = torch.cat(B_in_cams, dim=0).reshape(len(ob_in_cams),4,4)
|
||||
|
||||
B_in_cams_out = B_in_cams@torch.tensor(tf_to_center[None], device='cuda', dtype=torch.float)
|
||||
torch.cuda.empty_cache()
|
||||
self.last_trans_update = trans_delta
|
||||
self.last_rot_update = rot_mat_delta
|
||||
|
||||
if get_vis:
|
||||
logging.info("get_vis...")
|
||||
canvas = []
|
||||
padding = 2
|
||||
pose_data = make_crop_data_batch(self.cfg.input_resize, torch.as_tensor(ob_centered_in_cams), mesh_centered, rgb, depth, K, crop_ratio=crop_ratio, normal_map=normal_map, xyz_map=xyz_map_tensor, cfg=self.cfg, glctx=glctx, mesh_tensors=mesh_tensors, dataset=self.dataset, mesh_diameter=mesh_diameter)
|
||||
for id in range(0, len(B_in_cams)):
|
||||
rgbA_vis = (pose_data.rgbAs[id]*255).permute(1,2,0).data.cpu().numpy()
|
||||
rgbB_vis = (pose_data.rgbBs[id]*255).permute(1,2,0).data.cpu().numpy()
|
||||
row = [rgbA_vis, rgbB_vis]
|
||||
H,W = rgbA_vis.shape[:2]
|
||||
if pose_data.depthAs is not None:
|
||||
depthA = pose_data.depthAs[id].data.cpu().numpy().reshape(H,W)
|
||||
depthB = pose_data.depthBs[id].data.cpu().numpy().reshape(H,W)
|
||||
elif pose_data.xyz_mapAs is not None:
|
||||
depthA = pose_data.xyz_mapAs[id][2].data.cpu().numpy().reshape(H,W)
|
||||
depthB = pose_data.xyz_mapBs[id][2].data.cpu().numpy().reshape(H,W)
|
||||
zmin = min(depthA.min(), depthB.min())
|
||||
zmax = max(depthA.max(), depthB.max())
|
||||
depthA_vis = depth_to_vis(depthA, zmin=zmin, zmax=zmax, inverse=False)
|
||||
depthB_vis = depth_to_vis(depthB, zmin=zmin, zmax=zmax, inverse=False)
|
||||
row += [depthA_vis, depthB_vis]
|
||||
if pose_data.normalAs is not None:
|
||||
pass
|
||||
row = make_grid_image(row, nrow=len(row), padding=padding, pad_value=255)
|
||||
row = cv_draw_text(row, text=f'id:{id}', uv_top_left=(10,10), color=(0,255,0), fontScale=0.5)
|
||||
canvas.append(row)
|
||||
canvas = make_grid_image(canvas, nrow=1, padding=padding, pad_value=255)
|
||||
|
||||
pose_data = make_crop_data_batch(self.cfg.input_resize, B_in_cams, mesh_centered, rgb, depth, K, crop_ratio=crop_ratio, normal_map=normal_map, xyz_map=xyz_map_tensor, cfg=self.cfg, glctx=glctx, mesh_tensors=mesh_tensors, dataset=self.dataset, mesh_diameter=mesh_diameter)
|
||||
canvas_refined = []
|
||||
for id in range(0, len(B_in_cams)):
|
||||
rgbA_vis = (pose_data.rgbAs[id]*255).permute(1,2,0).data.cpu().numpy()
|
||||
rgbB_vis = (pose_data.rgbBs[id]*255).permute(1,2,0).data.cpu().numpy()
|
||||
row = [rgbA_vis, rgbB_vis]
|
||||
H,W = rgbA_vis.shape[:2]
|
||||
if pose_data.depthAs is not None:
|
||||
depthA = pose_data.depthAs[id].data.cpu().numpy().reshape(H,W)
|
||||
depthB = pose_data.depthBs[id].data.cpu().numpy().reshape(H,W)
|
||||
elif pose_data.xyz_mapAs is not None:
|
||||
depthA = pose_data.xyz_mapAs[id][2].data.cpu().numpy().reshape(H,W)
|
||||
depthB = pose_data.xyz_mapBs[id][2].data.cpu().numpy().reshape(H,W)
|
||||
zmin = min(depthA.min(), depthB.min())
|
||||
zmax = max(depthA.max(), depthB.max())
|
||||
depthA_vis = depth_to_vis(depthA, zmin=zmin, zmax=zmax, inverse=False)
|
||||
depthB_vis = depth_to_vis(depthB, zmin=zmin, zmax=zmax, inverse=False)
|
||||
row += [depthA_vis, depthB_vis]
|
||||
row = make_grid_image(row, nrow=len(row), padding=padding, pad_value=255)
|
||||
canvas_refined.append(row)
|
||||
|
||||
canvas_refined = make_grid_image(canvas_refined, nrow=1, padding=padding, pad_value=255)
|
||||
canvas = make_grid_image([canvas, canvas_refined], nrow=2, padding=padding, pad_value=255)
|
||||
torch.cuda.empty_cache()
|
||||
return B_in_cams_out, canvas
|
||||
|
||||
return B_in_cams_out, None
|
||||
|
||||
@@ -0,0 +1,227 @@
|
||||
# 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 functools
|
||||
import os,sys,kornia
|
||||
import time
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.distributed as dist
|
||||
from omegaconf import OmegaConf
|
||||
from tqdm import tqdm
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(f'{code_dir}/../../../')
|
||||
from learning.datasets.h5_dataset import *
|
||||
from learning.models.score_network import *
|
||||
from learning.datasets.pose_dataset import *
|
||||
from Utils import *
|
||||
from datareader import *
|
||||
|
||||
|
||||
def vis_batch_data_scores(pose_data, ids, scores, pad_margin=5):
|
||||
assert len(scores)==len(ids)
|
||||
canvas = []
|
||||
for id in ids:
|
||||
rgbA_vis = (pose_data.rgbAs[id]*255).permute(1,2,0).data.cpu().numpy()
|
||||
rgbB_vis = (pose_data.rgbBs[id]*255).permute(1,2,0).data.cpu().numpy()
|
||||
H,W = rgbA_vis.shape[:2]
|
||||
zmin = pose_data.depthAs[id].data.cpu().numpy().reshape(H,W).min()
|
||||
zmax = pose_data.depthAs[id].data.cpu().numpy().reshape(H,W).max()
|
||||
depthA_vis = depth_to_vis(pose_data.depthAs[id].data.cpu().numpy().reshape(H,W), zmin=zmin, zmax=zmax, inverse=False)
|
||||
depthB_vis = depth_to_vis(pose_data.depthBs[id].data.cpu().numpy().reshape(H,W), zmin=zmin, zmax=zmax, inverse=False)
|
||||
if pose_data.normalAs is not None:
|
||||
pass
|
||||
pad = np.ones((rgbA_vis.shape[0],pad_margin,3))*255
|
||||
if pose_data.normalAs is not None:
|
||||
pass
|
||||
else:
|
||||
row = np.concatenate([rgbA_vis, pad, depthA_vis, pad, rgbB_vis, pad, depthB_vis], axis=1)
|
||||
s = 100/row.shape[0]
|
||||
row = cv2.resize(row, fx=s, fy=s, dsize=None)
|
||||
row = cv_draw_text(row, text=f'id:{id}, score:{scores[id]:.3f}', uv_top_left=(10,10), color=(0,255,0), fontScale=0.5)
|
||||
canvas.append(row)
|
||||
pad = np.ones((pad_margin, row.shape[1], 3))*255
|
||||
canvas.append(pad)
|
||||
canvas = np.concatenate(canvas, axis=0).astype(np.uint8)
|
||||
return canvas
|
||||
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def make_crop_data_batch(render_size, ob_in_cams, mesh, rgb, depth, K, crop_ratio, normal_map=None, mesh_diameter=None, glctx=None, mesh_tensors=None, dataset:TripletH5Dataset=None, cfg=None):
|
||||
logging.info("Welcome make_crop_data_batch")
|
||||
H,W = depth.shape[:2]
|
||||
|
||||
args = []
|
||||
method = 'box_3d'
|
||||
tf_to_crops = compute_crop_window_tf_batch(pts=mesh.vertices, H=H, W=W, poses=ob_in_cams, K=K, crop_ratio=crop_ratio, out_size=(render_size[1], render_size[0]), method=method, mesh_diameter=mesh_diameter)
|
||||
logging.info("make tf_to_crops done")
|
||||
|
||||
B = len(ob_in_cams)
|
||||
poseAs = torch.as_tensor(ob_in_cams, dtype=torch.float, device='cuda')
|
||||
|
||||
bs = 512
|
||||
rgb_rs = []
|
||||
depth_rs = []
|
||||
xyz_map_rs = []
|
||||
|
||||
bbox2d_crop = torch.as_tensor(np.array([0, 0, cfg['input_resize'][0]-1, cfg['input_resize'][1]-1]).reshape(2,2), device='cuda', dtype=torch.float)
|
||||
bbox2d_ori = transform_pts(bbox2d_crop, tf_to_crops.inverse()[:,None]).reshape(-1,4)
|
||||
|
||||
for b in range(0,len(ob_in_cams),bs):
|
||||
extra = {}
|
||||
rgb_r, depth_r, normal_r = nvdiffrast_render(K=K, H=H, W=W, ob_in_cams=poseAs[b:b+bs], context='cuda', get_normal=cfg['use_normal'], glctx=glctx, mesh_tensors=mesh_tensors, output_size=cfg['input_resize'], bbox2d=bbox2d_ori[b:b+bs], use_light=True, extra=extra)
|
||||
rgb_rs.append(rgb_r)
|
||||
depth_rs.append(depth_r[...,None])
|
||||
xyz_map_rs.append(extra['xyz_map'])
|
||||
|
||||
rgb_rs = torch.cat(rgb_rs, dim=0).permute(0,3,1,2) * 255
|
||||
depth_rs = torch.cat(depth_rs, dim=0).permute(0,3,1,2)
|
||||
xyz_map_rs = torch.cat(xyz_map_rs, dim=0).permute(0,3,1,2) #(B,3,H,W)
|
||||
logging.info("render done")
|
||||
|
||||
rgbBs = kornia.geometry.transform.warp_perspective(torch.as_tensor(rgb, dtype=torch.float, device='cuda').permute(2,0,1)[None].expand(B,-1,-1,-1), tf_to_crops, dsize=render_size, mode='bilinear', align_corners=False)
|
||||
depthBs = kornia.geometry.transform.warp_perspective(torch.as_tensor(depth, dtype=torch.float, device='cuda')[None,None].expand(B,-1,-1,-1), tf_to_crops, dsize=render_size, mode='nearest', align_corners=False)
|
||||
if rgb_rs.shape[-2:]!=cfg['input_resize']:
|
||||
rgbAs = kornia.geometry.transform.warp_perspective(rgb_rs, tf_to_crops, dsize=render_size, mode='bilinear', align_corners=False)
|
||||
depthAs = kornia.geometry.transform.warp_perspective(depth_rs, tf_to_crops, dsize=render_size, mode='nearest', align_corners=False)
|
||||
else:
|
||||
rgbAs = rgb_rs
|
||||
depthAs = depth_rs
|
||||
|
||||
if xyz_map_rs.shape[-2:]!=cfg['input_resize']:
|
||||
xyz_mapAs = kornia.geometry.transform.warp_perspective(xyz_map_rs, tf_to_crops, dsize=render_size, mode='nearest', align_corners=False)
|
||||
else:
|
||||
xyz_mapAs = xyz_map_rs
|
||||
|
||||
normalAs = None
|
||||
normalBs = None
|
||||
|
||||
Ks = torch.as_tensor(K, dtype=torch.float).reshape(1,3,3).expand(len(rgbAs),3,3)
|
||||
mesh_diameters = torch.ones((len(rgbAs)), dtype=torch.float, device='cuda')*mesh_diameter
|
||||
|
||||
pose_data = BatchPoseData(rgbAs=rgbAs, rgbBs=rgbBs, depthAs=depthAs, depthBs=depthBs, normalAs=normalAs, normalBs=normalBs, poseA=poseAs, xyz_mapAs=xyz_mapAs, tf_to_crops=tf_to_crops, Ks=Ks, mesh_diameters=mesh_diameters)
|
||||
pose_data = dataset.transform_batch(pose_data, H_ori=H, W_ori=W, bound=1)
|
||||
|
||||
logging.info("pose batch data done")
|
||||
|
||||
return pose_data
|
||||
|
||||
|
||||
class ScorePredictor:
|
||||
def __init__(self, amp=True):
|
||||
self.amp = amp
|
||||
self.run_name = "2024-01-11-20-02-45"
|
||||
|
||||
model_name = 'model_best.pth'
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
ckpt_dir = f'{code_dir}/../../weights/{self.run_name}/{model_name}'
|
||||
|
||||
self.cfg = OmegaConf.load(f'{code_dir}/../../weights/{self.run_name}/config.yml')
|
||||
|
||||
self.cfg['ckpt_dir'] = ckpt_dir
|
||||
self.cfg['enable_amp'] = True
|
||||
|
||||
########## Defaults, to be backward compatible
|
||||
if 'use_normal' not in self.cfg:
|
||||
self.cfg['use_normal'] = False
|
||||
if 'use_BN' not in self.cfg:
|
||||
self.cfg['use_BN'] = False
|
||||
if 'zfar' not in self.cfg:
|
||||
self.cfg['zfar'] = np.inf
|
||||
if 'c_in' not in self.cfg:
|
||||
self.cfg['c_in'] = 4
|
||||
if 'normalize_xyz' not in self.cfg:
|
||||
self.cfg['normalize_xyz'] = False
|
||||
if 'crop_ratio' not in self.cfg or self.cfg['crop_ratio'] is None:
|
||||
self.cfg['crop_ratio'] = 1.2
|
||||
|
||||
logging.info(f"self.cfg: \n {OmegaConf.to_yaml(self.cfg)}")
|
||||
|
||||
self.dataset = ScoreMultiPairH5Dataset(cfg=self.cfg, mode='test', h5_file=None, max_num_key=1)
|
||||
self.model = ScoreNetMultiPair(cfg=self.cfg, c_in=self.cfg['c_in']).cuda()
|
||||
|
||||
logging.info(f"Using pretrained model from {ckpt_dir}")
|
||||
ckpt = torch.load(ckpt_dir)
|
||||
if 'model' in ckpt:
|
||||
ckpt = ckpt['model']
|
||||
self.model.load_state_dict(ckpt)
|
||||
|
||||
self.model.cuda().eval()
|
||||
logging.info("init done")
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def predict(self, rgb, depth, K, ob_in_cams, normal_map=None, get_vis=False, mesh=None, mesh_tensors=None, glctx=None, mesh_diameter=None):
|
||||
'''
|
||||
@rgb: np array (H,W,3)
|
||||
'''
|
||||
logging.info(f"ob_in_cams:{ob_in_cams.shape}")
|
||||
ob_in_cams = torch.as_tensor(ob_in_cams, dtype=torch.float, device='cuda')
|
||||
|
||||
logging.info(f'self.cfg.use_normal:{self.cfg.use_normal}')
|
||||
if not self.cfg.use_normal:
|
||||
normal_map = None
|
||||
|
||||
logging.info("making cropped data")
|
||||
|
||||
if mesh_tensors is None:
|
||||
mesh_tensors = make_mesh_tensors(mesh)
|
||||
|
||||
rgb = torch.as_tensor(rgb, device='cuda', dtype=torch.float)
|
||||
depth = torch.as_tensor(depth, device='cuda', dtype=torch.float)
|
||||
|
||||
pose_data = make_crop_data_batch(self.cfg.input_resize, ob_in_cams, mesh, rgb, depth, K, crop_ratio=self.cfg['crop_ratio'], glctx=glctx, mesh_tensors=mesh_tensors, dataset=self.dataset, cfg=self.cfg, mesh_diameter=mesh_diameter)
|
||||
|
||||
def find_best_among_pairs(pose_data:BatchPoseData):
|
||||
logging.info(f'pose_data.rgbAs.shape[0]: {pose_data.rgbAs.shape[0]}')
|
||||
ids = []
|
||||
scores = []
|
||||
bs = pose_data.rgbAs.shape[0]
|
||||
for b in range(0, pose_data.rgbAs.shape[0], bs):
|
||||
A = torch.cat([pose_data.rgbAs[b:b+bs].cuda(), pose_data.xyz_mapAs[b:b+bs].cuda()], dim=1).float()
|
||||
B = torch.cat([pose_data.rgbBs[b:b+bs].cuda(), pose_data.xyz_mapBs[b:b+bs].cuda()], dim=1).float()
|
||||
if pose_data.normalAs is not None:
|
||||
A = torch.cat([A, pose_data.normalAs.cuda().float()], dim=1)
|
||||
B = torch.cat([B, pose_data.normalBs.cuda().float()], dim=1)
|
||||
with torch.cuda.amp.autocast(enabled=self.amp):
|
||||
output = self.model(A, B, L=len(A))
|
||||
scores_cur = output["score_logit"].float().reshape(-1)
|
||||
ids.append(scores_cur.argmax()+b)
|
||||
scores.append(scores_cur)
|
||||
ids = torch.stack(ids, dim=0).reshape(-1)
|
||||
scores = torch.cat(scores, dim=0).reshape(-1)
|
||||
return ids, scores
|
||||
|
||||
pose_data_iter = pose_data
|
||||
global_ids = torch.arange(len(ob_in_cams), device='cuda', dtype=torch.long)
|
||||
scores_global = torch.zeros((len(ob_in_cams)), dtype=torch.float, device='cuda')
|
||||
|
||||
while 1:
|
||||
ids, scores = find_best_among_pairs(pose_data_iter)
|
||||
if len(ids)==1:
|
||||
scores_global[global_ids] = scores + 100
|
||||
break
|
||||
global_ids = global_ids[ids]
|
||||
pose_data_iter = pose_data.select_by_indices(global_ids)
|
||||
|
||||
scores = scores_global
|
||||
|
||||
logging.info(f'forward done')
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
if get_vis:
|
||||
logging.info("get_vis...")
|
||||
canvas = []
|
||||
ids = scores.argsort(descending=True)
|
||||
canvas = vis_batch_data_scores(pose_data, ids=ids, scores=scores)
|
||||
return scores, canvas
|
||||
|
||||
return scores, None
|
||||
|
||||
@@ -0,0 +1,102 @@
|
||||
# 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, field
|
||||
from typing import List, Optional, Tuple,Union
|
||||
import numpy as np
|
||||
import omegaconf
|
||||
import torch
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainingConfig(omegaconf.dictconfig.DictConfig):
|
||||
input_resize: tuple = (160, 160)
|
||||
normalize_xyz:Optional[bool] = True
|
||||
use_mask:Optional[bool] = False
|
||||
crop_ratio:Optional[float] = None
|
||||
split_objects_across_gpus: bool = True
|
||||
max_num_key: Optional[int] = None
|
||||
use_normal:bool = False
|
||||
n_view:int = 1
|
||||
zfar:float = np.inf
|
||||
c_in:int = 6
|
||||
train_num_pair:Optional[int] = None
|
||||
make_pair_online:Optional[bool] = False
|
||||
render_backend:Optional[str] = 'nvdiffrast'
|
||||
|
||||
# Run management
|
||||
run_id: Optional[str] = None
|
||||
exp_name:Optional[str] = None
|
||||
resume_run_id: Optional[str] = None
|
||||
save_dir: Optional[str] = None
|
||||
batch_size: int = 64
|
||||
epoch_size: int = 115200
|
||||
val_size: int = 1280
|
||||
n_epochs: int = 25
|
||||
save_epoch_interval: int = 100
|
||||
n_dataloader_workers: int = 20
|
||||
n_rendering_workers: int = 1
|
||||
gradient_max_norm:float = np.inf
|
||||
max_step_per_epoch: Optional[int] = 25000
|
||||
|
||||
# Network
|
||||
use_BN:bool = True
|
||||
loss_type:Optional[str] = 'pairwise_valid'
|
||||
|
||||
# Optimizer
|
||||
optimizer: str = "adam"
|
||||
weight_decay: float = 0.0
|
||||
clip_grad_norm: float = np.inf
|
||||
lr: float = 0.0001
|
||||
warmup_step: int = -1 # -1 means disable
|
||||
n_epochs_warmup: int = 1
|
||||
|
||||
# Visualization
|
||||
vis_interval: Optional[int] = 1000
|
||||
|
||||
debug: Optional[bool] = None
|
||||
|
||||
|
||||
|
||||
@dataclass
|
||||
class TrainRefinerConfig:
|
||||
# Datasets
|
||||
input_resize: tuple = (160, 160) #(W,H)
|
||||
crop_ratio:Optional[float] = None
|
||||
max_num_key: Optional[int] = None
|
||||
use_normal:bool = False
|
||||
use_mask:Optional[bool] = False
|
||||
normal_uint8:bool = False
|
||||
normalize_xyz:Optional[bool] = True
|
||||
trans_normalizer:Optional[list] = None
|
||||
rot_normalizer:Optional[float] = None
|
||||
c_in:int = 6
|
||||
n_view:int = 1
|
||||
zfar:float = np.inf
|
||||
trans_rep:str = 'tracknet' # tracknet/deepim
|
||||
rot_rep:Optional[str] = 'axis_angle' # 6d/axis_angle
|
||||
save_dir: Optional[str] = None
|
||||
|
||||
# Run management
|
||||
run_id: Optional[str] = None
|
||||
exp_name:Optional[str] = None
|
||||
batch_size: int = 64
|
||||
use_BN:bool = True
|
||||
optimizer: str = "adam"
|
||||
weight_decay: float = 0.0
|
||||
clip_grad_norm: float = np.inf
|
||||
lr: float = 0.0001
|
||||
warmup_step: int = -1
|
||||
loss_type:str = 'l2' # l1/l2/add
|
||||
|
||||
vis_interval: Optional[int] = 1000
|
||||
debug: Optional[bool] = None
|
||||
|
||||
|
||||
Reference in New Issue
Block a user