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
|
||||
|
||||
@@ -0,0 +1,138 @@
|
||||
# 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,copy,math,tqdm
|
||||
import numpy as np
|
||||
dir_path = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(dir_path)
|
||||
import torch.nn.functional as F
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import time
|
||||
import cv2
|
||||
sys.path.append(f'{dir_path}/../../../../')
|
||||
from Utils import *
|
||||
|
||||
|
||||
|
||||
class ConvBN(nn.Module):
|
||||
def __init__(self, C_in, C_out, kernel_size=3, stride=1, groups=1, bias=True,dilation=1,):
|
||||
super().__init__()
|
||||
padding = (kernel_size - 1) // 2
|
||||
self.net = nn.Sequential(
|
||||
nn.Conv2d(C_in, C_out, kernel_size, stride, padding, groups=groups, bias=bias,dilation=dilation),
|
||||
nn.BatchNorm2d(C_out),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class ConvBNReLU(nn.Module):
|
||||
def __init__(self, C_in, C_out, kernel_size=3, stride=1, groups=1, bias=True,dilation=1, norm_layer=nn.BatchNorm2d):
|
||||
super().__init__()
|
||||
padding = (kernel_size - 1) // 2
|
||||
layers = [
|
||||
nn.Conv2d(C_in, C_out, kernel_size, stride, padding, groups=groups, bias=bias,dilation=dilation),
|
||||
]
|
||||
if norm_layer is not None:
|
||||
layers.append(norm_layer(C_out))
|
||||
layers.append(nn.ReLU(inplace=True))
|
||||
self.net = nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
return self.net(x)
|
||||
|
||||
|
||||
class ConvPadding(nn.Module):
|
||||
def __init__(self,C_in, C_out, kernel_size=3, stride=1, groups=1, bias=True,dilation=1):
|
||||
super(ConvPadding, self).__init__()
|
||||
padding = (kernel_size - 1) // 2
|
||||
self.conv = nn.Conv2d(C_in, C_out, kernel_size, stride, padding, groups=groups, bias=bias,dilation=dilation)
|
||||
|
||||
def forward(self,x):
|
||||
return self.conv(x)
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1, bias=False):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=dilation, groups=groups, bias=bias, dilation=dilation)
|
||||
|
||||
|
||||
def conv1x1(in_planes, out_planes, stride=1, bias=False):
|
||||
"""1x1 convolution"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=bias)
|
||||
|
||||
class ResnetBasicBlock(nn.Module):
|
||||
__constants__ = ['downsample']
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1, base_width=64, dilation=1, norm_layer=nn.BatchNorm2d, bias=False):
|
||||
super().__init__()
|
||||
self.norm_layer = norm_layer
|
||||
if groups != 1 or base_width != 64:
|
||||
raise ValueError('BasicBlock only supports groups=1 and base_width=64')
|
||||
if dilation > 1:
|
||||
raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
|
||||
# Both self.conv1 and self.downsample layers downsample the input when stride != 1
|
||||
self.conv1 = conv3x3(inplanes, planes, stride,bias=bias)
|
||||
if self.norm_layer is not None:
|
||||
self.bn1 = norm_layer(planes)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes,bias=bias)
|
||||
if self.norm_layer is not None:
|
||||
self.bn2 = norm_layer(planes)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
identity = x
|
||||
|
||||
out = self.conv1(x)
|
||||
if self.norm_layer is not None:
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
if self.norm_layer is not None:
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
identity = self.downsample(x)
|
||||
out += identity
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
|
||||
class PositionalEmbedding(nn.Module):
|
||||
def __init__(self, d_model, max_len=512):
|
||||
super().__init__()
|
||||
|
||||
# Compute the positional encodings once in log space.
|
||||
pe = torch.zeros(max_len, d_model).float()
|
||||
pe.require_grad = False
|
||||
|
||||
position = torch.arange(0, max_len).float().unsqueeze(1) #(N,1)
|
||||
div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()[None]
|
||||
|
||||
pe[:, 0::2] = torch.sin(position * div_term) #(N, d_model/2)
|
||||
pe[:, 1::2] = torch.cos(position * div_term)
|
||||
|
||||
pe = pe.unsqueeze(0)
|
||||
self.register_buffer('pe', pe) #(1, max_len, D)
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
'''
|
||||
@x: (B,N,D)
|
||||
'''
|
||||
return x + self.pe[:, :x.size(1)]
|
||||
|
||||
@@ -0,0 +1,93 @@
|
||||
# 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
|
||||
import numpy as np
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(code_dir)
|
||||
sys.path.append(f'{code_dir}/../../../../')
|
||||
from Utils import *
|
||||
import torch.nn.functional as F
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import cv2
|
||||
from functools import partial
|
||||
from network_modules import *
|
||||
from Utils import *
|
||||
|
||||
|
||||
|
||||
class RefineNet(nn.Module):
|
||||
def __init__(self, cfg=None, c_in=4, n_view=1):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
if self.cfg.use_BN:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
norm_layer1d = nn.BatchNorm1d
|
||||
else:
|
||||
norm_layer = None
|
||||
norm_layer1d = None
|
||||
|
||||
self.encodeA = nn.Sequential(
|
||||
ConvBNReLU(C_in=c_in,C_out=64,kernel_size=7,stride=2, norm_layer=norm_layer),
|
||||
ConvBNReLU(C_in=64,C_out=128,kernel_size=3,stride=2, norm_layer=norm_layer),
|
||||
ResnetBasicBlock(128,128,bias=True, norm_layer=norm_layer),
|
||||
ResnetBasicBlock(128,128,bias=True, norm_layer=norm_layer),
|
||||
)
|
||||
|
||||
self.encodeAB = nn.Sequential(
|
||||
ResnetBasicBlock(256,256,bias=True, norm_layer=norm_layer),
|
||||
ResnetBasicBlock(256,256,bias=True, norm_layer=norm_layer),
|
||||
ConvBNReLU(256,512,kernel_size=3,stride=2, norm_layer=norm_layer),
|
||||
ResnetBasicBlock(512,512,bias=True, norm_layer=norm_layer),
|
||||
ResnetBasicBlock(512,512,bias=True, norm_layer=norm_layer),
|
||||
)
|
||||
|
||||
embed_dim = 512
|
||||
num_heads = 4
|
||||
self.pos_embed = PositionalEmbedding(d_model=embed_dim, max_len=400)
|
||||
|
||||
self.trans_head = nn.Sequential(
|
||||
nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dim_feedforward=512, batch_first=True),
|
||||
nn.Linear(512, 3),
|
||||
)
|
||||
|
||||
if self.cfg['rot_rep']=='axis_angle':
|
||||
rot_out_dim = 3
|
||||
elif self.cfg['rot_rep']=='6d':
|
||||
rot_out_dim = 6
|
||||
else:
|
||||
raise RuntimeError
|
||||
self.rot_head = nn.Sequential(
|
||||
nn.TransformerEncoderLayer(d_model=embed_dim, nhead=num_heads, dim_feedforward=512, batch_first=True),
|
||||
nn.Linear(512, rot_out_dim),
|
||||
)
|
||||
|
||||
|
||||
def forward(self, A, B):
|
||||
"""
|
||||
@A: (B,C,H,W)
|
||||
"""
|
||||
bs = len(A)
|
||||
output = {}
|
||||
|
||||
x = torch.cat([A,B], dim=0)
|
||||
x = self.encodeA(x)
|
||||
a = x[:bs]
|
||||
b = x[bs:]
|
||||
|
||||
ab = torch.cat((a,b),1).contiguous()
|
||||
ab = self.encodeAB(ab) #(B,C,H,W)
|
||||
|
||||
ab = self.pos_embed(ab.reshape(bs, ab.shape[1], -1).permute(0,2,1))
|
||||
|
||||
output['trans'] = self.trans_head(ab).mean(dim=1)
|
||||
output['rot'] = self.rot_head(ab).mean(dim=1)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,90 @@
|
||||
# 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
|
||||
import numpy as np
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(code_dir)
|
||||
sys.path.append(f'{code_dir}/../../../../')
|
||||
from Utils import *
|
||||
from functools import partial
|
||||
import torch.nn.functional as F
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import cv2
|
||||
from network_modules import *
|
||||
from Utils import *
|
||||
|
||||
|
||||
|
||||
|
||||
class ScoreNetMultiPair(nn.Module):
|
||||
def __init__(self, cfg=None, c_in=4):
|
||||
super().__init__()
|
||||
self.cfg = cfg
|
||||
if self.cfg.use_BN:
|
||||
norm_layer = nn.BatchNorm2d
|
||||
else:
|
||||
norm_layer = None
|
||||
|
||||
self.encoderA = nn.Sequential(
|
||||
ConvBNReLU(C_in=c_in,C_out=64,kernel_size=7,stride=2, norm_layer=norm_layer),
|
||||
ConvBNReLU(C_in=64,C_out=128,kernel_size=3,stride=2, norm_layer=norm_layer),
|
||||
ResnetBasicBlock(128,128,bias=True, norm_layer=norm_layer),
|
||||
ResnetBasicBlock(128,128,bias=True, norm_layer=norm_layer),
|
||||
)
|
||||
|
||||
self.encoderAB = nn.Sequential(
|
||||
ResnetBasicBlock(256,256,bias=True, norm_layer=norm_layer),
|
||||
ResnetBasicBlock(256,256,bias=True, norm_layer=norm_layer),
|
||||
ConvBNReLU(256,512,kernel_size=3,stride=2, norm_layer=norm_layer),
|
||||
ResnetBasicBlock(512,512,bias=True, norm_layer=norm_layer),
|
||||
ResnetBasicBlock(512,512,bias=True, norm_layer=norm_layer),
|
||||
)
|
||||
|
||||
embed_dim = 512
|
||||
num_heads = 4
|
||||
self.att = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, bias=True, batch_first=True)
|
||||
self.att_cross = nn.MultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, bias=True, batch_first=True)
|
||||
|
||||
self.pos_embed = PositionalEmbedding(d_model=embed_dim, max_len=400)
|
||||
self.linear = nn.Linear(embed_dim, 1)
|
||||
|
||||
|
||||
def extract_feat(self, A, B):
|
||||
"""
|
||||
@A: (B*L,C,H,W) L is num of pairs
|
||||
"""
|
||||
bs = A.shape[0] # B*L
|
||||
|
||||
x = torch.cat([A,B], dim=0)
|
||||
x = self.encoderA(x)
|
||||
a = x[:bs]
|
||||
b = x[bs:]
|
||||
ab = torch.cat((a,b), dim=1)
|
||||
ab = self.encoderAB(ab)
|
||||
ab = self.pos_embed(ab.reshape(bs, ab.shape[1], -1).permute(0,2,1))
|
||||
ab, _ = self.att(ab, ab, ab)
|
||||
return ab.mean(dim=1).reshape(bs,-1)
|
||||
|
||||
|
||||
def forward(self, A, B, L):
|
||||
"""
|
||||
@A: (B*L,C,H,W) L is num of pairs
|
||||
@L: num of pairs
|
||||
"""
|
||||
output = {}
|
||||
bs = A.shape[0]//L
|
||||
feats = self.extract_feat(A, B) #(B*L, C)
|
||||
x = feats.reshape(bs,L,-1)
|
||||
x, _ = self.att_cross(x, x, x)
|
||||
|
||||
output['score_logit'] = self.linear(x).reshape(bs,L) # (B,L)
|
||||
|
||||
return output
|
||||
@@ -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