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
+1
@@ -0,0 +1 @@
|
||||
../dyn-hamr/optim/BMC
|
||||
+1
@@ -0,0 +1 @@
|
||||
../../../hamer/_DATA/data/mano
|
||||
+1
@@ -0,0 +1 @@
|
||||
../../../hamer/_DATA/data/mano_mean_params.npz
|
||||
+1
@@ -0,0 +1 @@
|
||||
../../../hamer/_DATA/hamer_ckpts/checkpoints
|
||||
@@ -0,0 +1,42 @@
|
||||
COCOW-TRAIN:
|
||||
TYPE: ImageDataset
|
||||
URLS: hamer_training_data/dataset_tars/cocow-train/{000000..000036}.tar
|
||||
epoch_size: 78666
|
||||
DEX-TRAIN:
|
||||
TYPE: ImageDataset
|
||||
URLS: hamer_training_data/dataset_tars/dex-train/{000000..000406}.tar
|
||||
epoch_size: 406888
|
||||
FREIHAND-MOCAP:
|
||||
DATASET_FILE: hamer_training_data/freihand_mocap.npz
|
||||
FREIHAND-TRAIN:
|
||||
TYPE: ImageDataset
|
||||
URLS: hamer_training_data/dataset_tars/freihand-train/{000000..000130}.tar
|
||||
epoch_size: 130240
|
||||
H2O3D-TRAIN:
|
||||
TYPE: ImageDataset
|
||||
URLS: hamer_training_data/dataset_tars/h2o3d-train/{000000..000060}.tar
|
||||
epoch_size: 121996
|
||||
HALPE-TRAIN:
|
||||
TYPE: ImageDataset
|
||||
URLS: hamer_training_data/dataset_tars/halpe-train/{000000..000022}.tar
|
||||
epoch_size: 34289
|
||||
HO3D-TRAIN:
|
||||
TYPE: ImageDataset
|
||||
URLS: hamer_training_data/dataset_tars/ho3d-train/{000000..000083}.tar
|
||||
epoch_size: 83325
|
||||
INTERHAND26M-TRAIN:
|
||||
TYPE: ImageDataset
|
||||
URLS: hamer_training_data/dataset_tars/interhand26m-train/{000000..001056}.tar
|
||||
epoch_size: 1424632
|
||||
MPIINZSL-TRAIN:
|
||||
TYPE: ImageDataset
|
||||
URLS: hamer_training_data/dataset_tars/mpiinzsl-train/{000000..000015}.tar
|
||||
epoch_size: 15184
|
||||
MTC-TRAIN:
|
||||
TYPE: ImageDataset
|
||||
URLS: hamer_training_data/dataset_tars/mtc-train/{000000..000306}.tar
|
||||
epoch_size: 363947
|
||||
RHD-TRAIN:
|
||||
TYPE: ImageDataset
|
||||
URLS: hamer_training_data/dataset_tars/rhd-train/{000000..000041}.tar
|
||||
epoch_size: 61705
|
||||
@@ -0,0 +1,111 @@
|
||||
task_name: train
|
||||
tags:
|
||||
- dev
|
||||
train: true
|
||||
test: false
|
||||
ckpt_path: null
|
||||
seed: null
|
||||
DATASETS:
|
||||
TRAIN:
|
||||
FREIHAND-TRAIN:
|
||||
WEIGHT: 0.25
|
||||
INTERHAND26M-TRAIN:
|
||||
WEIGHT: 0.25
|
||||
MTC-TRAIN:
|
||||
WEIGHT: 0.1
|
||||
RHD-TRAIN:
|
||||
WEIGHT: 0.05
|
||||
COCOW-TRAIN:
|
||||
WEIGHT: 0.1
|
||||
HALPE-TRAIN:
|
||||
WEIGHT: 0.05
|
||||
MPIINZSL-TRAIN:
|
||||
WEIGHT: 0.05
|
||||
HO3D-TRAIN:
|
||||
WEIGHT: 0.05
|
||||
H2O3D-TRAIN:
|
||||
WEIGHT: 0.05
|
||||
DEX-TRAIN:
|
||||
WEIGHT: 0.05
|
||||
VAL:
|
||||
FREIHAND-TRAIN:
|
||||
WEIGHT: 1.0
|
||||
MOCAP: FREIHAND-MOCAP
|
||||
BETAS_REG: true
|
||||
CONFIG:
|
||||
SCALE_FACTOR: 0.3
|
||||
ROT_FACTOR: 30
|
||||
TRANS_FACTOR: 0.02
|
||||
COLOR_SCALE: 0.2
|
||||
ROT_AUG_RATE: 0.6
|
||||
TRANS_AUG_RATE: 0.5
|
||||
DO_FLIP: false
|
||||
FLIP_AUG_RATE: 0.0
|
||||
EXTREME_CROP_AUG_RATE: 0.0
|
||||
EXTREME_CROP_AUG_LEVEL: 1
|
||||
extras:
|
||||
ignore_warnings: false
|
||||
enforce_tags: true
|
||||
print_config: true
|
||||
exp_name: hamer
|
||||
MANO:
|
||||
DATA_DIR: _DATA/data/
|
||||
MODEL_PATH: data/mano
|
||||
GENDER: neutral
|
||||
NUM_HAND_JOINTS: 15
|
||||
MEAN_PARAMS: data/mano_mean_params.npz
|
||||
CREATE_BODY_POSE: false
|
||||
EXTRA:
|
||||
FOCAL_LENGTH: 5000
|
||||
NUM_LOG_IMAGES: 4
|
||||
NUM_LOG_SAMPLES_PER_IMAGE: 8
|
||||
PELVIS_IND: 0
|
||||
GENERAL:
|
||||
TOTAL_STEPS: 1000000
|
||||
LOG_STEPS: 1000
|
||||
VAL_STEPS: 1000
|
||||
CHECKPOINT_STEPS: 10000
|
||||
CHECKPOINT_SAVE_TOP_K: 1
|
||||
NUM_WORKERS: 8
|
||||
PREFETCH_FACTOR: 2
|
||||
TRAIN:
|
||||
LR: 1.0e-05
|
||||
WEIGHT_DECAY: 0.0001
|
||||
BATCH_SIZE: 32
|
||||
LOSS_REDUCTION: mean
|
||||
NUM_TRAIN_SAMPLES: 2
|
||||
NUM_TEST_SAMPLES: 64
|
||||
POSE_2D_NOISE_RATIO: 0.01
|
||||
SMPL_PARAM_NOISE_RATIO: 0.005
|
||||
MODEL:
|
||||
IMAGE_SIZE: 256
|
||||
IMAGE_MEAN:
|
||||
- 0.485
|
||||
- 0.456
|
||||
- 0.406
|
||||
IMAGE_STD:
|
||||
- 0.229
|
||||
- 0.224
|
||||
- 0.225
|
||||
BACKBONE:
|
||||
TYPE: vit
|
||||
PRETRAINED_WEIGHTS: hamer_training_data/vitpose_backbone.pth
|
||||
MANO_HEAD:
|
||||
TYPE: transformer_decoder
|
||||
IN_CHANNELS: 2048
|
||||
TRANSFORMER_DECODER:
|
||||
depth: 6
|
||||
heads: 8
|
||||
mlp_dim: 1024
|
||||
dim_head: 64
|
||||
dropout: 0.0
|
||||
emb_dropout: 0.0
|
||||
norm: layer
|
||||
context_dim: 1280
|
||||
LOSS_WEIGHTS:
|
||||
KEYPOINTS_3D: 0.05
|
||||
KEYPOINTS_2D: 0.01
|
||||
GLOBAL_ORIENT: 0.001
|
||||
HAND_POSE: 0.001
|
||||
BETAS: 0.0005
|
||||
ADVERSARIAL: 0.0005
|
||||
BIN
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
Binary file not shown.
Binary file not shown.
+6
-1
@@ -1037,6 +1037,11 @@ def batch_rodrigues(param):
|
||||
def smooth_global_rot_matrix(pred_rots, OE_filter):
|
||||
rot_mat = batch_rodrigues(pred_rots[None]).squeeze(0)
|
||||
smoothed_rot_mat = OE_filter.process(rot_mat)
|
||||
# Elementwise filtering leaves SO(3); project back before axis-angle conversion.
|
||||
u, _, vh = torch.linalg.svd(smoothed_rot_mat)
|
||||
correction = torch.eye(3, dtype=u.dtype, device=u.device)
|
||||
correction[-1, -1] = torch.det(u @ vh)
|
||||
smoothed_rot_mat = u @ correction @ vh
|
||||
smoothed_rot = rotation_matrix_to_angle_axis(smoothed_rot_mat.reshape(1,3,3)).reshape(-1)
|
||||
return smoothed_rot
|
||||
|
||||
@@ -1106,4 +1111,4 @@ class OneEuroFilter:
|
||||
cutoff = self.mincutoff + self.beta * torch.abs(edx)
|
||||
if print_inter:
|
||||
print(self.compute_alpha(cutoff))
|
||||
return self.x_filter.process(x, self.compute_alpha(cutoff))
|
||||
return self.x_filter.process(x, self.compute_alpha(cutoff))
|
||||
|
||||
+2
@@ -177,6 +177,8 @@ class AnimationBase(object):
|
||||
"""
|
||||
Add a single camera marker node that we move around when we animate
|
||||
"""
|
||||
if os.environ.get("DYNHAMR_HIDE_CAMERA_MARKERS") == "1":
|
||||
return
|
||||
print("ADDING CAMERA MARKERS")
|
||||
if self.cam_marker_node is None:
|
||||
cam_marker = make_camera_marker(up="y")
|
||||
|
||||
+1
@@ -0,0 +1 @@
|
||||
../../../../../weights/detector.pt
|
||||
Vendored
+94
@@ -0,0 +1,94 @@
|
||||
Copyright (c) 2022-Present, NVIDIA Corporation & affiliates. All rights reserved.
|
||||
|
||||
|
||||
=======================================================================
|
||||
|
||||
1. Definitions
|
||||
|
||||
"Licensor" means any person or entity that distributes its Work.
|
||||
|
||||
"Software" means the original work of authorship made available under
|
||||
this License.
|
||||
|
||||
"Work" means the Software and any additions to or derivative works of
|
||||
the Software that are made available under this License.
|
||||
|
||||
The terms "reproduce," "reproduction," "derivative works," and
|
||||
"distribution" have the meaning as provided under U.S. copyright law;
|
||||
provided, however, that for the purposes of this License, derivative
|
||||
works shall not include works that remain separable from, or merely
|
||||
link (or bind by name) to the interfaces of, the Work.
|
||||
|
||||
Works, including the Software, are "made available" under this License
|
||||
by including in or with the Work either (a) a copyright notice
|
||||
referencing the applicability of this License to the Work, or (b) a
|
||||
copy of this License.
|
||||
|
||||
2. License Grants
|
||||
|
||||
2.1 Copyright Grant. Subject to the terms and conditions of this
|
||||
License, each Licensor grants to you a perpetual, worldwide,
|
||||
non-exclusive, royalty-free, copyright license to reproduce,
|
||||
prepare derivative works of, publicly display, publicly perform,
|
||||
sublicense and distribute its Work and any resulting derivative
|
||||
works in any form.
|
||||
|
||||
3. Limitations
|
||||
|
||||
3.1 Redistribution. You may reproduce or distribute the Work only
|
||||
if (a) you do so under this License, (b) you include a complete
|
||||
copy of this License with your distribution, and (c) you retain
|
||||
without modification any copyright, patent, trademark, or
|
||||
attribution notices that are present in the Work.
|
||||
|
||||
3.2 Derivative Works. You may specify that additional or different
|
||||
terms apply to the use, reproduction, and distribution of your
|
||||
derivative works of the Work ("Your Terms") only if (a) Your Terms
|
||||
provide that the use limitation in Section 3.3 applies to your
|
||||
derivative works, and (b) you identify the specific derivative
|
||||
works that are subject to Your Terms. Notwithstanding Your Terms,
|
||||
this License (including the redistribution requirements in Section
|
||||
3.1) will continue to apply to the Work itself.
|
||||
|
||||
3.3 Use Limitation. The Work and any derivative works thereof only
|
||||
may be used or intended for use non-commercially. Notwithstanding
|
||||
the foregoing, NVIDIA and its affiliates may use the Work and any
|
||||
derivative works commercially. As used herein, "non-commercially"
|
||||
means for research or evaluation purposes only.
|
||||
|
||||
3.4 Patent Claims. If you bring or threaten to bring a patent claim
|
||||
against any Licensor (including any claim, cross-claim or
|
||||
counterclaim in a lawsuit) to enforce any patents that you allege
|
||||
are infringed by any Work, then your rights under this License from
|
||||
such Licensor (including the grant in Section 2.1) will terminate
|
||||
immediately.
|
||||
|
||||
3.5 Trademarks. This License does not grant any rights to use any
|
||||
Licensor�s or its affiliates� names, logos, or trademarks, except
|
||||
as necessary to reproduce the notices described in this License.
|
||||
|
||||
3.6 Termination. If you violate any term of this License, then your
|
||||
rights under this License (including the grant in Section 2.1) will
|
||||
terminate immediately.
|
||||
|
||||
4. Disclaimer of Warranty.
|
||||
|
||||
THE WORK IS PROVIDED "AS IS" WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
|
||||
MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR
|
||||
NON-INFRINGEMENT. YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER
|
||||
THIS LICENSE.
|
||||
|
||||
5. Limitation of Liability.
|
||||
|
||||
EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL
|
||||
THEORY, WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE
|
||||
SHALL ANY LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT,
|
||||
INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF
|
||||
OR RELATED TO THIS LICENSE, THE USE OR INABILITY TO USE THE WORK
|
||||
(INCLUDING BUT NOT LIMITED TO LOSS OF GOODWILL, BUSINESS INTERRUPTION,
|
||||
LOST PROFITS OR DATA, COMPUTER FAILURE OR MALFUNCTION, OR ANY OTHER
|
||||
COMMERCIAL DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN ADVISED OF
|
||||
THE POSSIBILITY OF SUCH DAMAGES.
|
||||
|
||||
=======================================================================
|
||||
Vendored
+1022
File diff suppressed because it is too large
Load Diff
+8
@@ -0,0 +1,8 @@
|
||||
DIR=$(pwd)
|
||||
|
||||
cd $DIR/mycpp/ && mkdir -p build && cd build && cmake .. -DPYTHON_EXECUTABLE=$(which python) && make -j11
|
||||
cd /kaolin && rm -rf build *egg* && pip install -e .
|
||||
# Optional: BundleSDF CUDA ops (required for model-free NeRF path only)
|
||||
# cd $DIR/bundlesdf/mycuda && rm -rf build *egg* && pip install -e .
|
||||
|
||||
cd "${DIR}"
|
||||
+23
@@ -0,0 +1,23 @@
|
||||
#!/usr/bin/env bash
|
||||
# Build native extensions for a local conda environment.
|
||||
# BundleSDF CUDA ops (bundlesdf/mycuda) are omitted; model-free NeRF path needs them separately.
|
||||
set -euo pipefail
|
||||
|
||||
PROJ_ROOT="$(cd -- "$(dirname "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
|
||||
|
||||
: "${CONDA_PREFIX:?Set CONDA_PREFIX by activating your conda env first.}"
|
||||
|
||||
export CMAKE_PREFIX_PATH="${CONDA_PREFIX}:${CMAKE_PREFIX_PATH:-}"
|
||||
|
||||
cd "${PROJ_ROOT}/mycpp"
|
||||
rm -rf build
|
||||
mkdir -p build
|
||||
cd build
|
||||
cmake .. \
|
||||
-DCMAKE_BUILD_TYPE=Release \
|
||||
-DCMAKE_PREFIX_PATH="${CMAKE_PREFIX_PATH}" \
|
||||
-DPython3_ROOT_DIR="${CONDA_PREFIX}" \
|
||||
-DPYBIND11_PYTHON_EXECUTABLE="${CONDA_PREFIX}/bin/python"
|
||||
cmake --build . -j"$(nproc)"
|
||||
|
||||
cd "${PROJ_ROOT}"
|
||||
@@ -0,0 +1,105 @@
|
||||
notes: ''
|
||||
n_step: 1000
|
||||
netdepth: 8
|
||||
netwidth: 256
|
||||
netdepth_fine: 8
|
||||
netwidth_fine: 256
|
||||
N_rand: 2048 # Batch number of rays
|
||||
first_frame_ray_in_batch: 0
|
||||
lrate: 0.01
|
||||
lrate_pose: 0.01
|
||||
pose_optimize_start: 0
|
||||
decay_rate: 0.1
|
||||
chunk: 99999999999
|
||||
netchunk: 6553600
|
||||
no_batching: 0
|
||||
amp: true
|
||||
|
||||
N_samples: 128 #number of coarse samples per ray
|
||||
N_samples_around_depth: 128
|
||||
N_importance: 0
|
||||
N_importance_iter: 1
|
||||
perturb: 1
|
||||
use_viewdirs: 1
|
||||
i_embed: 1 #set 1 for hashed embedding, 0 for default positional encoding, 2 for spherical; 3 for octree grid
|
||||
i_embed_views: 2 #set 1 for hashed embedding, 0 for default positional encoding, 2 for spherical
|
||||
multires: 8 #log2 of max freq for positional encoding (3D location)
|
||||
multires_views: 3 #log2 of max freq for positional encoding (2D direction)
|
||||
feature_grid_dim: 2
|
||||
raw_noise_std: 0
|
||||
white_bkgd: 0
|
||||
gradient_max_norm: 0.1
|
||||
gradient_pose_max_norm: 0.1
|
||||
|
||||
# logging/saving options
|
||||
i_print: 500
|
||||
i_img: 500
|
||||
i_weights: 500
|
||||
i_mesh: 500
|
||||
i_nerf_normals: 500
|
||||
i_save_ray: 500
|
||||
i_pose: 500
|
||||
save_octree_clouds: True
|
||||
|
||||
finest_res: 256
|
||||
base_res: 32
|
||||
num_levels: 16
|
||||
log2_hashmap_size: 22
|
||||
datadir: /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/github/HashNeRF-pytorch/data/bundlesdf_bundlesdf_e03000196
|
||||
n_train_image: 300
|
||||
use_octree: 1
|
||||
first_frame_weight: 1
|
||||
denoise_depth_use_octree_cloud: true
|
||||
octree_embed_base_voxel_size: 0.02
|
||||
octree_smallest_voxel_size: 0.02 # This determines the smallest feature vox size
|
||||
octree_raytracing_voxel_size: 0.02
|
||||
octree_dilate_size: 0.02 # meters
|
||||
down_scale_ratio: 1
|
||||
bounding_box: [[-1,-1,-1], [1,1,1]]
|
||||
farthest_pose_sampling: 0 # Sampling train images. This replace uniform skip
|
||||
use_mask: 1
|
||||
dilate_mask_size: 0
|
||||
rays_valid_depth_only: true
|
||||
near: 0.1
|
||||
far: 2
|
||||
rgb_weight: 1
|
||||
depth_weight: 0
|
||||
trunc: 0.01 #length of the truncation region in meters
|
||||
trunc_start: 0.01
|
||||
sdf_lambda: 5
|
||||
neg_trunc_ratio: 1 # -trunc distance ratio compared to +trunc
|
||||
trunc_decay_type: ''
|
||||
sdf_loss_type: l2
|
||||
fs_weight: 1000
|
||||
empty_weight: 1
|
||||
fs_rgb_weight: 0
|
||||
trunc_weight: 6000
|
||||
sparse_loss_weight: 0
|
||||
tv_loss_weight: 0
|
||||
frame_features: 2 #number of channels of the learnable per-frame features
|
||||
optimize_poses: 1 #optimize a pose refinement for the initial poses
|
||||
pose_reg_weight: 0
|
||||
point_cloud_loss_weight: 0
|
||||
point_cloud_loss_normal_weight: 0
|
||||
eikonal_weight: 0
|
||||
normal_loss_weight: 0
|
||||
feature_reg_weight: 0.1
|
||||
deformation_reg_weight: 0
|
||||
use_deformation_field: 0 #use a deformation field to account for inaccuracies in intrinsic parameters
|
||||
max_pix_shift: 0.02 # Pixel shift ratio relative to image size
|
||||
share_coarse_fine: 1
|
||||
mode: sdf
|
||||
fs_sdf: 1 # Uncertain free space
|
||||
crop: 0
|
||||
mesh_resolution: 0.003
|
||||
max_trans: 0.02 # meters
|
||||
max_rot: 10 # deg
|
||||
|
||||
continual: True
|
||||
|
||||
######### dbscan
|
||||
dbscan_eps: 0.01
|
||||
dbscan_eps_min_samples: 1
|
||||
|
||||
####### bundlenerf
|
||||
sync_max_delay: 0 # 0 for strict sync
|
||||
@@ -0,0 +1,105 @@
|
||||
notes: ''
|
||||
n_step: 1000
|
||||
netdepth: 8
|
||||
netwidth: 256
|
||||
netdepth_fine: 8
|
||||
netwidth_fine: 256
|
||||
N_rand: 2048 # Batch number of rays
|
||||
first_frame_ray_in_batch: 0
|
||||
lrate: 0.01
|
||||
lrate_pose: 0.01
|
||||
pose_optimize_start: 0
|
||||
decay_rate: 0.1
|
||||
chunk: 99999999999
|
||||
netchunk: 6553600
|
||||
no_batching: 0
|
||||
amp: true
|
||||
|
||||
N_samples: 128 #number of coarse samples per ray
|
||||
N_samples_around_depth: 128
|
||||
N_importance: 0
|
||||
N_importance_iter: 1
|
||||
perturb: 1
|
||||
use_viewdirs: 1
|
||||
i_embed: 1 #set 1 for hashed embedding, 0 for default positional encoding, 2 for spherical; 3 for octree grid
|
||||
i_embed_views: 2 #set 1 for hashed embedding, 0 for default positional encoding, 2 for spherical
|
||||
multires: 8 #log2 of max freq for positional encoding (3D location)
|
||||
multires_views: 3 #log2 of max freq for positional encoding (2D direction)
|
||||
feature_grid_dim: 2
|
||||
raw_noise_std: 0
|
||||
white_bkgd: 0
|
||||
gradient_max_norm: 0.1
|
||||
gradient_pose_max_norm: 0.1
|
||||
|
||||
# logging/saving options
|
||||
i_print: 500
|
||||
i_img: 500
|
||||
i_weights: 500
|
||||
i_mesh: 500
|
||||
i_nerf_normals: 500
|
||||
i_save_ray: 500
|
||||
i_pose: 500
|
||||
save_octree_clouds: True
|
||||
|
||||
finest_res: 512
|
||||
base_res: 32
|
||||
num_levels: 16
|
||||
log2_hashmap_size: 22
|
||||
datadir: /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/github/HashNeRF-pytorch/data/bundlesdf_bundlesdf_e03000196
|
||||
n_train_image: 300
|
||||
use_octree: 1
|
||||
first_frame_weight: 1
|
||||
denoise_depth_use_octree_cloud: true
|
||||
octree_embed_base_voxel_size: 0.02
|
||||
octree_smallest_voxel_size: 0.02 # This determines the smallest feature vox size
|
||||
octree_raytracing_voxel_size: 0.02
|
||||
octree_dilate_size: 0.02 # meters
|
||||
down_scale_ratio: 1
|
||||
bounding_box: [[-1,-1,-1], [1,1,1]]
|
||||
farthest_pose_sampling: 0 # Sampling train images. This replace uniform skip
|
||||
use_mask: 1
|
||||
dilate_mask_size: 0
|
||||
rays_valid_depth_only: true
|
||||
near: 0.1
|
||||
far: 2
|
||||
rgb_weight: 100
|
||||
depth_weight: 0
|
||||
trunc: 0.01 #length of the truncation region in meters
|
||||
trunc_start: 0.01
|
||||
sdf_lambda: 5
|
||||
neg_trunc_ratio: 1 # -trunc distance ratio compared to +trunc
|
||||
trunc_decay_type: ''
|
||||
sdf_loss_type: l2
|
||||
fs_weight: 100
|
||||
empty_weight: 1
|
||||
fs_rgb_weight: 0
|
||||
trunc_weight: 6000
|
||||
sparse_loss_weight: 0
|
||||
tv_loss_weight: 0
|
||||
frame_features: 2 #number of channels of the learnable per-frame features
|
||||
optimize_poses: 1 #optimize a pose refinement for the initial poses
|
||||
pose_reg_weight: 0
|
||||
point_cloud_loss_weight: 0
|
||||
point_cloud_loss_normal_weight: 0
|
||||
eikonal_weight: 0
|
||||
normal_loss_weight: 0
|
||||
feature_reg_weight: 0.1
|
||||
deformation_reg_weight: 0
|
||||
use_deformation_field: 0 #use a deformation field to account for inaccuracies in intrinsic parameters
|
||||
max_pix_shift: 0.02 # Pixel shift ratio relative to image size
|
||||
share_coarse_fine: 1
|
||||
mode: sdf
|
||||
fs_sdf: 1 # Uncertain free space
|
||||
crop: 0
|
||||
mesh_resolution: 0.003
|
||||
max_trans: 0.02 # meters
|
||||
max_rot: 10 # deg
|
||||
|
||||
continual: True
|
||||
|
||||
######### dbscan
|
||||
dbscan_eps: 0.01
|
||||
dbscan_eps_min_samples: 1
|
||||
|
||||
####### bundlenerf
|
||||
sync_max_delay: 0 # 0 for strict sync
|
||||
@@ -0,0 +1,20 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include "common.h"
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("sampleRaysUniformOccupiedVoxels", &sampleRaysUniformOccupiedVoxels);
|
||||
m.def("postprocessOctreeRayTracing", &postprocessOctreeRayTracing);
|
||||
m.def("rayColorToTextureImageCUDA", &rayColorToTextureImageCUDA);
|
||||
|
||||
}
|
||||
@@ -0,0 +1,313 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <bits/stdc++.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/torch.h>
|
||||
#include <stdint.h>
|
||||
#include <cstdio>
|
||||
#include <stdint.h>
|
||||
#include <stdexcept>
|
||||
#include <limits>
|
||||
#include <set>
|
||||
#include "common.h"
|
||||
#include "Eigen/Dense"
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
/**
|
||||
* @brief
|
||||
*
|
||||
* @tparam scalar_t
|
||||
* @param z_sampled
|
||||
* @param z_in_out
|
||||
* @param z_vals
|
||||
* @return __global__
|
||||
*/
|
||||
template <typename scalar_t>
|
||||
__global__ void sample_rays_uniform_occupied_voxels_kernel(const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> z_sampled, const torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> z_in_out, torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> z_vals)
|
||||
{
|
||||
const int i_ray = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
const int i_sample = blockIdx.y * blockDim.y + threadIdx.y;
|
||||
if (i_ray>=z_sampled.size(0)) return;
|
||||
if (i_sample>=z_sampled.size(1)) return;
|
||||
|
||||
int i_box = 0;
|
||||
float z_remain = z_sampled[i_ray][i_sample];
|
||||
auto z_in_out_cur_ray = z_in_out[i_ray];
|
||||
const float eps = 1e-4;
|
||||
const int max_n_box = z_in_out.size(1);
|
||||
|
||||
if (z_in_out_cur_ray[0][0]==0) return;
|
||||
|
||||
while (1)
|
||||
{
|
||||
if (i_box>=max_n_box)
|
||||
{
|
||||
if (z_remain<=eps)
|
||||
{
|
||||
z_vals[i_ray][i_sample] = z_in_out_cur_ray[max_n_box-1][1];
|
||||
}
|
||||
else
|
||||
{
|
||||
printf("ERROR sample_rays_uniform_occupied_voxels_kernel: z_remain=%f, i_ray=%d, i_sample=%d, i_box=%d, z_in_out_cur_ray=(%f,%f)\n",z_remain,i_ray,i_sample,i_box,z_in_out_cur_ray[i_box][0],z_in_out_cur_ray[i_box][1]);
|
||||
for (int i=0;i<z_in_out.size(1);i++)
|
||||
{
|
||||
printf("z_in_out_cur_ray[%d]=(%f,%f)\n",i,z_in_out_cur_ray[i][0],z_in_out_cur_ray[i][1]);
|
||||
}
|
||||
while (1){};
|
||||
}
|
||||
|
||||
|
||||
return;
|
||||
}
|
||||
|
||||
if (z_in_out_cur_ray[i_box][0]==0)
|
||||
{
|
||||
if (z_remain<=eps && i_box>=1)
|
||||
{
|
||||
z_vals[i_ray][i_sample] = z_in_out_cur_ray[i_box-1][1];
|
||||
return;
|
||||
}
|
||||
else
|
||||
{
|
||||
printf("ERROR sample_rays_uniform_occupied_voxels_kernel: z_remain=%f, i_ray=%d, i_sample=%d, i_box=%d, z_in_out_cur_ray=(%f,%f)\n",z_remain,i_ray,i_sample,i_box,z_in_out_cur_ray[i_box][0],z_in_out_cur_ray[i_box][1]);
|
||||
for (int i=0;i<z_in_out.size(1);i++)
|
||||
{
|
||||
printf("z_in_out_cur_ray[%d]=(%f,%f)\n",i,z_in_out_cur_ray[i][0],z_in_out_cur_ray[i][1]);
|
||||
}
|
||||
while (1){};
|
||||
}
|
||||
}
|
||||
|
||||
float box_len = z_in_out_cur_ray[i_box][1]-z_in_out_cur_ray[i_box][0];
|
||||
if (z_remain<=box_len)
|
||||
{
|
||||
z_vals[i_ray][i_sample] = z_in_out_cur_ray[i_box][0] + z_remain;
|
||||
return;
|
||||
}
|
||||
z_remain -= box_len;
|
||||
i_box++;
|
||||
}
|
||||
}
|
||||
|
||||
at::Tensor sampleRaysUniformOccupiedVoxels(const at::Tensor z_in_out, const at::Tensor z_sampled, at::Tensor z_vals)
|
||||
{
|
||||
CHECK_INPUT(z_in_out);
|
||||
CHECK_INPUT(z_sampled);
|
||||
CHECK_INPUT(z_vals);
|
||||
AT_ASSERTM(z_vals.sizes()==z_sampled.sizes());
|
||||
|
||||
const int N_rays = z_sampled.sizes()[0];
|
||||
const int N_samples = z_sampled.sizes()[1];
|
||||
const int threadx = 32;
|
||||
const int thready = 32;
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(z_in_out.type(), "sample_rays_uniform_occupied_voxels_kernel", ([&]
|
||||
{
|
||||
sample_rays_uniform_occupied_voxels_kernel<scalar_t><<<{divCeil(N_rays,threadx),divCeil(N_samples,thready)}, {threadx,thready}>>>(z_sampled.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(),z_in_out.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>(),z_vals.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>());
|
||||
}));
|
||||
|
||||
return z_vals;
|
||||
}
|
||||
|
||||
template<class scalar_t>
|
||||
__global__ void postprocessOctreeRayTracingKernel(const torch::PackedTensorAccessor32<long,1,torch::RestrictPtrTraits> ray_index, const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> depth_in_out, const torch::PackedTensorAccessor32<long,1,torch::RestrictPtrTraits> unique_intersect_ray_ids, const torch::PackedTensorAccessor32<long,1,torch::RestrictPtrTraits> start_poss, torch::PackedTensorAccessor32<scalar_t,3,torch::RestrictPtrTraits> depths_in_out_padded)
|
||||
{
|
||||
const int unique_id_pos = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (unique_id_pos>=unique_intersect_ray_ids.size(0)) return;
|
||||
const int i_ray = unique_intersect_ray_ids[unique_id_pos];
|
||||
|
||||
int i_intersect = 0;
|
||||
auto cur_depths_in_out_padded = depths_in_out_padded[i_ray];
|
||||
for (int i=start_poss[unique_id_pos];i<ray_index.size(0);i++)
|
||||
{
|
||||
if (ray_index[i]!=i_ray) break;
|
||||
if (depth_in_out[i][0]==0 || depth_in_out[i][1]==0) break;
|
||||
if (depth_in_out[i][0]>depth_in_out[i][1]) continue;
|
||||
if (abs(depth_in_out[i][1]-depth_in_out[i][0])<1e-4) continue;
|
||||
|
||||
cur_depths_in_out_padded[i_intersect][0] = depth_in_out[i][0];
|
||||
cur_depths_in_out_padded[i_intersect][1] = depth_in_out[i][1];
|
||||
|
||||
i_intersect++;
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
at::Tensor postprocessOctreeRayTracing(const at::Tensor ray_index, const at::Tensor depth_in_out, const at::Tensor unique_intersect_ray_ids, const at::Tensor start_poss, const int max_intersections, const int N_rays)
|
||||
{
|
||||
CHECK_INPUT(ray_index);
|
||||
CHECK_INPUT(depth_in_out);
|
||||
CHECK_INPUT(start_poss);
|
||||
|
||||
const int n_unique_ids = unique_intersect_ray_ids.sizes()[0];
|
||||
at::Tensor depths_in_out_padded = at::zeros({N_rays,max_intersections,2}, torch::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA, 0).requires_grad(false));
|
||||
dim3 threads = {256};
|
||||
dim3 blocks = {divCeil(n_unique_ids,threads.x)};
|
||||
AT_DISPATCH_FLOATING_TYPES(depth_in_out.type(), "postprocessOctreeRayTracingKernel", ([&]
|
||||
{
|
||||
postprocessOctreeRayTracingKernel<scalar_t><<<blocks,threads>>>(ray_index.packed_accessor32<long,1,torch::RestrictPtrTraits>(), depth_in_out.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), unique_intersect_ray_ids.packed_accessor32<long,1,torch::RestrictPtrTraits>(), start_poss.packed_accessor32<long,1,torch::RestrictPtrTraits>(), depths_in_out_padded.packed_accessor32<scalar_t,3,torch::RestrictPtrTraits>());
|
||||
}));
|
||||
|
||||
return depths_in_out_padded;
|
||||
}
|
||||
|
||||
|
||||
at::Tensor calculateBarycentricCoordinate3D(const at::Tensor &triangle, const at::Tensor &p)
|
||||
{
|
||||
auto vector_A = triangle[1]-triangle[2];
|
||||
auto vector_B = triangle[1]-triangle[0];
|
||||
auto normal = vector_A.cross(vector_B);
|
||||
auto areaABC = (normal * at::cross(triangle[1]-triangle[0], triangle[2]-triangle[0])).sum();
|
||||
auto areaPBC = (normal * at::cross(triangle[1]-p, triangle[2]-p)).sum();
|
||||
auto areaPCA = (normal * at::cross(triangle[2]-p, triangle[0]-p)).sum();
|
||||
at::Tensor ws = at::zeros({3}).to(torch::kFloat32);
|
||||
ws[0] = areaPBC / areaABC;
|
||||
ws[1] = areaPCA / areaABC;
|
||||
ws[2] = 1-ws[0]-ws[1];
|
||||
return ws;
|
||||
}
|
||||
|
||||
|
||||
__device__ Eigen::Vector3f calculateBarycentricCoordinate3DKernel(const Eigen::Matrix<float,3,3> &triangle, const Eigen::Vector3f &p)
|
||||
{
|
||||
Eigen::Vector3f vector_A = triangle.row(1)-triangle.row(2);
|
||||
Eigen::Vector3f vector_B = triangle.row(1)-triangle.row(0);
|
||||
Eigen::Vector3f normal = vector_A.cross(vector_B);
|
||||
float areaABC = (normal.array() * (triangle.row(1)-triangle.row(0)).cross(triangle.row(2)-triangle.row(0)).transpose().array()).sum();
|
||||
float areaPBC = (normal.array() * (triangle.row(1).transpose()-p).cross(triangle.row(2).transpose()-p).array()).sum();
|
||||
float areaPCA = (normal.array() * (triangle.row(2).transpose()-p).cross(triangle.row(0).transpose()-p).array()).sum();
|
||||
Eigen::Vector3f w;
|
||||
w[0] = areaPBC / areaABC;
|
||||
w[1] = areaPCA / areaABC;
|
||||
w[2] = 1-w[0]-w[1];
|
||||
return w;
|
||||
|
||||
}
|
||||
|
||||
|
||||
__device__ void calculateBarycentricCoordinate2DKernel(const Eigen::Matrix<float,3,2> &triangle, const Eigen::Vector2f &p, Eigen::Vector3f &w)
|
||||
{
|
||||
Eigen::Vector2f CA = triangle.row(0)-triangle.row(2);
|
||||
Eigen::Vector2f AC = -CA;
|
||||
Eigen::Vector2f CP = p-triangle.row(2).transpose();
|
||||
Eigen::Vector2f AB = triangle.row(1)-triangle.row(0);
|
||||
Eigen::Vector2f AP = p-triangle.row(0).transpose();
|
||||
Eigen::Matrix2f numerator, denominator;
|
||||
numerator << CA, CP;
|
||||
denominator<<AB, AC;
|
||||
w(1) = numerator.determinant()/denominator.determinant();
|
||||
numerator << AB, AP;
|
||||
w(2) = numerator.determinant()/denominator.determinant();
|
||||
w(0) = 1-w(1)-w(2);
|
||||
}
|
||||
|
||||
|
||||
|
||||
template<class scalar_t>
|
||||
__global__ void rayColorToTextureImageKernel(const torch::PackedTensorAccessor32<long,2,torch::RestrictPtrTraits> F, const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> V, const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> hit_locations, const torch::PackedTensorAccessor32<long,1,torch::RestrictPtrTraits> hit_face_ids, const torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> uvs_tex, torch::PackedTensorAccessor32<scalar_t,2,torch::RestrictPtrTraits> uvs)
|
||||
{
|
||||
const int i_hit = blockIdx.x*blockDim.x + threadIdx.x;
|
||||
if (i_hit>=hit_locations.size(0)) return;
|
||||
|
||||
auto face = F[hit_face_ids[i_hit]];
|
||||
Eigen::Matrix3f tri_v = Eigen::Matrix3f::Zero();
|
||||
for (int r=0;r<3;r++)
|
||||
{
|
||||
for (int c=0;c<3;c++)
|
||||
{
|
||||
tri_v(r,c) = V[face[r]][c];
|
||||
}
|
||||
}
|
||||
auto location = hit_locations[i_hit];
|
||||
Eigen::Vector3f p(location[0], location[1], location[2]);
|
||||
auto w = calculateBarycentricCoordinate3DKernel(tri_v, p);
|
||||
|
||||
Eigen::Matrix<float,3,2> tri_uvs;
|
||||
for (int i=0;i<3;i++)
|
||||
{
|
||||
for (int j=0;j<2;j++)
|
||||
{
|
||||
tri_uvs(i,j) = uvs_tex[face[i]][j];
|
||||
}
|
||||
}
|
||||
|
||||
Eigen::Vector2f cur_uv = tri_uvs.transpose() * w;
|
||||
uvs[i_hit][0] = cur_uv(0);
|
||||
uvs[i_hit][1] = cur_uv(1);
|
||||
}
|
||||
|
||||
|
||||
|
||||
void rayColorToTextureImageCUDA(const at::Tensor &F, const at::Tensor &V, const at::Tensor &hit_locations, const at::Tensor &hit_face_ids, const at::Tensor &uvs_tex, at::Tensor &uvs)
|
||||
{
|
||||
CHECK_CONTIGUOUS(F);
|
||||
CHECK_CONTIGUOUS(V);
|
||||
CHECK_CONTIGUOUS(hit_locations);
|
||||
CHECK_CONTIGUOUS(hit_face_ids);
|
||||
CHECK_CONTIGUOUS(uvs_tex);
|
||||
|
||||
dim3 threads = {512};
|
||||
dim3 blocks = {divCeil(int(hit_locations.sizes()[0]),threads.x)};
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES(V.type(), "rayColorToTextureImageKernel", ([&]
|
||||
{
|
||||
rayColorToTextureImageKernel<scalar_t><<<blocks,threads>>>(F.packed_accessor32<long,2,torch::RestrictPtrTraits>(), V.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), hit_locations.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), hit_face_ids.packed_accessor32<long,1,torch::RestrictPtrTraits>(), uvs_tex.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>(), uvs.packed_accessor32<scalar_t,2,torch::RestrictPtrTraits>());
|
||||
}));
|
||||
}
|
||||
|
||||
|
||||
|
||||
__device__ bool isLineIntersectLine(Eigen::Vector2f l1p1, Eigen::Vector2f l1p2, Eigen::Vector2f l2p1, Eigen::Vector2f l2p2)
|
||||
{
|
||||
float q = (l1p1(1) - l2p1(1)) * (l2p2(0) - l2p1(0)) - (l1p1(0) - l2p1(0)) * (l2p2(1) - l2p1(1));
|
||||
float d = (l1p2(0) - l1p1(0)) * (l2p2(1) - l2p1(1)) - (l1p2(1) - l1p1(1)) * (l2p2(0) - l2p1(0));
|
||||
|
||||
if ( d == 0 )
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
float r = q / d;
|
||||
|
||||
q = (l1p1(1) - l2p1(1)) * (l1p2(0) - l1p1(0)) - (l1p1(0) - l2p1(0)) * (l1p2(1) - l1p1(1));
|
||||
float s = q / d;
|
||||
|
||||
if( r < 0 || r > 1 || s < 0 || s > 1 )
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
__device__ bool isLineIntersectSquare(const Eigen::Vector2f &la, const Eigen::Vector2f &lb, const Eigen::Vector2f &p0, const Eigen::Vector2f &p1, const Eigen::Vector2f &p2, const Eigen::Vector2f &p3)
|
||||
{
|
||||
return isLineIntersectLine(la, lb, p0, p1) || isLineIntersectLine(la, lb, p1, p2) || isLineIntersectLine(la, lb, p2, p3) || isLineIntersectLine(la, lb, p3, p0);
|
||||
}
|
||||
|
||||
|
||||
__device__ bool isPixelInsideTriangle(const Eigen::Matrix<float,3,2> &triangle, Eigen::Vector2f p, Eigen::Vector3f &w)
|
||||
{
|
||||
calculateBarycentricCoordinate2DKernel(triangle, p, w);
|
||||
for (int j=0;j<3;j++)
|
||||
{
|
||||
if (w(j)<0) return false;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
/*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
|
||||
#define CHECK_CUDA(x) AT_ASSERTM(x.is_cuda(), #x " must be a CUDA tensor")
|
||||
#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
|
||||
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
|
||||
|
||||
|
||||
inline int divCeil(int a, int b)
|
||||
{
|
||||
return (a+b-1)/b;
|
||||
};
|
||||
|
||||
|
||||
at::Tensor sampleRaysUniformOccupiedVoxels(const at::Tensor z_in_out, const at::Tensor z_sampled, at::Tensor z_vals);
|
||||
at::Tensor postprocessOctreeRayTracing(const at::Tensor ray_index, const at::Tensor depth_in_out, const at::Tensor unique_intersect_ray_ids, const at::Tensor start_poss, const int max_intersections, const int N_rays);
|
||||
void rayColorToTextureImageCUDA(const at::Tensor &F, const at::Tensor &V, const at::Tensor &hit_locations, const at::Tensor &hit_face_ids, const at::Tensor &uvs_tex, at::Tensor &uvs);
|
||||
@@ -0,0 +1,41 @@
|
||||
# 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.
|
||||
|
||||
|
||||
from setuptools import setup
|
||||
import os,sys
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
from torch.utils.cpp_extension import load
|
||||
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
|
||||
|
||||
nvcc_flags = ['-Xcompiler', '-O3', '-std=c++14', '-U__CUDA_NO_HALF_OPERATORS__', '-U__CUDA_NO_HALF_CONVERSIONS__', '-U__CUDA_NO_HALF2_OPERATORS__']
|
||||
c_flags = ['-O3', '-std=c++14']
|
||||
|
||||
setup(
|
||||
name='common',
|
||||
extra_cflags=c_flags,
|
||||
extra_cuda_cflags=nvcc_flags,
|
||||
ext_modules=[
|
||||
CUDAExtension('common', [
|
||||
'bindings.cpp',
|
||||
'common.cu',
|
||||
],extra_compile_args={'gcc': c_flags, 'nvcc': nvcc_flags}),
|
||||
CUDAExtension('gridencoder', [
|
||||
f"{code_dir}/torch_ngp_grid_encoder/gridencoder.cu",
|
||||
f"{code_dir}/torch_ngp_grid_encoder/bindings.cpp",
|
||||
],extra_compile_args={'gcc': c_flags, 'nvcc': nvcc_flags}),
|
||||
],
|
||||
include_dirs=[
|
||||
"/usr/local/include/eigen3",
|
||||
"/usr/include/eigen3",
|
||||
],
|
||||
cmdclass={
|
||||
'build_ext': BuildExtension
|
||||
})
|
||||
@@ -0,0 +1,21 @@
|
||||
MIT License
|
||||
|
||||
Copyright (c) 2022 hawkey
|
||||
|
||||
Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
of this software and associated documentation files (the "Software"), to deal
|
||||
in the Software without restriction, including without limitation the rights
|
||||
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
copies of the Software, and to permit persons to whom the Software is
|
||||
furnished to do so, subject to the following conditions:
|
||||
|
||||
The above copyright notice and this permission notice shall be included in all
|
||||
copies or substantial portions of the Software.
|
||||
|
||||
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
SOFTWARE.
|
||||
+8
@@ -0,0 +1,8 @@
|
||||
#include <torch/extension.h>
|
||||
|
||||
#include "gridencoder.h"
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("grid_encode_forward", &grid_encode_forward, "grid_encode_forward (CUDA)");
|
||||
m.def("grid_encode_backward", &grid_encode_backward, "grid_encode_backward (CUDA)");
|
||||
}
|
||||
@@ -0,0 +1,158 @@
|
||||
import numpy as np
|
||||
import os,sys,pdb
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(code_dir)
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.autograd import Function
|
||||
from torch.autograd.function import once_differentiable
|
||||
from torch.cuda.amp import custom_bwd, custom_fwd
|
||||
from setuptools import setup
|
||||
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
|
||||
from torch.utils.cpp_extension import load
|
||||
import os,sys
|
||||
import gridencoder
|
||||
|
||||
|
||||
_gridtype_to_id = {
|
||||
'hash': 0,
|
||||
'tiled': 1,
|
||||
}
|
||||
|
||||
class _grid_encode(Function):
|
||||
@staticmethod
|
||||
@custom_fwd
|
||||
def forward(ctx, inputs, embeddings, offsets, per_level_scale, base_resolution, calc_grad_inputs=False, gridtype=0, align_corners=False):
|
||||
# inputs: [B, D], float in [0, 1]
|
||||
# embeddings: [sO, C], float
|
||||
# offsets: [L + 1], int
|
||||
# RETURN: [B, F], float
|
||||
|
||||
inputs = inputs.contiguous()
|
||||
|
||||
B, D = inputs.shape
|
||||
L = offsets.shape[0] - 1
|
||||
C = embeddings.shape[1]
|
||||
S = np.log2(per_level_scale)
|
||||
H = base_resolution
|
||||
|
||||
if torch.is_autocast_enabled() and C % 2 == 0:
|
||||
embeddings = embeddings.to(torch.half)
|
||||
|
||||
outputs = torch.empty(L, B, C, device=inputs.device, dtype=embeddings.dtype)
|
||||
|
||||
if calc_grad_inputs:
|
||||
dy_dx = torch.empty(B, L * D * C, device=inputs.device, dtype=embeddings.dtype)
|
||||
else:
|
||||
dy_dx = torch.empty(1, device=inputs.device, dtype=embeddings.dtype) # placeholder... TODO: a better way?
|
||||
|
||||
gridencoder.grid_encode_forward(inputs, embeddings, offsets, outputs, B, D, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners)
|
||||
|
||||
outputs = outputs.permute(1, 0, 2).reshape(B, L * C)
|
||||
|
||||
ctx.save_for_backward(inputs, embeddings, offsets, dy_dx)
|
||||
ctx.dims = [B, D, C, L, S, H, gridtype]
|
||||
ctx.calc_grad_inputs = calc_grad_inputs
|
||||
ctx.align_corners = align_corners
|
||||
|
||||
return outputs
|
||||
|
||||
@staticmethod
|
||||
@custom_bwd
|
||||
def backward(ctx, grad):
|
||||
|
||||
inputs, embeddings, offsets, dy_dx = ctx.saved_tensors
|
||||
B, D, C, L, S, H, gridtype = ctx.dims
|
||||
calc_grad_inputs = ctx.calc_grad_inputs
|
||||
align_corners = ctx.align_corners
|
||||
|
||||
# grad: [B, L * C] --> [L, B, C]
|
||||
grad = grad.view(B, L, C).permute(1, 0, 2).contiguous()
|
||||
|
||||
grad_embeddings = torch.zeros_like(embeddings)
|
||||
|
||||
if calc_grad_inputs:
|
||||
grad_inputs = torch.zeros_like(inputs, dtype=embeddings.dtype)
|
||||
else:
|
||||
grad_inputs = torch.zeros(1, device=inputs.device, dtype=embeddings.dtype)
|
||||
|
||||
gridencoder.grid_encode_backward(grad, inputs, embeddings, offsets, grad_embeddings, B, D, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners)
|
||||
|
||||
if calc_grad_inputs:
|
||||
grad_inputs = grad_inputs.to(inputs.dtype)
|
||||
return grad_inputs, grad_embeddings, None, None, None, None, None, None
|
||||
else:
|
||||
return None, grad_embeddings, None, None, None, None, None, None
|
||||
|
||||
|
||||
grid_encode = _grid_encode.apply
|
||||
|
||||
|
||||
# https://github.com/ashawkey/torch-ngp/blob/main/gridencoder/grid.py
|
||||
class GridEncoder(nn.Module):
|
||||
def __init__(self, input_dim=3, n_levels=16, level_dim=2, base_resolution=16, log2_hashmap_size=19, desired_resolution=None, gridtype='hash', align_corners=False):
|
||||
super().__init__()
|
||||
|
||||
per_level_scale = np.exp2(np.log2(desired_resolution / base_resolution) / (n_levels - 1))
|
||||
|
||||
self.input_dim = input_dim # coord dims, 2 or 3
|
||||
self.n_levels = n_levels # num levels, each level multiply resolution by 2
|
||||
self.level_dim = level_dim # encode channels per level
|
||||
self.per_level_scale = per_level_scale # multiply resolution by this scale at each level.
|
||||
self.log2_hashmap_size = log2_hashmap_size
|
||||
self.base_resolution = base_resolution
|
||||
self.out_dim = n_levels * level_dim
|
||||
self.gridtype = gridtype
|
||||
self.gridtype_id = _gridtype_to_id[gridtype] # "tiled" or "hash"
|
||||
self.align_corners = align_corners
|
||||
|
||||
# allocate parameters
|
||||
offsets = []
|
||||
offset = 0
|
||||
self.max_params = 2 ** log2_hashmap_size
|
||||
for i in range(n_levels):
|
||||
resolution = int(np.ceil(base_resolution * per_level_scale ** i))
|
||||
params_in_level = min(self.max_params, (resolution if align_corners else resolution + 1) ** input_dim) # limit max number
|
||||
params_in_level = int(np.ceil(params_in_level / 8) * 8) # make divisible
|
||||
print(f"level {i}, resolution: {resolution}")
|
||||
offsets.append(offset)
|
||||
offset += params_in_level
|
||||
offsets.append(offset)
|
||||
offsets = torch.from_numpy(np.array(offsets, dtype=np.int32))
|
||||
self.register_buffer('offsets', offsets)
|
||||
|
||||
self.n_params = offsets[-1] * level_dim
|
||||
|
||||
# parameters
|
||||
self.embeddings = nn.Parameter(torch.empty(offset, level_dim))
|
||||
# self.embeddings = nn.Embedding(offset, level_dim, sparse=True)
|
||||
|
||||
self.reset_parameters()
|
||||
|
||||
def reset_parameters(self):
|
||||
std = 1e-4
|
||||
self.embeddings.data.uniform_(-std, std)
|
||||
# self.embeddings.weight.data.uniform_(-std, std)
|
||||
|
||||
|
||||
def __repr__(self):
|
||||
return f"GridEncoder: input_dim={self.input_dim} n_levels={self.n_levels} level_dim={self.level_dim} resolution={self.base_resolution} -> {int(round(self.base_resolution * self.per_level_scale ** (self.n_levels - 1)))} per_level_scale={self.per_level_scale:.4f} params={tuple(self.embeddings.shape)} gridtype={self.gridtype} align_corners={self.align_corners}"
|
||||
|
||||
|
||||
def forward(self, inputs, bound=1):
|
||||
# inputs: [..., input_dim], normalized real world positions in [-bound, bound]
|
||||
# return: [..., num_levels * level_dim]
|
||||
|
||||
inputs = (inputs + bound) / (2 * bound) # map to [0, 1]
|
||||
|
||||
#print('inputs', inputs.shape, inputs.dtype, inputs.min().item(), inputs.max().item())
|
||||
|
||||
prefix_shape = list(inputs.shape[:-1])
|
||||
inputs = inputs.view(-1, self.input_dim)
|
||||
|
||||
outputs = grid_encode(inputs, self.embeddings, self.offsets, self.per_level_scale, self.base_resolution, inputs.requires_grad, self.gridtype_id, self.align_corners)
|
||||
outputs = outputs.view(prefix_shape + [self.out_dim])
|
||||
|
||||
#print('outputs', outputs.shape, outputs.dtype, outputs.min().item(), outputs.max().item())
|
||||
|
||||
return outputs
|
||||
+500
@@ -0,0 +1,500 @@
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <algorithm>
|
||||
#include <stdexcept>
|
||||
|
||||
#include <stdint.h>
|
||||
#include <cstdio>
|
||||
|
||||
|
||||
#define CHECK_CUDA(x) TORCH_CHECK(x.device().is_cuda(), #x " must be a CUDA tensor")
|
||||
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor")
|
||||
#define CHECK_IS_INT(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, #x " must be an int tensor")
|
||||
#define CHECK_IS_FLOATING(x) TORCH_CHECK(x.scalar_type() == at::ScalarType::Float || x.scalar_type() == at::ScalarType::Half || x.scalar_type() == at::ScalarType::Double, #x " must be a floating tensor")
|
||||
|
||||
|
||||
// just for compatability of half precision in AT_DISPATCH_FLOATING_TYPES_AND_HALF...
|
||||
static inline __device__ at::Half atomicAdd(at::Half *address, at::Half val) {
|
||||
// requires CUDA >= 10 and ARCH >= 70
|
||||
// this is very slow compared to float or __half2, and never used.
|
||||
//return atomicAdd(reinterpret_cast<__half*>(address), val);
|
||||
}
|
||||
|
||||
|
||||
template <typename T>
|
||||
static inline __host__ __device__ T div_round_up(T val, T divisor) {
|
||||
return (val + divisor - 1) / divisor;
|
||||
}
|
||||
|
||||
|
||||
template <uint32_t D>
|
||||
__device__ uint32_t fast_hash(const uint32_t pos_grid[D]) {
|
||||
static_assert(D <= 7, "fast_hash can only hash up to 7 dimensions.");
|
||||
|
||||
// While 1 is technically not a good prime for hashing (or a prime at all), it helps memory coherence
|
||||
// and is sufficient for our use case of obtaining a uniformly colliding index from high-dimensional
|
||||
// coordinates.
|
||||
constexpr uint32_t primes[7] = { 1, 2654435761, 805459861, 3674653429, 2097192037, 1434869437, 2165219737 };
|
||||
|
||||
uint32_t result = 0;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < D; ++i) {
|
||||
result ^= pos_grid[i] * primes[i];
|
||||
}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
|
||||
template <uint32_t D, uint32_t C>
|
||||
__device__ uint32_t get_grid_index(const uint32_t gridtype, const bool align_corners, const uint32_t ch, const uint32_t hashmap_size, const uint32_t resolution, const uint32_t pos_grid[D]) {
|
||||
uint32_t stride = 1;
|
||||
uint32_t index = 0;
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D && stride <= hashmap_size; d++) {
|
||||
index += pos_grid[d] * stride;
|
||||
stride *= align_corners ? resolution: (resolution + 1);
|
||||
}
|
||||
|
||||
// NOTE: for NeRF, the hash is in fact not necessary. Check https://github.com/NVlabs/instant-ngp/issues/97.
|
||||
// gridtype: 0 == hash, 1 == tiled
|
||||
if (gridtype == 0 && stride > hashmap_size) {
|
||||
index = fast_hash<D>(pos_grid);
|
||||
}
|
||||
|
||||
return (index % hashmap_size) * C + ch;
|
||||
}
|
||||
|
||||
|
||||
/**
|
||||
* @brief
|
||||
*
|
||||
* @tparam scalar_t
|
||||
* @tparam D : point dimension usually 3D
|
||||
* @tparam C : embedding dim 1/2/4
|
||||
* @param inputs
|
||||
* @param grid
|
||||
* @param offsets
|
||||
* @param outputs
|
||||
* @param B
|
||||
* @param L
|
||||
* @param S
|
||||
* @param H
|
||||
* @param calc_grad_inputs
|
||||
* @param dy_dx
|
||||
* @param gridtype
|
||||
* @param align_corners
|
||||
* @return __global__
|
||||
*/
|
||||
template <typename scalar_t, uint32_t D, uint32_t C>
|
||||
__global__ void kernel_grid(
|
||||
const float * __restrict__ inputs,
|
||||
const scalar_t * __restrict__ grid,
|
||||
const int * __restrict__ offsets,
|
||||
scalar_t * __restrict__ outputs,
|
||||
const uint32_t B, const uint32_t L, const float S, const uint32_t H,
|
||||
const bool calc_grad_inputs,
|
||||
scalar_t * __restrict__ dy_dx,
|
||||
const uint32_t gridtype,
|
||||
const bool align_corners
|
||||
) {
|
||||
const uint32_t b = blockIdx.x * blockDim.x + threadIdx.x; // batch id
|
||||
|
||||
if (b >= B) return;
|
||||
|
||||
const uint32_t level = blockIdx.y;
|
||||
|
||||
// locate
|
||||
grid += (uint32_t)offsets[level] * C;
|
||||
inputs += b * D;
|
||||
outputs += level * B * C + b * C;
|
||||
|
||||
// check input range (should be in [0, 1])
|
||||
bool flag_oob = false;
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
if (inputs[d] < 0 || inputs[d] > 1) {
|
||||
flag_oob = true;
|
||||
}
|
||||
}
|
||||
// if input out of bound, just set output to 0
|
||||
if (flag_oob) {
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
outputs[ch] = 0;
|
||||
}
|
||||
if (calc_grad_inputs) {
|
||||
dy_dx += b * D * L * C + level * D * C; // B L D C
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
dy_dx[d * C + ch] = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const uint32_t hashmap_size = offsets[level + 1] - offsets[level];
|
||||
const float scale = exp2f(level * S) * H - 1.0f;
|
||||
const uint32_t resolution = (uint32_t)ceil(scale) + 1;
|
||||
|
||||
// calculate coordinate
|
||||
float pos[D];
|
||||
uint32_t pos_grid[D];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
pos[d] = inputs[d] * scale + (align_corners ? 0.0f : 0.5f);
|
||||
pos_grid[d] = floorf(pos[d]);
|
||||
pos[d] -= (float)pos_grid[d];
|
||||
}
|
||||
|
||||
//printf("[b=%d, l=%d] pos=(%f, %f)+(%d, %d)\n", b, level, pos[0], pos[1], pos_grid[0], pos_grid[1]);
|
||||
|
||||
// interpolate
|
||||
scalar_t results[C] = {0}; // temp results in register
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t idx = 0; idx < (1 << D); idx++) {
|
||||
float w = 1;
|
||||
uint32_t pos_grid_local[D];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
if ((idx & (1 << d)) == 0) {
|
||||
w *= 1 - pos[d];
|
||||
pos_grid_local[d] = pos_grid[d];
|
||||
} else {
|
||||
w *= pos[d];
|
||||
pos_grid_local[d] = pos_grid[d] + 1;
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t index = get_grid_index<D, C>(gridtype, align_corners, 0, hashmap_size, resolution, pos_grid_local); // Hashing
|
||||
|
||||
// writing to register (fast)
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
results[ch] += w * grid[index + ch];
|
||||
}
|
||||
|
||||
//printf("[b=%d, l=%d] int %d, idx %d, w %f, val %f\n", b, level, idx, index, w, grid[index]);
|
||||
}
|
||||
|
||||
// writing to global memory (slow)
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
outputs[ch] = results[ch];
|
||||
}
|
||||
|
||||
// prepare dy_dx for calc_grad_inputs
|
||||
// differentiable (soft) indexing: https://discuss.pytorch.org/t/differentiable-indexing/17647/9
|
||||
if (calc_grad_inputs) {
|
||||
|
||||
dy_dx += b * D * L * C + level * D * C; // B L D C
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t gd = 0; gd < D; gd++) {
|
||||
|
||||
scalar_t results_grad[C] = {0};
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t idx = 0; idx < (1 << (D - 1)); idx++) {
|
||||
float w = scale;
|
||||
uint32_t pos_grid_local[D];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t nd = 0; nd < D - 1; nd++) {
|
||||
const uint32_t d = (nd >= gd) ? (nd + 1) : nd;
|
||||
|
||||
if ((idx & (1 << nd)) == 0) {
|
||||
w *= 1 - pos[d];
|
||||
pos_grid_local[d] = pos_grid[d];
|
||||
} else {
|
||||
w *= pos[d];
|
||||
pos_grid_local[d] = pos_grid[d] + 1;
|
||||
}
|
||||
}
|
||||
|
||||
pos_grid_local[gd] = pos_grid[gd];
|
||||
uint32_t index_left = get_grid_index<D, C>(gridtype, align_corners, 0, hashmap_size, resolution, pos_grid_local);
|
||||
pos_grid_local[gd] = pos_grid[gd] + 1;
|
||||
uint32_t index_right = get_grid_index<D, C>(gridtype, align_corners, 0, hashmap_size, resolution, pos_grid_local);
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
results_grad[ch] += w * (grid[index_right + ch] - grid[index_left + ch]);
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t ch = 0; ch < C; ch++) {
|
||||
dy_dx[gd * C + ch] = results_grad[ch];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename scalar_t, uint32_t D, uint32_t C, uint32_t N_C>
|
||||
__global__ void kernel_grid_backward(
|
||||
const scalar_t * __restrict__ grad,
|
||||
const float * __restrict__ inputs,
|
||||
const scalar_t * __restrict__ grid,
|
||||
const int * __restrict__ offsets,
|
||||
scalar_t * __restrict__ grad_grid,
|
||||
const uint32_t B, const uint32_t L, const float S, const uint32_t H,
|
||||
const uint32_t gridtype,
|
||||
const bool align_corners
|
||||
) {
|
||||
const uint32_t b = (blockIdx.x * blockDim.x + threadIdx.x) * N_C / C;
|
||||
if (b >= B) return;
|
||||
|
||||
const uint32_t level = blockIdx.y;
|
||||
const uint32_t ch = (blockIdx.x * blockDim.x + threadIdx.x) * N_C - b * C;
|
||||
|
||||
// locate
|
||||
grad_grid += offsets[level] * C;
|
||||
inputs += b * D;
|
||||
grad += level * B * C + b * C + ch; // L, B, C
|
||||
|
||||
const uint32_t hashmap_size = offsets[level + 1] - offsets[level];
|
||||
const float scale = exp2f(level * S) * H - 1.0f;
|
||||
const uint32_t resolution = (uint32_t)ceil(scale) + 1;
|
||||
|
||||
// check input range (should be in [0, 1])
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
if (inputs[d] < 0 || inputs[d] > 1) {
|
||||
return; // grad is init as 0, so we simply return.
|
||||
}
|
||||
}
|
||||
|
||||
// calculate coordinate
|
||||
float pos[D];
|
||||
uint32_t pos_grid[D];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
pos[d] = inputs[d] * scale + (align_corners ? 0.0f : 0.5f);
|
||||
pos_grid[d] = floorf(pos[d]);
|
||||
pos[d] -= (float)pos_grid[d];
|
||||
}
|
||||
|
||||
scalar_t grad_cur[N_C] = {0}; // fetch to register
|
||||
#pragma unroll
|
||||
for (uint32_t c = 0; c < N_C; c++) {
|
||||
grad_cur[c] = grad[c];
|
||||
}
|
||||
|
||||
// interpolate
|
||||
#pragma unroll
|
||||
for (uint32_t idx = 0; idx < (1 << D); idx++) {
|
||||
float w = 1;
|
||||
uint32_t pos_grid_local[D];
|
||||
|
||||
#pragma unroll
|
||||
for (uint32_t d = 0; d < D; d++) {
|
||||
if ((idx & (1 << d)) == 0) {
|
||||
w *= 1 - pos[d];
|
||||
pos_grid_local[d] = pos_grid[d];
|
||||
} else {
|
||||
w *= pos[d];
|
||||
pos_grid_local[d] = pos_grid[d] + 1;
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t index = get_grid_index<D, C>(gridtype, align_corners, ch, hashmap_size, resolution, pos_grid_local);
|
||||
|
||||
// atomicAdd for __half is slow (especially for large values), so we use __half2 if N_C % 2 == 0
|
||||
// TODO: use float which is better than __half, if N_C % 2 != 0
|
||||
if (std::is_same<scalar_t, at::Half>::value && N_C % 2 == 0) {
|
||||
#pragma unroll
|
||||
for (uint32_t c = 0; c < N_C; c += 2) {
|
||||
// process two __half at once (by interpreting as a __half2)
|
||||
__half2 v = {(__half)(w * grad_cur[c]), (__half)(w * grad_cur[c + 1])};
|
||||
atomicAdd((__half2*)&grad_grid[index + c], v);
|
||||
}
|
||||
// float, or __half when N_C % 2 != 0 (which means C == 1)
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (uint32_t c = 0; c < N_C; c++) {
|
||||
atomicAdd(&grad_grid[index + c], w * grad_cur[c]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
template <typename scalar_t, uint32_t D, uint32_t C>
|
||||
__global__ void kernel_input_backward(
|
||||
const scalar_t * __restrict__ grad,
|
||||
const scalar_t * __restrict__ dy_dx,
|
||||
scalar_t * __restrict__ grad_inputs,
|
||||
uint32_t B, uint32_t L
|
||||
) {
|
||||
const uint32_t t = threadIdx.x + blockIdx.x * blockDim.x;
|
||||
if (t >= B * D) return;
|
||||
|
||||
const uint32_t b = t / D;
|
||||
const uint32_t d = t - b * D;
|
||||
|
||||
dy_dx += b * L * D * C;
|
||||
|
||||
scalar_t result = 0;
|
||||
|
||||
# pragma unroll
|
||||
for (int l = 0; l < L; l++) {
|
||||
# pragma unroll
|
||||
for (int ch = 0; ch < C; ch++) {
|
||||
result += grad[l * B * C + b * C + ch] * dy_dx[l * D * C + d * C + ch];
|
||||
}
|
||||
}
|
||||
|
||||
grad_inputs[t] = result;
|
||||
}
|
||||
|
||||
|
||||
template <typename scalar_t, uint32_t D>
|
||||
void kernel_grid_wrapper(const float *inputs, const scalar_t *embeddings, const int *offsets, scalar_t *outputs, const uint32_t B, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, scalar_t *dy_dx, const uint32_t gridtype, const bool align_corners) {
|
||||
static constexpr uint32_t N_THREAD = 512;
|
||||
const dim3 blocks_hashgrid = { div_round_up(B, N_THREAD), L, 1 };
|
||||
switch (C) {
|
||||
case 1: kernel_grid<scalar_t, D, 1><<<blocks_hashgrid, N_THREAD>>>(inputs, embeddings, offsets, outputs, B, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 2: kernel_grid<scalar_t, D, 2><<<blocks_hashgrid, N_THREAD>>>(inputs, embeddings, offsets, outputs, B, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 4: kernel_grid<scalar_t, D, 4><<<blocks_hashgrid, N_THREAD>>>(inputs, embeddings, offsets, outputs, B, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 8: kernel_grid<scalar_t, D, 8><<<blocks_hashgrid, N_THREAD>>>(inputs, embeddings, offsets, outputs, B, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
default: throw std::runtime_error{"GridEncoding: C must be 1, 2, 4, or 8."};
|
||||
}
|
||||
}
|
||||
|
||||
// inputs: [B, D], float, in [0, 1]
|
||||
// embeddings: [sO, C], float
|
||||
// offsets: [L + 1], uint32_t
|
||||
// outputs: [L, B, C], float (L first, so only one level of hashmap needs to fit into cache at a time.)
|
||||
// H: base resolution
|
||||
// dy_dx: [B, L * D * C]
|
||||
template <typename scalar_t>
|
||||
void grid_encode_forward_cuda(const float *inputs, const scalar_t *embeddings, const int *offsets, scalar_t *outputs, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, scalar_t *dy_dx, const uint32_t gridtype, const bool align_corners) {
|
||||
switch (D) {
|
||||
case 1: kernel_grid_wrapper<scalar_t, 1>(inputs, embeddings, offsets, outputs, B, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 2: kernel_grid_wrapper<scalar_t, 2>(inputs, embeddings, offsets, outputs, B, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 3: kernel_grid_wrapper<scalar_t, 3>(inputs, embeddings, offsets, outputs, B, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 4: kernel_grid_wrapper<scalar_t, 4>(inputs, embeddings, offsets, outputs, B, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
case 5: kernel_grid_wrapper<scalar_t, 5>(inputs, embeddings, offsets, outputs, B, C, L, S, H, calc_grad_inputs, dy_dx, gridtype, align_corners); break;
|
||||
default: throw std::runtime_error{"GridEncoding: D must be 1, 2, 3, 4, or 5."};
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
template <typename scalar_t, uint32_t D>
|
||||
void kernel_grid_backward_wrapper(const scalar_t *grad, const float *inputs, const scalar_t *embeddings, const int *offsets, scalar_t *grad_embeddings, const uint32_t B, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, scalar_t *dy_dx, scalar_t *grad_inputs, const uint32_t gridtype, const bool align_corners) {
|
||||
static constexpr uint32_t N_THREAD = 256;
|
||||
const uint32_t N_C = std::min(2u, C); // n_features_per_thread
|
||||
const dim3 blocks_hashgrid = { div_round_up(B * C / N_C, N_THREAD), L, 1 };
|
||||
switch (C) {
|
||||
case 1:
|
||||
kernel_grid_backward<scalar_t, D, 1, 1><<<blocks_hashgrid, N_THREAD>>>(grad, inputs, embeddings, offsets, grad_embeddings, B, L, S, H, gridtype, align_corners);
|
||||
if (calc_grad_inputs) kernel_input_backward<scalar_t, D, 1><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(grad, dy_dx, grad_inputs, B, L);
|
||||
break;
|
||||
case 2:
|
||||
kernel_grid_backward<scalar_t, D, 2, 2><<<blocks_hashgrid, N_THREAD>>>(grad, inputs, embeddings, offsets, grad_embeddings, B, L, S, H, gridtype, align_corners);
|
||||
if (calc_grad_inputs) kernel_input_backward<scalar_t, D, 2><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(grad, dy_dx, grad_inputs, B, L);
|
||||
break;
|
||||
case 4:
|
||||
kernel_grid_backward<scalar_t, D, 4, 2><<<blocks_hashgrid, N_THREAD>>>(grad, inputs, embeddings, offsets, grad_embeddings, B, L, S, H, gridtype, align_corners);
|
||||
if (calc_grad_inputs) kernel_input_backward<scalar_t, D, 4><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(grad, dy_dx, grad_inputs, B, L);
|
||||
break;
|
||||
case 8:
|
||||
kernel_grid_backward<scalar_t, D, 8, 2><<<blocks_hashgrid, N_THREAD>>>(grad, inputs, embeddings, offsets, grad_embeddings, B, L, S, H, gridtype, align_corners);
|
||||
if (calc_grad_inputs) kernel_input_backward<scalar_t, D, 8><<<div_round_up(B * D, N_THREAD), N_THREAD>>>(grad, dy_dx, grad_inputs, B, L);
|
||||
break;
|
||||
default: throw std::runtime_error{"GridEncoding: C must be 1, 2, 4, or 8."};
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
// grad: [L, B, C], float
|
||||
// inputs: [B, D], float, in [0, 1]
|
||||
// embeddings: [sO, C], float
|
||||
// offsets: [L + 1], uint32_t
|
||||
// grad_embeddings: [sO, C]
|
||||
// H: base resolution
|
||||
template <typename scalar_t>
|
||||
void grid_encode_backward_cuda(const scalar_t *grad, const float *inputs, const scalar_t *embeddings, const int *offsets, scalar_t *grad_embeddings, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, scalar_t *dy_dx, scalar_t *grad_inputs, const uint32_t gridtype, const bool align_corners) {
|
||||
switch (D) {
|
||||
case 1: kernel_grid_backward_wrapper<scalar_t, 1>(grad, inputs, embeddings, offsets, grad_embeddings, B, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners); break;
|
||||
case 2: kernel_grid_backward_wrapper<scalar_t, 2>(grad, inputs, embeddings, offsets, grad_embeddings, B, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners); break;
|
||||
case 3: kernel_grid_backward_wrapper<scalar_t, 3>(grad, inputs, embeddings, offsets, grad_embeddings, B, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners); break;
|
||||
case 4: kernel_grid_backward_wrapper<scalar_t, 4>(grad, inputs, embeddings, offsets, grad_embeddings, B, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners); break;
|
||||
case 5: kernel_grid_backward_wrapper<scalar_t, 5>(grad, inputs, embeddings, offsets, grad_embeddings, B, C, L, S, H, calc_grad_inputs, dy_dx, grad_inputs, gridtype, align_corners); break;
|
||||
default: throw std::runtime_error{"GridEncoding: D must be 1, 2, 3, 4, or 5."};
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
void grid_encode_forward(const at::Tensor inputs, const at::Tensor embeddings, const at::Tensor offsets, at::Tensor outputs, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, at::Tensor dy_dx, const uint32_t gridtype, const bool align_corners) {
|
||||
CHECK_CUDA(inputs);
|
||||
CHECK_CUDA(embeddings);
|
||||
CHECK_CUDA(offsets);
|
||||
CHECK_CUDA(outputs);
|
||||
CHECK_CUDA(dy_dx);
|
||||
|
||||
CHECK_CONTIGUOUS(inputs);
|
||||
CHECK_CONTIGUOUS(embeddings);
|
||||
CHECK_CONTIGUOUS(offsets);
|
||||
CHECK_CONTIGUOUS(outputs);
|
||||
CHECK_CONTIGUOUS(dy_dx);
|
||||
|
||||
CHECK_IS_FLOATING(inputs);
|
||||
CHECK_IS_FLOATING(embeddings);
|
||||
CHECK_IS_INT(offsets);
|
||||
CHECK_IS_FLOATING(outputs);
|
||||
CHECK_IS_FLOATING(dy_dx);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
|
||||
embeddings.scalar_type(), "grid_encode_forward", ([&] {
|
||||
grid_encode_forward_cuda<scalar_t>(inputs.data_ptr<float>(), embeddings.data_ptr<scalar_t>(), offsets.data_ptr<int>(), outputs.data_ptr<scalar_t>(), B, D, C, L, S, H, calc_grad_inputs, dy_dx.data_ptr<scalar_t>(), gridtype, align_corners);
|
||||
}));
|
||||
}
|
||||
|
||||
void grid_encode_backward(const at::Tensor grad, const at::Tensor inputs, const at::Tensor embeddings, const at::Tensor offsets, at::Tensor grad_embeddings, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, const at::Tensor dy_dx, at::Tensor grad_inputs, const uint32_t gridtype, const bool align_corners) {
|
||||
CHECK_CUDA(grad);
|
||||
CHECK_CUDA(inputs);
|
||||
CHECK_CUDA(embeddings);
|
||||
CHECK_CUDA(offsets);
|
||||
CHECK_CUDA(grad_embeddings);
|
||||
CHECK_CUDA(dy_dx);
|
||||
CHECK_CUDA(grad_inputs);
|
||||
|
||||
CHECK_CONTIGUOUS(grad);
|
||||
CHECK_CONTIGUOUS(inputs);
|
||||
CHECK_CONTIGUOUS(embeddings);
|
||||
CHECK_CONTIGUOUS(offsets);
|
||||
CHECK_CONTIGUOUS(grad_embeddings);
|
||||
CHECK_CONTIGUOUS(dy_dx);
|
||||
CHECK_CONTIGUOUS(grad_inputs);
|
||||
|
||||
CHECK_IS_FLOATING(grad);
|
||||
CHECK_IS_FLOATING(inputs);
|
||||
CHECK_IS_FLOATING(embeddings);
|
||||
CHECK_IS_INT(offsets);
|
||||
CHECK_IS_FLOATING(grad_embeddings);
|
||||
CHECK_IS_FLOATING(dy_dx);
|
||||
CHECK_IS_FLOATING(grad_inputs);
|
||||
|
||||
AT_DISPATCH_FLOATING_TYPES_AND_HALF(
|
||||
grad.scalar_type(), "grid_encode_backward", ([&] {
|
||||
grid_encode_backward_cuda<scalar_t>(grad.data_ptr<scalar_t>(), inputs.data_ptr<float>(), embeddings.data_ptr<scalar_t>(), offsets.data_ptr<int>(), grad_embeddings.data_ptr<scalar_t>(), B, D, C, L, S, H, calc_grad_inputs, dy_dx.data_ptr<scalar_t>(), grad_inputs.data_ptr<scalar_t>(), gridtype, align_corners);
|
||||
}));
|
||||
|
||||
}
|
||||
+15
@@ -0,0 +1,15 @@
|
||||
#ifndef _HASH_ENCODE_H
|
||||
#define _HASH_ENCODE_H
|
||||
|
||||
#include <stdint.h>
|
||||
#include <torch/torch.h>
|
||||
|
||||
// inputs: [B, D], float, in [0, 1]
|
||||
// embeddings: [sO, C], float
|
||||
// offsets: [L + 1], uint32_t
|
||||
// outputs: [B, L * C], float
|
||||
// H: base resolution
|
||||
void grid_encode_forward(const at::Tensor inputs, const at::Tensor embeddings, const at::Tensor offsets, at::Tensor outputs, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, at::Tensor dy_dx, const uint32_t gridtype, const bool align_corners);
|
||||
void grid_encode_backward(const at::Tensor grad, const at::Tensor inputs, const at::Tensor embeddings, const at::Tensor offsets, at::Tensor grad_embeddings, const uint32_t B, const uint32_t D, const uint32_t C, const uint32_t L, const float S, const uint32_t H, const bool calc_grad_inputs, const at::Tensor dy_dx, at::Tensor grad_inputs, const uint32_t gridtype, const bool align_corners);
|
||||
|
||||
#endif
|
||||
@@ -0,0 +1,475 @@
|
||||
# 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 torch,pdb,os,sys
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(f'{code_dir}/../')
|
||||
from Utils import *
|
||||
from pytorch3d.transforms import so3_log_map,so3_exp_map,se3_exp_map
|
||||
|
||||
img2mse = lambda x, y : torch.mean((x - y) ** 2)
|
||||
img2mae = lambda x, y: (torch.abs(x - y)).mean()
|
||||
mse2psnr = lambda x : -10. * torch.log(x) / torch.log(torch.Tensor([10.]))
|
||||
to8b = lambda x : (255*np.clip(x,0,1)).astype(np.uint8)
|
||||
|
||||
|
||||
class FeatureArray(nn.Module):
|
||||
"""
|
||||
Per-frame corrective latent code.
|
||||
"""
|
||||
|
||||
def __init__(self, num_frames, num_channels):
|
||||
super().__init__()
|
||||
|
||||
self.num_frames = num_frames
|
||||
self.num_channels = num_channels
|
||||
|
||||
self.data = nn.parameter.Parameter(torch.normal(0,1,size=[num_frames, num_channels]).float(), requires_grad=True)
|
||||
self.register_parameter('data',self.data)
|
||||
|
||||
|
||||
def __call__(self, ids):
|
||||
return self.data[ids]
|
||||
|
||||
|
||||
class PoseArray(nn.Module):
|
||||
def __init__(self, num_frames,max_trans,max_rot):
|
||||
super().__init__()
|
||||
self.num_frames = num_frames
|
||||
self.max_trans = max_trans
|
||||
self.max_rot = max_rot
|
||||
self.data = nn.parameter.Parameter(torch.zeros([num_frames, 6]).float(), requires_grad=True)
|
||||
self.register_parameter('data',self.data)
|
||||
|
||||
|
||||
def get_matrices(self,ids):
|
||||
if not torch.is_tensor(ids):
|
||||
ids = torch.tensor(ids).long()
|
||||
theta = torch.tanh(self.data)
|
||||
trans = theta[:,:3] * self.max_trans
|
||||
rot = theta[:,3:6] * self.max_rot/180.0*np.pi
|
||||
Ts_data = se3_exp_map(torch.cat((trans,rot),dim=-1)).permute(0,2,1)
|
||||
Ts = torch.eye(4, device=self.data.device).reshape(1,4,4).repeat(len(ids),1,1)
|
||||
mask = ids!=0
|
||||
Ts[mask] = Ts_data[ids[mask]]
|
||||
return Ts
|
||||
|
||||
|
||||
|
||||
class SHEncoder(nn.Module):
|
||||
'''Spherical encoding
|
||||
'''
|
||||
def __init__(self, input_dim=3, degree=4):
|
||||
|
||||
super().__init__()
|
||||
|
||||
self.input_dim = input_dim
|
||||
self.degree = degree
|
||||
|
||||
assert self.input_dim == 3
|
||||
assert self.degree >= 1 and self.degree <= 5
|
||||
|
||||
self.out_dim = degree ** 2
|
||||
|
||||
self.C0 = 0.28209479177387814
|
||||
self.C1 = 0.4886025119029199
|
||||
self.C2 = [
|
||||
1.0925484305920792,
|
||||
-1.0925484305920792,
|
||||
0.31539156525252005,
|
||||
-1.0925484305920792,
|
||||
0.5462742152960396
|
||||
]
|
||||
self.C3 = [
|
||||
-0.5900435899266435,
|
||||
2.890611442640554,
|
||||
-0.4570457994644658,
|
||||
0.3731763325901154,
|
||||
-0.4570457994644658,
|
||||
1.445305721320277,
|
||||
-0.5900435899266435
|
||||
]
|
||||
self.C4 = [
|
||||
2.5033429417967046,
|
||||
-1.7701307697799304,
|
||||
0.9461746957575601,
|
||||
-0.6690465435572892,
|
||||
0.10578554691520431,
|
||||
-0.6690465435572892,
|
||||
0.47308734787878004,
|
||||
-1.7701307697799304,
|
||||
0.6258357354491761
|
||||
]
|
||||
|
||||
def forward(self, input, **kwargs):
|
||||
|
||||
result = torch.empty((*input.shape[:-1], self.out_dim), dtype=input.dtype, device=input.device)
|
||||
x, y, z = input.unbind(-1)
|
||||
|
||||
result[..., 0] = self.C0
|
||||
if self.degree > 1:
|
||||
result[..., 1] = -self.C1 * y
|
||||
result[..., 2] = self.C1 * z
|
||||
result[..., 3] = -self.C1 * x
|
||||
if self.degree > 2:
|
||||
xx, yy, zz = x * x, y * y, z * z
|
||||
xy, yz, xz = x * y, y * z, x * z
|
||||
result[..., 4] = self.C2[0] * xy
|
||||
result[..., 5] = self.C2[1] * yz
|
||||
result[..., 6] = self.C2[2] * (2.0 * zz - xx - yy)
|
||||
#result[..., 6] = self.C2[2] * (3.0 * zz - 1) # xx + yy + zz == 1, but this will lead to different backward gradients, interesting...
|
||||
result[..., 7] = self.C2[3] * xz
|
||||
result[..., 8] = self.C2[4] * (xx - yy)
|
||||
if self.degree > 3:
|
||||
result[..., 9] = self.C3[0] * y * (3 * xx - yy)
|
||||
result[..., 10] = self.C3[1] * xy * z
|
||||
result[..., 11] = self.C3[2] * y * (4 * zz - xx - yy)
|
||||
result[..., 12] = self.C3[3] * z * (2 * zz - 3 * xx - 3 * yy)
|
||||
result[..., 13] = self.C3[4] * x * (4 * zz - xx - yy)
|
||||
result[..., 14] = self.C3[5] * z * (xx - yy)
|
||||
result[..., 15] = self.C3[6] * x * (xx - 3 * yy)
|
||||
if self.degree > 4:
|
||||
result[..., 16] = self.C4[0] * xy * (xx - yy)
|
||||
result[..., 17] = self.C4[1] * yz * (3 * xx - yy)
|
||||
result[..., 18] = self.C4[2] * xy * (7 * zz - 1)
|
||||
result[..., 19] = self.C4[3] * yz * (7 * zz - 3)
|
||||
result[..., 20] = self.C4[4] * (zz * (35 * zz - 30) + 3)
|
||||
result[..., 21] = self.C4[5] * xz * (7 * zz - 3)
|
||||
result[..., 22] = self.C4[6] * (xx - yy) * (7 * zz - 1)
|
||||
result[..., 23] = self.C4[7] * xz * (xx - 3 * yy)
|
||||
result[..., 24] = self.C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy))
|
||||
|
||||
return result
|
||||
|
||||
|
||||
class Embedder(nn.Module):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__()
|
||||
self.kwargs = kwargs
|
||||
self.create_embedding_fn()
|
||||
|
||||
def create_embedding_fn(self):
|
||||
embed_fns = []
|
||||
d = self.kwargs['input_dims']
|
||||
out_dim = 0
|
||||
if self.kwargs['include_input']:
|
||||
embed_fns.append(lambda x : x)
|
||||
out_dim += d
|
||||
|
||||
max_freq = self.kwargs['max_freq_log2']
|
||||
N_freqs = self.kwargs['num_freqs']
|
||||
|
||||
if self.kwargs['log_sampling']:
|
||||
freq_bands = 2.**torch.linspace(0., max_freq, steps=N_freqs)
|
||||
else:
|
||||
freq_bands = torch.linspace(2.**0., 2.**max_freq, steps=N_freqs)
|
||||
|
||||
for freq in freq_bands:
|
||||
for p_fn in self.kwargs['periodic_fns']:
|
||||
embed_fns.append(lambda x, p_fn=p_fn, freq=freq : p_fn(x * freq))
|
||||
out_dim += d
|
||||
|
||||
self.embed_fns = embed_fns
|
||||
self.out_dim = out_dim
|
||||
|
||||
def forward(self, inputs):
|
||||
return torch.cat([fn(inputs) for fn in self.embed_fns], -1)
|
||||
|
||||
|
||||
def get_embedder(multires, cfg, i=0, octree_m=None):
|
||||
if i == -1:
|
||||
return nn.Identity(), 3
|
||||
elif i==0:
|
||||
embed_kwargs = {
|
||||
'include_input' : True,
|
||||
'input_dims' : 3,
|
||||
'max_freq_log2' : multires-1,
|
||||
'num_freqs' : multires,
|
||||
'log_sampling' : True,
|
||||
'periodic_fns' : [torch.sin, torch.cos],
|
||||
}
|
||||
|
||||
embed = Embedder(**embed_kwargs)
|
||||
out_dim = embed.out_dim
|
||||
elif i==1:
|
||||
from mycuda.torch_ngp_grid_encoder.grid import GridEncoder
|
||||
embed = GridEncoder(input_dim=3, n_levels=cfg['num_levels'], log2_hashmap_size=cfg['log2_hashmap_size'], desired_resolution=cfg['finest_res'], base_resolution=cfg['base_res'], level_dim=cfg['feature_grid_dim'])
|
||||
print(embed)
|
||||
out_dim = embed.out_dim
|
||||
elif i==2:
|
||||
embed = SHEncoder(degree=cfg['multires_views'])
|
||||
out_dim = embed.out_dim
|
||||
return embed, out_dim
|
||||
|
||||
|
||||
|
||||
def mesh_to_real_world(mesh,pose_offset,translation,sc_factor):
|
||||
'''
|
||||
@pose_offset: optimized delta pose of the first frame. Usually it's identity
|
||||
'''
|
||||
mesh.vertices = mesh.vertices/sc_factor - np.array(translation).reshape(1,3)
|
||||
mesh.apply_transform(pose_offset)
|
||||
return mesh
|
||||
|
||||
|
||||
def get_optimized_poses_in_real_world(poses_normalized, pose_array, sc_factor, translation):
|
||||
'''
|
||||
@poses_normalized: np array, cam_in_ob (opengl convention), normalized to [-1,1] and centered
|
||||
@pose_array: PoseArray, delta poses
|
||||
Return:
|
||||
cam_in_ob, real-world unit, opencv convention
|
||||
'''
|
||||
original_poses = poses_normalized.copy()
|
||||
original_poses[:, :3, 3] /= sc_factor # To true world scale
|
||||
original_poses[:, :3, 3] -= translation
|
||||
|
||||
# Apply pose transformation
|
||||
tf = pose_array.get_matrices(np.arange(len(poses_normalized))).reshape(-1,4,4).data.cpu().numpy()
|
||||
optimized_poses = tf@poses_normalized
|
||||
|
||||
optimized_poses = np.array(optimized_poses).astype(np.float32)
|
||||
optimized_poses[:, :3, 3] /= sc_factor
|
||||
optimized_poses[:, :3, 3] -= translation
|
||||
|
||||
original_init_ob_in_cam = optimized_poses[0].copy()
|
||||
offset = np.linalg.inv(original_init_ob_in_cam)@original_poses[0]
|
||||
for i in range(len(optimized_poses)):
|
||||
new_ob_in_cam = optimized_poses[i]@offset
|
||||
optimized_poses[i] = new_ob_in_cam
|
||||
optimized_poses[i] = optimized_poses[i]@glcam_in_cvcam
|
||||
|
||||
return optimized_poses,offset
|
||||
|
||||
def preprocess_data(rgbs,depths,masks,normal_maps,poses,sc_factor,translation):
|
||||
'''
|
||||
@rgbs: np array (N,H,W,3)
|
||||
@depths: (N,H,W)
|
||||
@masks: (N,H,W)
|
||||
@normal_maps: (N,H,W,3)
|
||||
@poses: (N,4,4)
|
||||
'''
|
||||
depths[depths<0.001] = BAD_DEPTH
|
||||
if masks is not None:
|
||||
rgbs[masks==0] = BAD_COLOR
|
||||
depths[masks==0] = BAD_DEPTH
|
||||
if normal_maps is not None:
|
||||
normal_maps[...,[1,2]] *= -1 # To OpenGL
|
||||
normal_maps[masks==0] = 0
|
||||
masks = masks[...,None]
|
||||
|
||||
rgbs = (rgbs / 255.0).astype(np.float32)
|
||||
depths *= sc_factor
|
||||
depths = depths[...,None]
|
||||
poses[:, :3, 3] += translation
|
||||
poses[:, :3, 3] *= sc_factor
|
||||
return rgbs,depths,masks,normal_maps,poses
|
||||
|
||||
|
||||
class NeRFSmall(nn.Module):
|
||||
def __init__(self,num_layers=3,hidden_dim=64,geo_feat_dim=15,num_layers_color=4,hidden_dim_color=64,input_ch=3, input_ch_views=3):
|
||||
super(NeRFSmall, self).__init__()
|
||||
|
||||
self.input_ch = input_ch
|
||||
self.input_ch_views = input_ch_views
|
||||
|
||||
# sigma network
|
||||
self.num_layers = num_layers
|
||||
self.hidden_dim = hidden_dim
|
||||
self.geo_feat_dim = geo_feat_dim
|
||||
|
||||
sigma_net = []
|
||||
for l in range(num_layers):
|
||||
if l == 0:
|
||||
in_dim = self.input_ch
|
||||
else:
|
||||
in_dim = hidden_dim
|
||||
|
||||
if l == num_layers - 1:
|
||||
out_dim = 1 + self.geo_feat_dim # 1 sigma + 15 SH features for color
|
||||
else:
|
||||
out_dim = hidden_dim
|
||||
|
||||
sigma_net.append(nn.Linear(in_dim, out_dim, bias=True))
|
||||
if l!=num_layers-1:
|
||||
sigma_net.append(nn.ReLU(inplace=True))
|
||||
|
||||
self.sigma_net = nn.Sequential(*sigma_net)
|
||||
torch.nn.init.constant_(self.sigma_net[-1].bias, 0.1) # Encourage last layer predict positive SDF
|
||||
|
||||
# color network
|
||||
self.num_layers_color = num_layers_color
|
||||
self.hidden_dim_color = hidden_dim_color
|
||||
|
||||
color_net = []
|
||||
for l in range(num_layers_color):
|
||||
if l == 0:
|
||||
in_dim = self.input_ch_views + self.geo_feat_dim
|
||||
else:
|
||||
in_dim = hidden_dim
|
||||
|
||||
if l == num_layers_color - 1:
|
||||
out_dim = 3 # 3 rgb
|
||||
else:
|
||||
out_dim = hidden_dim
|
||||
|
||||
color_net.append(nn.Linear(in_dim, out_dim, bias=True))
|
||||
if l!=num_layers_color-1:
|
||||
color_net.append(nn.ReLU(inplace=True))
|
||||
|
||||
self.color_net = nn.Sequential(*color_net)
|
||||
|
||||
def forward_sdf(self,x):
|
||||
'''
|
||||
@x: embedded positions
|
||||
'''
|
||||
h = self.sigma_net(x)
|
||||
sigma, geo_feat = h[..., 0], h[..., 1:]
|
||||
return sigma
|
||||
|
||||
|
||||
def forward(self, x):
|
||||
x = x.float()
|
||||
input_pts, input_views = torch.split(x, [self.input_ch, self.input_ch_views], dim=-1)
|
||||
|
||||
# sigma
|
||||
h = input_pts
|
||||
h = self.sigma_net(h)
|
||||
|
||||
sigma, geo_feat = h[..., 0], h[..., 1:]
|
||||
|
||||
# color
|
||||
h = torch.cat([input_views, geo_feat], dim=-1)
|
||||
color = self.color_net(h)
|
||||
|
||||
outputs = torch.cat([color, sigma.unsqueeze(dim=-1)], -1)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
def sample_pdf(bins, weights, N_samples, det=False):
|
||||
weights = weights + 1e-5 # prevent nans
|
||||
pdf = weights / torch.sum(weights, -1, keepdim=True)
|
||||
cdf = torch.cumsum(pdf, -1)
|
||||
cdf = torch.cat([torch.zeros_like(cdf[...,:1]), cdf], -1) # (batch, len(bins))
|
||||
|
||||
if det:
|
||||
u = torch.linspace(0., 1., steps=N_samples)
|
||||
u = u.expand(list(cdf.shape[:-1]) + [N_samples])
|
||||
else:
|
||||
u = torch.rand(list(cdf.shape[:-1]) + [N_samples])
|
||||
|
||||
u = u.contiguous()
|
||||
inds = torch.searchsorted(cdf, u, right=True)
|
||||
below = torch.max(torch.zeros_like(inds-1), inds-1)
|
||||
above = torch.min((cdf.shape[-1]-1) * torch.ones_like(inds), inds)
|
||||
inds_g = torch.stack([below, above], -1) # (batch, N_samples, 2)
|
||||
|
||||
matched_shape = [inds_g.shape[0], inds_g.shape[1], cdf.shape[-1]]
|
||||
cdf_g = torch.gather(cdf.unsqueeze(1).expand(matched_shape), 2, inds_g)
|
||||
bins_g = torch.gather(bins.unsqueeze(1).expand(matched_shape), 2, inds_g)
|
||||
|
||||
denom = (cdf_g[...,1]-cdf_g[...,0])
|
||||
denom = torch.where(denom<1e-5, torch.ones_like(denom), denom)
|
||||
t = (u-cdf_g[...,0])/denom
|
||||
samples = bins_g[...,0] + t * (bins_g[...,1]-bins_g[...,0])
|
||||
|
||||
return samples
|
||||
|
||||
|
||||
|
||||
def get_camera_rays_np(H, W, K):
|
||||
"""Get ray origins, directions from a pinhole camera."""
|
||||
i, j = np.meshgrid(np.arange(W, dtype=np.float32),
|
||||
np.arange(H, dtype=np.float32), indexing='xy')
|
||||
dirs = np.stack([(i - K[0,2])/K[0,0], -(j - K[1,2])/K[1,1], -np.ones_like(i)], axis=-1)
|
||||
return dirs
|
||||
|
||||
|
||||
|
||||
def get_masks(z_vals, target_d, truncation, cfg, dir_norm=None):
|
||||
valid_depth_mask = (target_d>=cfg['near']*cfg['sc_factor']) & (target_d<=cfg['far']*cfg['sc_factor'])
|
||||
front_mask = (z_vals < target_d - truncation)
|
||||
back_mask = (z_vals > target_d + truncation*cfg['neg_trunc_ratio'])
|
||||
|
||||
sdf_mask = (1.0 - front_mask.float()) * (1.0 - back_mask.float()) * valid_depth_mask
|
||||
|
||||
num_fs_samples = front_mask.sum()
|
||||
num_sdf_samples = sdf_mask.sum()
|
||||
num_samples = num_sdf_samples + num_fs_samples
|
||||
fs_weight = 0.5
|
||||
sdf_weight = 1.0 - fs_weight
|
||||
return front_mask.bool(), sdf_mask.bool(), fs_weight, sdf_weight
|
||||
|
||||
|
||||
def get_sdf_loss(z_vals, target_d, predicted_sdf, truncation, cfg, return_mask=False, sample_weights=None, rays_d=None):
|
||||
dir_norm = rays_d.norm(dim=-1,keepdim=True)
|
||||
front_mask, sdf_mask, fs_weight, sdf_weight = get_masks(z_vals, target_d, truncation, cfg, dir_norm=dir_norm)
|
||||
front_mask = front_mask.bool()
|
||||
|
||||
mask = (target_d>cfg['far']*cfg['sc_factor']) & (predicted_sdf<cfg['fs_sdf'])
|
||||
|
||||
fs_loss = torch.mean(((predicted_sdf-cfg['fs_sdf']) * mask)**2 * sample_weights) * fs_weight
|
||||
|
||||
mask = front_mask & (target_d<=cfg['far']*cfg['sc_factor']) & (predicted_sdf<1)
|
||||
empty_loss = torch.mean(torch.abs(predicted_sdf-1) * mask * sample_weights)
|
||||
sdf_loss = torch.mean(((z_vals + predicted_sdf * truncation) * sdf_mask - target_d * sdf_mask)**2 * sample_weights) * sdf_weight
|
||||
|
||||
if return_mask:
|
||||
return fs_loss,sdf_loss,empty_loss, front_mask,sdf_mask
|
||||
return fs_loss, sdf_loss, empty_loss
|
||||
|
||||
|
||||
|
||||
def ray_box_intersection_batch(origins, dirs, bounds):
|
||||
'''
|
||||
@origins: (N,3) origin and directions. In the same coordinate frame as the bounding box
|
||||
@bounds: (2,3) xyz_min and max
|
||||
'''
|
||||
if not torch.is_tensor(origins):
|
||||
origins = torch.tensor(origins)
|
||||
dirs = torch.tensor(dirs)
|
||||
if not torch.is_tensor(bounds):
|
||||
bounds = torch.tensor(bounds)
|
||||
|
||||
dirs = dirs/(torch.norm(dirs,dim=-1,keepdim=True)+1e-10)
|
||||
inv_dirs = 1/dirs
|
||||
bounds = bounds[None].expand(len(dirs),-1,-1) #(N,2,3)
|
||||
|
||||
sign = torch.zeros((len(dirs),3)).long().to(dirs.device) #(N,3)
|
||||
sign[:,0] = (inv_dirs[:,0] < 0)
|
||||
sign[:,1] = (inv_dirs[:,1] < 0)
|
||||
sign[:,2] = (inv_dirs[:,2] < 0)
|
||||
|
||||
tmin = (torch.gather(bounds[...,0],dim=1,index=sign[:,0].reshape(-1,1)).reshape(-1) - origins[:,0]) * inv_dirs[:,0] #(N)
|
||||
tmin[tmin<0] = 0
|
||||
tmax = (torch.gather(bounds[...,0],dim=1,index=1-sign[:,0].reshape(-1,1)).reshape(-1) - origins[:,0]) * inv_dirs[:,0]
|
||||
tymin = (torch.gather(bounds[...,1],dim=1,index=sign[:,1].reshape(-1,1)).reshape(-1) - origins[:,1]) * inv_dirs[:,1]
|
||||
tymin[tymin<0] = 0
|
||||
tymax = (torch.gather(bounds[...,1],dim=1,index=1-sign[:,1].reshape(-1,1)).reshape(-1) - origins[:,1]) * inv_dirs[:,1]
|
||||
|
||||
ishit = torch.ones(len(dirs)).bool().to(dirs.device)
|
||||
ishit[(tmin > tymax) | (tymin > tmax)] = 0
|
||||
tmin[tymin>tmin] = tymin[tymin>tmin]
|
||||
tmax[tymax<tmax] = tymax[tymax<tmax]
|
||||
|
||||
tzmin = (torch.gather(bounds[...,2],dim=1,index=sign[:,2].reshape(-1,1)).reshape(-1) - origins[:,2]) * inv_dirs[:,2]
|
||||
tzmin[tzmin<0] = 0
|
||||
tzmax = (torch.gather(bounds[...,2],dim=1,index=1-sign[:,2].reshape(-1,1)).reshape(-1) - origins[:,2]) * inv_dirs[:,2]
|
||||
|
||||
ishit[(tmin > tzmax) | (tzmin > tmax)] = 0
|
||||
tmin[tzmin>tmin] = tzmin[tzmin>tmin] #(N)
|
||||
tmax[tzmax<tmax] = tzmax[tzmax<tmax]
|
||||
|
||||
tmin[ishit==0] = -1
|
||||
tmax[ishit==0] = -1
|
||||
|
||||
return tmin, tmax
|
||||
+1232
File diff suppressed because it is too large
Load Diff
+115
@@ -0,0 +1,115 @@
|
||||
# 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.
|
||||
|
||||
|
||||
from nerf_runner import *
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(f'{code_dir}/../')
|
||||
from datareader import *
|
||||
from bundlesdf.tool import *
|
||||
import yaml,argparse
|
||||
|
||||
|
||||
def run_neural_object_field(cfg, K, rgbs, depths, masks, cam_in_obs, debug=0, save_dir='/home/bowen/debug/foundationpose_bundlesdf'):
|
||||
rgbs = np.asarray(rgbs)
|
||||
depths = np.asarray(depths)
|
||||
masks = np.asarray(masks)
|
||||
cam_in_obs = np.asarray(cam_in_obs)
|
||||
glcam_in_obs = cam_in_obs@glcam_in_cvcam
|
||||
|
||||
cfg['save_dir'] = save_dir
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
|
||||
for i,rgb in enumerate(rgbs):
|
||||
imageio.imwrite(f'{save_dir}/rgb_{i:07d}.png', rgb)
|
||||
|
||||
sc_factor,translation,pcd_real_scale, pcd_normalized = compute_scene_bounds(None,glcam_in_obs, K, use_mask=True,base_dir=save_dir,rgbs=rgbs,depths=depths,masks=masks, eps=cfg['dbscan_eps'], min_samples=cfg['dbscan_eps_min_samples'])
|
||||
cfg['sc_factor'] = sc_factor
|
||||
cfg['translation'] = translation
|
||||
|
||||
o3d.io.write_point_cloud(f'{save_dir}/pcd_normalized.ply', pcd_normalized)
|
||||
|
||||
rgbs_, depths_, masks_, normal_maps,poses = preprocess_data(rgbs, depths, masks,normal_maps=None,poses=glcam_in_obs,sc_factor=cfg['sc_factor'],translation=cfg['translation'])
|
||||
|
||||
nerf = NerfRunner(cfg, rgbs_, depths_, masks_, normal_maps=None, poses=poses, K=K, occ_masks=None, build_octree_pcd=pcd_normalized)
|
||||
nerf.train()
|
||||
|
||||
mesh = nerf.extract_mesh(isolevel=0,voxel_size=cfg['mesh_resolution'])
|
||||
mesh = nerf.mesh_texture_from_train_images(mesh, rgbs_raw=rgbs, tex_res=1028)
|
||||
optimized_cvcam_in_obs,offset = get_optimized_poses_in_real_world(poses,nerf.models['pose_array'],cfg['sc_factor'],cfg['translation'])
|
||||
mesh = mesh_to_real_world(mesh, pose_offset=offset, translation=nerf.cfg['translation'], sc_factor=nerf.cfg['sc_factor'])
|
||||
return mesh
|
||||
|
||||
|
||||
def run_one_ob(base_dir, cfg, use_refined_mask=False):
|
||||
save_dir = f'{base_dir}/nerf'
|
||||
os.system(f'rm -rf {save_dir} && mkdir -p {save_dir}')
|
||||
with open(f'{base_dir}/select_frames.yml','r') as ff:
|
||||
info = yaml.safe_load(ff)
|
||||
rgbs = []
|
||||
depths = []
|
||||
masks = []
|
||||
cam_in_obs = []
|
||||
color_files = sorted(glob.glob(f'{base_dir}/rgb/*.png'))
|
||||
K = np.loadtxt(f'{base_dir}/K.txt')
|
||||
for i,color_file in enumerate(color_files):
|
||||
rgb = imageio.imread(color_file)
|
||||
depth = cv2.imread(color_file.replace('rgb','depth_enhanced'), -1)/1e3
|
||||
if use_refined_mask:
|
||||
mask = cv2.imread(color_file.replace('rgb','mask_refined'), -1)
|
||||
else:
|
||||
mask = cv2.imread(color_file.replace('rgb','mask'), -1)
|
||||
cam_in_ob = np.loadtxt(color_file.replace('rgb','cam_in_ob').replace('.png','.txt')).reshape(4,4)
|
||||
rgbs.append(rgb)
|
||||
depths.append(depth)
|
||||
masks.append(mask)
|
||||
cam_in_obs.append(cam_in_ob)
|
||||
|
||||
mesh = run_neural_object_field(cfg, K, rgbs, depths, masks, cam_in_obs, save_dir=save_dir, debug=0)
|
||||
return mesh
|
||||
|
||||
|
||||
def run_ycbv():
|
||||
ob_ids = np.arange(1,22)
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
with open(f'{code_dir}/config_ycbv.yml','r') as ff:
|
||||
cfg = yaml.safe_load(ff)
|
||||
|
||||
for ob_id in ob_ids:
|
||||
base_dir = f'{args.ref_view_dir}/ob_{ob_id:07d}'
|
||||
mesh = run_one_ob(base_dir=base_dir, cfg=cfg)
|
||||
out_file = f'{base_dir}/model/model.obj'
|
||||
os.makedirs(os.path.dirname(out_file), exist_ok=True)
|
||||
mesh.export(out_file)
|
||||
|
||||
|
||||
def run_linemod():
|
||||
ob_ids = np.setdiff1d(np.arange(1,16), np.array([7,3])).tolist()
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
with open(f'{code_dir}/config_linemod.yml','r') as ff:
|
||||
cfg = yaml.safe_load(ff)
|
||||
for ob_id in ob_ids:
|
||||
base_dir = f'{args.ref_view_dir}/ob_{ob_id:07d}'
|
||||
mesh = run_one_ob(base_dir=base_dir, cfg=cfg, use_refined_mask=True)
|
||||
out_file = f'{base_dir}/model/model.obj'
|
||||
os.makedirs(os.path.dirname(out_file), exist_ok=True)
|
||||
mesh.export(out_file)
|
||||
logging.info(f"saved to {out_file}")
|
||||
|
||||
|
||||
if __name__=="__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
parser.add_argument('--ref_view_dir', type=str, default=f'/mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video/bowen_addon/ref_views_16')
|
||||
parser.add_argument('--dataset', type=str, default=f'ycbv', help='one of [ycbv/linemod]')
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.dataset=='ycbv':
|
||||
run_ycbv()
|
||||
else:
|
||||
run_linemod()
|
||||
+132
@@ -0,0 +1,132 @@
|
||||
# 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 joblib,json,gzip,pickle
|
||||
from sklearn.cluster import DBSCAN
|
||||
import shutil,re,imageio,pdb,os,sys
|
||||
from Utils import *
|
||||
import pandas as pd
|
||||
|
||||
|
||||
def find_biggest_cluster(pts, eps=0.06, min_samples=1):
|
||||
dbscan = DBSCAN(eps=eps,min_samples=min_samples,n_jobs=-1)
|
||||
dbscan.fit(pts)
|
||||
ids, cnts = np.unique(dbscan.labels_, return_counts=True)
|
||||
best_id = ids[cnts.argsort()[-1]]
|
||||
keep_mask = dbscan.labels_==best_id
|
||||
pts_cluster = pts[keep_mask]
|
||||
return pts_cluster, keep_mask
|
||||
|
||||
|
||||
def compute_translation_scales(pts,max_dim=2,cluster=True, eps=0.06, min_samples=1):
|
||||
if cluster:
|
||||
pts, keep_mask = find_biggest_cluster(pts, eps, min_samples)
|
||||
else:
|
||||
keep_mask = np.ones((len(pts)), dtype=bool)
|
||||
max_xyz = pts.max(axis=0)
|
||||
min_xyz = pts.min(axis=0)
|
||||
center = (max_xyz+min_xyz)/2
|
||||
sc_factor = max_dim/(max_xyz-min_xyz).max() #Normalize to [-1,1]
|
||||
sc_factor *= 0.9
|
||||
translation_cvcam = -center
|
||||
return translation_cvcam, sc_factor, keep_mask
|
||||
|
||||
|
||||
def compute_scene_bounds_worker(color_file,K,glcam_in_world,use_mask,rgb=None,depth=None,mask=None):
|
||||
if rgb is None:
|
||||
depth_file = color_file.replace('images','depth_filtered')
|
||||
mask_file = color_file.replace('images','masks')
|
||||
rgb = np.array(Image.open(color_file))[...,:3]
|
||||
depth = cv2.imread(depth_file,-1)/1e3
|
||||
xyz_map = depth2xyzmap(depth,K)
|
||||
valid = depth>=0.1
|
||||
if use_mask:
|
||||
if mask is None:
|
||||
mask = cv2.imread(mask_file,-1)
|
||||
valid = valid & (mask>0)
|
||||
pts = xyz_map[valid].reshape(-1,3)
|
||||
if len(pts)==0:
|
||||
return None
|
||||
colors = rgb[valid].reshape(-1,3)
|
||||
pcd = toOpen3dCloud(pts,colors)
|
||||
pcd = pcd.voxel_down_sample(0.01)
|
||||
pcd, ind = pcd.remove_statistical_outlier(nb_neighbors=30,std_ratio=2.0)
|
||||
cam_in_world = glcam_in_world@glcam_in_cvcam
|
||||
pcd.transform(cam_in_world)
|
||||
return np.asarray(pcd.points).copy(), np.asarray(pcd.colors).copy()
|
||||
|
||||
|
||||
def compute_scene_bounds(color_files,glcam_in_worlds,K,use_mask=True,base_dir=None,rgbs=None,depths=None,masks=None,cluster=True, translation_cvcam=None, sc_factor=None, eps=0.06, min_samples=1):
|
||||
assert color_files is None or rgbs is None
|
||||
|
||||
if base_dir is None:
|
||||
base_dir = os.path.dirname(color_files[0])+'/../'
|
||||
|
||||
args = []
|
||||
if rgbs is not None:
|
||||
for i in range(len(rgbs)):
|
||||
args.append((None,K,glcam_in_worlds[i],use_mask,rgbs[i],depths[i],masks[i]))
|
||||
else:
|
||||
for i in range(len(color_files)):
|
||||
args.append((color_files[i],K,glcam_in_worlds[i],use_mask))
|
||||
|
||||
logging.info(f"compute_scene_bounds_worker start")
|
||||
ret = joblib.Parallel(n_jobs=10, prefer="threads")(joblib.delayed(compute_scene_bounds_worker)(*arg) for arg in args)
|
||||
logging.info(f"compute_scene_bounds_worker done")
|
||||
|
||||
pcd_all = None
|
||||
for r in ret:
|
||||
if r is None:
|
||||
continue
|
||||
if pcd_all is None:
|
||||
pcd_all = toOpen3dCloud(r[0],r[1])
|
||||
else:
|
||||
pcd_all += toOpen3dCloud(r[0],r[1])
|
||||
pcd = pcd_all.voxel_down_sample(eps/5)
|
||||
|
||||
logging.info(f"merge pcd")
|
||||
|
||||
o3d.io.write_point_cloud(f'{base_dir}/naive_fusion.ply',pcd)
|
||||
pts = np.asarray(pcd.points).copy()
|
||||
|
||||
def make_tf(translation_cvcam, sc_factor):
|
||||
tf = np.eye(4)
|
||||
tf[:3,3] = translation_cvcam
|
||||
tf1 = np.eye(4)
|
||||
tf1[:3,:3] *= sc_factor
|
||||
tf = tf1@tf
|
||||
return tf
|
||||
|
||||
if translation_cvcam is None:
|
||||
translation_cvcam, sc_factor, keep_mask = compute_translation_scales(pts, cluster=cluster, eps=eps, min_samples=min_samples)
|
||||
tf = make_tf(translation_cvcam, sc_factor)
|
||||
else:
|
||||
tf = make_tf(translation_cvcam, sc_factor)
|
||||
tmp = copy.deepcopy(pcd)
|
||||
tmp.transform(tf)
|
||||
tmp_pts = np.asarray(tmp.points)
|
||||
keep_mask = (np.abs(tmp_pts)<1).all(axis=-1)
|
||||
|
||||
logging.info(f"compute_translation_scales done")
|
||||
|
||||
pcd = toOpen3dCloud(pts[keep_mask],np.asarray(pcd.colors)[keep_mask])
|
||||
o3d.io.write_point_cloud(f"{base_dir}/naive_fusion_biggest_cluster.ply",pcd)
|
||||
pcd_real_scale = copy.deepcopy(pcd)
|
||||
print(f'translation_cvcam={translation_cvcam}, sc_factor={sc_factor}')
|
||||
with open(f'{base_dir}/normalization.yml','w') as ff:
|
||||
tmp = {
|
||||
'translation_cvcam':translation_cvcam.tolist(),
|
||||
'sc_factor':float(sc_factor),
|
||||
}
|
||||
yaml.dump(tmp,ff)
|
||||
|
||||
pcd.transform(tf)
|
||||
return sc_factor, translation_cvcam, pcd_real_scale, pcd
|
||||
|
||||
|
||||
+122
@@ -0,0 +1,122 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Quick sanity checks for a local FoundationPose install (no inference)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--demo-data",
|
||||
action="store_true",
|
||||
help="Also check default paths used by run_demo.py (demo_data/mustard0).",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
repo_root = Path(__file__).resolve().parent
|
||||
os.chdir(repo_root)
|
||||
if str(repo_root) not in sys.path:
|
||||
sys.path.insert(0, str(repo_root))
|
||||
|
||||
errors: list[str] = []
|
||||
warnings: list[str] = []
|
||||
|
||||
print(f"Repository: {repo_root}")
|
||||
print(f"Python: {sys.version.split()[0]}")
|
||||
|
||||
# --- PyTorch / CUDA ---
|
||||
try:
|
||||
import torch
|
||||
|
||||
print(f"PyTorch: {torch.__version__}")
|
||||
print(f"CUDA available: {torch.cuda.is_available()}")
|
||||
if torch.cuda.is_available():
|
||||
print(f"CUDA device: {torch.cuda.get_device_name(0)}")
|
||||
except Exception as e:
|
||||
errors.append(f"torch: {e}")
|
||||
|
||||
# --- Core imports (same chain as run_demo / estimater) ---
|
||||
modules = [
|
||||
"pytorch3d",
|
||||
"nvdiffrast.torch",
|
||||
"trimesh",
|
||||
"open3d",
|
||||
"cv2",
|
||||
"warp",
|
||||
]
|
||||
for mod in modules:
|
||||
try:
|
||||
__import__(mod)
|
||||
print(f"import {mod}: ok")
|
||||
except Exception as e:
|
||||
errors.append(f"import {mod}: {e}")
|
||||
|
||||
# --- Native extension ---
|
||||
build_dir = repo_root / "mycpp" / "build"
|
||||
so_files = list(build_dir.glob("mycpp*.so"))
|
||||
if not so_files:
|
||||
warnings.append(f"No mycpp extension found under {build_dir} (run bash build_all_conda.sh).")
|
||||
else:
|
||||
print(f"mycpp extension: {so_files[0].name}")
|
||||
|
||||
try:
|
||||
sys.path.insert(0, str(build_dir))
|
||||
import mycpp
|
||||
|
||||
print(f"import mycpp: ok ({mycpp})")
|
||||
except Exception as e:
|
||||
errors.append(f"import mycpp: {e}")
|
||||
|
||||
try:
|
||||
import estimater
|
||||
|
||||
print("import estimater: ok")
|
||||
except Exception as e:
|
||||
errors.append(f"import estimater: {e}")
|
||||
|
||||
# --- Weights (hard-coded run names in learning/training/predict_*.py) ---
|
||||
weight_sets = [
|
||||
("scorer", "2024-01-11-20-02-45"),
|
||||
("refiner", "2023-10-28-18-33-37"),
|
||||
]
|
||||
weights_root = repo_root / "weights"
|
||||
for label, stamp in weight_sets:
|
||||
d = weights_root / stamp
|
||||
ckpt = d / "model_best.pth"
|
||||
cfg = d / "config.yml"
|
||||
missing = [p.name for p in (ckpt, cfg) if not p.is_file()]
|
||||
if missing:
|
||||
warnings.append(
|
||||
f"Weights ({label}, {stamp}): missing {', '.join(missing)} — download per readme Data prepare."
|
||||
)
|
||||
else:
|
||||
print(f"Weights ({label}, {stamp}): ok")
|
||||
|
||||
# --- Optional demo bundle ---
|
||||
if args.demo_data:
|
||||
demo_scene = repo_root / "demo_data" / "mustard0"
|
||||
demo_mesh = demo_scene / "mesh" / "textured_simple.obj"
|
||||
if not demo_mesh.is_file():
|
||||
warnings.append(f"run_demo demo scene: missing {demo_mesh}")
|
||||
else:
|
||||
print(f"run_demo demo_data: ok ({demo_mesh.relative_to(repo_root)})")
|
||||
|
||||
print()
|
||||
for w in warnings:
|
||||
print(f"WARNING: {w}")
|
||||
for e in errors:
|
||||
print(f"ERROR: {e}")
|
||||
|
||||
# Non-zero only when imports or core checks fail; missing weights/data still exit 0.
|
||||
if errors:
|
||||
return 1
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
+613
@@ -0,0 +1,613 @@
|
||||
# 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.
|
||||
|
||||
|
||||
from Utils import *
|
||||
import json,os,sys
|
||||
|
||||
|
||||
BOP_LIST = ['lmo','tless','ycbv','hb','tudl','icbin','itodd']
|
||||
BOP_DIR = os.getenv('BOP_DIR')
|
||||
|
||||
def get_bop_reader(video_dir, zfar=np.inf):
|
||||
if 'ycbv' in video_dir or 'YCB' in video_dir:
|
||||
return YcbVideoReader(video_dir, zfar=zfar)
|
||||
if 'lmo' in video_dir or 'LINEMOD-O' in video_dir:
|
||||
return LinemodOcclusionReader(video_dir, zfar=zfar)
|
||||
if 'tless' in video_dir or 'TLESS' in video_dir:
|
||||
return TlessReader(video_dir, zfar=zfar)
|
||||
if 'hb' in video_dir:
|
||||
return HomebrewedReader(video_dir, zfar=zfar)
|
||||
if 'tudl' in video_dir:
|
||||
return TudlReader(video_dir, zfar=zfar)
|
||||
if 'icbin' in video_dir:
|
||||
return IcbinReader(video_dir, zfar=zfar)
|
||||
if 'itodd' in video_dir:
|
||||
return ItoddReader(video_dir, zfar=zfar)
|
||||
else:
|
||||
raise RuntimeError
|
||||
|
||||
|
||||
def get_bop_video_dirs(dataset):
|
||||
if dataset=='ycbv':
|
||||
video_dirs = sorted(glob.glob(f'{BOP_DIR}/ycbv/test/*'))
|
||||
elif dataset=='lmo':
|
||||
video_dirs = sorted(glob.glob(f'{BOP_DIR}/lmo/lmo_test_bop19/test/*'))
|
||||
elif dataset=='tless':
|
||||
video_dirs = sorted(glob.glob(f'{BOP_DIR}/tless/tless_test_primesense_bop19/test_primesense/*'))
|
||||
elif dataset=='hb':
|
||||
video_dirs = sorted(glob.glob(f'{BOP_DIR}/hb/hb_test_primesense_bop19/test_primesense/*'))
|
||||
elif dataset=='tudl':
|
||||
video_dirs = sorted(glob.glob(f'{BOP_DIR}/tudl/tudl_test_bop19/test/*'))
|
||||
elif dataset=='icbin':
|
||||
video_dirs = sorted(glob.glob(f'{BOP_DIR}/icbin/icbin_test_bop19/test/*'))
|
||||
elif dataset=='itodd':
|
||||
video_dirs = sorted(glob.glob(f'{BOP_DIR}/itodd/itodd_test_bop19/test/*'))
|
||||
else:
|
||||
raise RuntimeError
|
||||
return video_dirs
|
||||
|
||||
|
||||
|
||||
class YcbineoatReader:
|
||||
def __init__(self,video_dir, downscale=1, shorter_side=None, zfar=np.inf):
|
||||
self.video_dir = video_dir
|
||||
self.downscale = downscale
|
||||
self.zfar = zfar
|
||||
self.color_files = sorted(glob.glob(f"{self.video_dir}/rgb/*.png"))
|
||||
self.K = np.loadtxt(f'{video_dir}/cam_K.txt').reshape(3,3)
|
||||
self.id_strs = []
|
||||
for color_file in self.color_files:
|
||||
id_str = os.path.basename(color_file).replace('.png','')
|
||||
self.id_strs.append(id_str)
|
||||
self.H,self.W = cv2.imread(self.color_files[0]).shape[:2]
|
||||
|
||||
if shorter_side is not None:
|
||||
self.downscale = shorter_side/min(self.H, self.W)
|
||||
|
||||
self.H = int(self.H*self.downscale)
|
||||
self.W = int(self.W*self.downscale)
|
||||
self.K[:2] *= self.downscale
|
||||
|
||||
self.gt_pose_files = sorted(glob.glob(f'{self.video_dir}/annotated_poses/*'))
|
||||
|
||||
self.videoname_to_object = {
|
||||
'bleach0': "021_bleach_cleanser",
|
||||
'bleach_hard_00_03_chaitanya': "021_bleach_cleanser",
|
||||
'cracker_box_reorient': '003_cracker_box',
|
||||
'cracker_box_yalehand0': '003_cracker_box',
|
||||
'mustard0': '006_mustard_bottle',
|
||||
'mustard_easy_00_02': '006_mustard_bottle',
|
||||
'sugar_box1': '004_sugar_box',
|
||||
'sugar_box_yalehand0': '004_sugar_box',
|
||||
'tomato_soup_can_yalehand0': '005_tomato_soup_can',
|
||||
}
|
||||
|
||||
|
||||
def get_video_name(self):
|
||||
return self.video_dir.split('/')[-1]
|
||||
|
||||
def __len__(self):
|
||||
return len(self.color_files)
|
||||
|
||||
def get_gt_pose(self,i):
|
||||
try:
|
||||
pose = np.loadtxt(self.gt_pose_files[i]).reshape(4,4)
|
||||
return pose
|
||||
except:
|
||||
logging.info("GT pose not found, return None")
|
||||
return None
|
||||
|
||||
|
||||
def get_color(self,i):
|
||||
color = imageio.imread(self.color_files[i])[...,:3]
|
||||
color = cv2.resize(color, (self.W,self.H), interpolation=cv2.INTER_NEAREST)
|
||||
return color
|
||||
|
||||
def get_mask(self,i):
|
||||
mask = cv2.imread(self.color_files[i].replace('rgb','masks'),-1)
|
||||
if len(mask.shape)==3:
|
||||
for c in range(3):
|
||||
if mask[...,c].sum()>0:
|
||||
mask = mask[...,c]
|
||||
break
|
||||
mask = cv2.resize(mask, (self.W,self.H), interpolation=cv2.INTER_NEAREST).astype(bool).astype(np.uint8)
|
||||
return mask
|
||||
|
||||
def get_depth(self,i):
|
||||
depth = cv2.imread(self.color_files[i].replace('rgb','depth'),-1)/1e3
|
||||
depth = cv2.resize(depth, (self.W,self.H), interpolation=cv2.INTER_NEAREST)
|
||||
depth[(depth<0.001) | (depth>=self.zfar)] = 0
|
||||
return depth
|
||||
|
||||
|
||||
def get_xyz_map(self,i):
|
||||
depth = self.get_depth(i)
|
||||
xyz_map = depth2xyzmap(depth, self.K)
|
||||
return xyz_map
|
||||
|
||||
def get_occ_mask(self,i):
|
||||
hand_mask_file = self.color_files[i].replace('rgb','masks_hand')
|
||||
occ_mask = np.zeros((self.H,self.W), dtype=bool)
|
||||
if os.path.exists(hand_mask_file):
|
||||
occ_mask = occ_mask | (cv2.imread(hand_mask_file,-1)>0)
|
||||
|
||||
right_hand_mask_file = self.color_files[i].replace('rgb','masks_hand_right')
|
||||
if os.path.exists(right_hand_mask_file):
|
||||
occ_mask = occ_mask | (cv2.imread(right_hand_mask_file,-1)>0)
|
||||
|
||||
occ_mask = cv2.resize(occ_mask, (self.W,self.H), interpolation=cv2.INTER_NEAREST)
|
||||
|
||||
return occ_mask.astype(np.uint8)
|
||||
|
||||
def get_gt_mesh(self):
|
||||
ob_name = self.videoname_to_object[self.get_video_name()]
|
||||
YCB_VIDEO_DIR = os.getenv('YCB_VIDEO_DIR')
|
||||
mesh = trimesh.load(f'{YCB_VIDEO_DIR}/models/{ob_name}/textured_simple.obj')
|
||||
return mesh
|
||||
|
||||
|
||||
class BopBaseReader:
|
||||
def __init__(self, base_dir, zfar=np.inf, resize=1):
|
||||
self.base_dir = base_dir
|
||||
self.resize = resize
|
||||
self.dataset_name = None
|
||||
self.color_files = sorted(glob.glob(f"{self.base_dir}/rgb/*"))
|
||||
if len(self.color_files)==0:
|
||||
self.color_files = sorted(glob.glob(f"{self.base_dir}/gray/*"))
|
||||
self.zfar = zfar
|
||||
|
||||
self.K_table = {}
|
||||
with open(f'{self.base_dir}/scene_camera.json','r') as ff:
|
||||
info = json.load(ff)
|
||||
for k in info:
|
||||
self.K_table[f'{int(k):06d}'] = np.array(info[k]['cam_K']).reshape(3,3)
|
||||
self.bop_depth_scale = info[k]['depth_scale']
|
||||
|
||||
if os.path.exists(f'{self.base_dir}/scene_gt.json'):
|
||||
with open(f'{self.base_dir}/scene_gt.json','r') as ff:
|
||||
self.scene_gt = json.load(ff)
|
||||
self.scene_gt = copy.deepcopy(self.scene_gt) # Release file handle to be pickle-able by joblib
|
||||
assert len(self.scene_gt)==len(self.color_files)
|
||||
else:
|
||||
self.scene_gt = None
|
||||
|
||||
self.make_id_strs()
|
||||
|
||||
|
||||
def make_scene_ob_ids_dict(self):
|
||||
with open(f'{BOP_DIR}/{self.dataset_name}/test_targets_bop19.json','r') as ff:
|
||||
self.scene_ob_ids_dict = {}
|
||||
data = json.load(ff)
|
||||
for d in data:
|
||||
if d['scene_id']==self.get_video_id():
|
||||
id_str = f"{d['im_id']:06d}"
|
||||
if id_str not in self.scene_ob_ids_dict:
|
||||
self.scene_ob_ids_dict[id_str] = []
|
||||
self.scene_ob_ids_dict[id_str] += [d['obj_id']]*d['inst_count']
|
||||
|
||||
|
||||
def get_K(self, i_frame):
|
||||
K = self.K_table[self.id_strs[i_frame]]
|
||||
if self.resize!=1:
|
||||
K[:2,:2] *= self.resize
|
||||
return K
|
||||
|
||||
|
||||
def get_video_dir(self):
|
||||
video_id = int(self.base_dir.rstrip('/').split('/')[-1])
|
||||
return video_id
|
||||
|
||||
def make_id_strs(self):
|
||||
self.id_strs = []
|
||||
for i in range(len(self.color_files)):
|
||||
name = os.path.basename(self.color_files[i]).split('.')[0]
|
||||
self.id_strs.append(name)
|
||||
|
||||
|
||||
def get_instance_ids_in_image(self, i_frame:int):
|
||||
ob_ids = []
|
||||
if self.scene_gt is not None:
|
||||
name = int(os.path.basename(self.color_files[i_frame]).split('.')[0])
|
||||
for k in self.scene_gt[str(name)]:
|
||||
ob_ids.append(k['obj_id'])
|
||||
elif self.scene_ob_ids_dict is not None:
|
||||
return np.array(self.scene_ob_ids_dict[self.id_strs[i_frame]])
|
||||
else:
|
||||
mask_dir = os.path.dirname(self.color_files[0]).replace('rgb','mask_visib')
|
||||
id_str = self.id_strs[i_frame]
|
||||
mask_files = sorted(glob.glob(f'{mask_dir}/{id_str}_*.png'))
|
||||
ob_ids = []
|
||||
for mask_file in mask_files:
|
||||
ob_id = int(os.path.basename(mask_file).split('.')[0].split('_')[1])
|
||||
ob_ids.append(ob_id)
|
||||
ob_ids = np.asarray(ob_ids)
|
||||
return ob_ids
|
||||
|
||||
|
||||
def get_gt_mesh_file(self, ob_id):
|
||||
raise RuntimeError("You should override this")
|
||||
|
||||
|
||||
def get_color(self,i):
|
||||
color = imageio.imread(self.color_files[i])
|
||||
if len(color.shape)==2:
|
||||
color = np.tile(color[...,None], (1,1,3)) # Gray to RGB
|
||||
if self.resize!=1:
|
||||
color = cv2.resize(color, fx=self.resize, fy=self.resize, dsize=None)
|
||||
return color
|
||||
|
||||
|
||||
def get_depth(self,i, filled=False):
|
||||
if filled:
|
||||
depth_file = self.color_files[i].replace('rgb','depth_filled')
|
||||
depth_file = f'{os.path.dirname(depth_file)}/0{os.path.basename(depth_file)}'
|
||||
depth = cv2.imread(depth_file,-1)/1e3
|
||||
else:
|
||||
depth_file = self.color_files[i].replace('rgb','depth').replace('gray','depth')
|
||||
depth = cv2.imread(depth_file,-1)*1e-3*self.bop_depth_scale
|
||||
if self.resize!=1:
|
||||
depth = cv2.resize(depth, fx=self.resize, fy=self.resize, dsize=None, interpolation=cv2.INTER_NEAREST)
|
||||
depth[depth<0.001] = 0
|
||||
depth[depth>self.zfar] = 0
|
||||
return depth
|
||||
|
||||
def get_xyz_map(self,i):
|
||||
depth = self.get_depth(i)
|
||||
xyz_map = depth2xyzmap(depth, self.get_K(i))
|
||||
return xyz_map
|
||||
|
||||
|
||||
def get_mask(self, i_frame:int, ob_id:int, type='mask_visib'):
|
||||
'''
|
||||
@type: mask_visib (only visible part) / mask (projected mask from whole model)
|
||||
'''
|
||||
pos = 0
|
||||
name = int(os.path.basename(self.color_files[i_frame]).split('.')[0])
|
||||
if self.scene_gt is not None:
|
||||
for k in self.scene_gt[str(name)]:
|
||||
if k['obj_id']==ob_id:
|
||||
break
|
||||
pos += 1
|
||||
mask_file = f'{self.base_dir}/{type}/{name:06d}_{pos:06d}.png'
|
||||
if not os.path.exists(mask_file):
|
||||
logging.info(f'{mask_file} not found')
|
||||
return None
|
||||
else:
|
||||
# mask_dir = os.path.dirname(self.color_files[0]).replace('rgb',type)
|
||||
# mask_file = f'{mask_dir}/{self.id_strs[i_frame]}_{ob_id:06d}.png'
|
||||
raise RuntimeError
|
||||
mask = cv2.imread(mask_file, -1)
|
||||
if self.resize!=1:
|
||||
mask = cv2.resize(mask, fx=self.resize, fy=self.resize, dsize=None, interpolation=cv2.INTER_NEAREST)
|
||||
return mask>0
|
||||
|
||||
|
||||
def get_gt_mesh(self, ob_id:int):
|
||||
mesh_file = self.get_gt_mesh_file(ob_id)
|
||||
mesh = trimesh.load(mesh_file)
|
||||
mesh.vertices *= 1e-3
|
||||
return mesh
|
||||
|
||||
|
||||
def get_model_diameter(self, ob_id):
|
||||
dir = os.path.dirname(self.get_gt_mesh_file(self.ob_ids[0]))
|
||||
info_file = f'{dir}/models_info.json'
|
||||
with open(info_file,'r') as ff:
|
||||
info = json.load(ff)
|
||||
return info[str(ob_id)]['diameter']/1e3
|
||||
|
||||
|
||||
|
||||
def get_gt_poses(self, i_frame, ob_id):
|
||||
gt_poses = []
|
||||
name = int(self.id_strs[i_frame])
|
||||
for i_k, k in enumerate(self.scene_gt[str(name)]):
|
||||
if k['obj_id']==ob_id:
|
||||
cur = np.eye(4)
|
||||
cur[:3,:3] = np.array(k['cam_R_m2c']).reshape(3,3)
|
||||
cur[:3,3] = np.array(k['cam_t_m2c'])/1e3
|
||||
gt_poses.append(cur)
|
||||
return np.asarray(gt_poses).reshape(-1,4,4)
|
||||
|
||||
|
||||
def get_gt_pose(self, i_frame:int, ob_id, mask=None, use_my_correction=False):
|
||||
ob_in_cam = np.eye(4)
|
||||
best_iou = -np.inf
|
||||
best_gt_mask = None
|
||||
name = int(self.id_strs[i_frame])
|
||||
for i_k, k in enumerate(self.scene_gt[str(name)]):
|
||||
if k['obj_id']==ob_id:
|
||||
cur = np.eye(4)
|
||||
cur[:3,:3] = np.array(k['cam_R_m2c']).reshape(3,3)
|
||||
cur[:3,3] = np.array(k['cam_t_m2c'])/1e3
|
||||
if mask is not None: # When multi-instance exists, use mask to determine which one
|
||||
gt_mask = cv2.imread(f'{self.base_dir}/mask_visib/{self.id_strs[i_frame]}_{i_k:06d}.png', -1).astype(bool)
|
||||
intersect = (gt_mask*mask).astype(bool)
|
||||
union = (gt_mask+mask).astype(bool)
|
||||
iou = float(intersect.sum())/union.sum()
|
||||
if iou>best_iou:
|
||||
best_iou = iou
|
||||
best_gt_mask = gt_mask
|
||||
ob_in_cam = cur
|
||||
else:
|
||||
ob_in_cam = cur
|
||||
break
|
||||
|
||||
|
||||
if use_my_correction:
|
||||
if 'ycb' in self.base_dir.lower() and 'train_real' in self.color_files[i_frame]:
|
||||
video_id = self.get_video_id()
|
||||
if ob_id==1:
|
||||
if video_id in [12,13,14,17,24]:
|
||||
ob_in_cam = ob_in_cam@self.symmetry_tfs[ob_id][1]
|
||||
return ob_in_cam
|
||||
|
||||
|
||||
def load_symmetry_tfs(self):
|
||||
dir = os.path.dirname(self.get_gt_mesh_file(self.ob_ids[0]))
|
||||
info_file = f'{dir}/models_info.json'
|
||||
with open(info_file,'r') as ff:
|
||||
info = json.load(ff)
|
||||
self.symmetry_tfs = {}
|
||||
self.symmetry_info_table = {}
|
||||
for ob_id in self.ob_ids:
|
||||
self.symmetry_info_table[ob_id] = info[str(ob_id)]
|
||||
self.symmetry_tfs[ob_id] = symmetry_tfs_from_info(info[str(ob_id)], rot_angle_discrete=5)
|
||||
self.geometry_symmetry_info_table = copy.deepcopy(self.symmetry_info_table)
|
||||
|
||||
|
||||
def get_video_id(self):
|
||||
return int(self.base_dir.split('/')[-1])
|
||||
|
||||
|
||||
class LinemodOcclusionReader(BopBaseReader):
|
||||
def __init__(self,base_dir='/mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/LINEMOD-O/lmo_test_all/test/000002', zfar=np.inf):
|
||||
super().__init__(base_dir, zfar=zfar)
|
||||
self.dataset_name = 'lmo'
|
||||
self.K = list(self.K_table.values())[0]
|
||||
self.ob_ids = [1,5,6,8,9,10,11,12]
|
||||
self.ob_id_to_names = {
|
||||
1: 'ape',
|
||||
2: 'benchvise',
|
||||
3: 'bowl',
|
||||
4: 'camera',
|
||||
5: 'water_pour',
|
||||
6: 'cat',
|
||||
7: 'cup',
|
||||
8: 'driller',
|
||||
9: 'duck',
|
||||
10: 'eggbox',
|
||||
11: 'glue',
|
||||
12: 'holepuncher',
|
||||
13: 'iron',
|
||||
14: 'lamp',
|
||||
15: 'phone',
|
||||
}
|
||||
self.load_symmetry_tfs()
|
||||
|
||||
def get_gt_mesh_file(self, ob_id):
|
||||
mesh_dir = f'{BOP_DIR}/{self.dataset_name}/models/obj_{ob_id:06d}.ply'
|
||||
return mesh_dir
|
||||
|
||||
|
||||
|
||||
class LinemodReader(LinemodOcclusionReader):
|
||||
def __init__(self, base_dir='/mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/LINEMOD/lm_test_all/test/000001', zfar=np.inf, split=None):
|
||||
super().__init__(base_dir, zfar=zfar)
|
||||
self.dataset_name = 'lm'
|
||||
if split is not None: # train/test
|
||||
with open(f'/mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/LINEMOD/Linemod_preprocessed/data/{self.get_video_id():02d}/{split}.txt','r') as ff:
|
||||
lines = ff.read().splitlines()
|
||||
self.color_files = []
|
||||
for line in lines:
|
||||
id = int(line)
|
||||
self.color_files.append(f'{self.base_dir}/rgb/{id:06d}.png')
|
||||
self.make_id_strs()
|
||||
|
||||
self.ob_ids = np.setdiff1d(np.arange(1,16), np.array([7,3])).tolist() # Exclude bowl and mug
|
||||
self.load_symmetry_tfs()
|
||||
|
||||
|
||||
def get_gt_mesh_file(self, ob_id):
|
||||
root = self.base_dir
|
||||
while 1:
|
||||
if os.path.exists(f'{root}/lm_models'):
|
||||
mesh_dir = f'{root}/lm_models/models/obj_{ob_id:06d}.ply'
|
||||
break
|
||||
else:
|
||||
root = os.path.abspath(f'{root}/../')
|
||||
return mesh_dir
|
||||
|
||||
|
||||
def get_reconstructed_mesh(self, ob_id, ref_view_dir):
|
||||
mesh = trimesh.load(os.path.abspath(f'{ref_view_dir}/ob_{ob_id:07d}/model/model.obj'))
|
||||
return mesh
|
||||
|
||||
|
||||
class YcbVideoReader(BopBaseReader):
|
||||
def __init__(self, base_dir, zfar=np.inf):
|
||||
super().__init__(base_dir, zfar=zfar)
|
||||
self.dataset_name = 'ycbv'
|
||||
self.K = list(self.K_table.values())[0]
|
||||
|
||||
self.make_id_strs()
|
||||
|
||||
self.ob_ids = np.arange(1,22).astype(int).tolist()
|
||||
YCB_VIDEO_DIR = os.getenv('YCB_VIDEO_DIR')
|
||||
names = sorted(os.listdir(f'{YCB_VIDEO_DIR}/models/'))
|
||||
self.ob_id_to_names = {}
|
||||
self.name_to_ob_id = {}
|
||||
for i,ob_id in enumerate(self.ob_ids):
|
||||
self.ob_id_to_names[ob_id] = names[i]
|
||||
self.name_to_ob_id[names[i]] = ob_id
|
||||
|
||||
if 'BOP' not in self.base_dir:
|
||||
with open(f'{self.base_dir}/../../keyframe.txt','r') as ff:
|
||||
self.keyframe_lines = ff.read().splitlines()
|
||||
|
||||
self.load_symmetry_tfs()
|
||||
for ob_id in self.ob_ids:
|
||||
if ob_id in [1,4,6,18]: # Cylinder
|
||||
self.geometry_symmetry_info_table[ob_id] = {
|
||||
'symmetries_continuous': [
|
||||
{'axis':[0,0,1], 'offset':[0,0,0]},
|
||||
],
|
||||
'symmetries_discrete': euler_matrix(0, np.pi, 0).reshape(1,4,4).tolist(),
|
||||
}
|
||||
elif ob_id in [13]:
|
||||
self.geometry_symmetry_info_table[ob_id] = {
|
||||
'symmetries_continuous': [
|
||||
{'axis':[0,0,1], 'offset':[0,0,0]},
|
||||
],
|
||||
}
|
||||
elif ob_id in [2,3,9,21]: # Rectangle box
|
||||
tfs = []
|
||||
for rz in [0, np.pi]:
|
||||
for rx in [0,np.pi]:
|
||||
for ry in [0,np.pi]:
|
||||
tfs.append(euler_matrix(rx, ry, rz))
|
||||
self.geometry_symmetry_info_table[ob_id] = {
|
||||
'symmetries_discrete': np.asarray(tfs).reshape(-1,4,4).tolist(),
|
||||
}
|
||||
else:
|
||||
pass
|
||||
|
||||
def get_gt_mesh_file(self, ob_id):
|
||||
if 'BOP' in self.base_dir:
|
||||
mesh_file = os.path.abspath(f'{self.base_dir}/../../ycbv_models/models/obj_{ob_id:06d}.ply')
|
||||
else:
|
||||
mesh_file = f'{self.base_dir}/../../ycbv_models/models/obj_{ob_id:06d}.ply'
|
||||
return mesh_file
|
||||
|
||||
|
||||
def get_gt_mesh(self, ob_id:int, get_posecnn_version=False):
|
||||
if get_posecnn_version:
|
||||
YCB_VIDEO_DIR = os.getenv('YCB_VIDEO_DIR')
|
||||
mesh = trimesh.load(f'{YCB_VIDEO_DIR}/models/{self.ob_id_to_names[ob_id]}/textured_simple.obj')
|
||||
return mesh
|
||||
mesh_file = self.get_gt_mesh_file(ob_id)
|
||||
mesh = trimesh.load(mesh_file, process=False)
|
||||
mesh.vertices *= 1e-3
|
||||
tex_file = mesh_file.replace('.ply','.png')
|
||||
if os.path.exists(tex_file):
|
||||
from PIL import Image
|
||||
im = Image.open(tex_file)
|
||||
uv = mesh.visual.uv
|
||||
material = trimesh.visual.texture.SimpleMaterial(image=im)
|
||||
color_visuals = trimesh.visual.TextureVisuals(uv=uv, image=im, material=material)
|
||||
mesh.visual = color_visuals
|
||||
return mesh
|
||||
|
||||
|
||||
def get_reconstructed_mesh(self, ob_id, ref_view_dir):
|
||||
mesh = trimesh.load(os.path.abspath(f'{ref_view_dir}/ob_{ob_id:07d}/model/model.obj'))
|
||||
return mesh
|
||||
|
||||
|
||||
def get_transform_reconstructed_to_gt_model(self, ob_id):
|
||||
out = np.eye(4)
|
||||
return out
|
||||
|
||||
|
||||
def get_visible_cloud(self, ob_id):
|
||||
file = os.path.abspath(f'{self.base_dir}/../../models/{self.ob_id_to_names[ob_id]}/visible_cloud.ply')
|
||||
pcd = o3d.io.read_point_cloud(file)
|
||||
return pcd
|
||||
|
||||
|
||||
def is_keyframe(self, i):
|
||||
color_file = self.color_files[i]
|
||||
video_id = self.get_video_id()
|
||||
frame_id = int(os.path.basename(color_file).split('.')[0])
|
||||
key = f'{video_id:04d}/{frame_id:06d}'
|
||||
return (key in self.keyframe_lines)
|
||||
|
||||
|
||||
|
||||
class TlessReader(BopBaseReader):
|
||||
def __init__(self, base_dir, zfar=np.inf):
|
||||
super().__init__(base_dir, zfar=zfar)
|
||||
self.dataset_name = 'tless'
|
||||
|
||||
self.ob_ids = np.arange(1,31).astype(int).tolist()
|
||||
self.load_symmetry_tfs()
|
||||
|
||||
|
||||
def get_gt_mesh_file(self, ob_id):
|
||||
mesh_file = f'{self.base_dir}/../../../models_cad/obj_{ob_id:06d}.ply'
|
||||
return mesh_file
|
||||
|
||||
|
||||
def get_gt_mesh(self, ob_id):
|
||||
mesh = trimesh.load(self.get_gt_mesh_file(ob_id))
|
||||
mesh.vertices *= 1e-3
|
||||
mesh = trimesh_add_pure_colored_texture(mesh, color=np.ones((3))*200)
|
||||
return mesh
|
||||
|
||||
|
||||
class HomebrewedReader(BopBaseReader):
|
||||
def __init__(self, base_dir, zfar=np.inf):
|
||||
super().__init__(base_dir, zfar=zfar)
|
||||
self.dataset_name = 'hb'
|
||||
self.ob_ids = np.arange(1,34).astype(int).tolist()
|
||||
self.load_symmetry_tfs()
|
||||
self.make_scene_ob_ids_dict()
|
||||
|
||||
|
||||
def get_gt_mesh_file(self, ob_id):
|
||||
mesh_file = f'{self.base_dir}/../../../hb_models/models/obj_{ob_id:06d}.ply'
|
||||
return mesh_file
|
||||
|
||||
|
||||
def get_gt_pose(self, i_frame:int, ob_id, use_my_correction=False):
|
||||
logging.info("WARN HomeBrewed doesn't have GT pose")
|
||||
return np.eye(4)
|
||||
|
||||
|
||||
|
||||
class ItoddReader(BopBaseReader):
|
||||
def __init__(self, base_dir, zfar=np.inf):
|
||||
super().__init__(base_dir, zfar=zfar)
|
||||
self.dataset_name = 'itodd'
|
||||
self.make_id_strs()
|
||||
|
||||
self.ob_ids = np.arange(1,29).astype(int).tolist()
|
||||
self.load_symmetry_tfs()
|
||||
self.make_scene_ob_ids_dict()
|
||||
|
||||
|
||||
def get_gt_mesh_file(self, ob_id):
|
||||
mesh_file = f'{self.base_dir}/../../../itodd_models/models/obj_{ob_id:06d}.ply'
|
||||
return mesh_file
|
||||
|
||||
|
||||
class IcbinReader(BopBaseReader):
|
||||
def __init__(self, base_dir, zfar=np.inf):
|
||||
super().__init__(base_dir, zfar=zfar)
|
||||
self.dataset_name = 'icbin'
|
||||
self.ob_ids = np.arange(1,3).astype(int).tolist()
|
||||
self.load_symmetry_tfs()
|
||||
|
||||
def get_gt_mesh_file(self, ob_id):
|
||||
mesh_file = f'{self.base_dir}/../../../icbin_models/models/obj_{ob_id:06d}.ply'
|
||||
return mesh_file
|
||||
|
||||
|
||||
class TudlReader(BopBaseReader):
|
||||
def __init__(self, base_dir, zfar=np.inf):
|
||||
super().__init__(base_dir, zfar=zfar)
|
||||
self.dataset_name = 'tudl'
|
||||
self.ob_ids = np.arange(1,4).astype(int).tolist()
|
||||
self.load_symmetry_tfs()
|
||||
|
||||
def get_gt_mesh_file(self, ob_id):
|
||||
mesh_file = f'{self.base_dir}/../../../tudl_models/models/obj_{ob_id:06d}.ply'
|
||||
return mesh_file
|
||||
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
docker rm -f foundationpose
|
||||
DIR=$(pwd)/../
|
||||
xhost + && docker run --gpus all --env NVIDIA_DISABLE_REQUIRE=1 -it --network=host --name foundationpose --cap-add=SYS_PTRACE --security-opt seccomp=unconfined -v $DIR:$DIR -v /home:/home -v /mnt:/mnt -v /tmp/.X11-unix:/tmp/.X11-unix -v /tmp:/tmp --ipc=host -e DISPLAY=${DISPLAY} -e GIT_INDEX_FILE foundationpose:latest bash -c "cd $DIR && bash"
|
||||
+17
@@ -0,0 +1,17 @@
|
||||
# Conda-only dependencies (compilers, CMake, Eigen, Boost, pybind11).
|
||||
# After creating the env, install PyTorch and GPU extensions as described in readme.md
|
||||
# (PyTorch index URL may need to match your CUDA driver), then:
|
||||
# pip install -r requirements.txt
|
||||
# bash build_all_conda.sh
|
||||
name: foundationpose
|
||||
channels:
|
||||
- conda-forge
|
||||
dependencies:
|
||||
- python=3.11
|
||||
- pip
|
||||
- cmake
|
||||
- ninja
|
||||
- eigen
|
||||
- boost-cpp
|
||||
- pybind11
|
||||
- cxx-compiler
|
||||
+270
@@ -0,0 +1,270 @@
|
||||
# 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.
|
||||
|
||||
|
||||
from Utils import *
|
||||
from datareader import *
|
||||
import itertools
|
||||
from learning.training.predict_score import *
|
||||
from learning.training.predict_pose_refine import *
|
||||
import yaml
|
||||
|
||||
|
||||
class FoundationPose:
|
||||
def __init__(self, model_pts, model_normals, symmetry_tfs=None, mesh=None, scorer:ScorePredictor=None, refiner:PoseRefinePredictor=None, glctx=None, debug=0, debug_dir='/home/bowen/debug/novel_pose_debug/'):
|
||||
self.gt_pose = None
|
||||
self.ignore_normal_flip = True
|
||||
self.debug = debug
|
||||
self.debug_dir = debug_dir
|
||||
os.makedirs(debug_dir, exist_ok=True)
|
||||
|
||||
self.reset_object(model_pts, model_normals, symmetry_tfs=symmetry_tfs, mesh=mesh)
|
||||
self.make_rotation_grid(min_n_views=40, inplane_step=60)
|
||||
|
||||
self.glctx = glctx
|
||||
|
||||
if scorer is not None:
|
||||
self.scorer = scorer
|
||||
else:
|
||||
self.scorer = ScorePredictor()
|
||||
|
||||
if refiner is not None:
|
||||
self.refiner = refiner
|
||||
else:
|
||||
self.refiner = PoseRefinePredictor()
|
||||
|
||||
self.pose_last = None # Used for tracking; per the centered mesh
|
||||
|
||||
|
||||
def reset_object(self, model_pts, model_normals, symmetry_tfs=None, mesh=None):
|
||||
max_xyz = mesh.vertices.max(axis=0)
|
||||
min_xyz = mesh.vertices.min(axis=0)
|
||||
self.model_center = (min_xyz+max_xyz)/2
|
||||
if mesh is not None:
|
||||
self.mesh_ori = mesh.copy()
|
||||
mesh = mesh.copy()
|
||||
mesh.vertices = mesh.vertices - self.model_center.reshape(1,3)
|
||||
|
||||
model_pts = mesh.vertices
|
||||
self.diameter = compute_mesh_diameter(model_pts=mesh.vertices, n_sample=10000)
|
||||
self.vox_size = max(self.diameter/20.0, 0.003)
|
||||
logging.info(f'self.diameter:{self.diameter}, vox_size:{self.vox_size}')
|
||||
self.dist_bin = self.vox_size/2
|
||||
self.angle_bin = 20 # Deg
|
||||
pcd = toOpen3dCloud(model_pts, normals=model_normals)
|
||||
pcd = pcd.voxel_down_sample(self.vox_size)
|
||||
self.max_xyz = np.asarray(pcd.points).max(axis=0)
|
||||
self.min_xyz = np.asarray(pcd.points).min(axis=0)
|
||||
self.pts = torch.tensor(np.asarray(pcd.points), dtype=torch.float32, device='cuda')
|
||||
self.normals = F.normalize(torch.tensor(np.asarray(pcd.normals), dtype=torch.float32, device='cuda'), dim=-1)
|
||||
logging.info(f'self.pts:{self.pts.shape}')
|
||||
self.mesh_path = None
|
||||
self.mesh = mesh
|
||||
if self.mesh is not None:
|
||||
self.mesh_path = f'/tmp/{uuid.uuid4()}.obj'
|
||||
self.mesh.export(self.mesh_path)
|
||||
self.mesh_tensors = make_mesh_tensors(self.mesh)
|
||||
|
||||
if symmetry_tfs is None:
|
||||
self.symmetry_tfs = torch.eye(4).float().cuda()[None]
|
||||
else:
|
||||
self.symmetry_tfs = torch.as_tensor(symmetry_tfs, device='cuda', dtype=torch.float)
|
||||
|
||||
logging.info("reset done")
|
||||
|
||||
|
||||
|
||||
def get_tf_to_centered_mesh(self):
|
||||
tf_to_center = torch.eye(4, dtype=torch.float, device='cuda')
|
||||
tf_to_center[:3,3] = -torch.as_tensor(self.model_center, device='cuda', dtype=torch.float)
|
||||
return tf_to_center
|
||||
|
||||
|
||||
def to_device(self, s='cuda:0'):
|
||||
for k in self.__dict__:
|
||||
self.__dict__[k] = self.__dict__[k]
|
||||
if torch.is_tensor(self.__dict__[k]) or isinstance(self.__dict__[k], nn.Module):
|
||||
logging.info(f"Moving {k} to device {s}")
|
||||
self.__dict__[k] = self.__dict__[k].to(s)
|
||||
for k in self.mesh_tensors:
|
||||
logging.info(f"Moving {k} to device {s}")
|
||||
self.mesh_tensors[k] = self.mesh_tensors[k].to(s)
|
||||
if self.refiner is not None:
|
||||
self.refiner.model.to(s)
|
||||
if self.scorer is not None:
|
||||
self.scorer.model.to(s)
|
||||
if self.glctx is not None:
|
||||
self.glctx = dr.RasterizeCudaContext(s)
|
||||
|
||||
|
||||
|
||||
def make_rotation_grid(self, min_n_views=40, inplane_step=60):
|
||||
cam_in_obs = sample_views_icosphere(n_views=min_n_views)
|
||||
logging.info(f'cam_in_obs:{cam_in_obs.shape}')
|
||||
rot_grid = []
|
||||
for i in range(len(cam_in_obs)):
|
||||
for inplane_rot in np.deg2rad(np.arange(0, 360, inplane_step)):
|
||||
cam_in_ob = cam_in_obs[i]
|
||||
R_inplane = euler_matrix(0,0,inplane_rot)
|
||||
cam_in_ob = cam_in_ob@R_inplane
|
||||
ob_in_cam = np.linalg.inv(cam_in_ob)
|
||||
rot_grid.append(ob_in_cam)
|
||||
|
||||
rot_grid = np.asarray(rot_grid)
|
||||
logging.info(f"rot_grid:{rot_grid.shape}")
|
||||
rot_grid = mycpp.cluster_poses(30, 99999, rot_grid, self.symmetry_tfs.data.cpu().numpy())
|
||||
rot_grid = np.asarray(rot_grid)
|
||||
logging.info(f"after cluster, rot_grid:{rot_grid.shape}")
|
||||
self.rot_grid = torch.as_tensor(rot_grid, device='cuda', dtype=torch.float)
|
||||
logging.info(f"self.rot_grid: {self.rot_grid.shape}")
|
||||
|
||||
|
||||
def generate_random_pose_hypo(self, K, rgb, depth, mask, scene_pts=None):
|
||||
'''
|
||||
@scene_pts: torch tensor (N,3)
|
||||
'''
|
||||
ob_in_cams = self.rot_grid.clone()
|
||||
center = self.guess_translation(depth=depth, mask=mask, K=K)
|
||||
ob_in_cams[:,:3,3] = torch.tensor(center, device='cuda', dtype=torch.float).reshape(1,3)
|
||||
return ob_in_cams
|
||||
|
||||
|
||||
def guess_translation(self, depth, mask, K):
|
||||
vs,us = np.where(mask>0)
|
||||
if len(us)==0:
|
||||
logging.info(f'mask is all zero')
|
||||
return np.zeros((3))
|
||||
uc = (us.min()+us.max())/2.0
|
||||
vc = (vs.min()+vs.max())/2.0
|
||||
valid = mask.astype(bool) & (depth>=0.001)
|
||||
if not valid.any():
|
||||
logging.info(f"valid is empty")
|
||||
return np.zeros((3))
|
||||
|
||||
zc = np.median(depth[valid])
|
||||
center = (np.linalg.inv(K)@np.asarray([uc,vc,1]).reshape(3,1))*zc
|
||||
|
||||
if self.debug>=2:
|
||||
pcd = toOpen3dCloud(center.reshape(1,3))
|
||||
o3d.io.write_point_cloud(f'{self.debug_dir}/init_center.ply', pcd)
|
||||
|
||||
return center.reshape(3)
|
||||
|
||||
|
||||
def register(self, K, rgb, depth, ob_mask, ob_id=None, glctx=None, iteration=5):
|
||||
'''Copmute pose from given pts to self.pcd
|
||||
@pts: (N,3) np array, downsampled scene points
|
||||
'''
|
||||
set_seed(0)
|
||||
logging.info('Welcome')
|
||||
|
||||
if self.glctx is None:
|
||||
if glctx is None:
|
||||
self.glctx = dr.RasterizeCudaContext()
|
||||
# self.glctx = dr.RasterizeGLContext()
|
||||
else:
|
||||
self.glctx = glctx
|
||||
|
||||
depth = erode_depth(depth, radius=2, device='cuda')
|
||||
depth = bilateral_filter_depth(depth, radius=2, device='cuda')
|
||||
|
||||
if self.debug>=2:
|
||||
xyz_map = depth2xyzmap(depth, K)
|
||||
valid = xyz_map[...,2]>=0.001
|
||||
pcd = toOpen3dCloud(xyz_map[valid], rgb[valid])
|
||||
o3d.io.write_point_cloud(f'{self.debug_dir}/scene_raw.ply',pcd)
|
||||
cv2.imwrite(f'{self.debug_dir}/ob_mask.png', (ob_mask*255.0).clip(0,255))
|
||||
|
||||
normal_map = None
|
||||
valid = (depth>=0.001) & (ob_mask>0)
|
||||
if valid.sum()<4:
|
||||
logging.info(f'valid too small, return')
|
||||
pose = np.eye(4)
|
||||
pose[:3,3] = self.guess_translation(depth=depth, mask=ob_mask, K=K)
|
||||
return pose
|
||||
|
||||
if self.debug>=2:
|
||||
imageio.imwrite(f'{self.debug_dir}/color.png', rgb)
|
||||
cv2.imwrite(f'{self.debug_dir}/depth.png', (depth*1000).astype(np.uint16))
|
||||
valid = xyz_map[...,2]>=0.001
|
||||
pcd = toOpen3dCloud(xyz_map[valid], rgb[valid])
|
||||
o3d.io.write_point_cloud(f'{self.debug_dir}/scene_complete.ply',pcd)
|
||||
|
||||
self.H, self.W = depth.shape[:2]
|
||||
self.K = K
|
||||
self.ob_id = ob_id
|
||||
self.ob_mask = ob_mask
|
||||
|
||||
poses = self.generate_random_pose_hypo(K=K, rgb=rgb, depth=depth, mask=ob_mask, scene_pts=None)
|
||||
poses = poses.data.cpu().numpy()
|
||||
logging.info(f'poses:{poses.shape}')
|
||||
center = self.guess_translation(depth=depth, mask=ob_mask, K=K)
|
||||
|
||||
poses = torch.as_tensor(poses, device='cuda', dtype=torch.float)
|
||||
poses[:,:3,3] = torch.as_tensor(center.reshape(1,3), device='cuda')
|
||||
|
||||
add_errs = self.compute_add_err_to_gt_pose(poses)
|
||||
logging.info(f"after viewpoint, add_errs min:{add_errs.min()}")
|
||||
|
||||
xyz_map = depth2xyzmap(depth, K)
|
||||
poses, vis = self.refiner.predict(mesh=self.mesh, mesh_tensors=self.mesh_tensors, rgb=rgb, depth=depth, K=K, ob_in_cams=poses.data.cpu().numpy(), normal_map=normal_map, xyz_map=xyz_map, glctx=self.glctx, mesh_diameter=self.diameter, iteration=iteration, get_vis=self.debug>=2)
|
||||
if vis is not None:
|
||||
imageio.imwrite(f'{self.debug_dir}/vis_refiner.png', vis)
|
||||
|
||||
scores, vis = self.scorer.predict(mesh=self.mesh, rgb=rgb, depth=depth, K=K, ob_in_cams=poses.data.cpu().numpy(), normal_map=normal_map, mesh_tensors=self.mesh_tensors, glctx=self.glctx, mesh_diameter=self.diameter, get_vis=self.debug>=2)
|
||||
if vis is not None:
|
||||
imageio.imwrite(f'{self.debug_dir}/vis_score.png', vis)
|
||||
|
||||
add_errs = self.compute_add_err_to_gt_pose(poses)
|
||||
logging.info(f"final, add_errs min:{add_errs.min()}")
|
||||
|
||||
ids = torch.as_tensor(scores).argsort(descending=True)
|
||||
logging.info(f'sort ids:{ids}')
|
||||
scores = scores[ids]
|
||||
poses = poses[ids]
|
||||
|
||||
logging.info(f'sorted scores:{scores}')
|
||||
|
||||
best_pose = poses[0]@self.get_tf_to_centered_mesh()
|
||||
self.pose_last = poses[0]
|
||||
self.best_id = ids[0]
|
||||
|
||||
self.poses = poses
|
||||
self.scores = scores
|
||||
|
||||
return best_pose.data.cpu().numpy()
|
||||
|
||||
|
||||
def compute_add_err_to_gt_pose(self, poses):
|
||||
'''
|
||||
@poses: wrt. the centered mesh
|
||||
'''
|
||||
return -torch.ones(len(poses), device='cuda', dtype=torch.float)
|
||||
|
||||
|
||||
def track_one(self, rgb, depth, K, iteration, extra={}):
|
||||
if self.pose_last is None:
|
||||
logging.info("Please init pose by register first")
|
||||
raise RuntimeError
|
||||
logging.info("Welcome")
|
||||
|
||||
depth = torch.as_tensor(depth, device='cuda', dtype=torch.float)
|
||||
depth = erode_depth(depth, radius=2, device='cuda')
|
||||
depth = bilateral_filter_depth(depth, radius=2, device='cuda')
|
||||
logging.info("depth processing done")
|
||||
|
||||
xyz_map = depth2xyzmap_batch(depth[None], torch.as_tensor(K, dtype=torch.float, device='cuda')[None], zfar=np.inf)[0]
|
||||
|
||||
pose, vis = self.refiner.predict(mesh=self.mesh, mesh_tensors=self.mesh_tensors, rgb=rgb, depth=depth, K=K, ob_in_cams=self.pose_last.reshape(1,4,4).data.cpu().numpy(), normal_map=None, xyz_map=xyz_map, mesh_diameter=self.diameter, glctx=self.glctx, iteration=iteration, get_vis=self.debug>=2)
|
||||
logging.info("pose done")
|
||||
if self.debug>=2:
|
||||
extra['vis'] = vis
|
||||
self.pose_last = pose
|
||||
return (pose@self.get_tf_to_centered_mesh()).data.cpu().numpy().reshape(4,4)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
+25
@@ -0,0 +1,25 @@
|
||||
cmake_minimum_required(VERSION 3.15)
|
||||
project(mycpp)
|
||||
|
||||
|
||||
set(CMAKE_BUILD_TYPE Release)
|
||||
set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++14 -fopenmp -g3 -O3")
|
||||
|
||||
|
||||
find_package(Boost REQUIRED COMPONENTS system program_options)
|
||||
find_package(OpenMP REQUIRED)
|
||||
find_package(Eigen3 REQUIRED)
|
||||
find_package(pybind11 REQUIRED)
|
||||
|
||||
include_directories(
|
||||
include
|
||||
${BLAS_INCLUDE_DIR}
|
||||
)
|
||||
|
||||
file(GLOB MY_SRC ${PROJECT_SOURCE_DIR}/src/*.cpp)
|
||||
|
||||
set(PYBIND11_CPP_STANDARD -std=c++14)
|
||||
|
||||
pybind11_add_module(mycpp src/app/pybind_api.cpp ${MY_SRC})
|
||||
target_link_libraries(mycpp PRIVATE ${Boost_LIBRARIES} ${OpenMP_CXX_FLAGS} Eigen3::Eigen)
|
||||
@@ -0,0 +1,46 @@
|
||||
/*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.
|
||||
*/
|
||||
|
||||
|
||||
#pragma once
|
||||
|
||||
// STL
|
||||
#include <iostream>
|
||||
#include <string>
|
||||
#include <map>
|
||||
#include <vector>
|
||||
#include <cmath>
|
||||
#include <time.h>
|
||||
#include <queue>
|
||||
#include <climits>
|
||||
#include <tuple>
|
||||
#include <unordered_map>
|
||||
#include <boost/assign.hpp>
|
||||
#include <thread>
|
||||
#include <Eigen/Dense>
|
||||
#include <Eigen/Geometry>
|
||||
#include <string>
|
||||
#include <limits.h>
|
||||
#include <unistd.h>
|
||||
#include <memory>
|
||||
#include <math.h>
|
||||
#include <boost/format.hpp>
|
||||
#include <numeric>
|
||||
#include <thread>
|
||||
#include <omp.h>
|
||||
|
||||
|
||||
using vectorMatrix4f = std::vector<Eigen::Matrix4f,Eigen::aligned_allocator<Eigen::Matrix4f>>;
|
||||
|
||||
|
||||
namespace Utils
|
||||
{
|
||||
float rotationGeodesicDistance(const Eigen::Matrix3f &R1, const Eigen::Matrix3f &R2);
|
||||
|
||||
} // namespace Utils
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
/*
|
||||
# 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.
|
||||
*/
|
||||
|
||||
|
||||
#include "Utils.h"
|
||||
|
||||
|
||||
|
||||
namespace Utils
|
||||
{
|
||||
|
||||
|
||||
// Difference angle in radian
|
||||
float rotationGeodesicDistance(const Eigen::Matrix3f &R1, const Eigen::Matrix3f &R2)
|
||||
{
|
||||
float cos = ((R1 * R2.transpose()).trace()-1) / 2.0;
|
||||
cos = std::max(std::min(cos, 1.0f), -1.0f);
|
||||
return std::acos(cos);
|
||||
}
|
||||
|
||||
|
||||
} // namespace Utils
|
||||
@@ -0,0 +1,77 @@
|
||||
/*
|
||||
# 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.
|
||||
*/
|
||||
|
||||
|
||||
#include "Utils.h"
|
||||
#include <boost/algorithm/string.hpp>
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <pybind11/stl.h>
|
||||
#include <pybind11/eigen.h>
|
||||
|
||||
namespace py = pybind11;
|
||||
|
||||
|
||||
|
||||
//@angle_diff: unit is degree
|
||||
//@dist_diff: unit is meter
|
||||
vectorMatrix4f cluster_poses(float angle_diff, float dist_diff, const vectorMatrix4f &poses_in, const vectorMatrix4f &symmetry_tfs)
|
||||
{
|
||||
printf("num original candidates = %d\n",poses_in.size());
|
||||
vectorMatrix4f poses_out;
|
||||
poses_out.push_back(poses_in[0]);
|
||||
|
||||
const float radian_thres = angle_diff/180.0*M_PI;
|
||||
|
||||
for (int i=1;i<poses_in.size();i++)
|
||||
{
|
||||
bool isnew = true;
|
||||
Eigen::Matrix4f cur_pose = poses_in[i];
|
||||
for (const auto &cluster:poses_out)
|
||||
{
|
||||
Eigen::Vector3f t0 = cluster.block(0,3,3,1);
|
||||
Eigen::Vector3f t1 = cur_pose.block(0,3,3,1);
|
||||
|
||||
if ((t0-t1).norm()>=dist_diff)
|
||||
{
|
||||
continue;
|
||||
}
|
||||
|
||||
for (const auto &tf: symmetry_tfs)
|
||||
{
|
||||
Eigen::Matrix4f cur_pose_tmp = cur_pose*tf;
|
||||
float rot_diff = Utils::rotationGeodesicDistance(cur_pose_tmp.block(0,0,3,3), cluster.block(0,0,3,3));
|
||||
if (rot_diff < radian_thres)
|
||||
{
|
||||
isnew = false;
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
if (!isnew) break;
|
||||
}
|
||||
|
||||
if (isnew)
|
||||
{
|
||||
poses_out.push_back(poses_in[i]);
|
||||
}
|
||||
}
|
||||
|
||||
printf("num of pose after clustering: %d\n",poses_out.size());
|
||||
return poses_out;
|
||||
}
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
PYBIND11_MODULE(mycpp, m)
|
||||
{
|
||||
m.def("cluster_poses", &cluster_poses, py::call_guard<py::gil_scoped_release>());
|
||||
}
|
||||
@@ -0,0 +1,79 @@
|
||||
# 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,time,socket
|
||||
code_path = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(code_path)
|
||||
import open3d as o3d
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import cv2,imageio
|
||||
import time
|
||||
import trimesh
|
||||
import pyrender
|
||||
from transformations import *
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import cv2
|
||||
import time
|
||||
import argparse,pickle
|
||||
from Utils import *
|
||||
|
||||
|
||||
cvcam_in_glcam = np.array([[1,0,0,0],
|
||||
[0,-1,0,0],
|
||||
[0,0,-1,0],
|
||||
[0,0,0,1]])
|
||||
|
||||
class ModelRendererOffscreen:
|
||||
def __init__(self, cam_K, H,W, zfar=100):
|
||||
'''
|
||||
@window_sizes: H,W
|
||||
'''
|
||||
self.K = cam_K
|
||||
self.scene = pyrender.Scene(ambient_light=[1., 1., 1.],bg_color=[0,0,0])
|
||||
self.camera = pyrender.IntrinsicsCamera(fx=cam_K[0,0],fy=cam_K[1,1],cx=cam_K[0,2],cy=cam_K[1,2],znear=0.001,zfar=zfar)
|
||||
self.cam_node = self.scene.add(self.camera, pose=np.eye(4), name='cam')
|
||||
self.mesh_nodes = []
|
||||
|
||||
self.H = H
|
||||
self.W = W
|
||||
self.r = pyrender.OffscreenRenderer(self.W, self.H)
|
||||
|
||||
|
||||
def set_cam_pose(self, cam_pose):
|
||||
self.cam_node.matrix = cam_pose
|
||||
|
||||
def add_mesh(self, mesh):
|
||||
mesh = pyrender.Mesh.from_trimesh(mesh, smooth=False)
|
||||
mesh_node = self.scene.add(mesh,pose=np.eye(4), name='ob') # Object pose parent is cam
|
||||
self.mesh_nodes.append(mesh_node)
|
||||
|
||||
|
||||
def add_point_light(self, intensity=3):
|
||||
light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=intensity)
|
||||
self.scene.add(light, pose=np.eye(4)) # Same as camera position
|
||||
|
||||
|
||||
def clear_mesh_nodes(self):
|
||||
for n in self.mesh_nodes:
|
||||
self.scene.remove_node(n)
|
||||
|
||||
|
||||
def render(self,mesh=None,ob_in_cvcam=None, get_normal=False):
|
||||
if mesh is not None:
|
||||
mesh = mesh.copy()
|
||||
mesh.apply_transform(cvcam_in_glcam@ob_in_cvcam)
|
||||
mesh = pyrender.Mesh.from_trimesh(mesh, smooth=False)
|
||||
mesh_node = self.scene.add(mesh, pose=np.eye(4), name='ob') # Object pose parent is cam
|
||||
color, depth = self.r.render(self.scene) # depth: float
|
||||
if mesh is not None:
|
||||
self.scene.remove_node(mesh_node)
|
||||
|
||||
return color, depth
|
||||
Vendored
+226
@@ -0,0 +1,226 @@
|
||||
# FoundationPose: Unified 6D Pose Estimation and Tracking of Novel Objects
|
||||
[[Paper]](https://arxiv.org/abs/2312.08344) [[Website]](https://nvlabs.github.io/FoundationPose/)
|
||||
|
||||
This is the official implementation of our paper to be appeared in CVPR 2024 (Highlight)
|
||||
|
||||
Contributors: Bowen Wen, Wei Yang, Jan Kautz, Stan Birchfield
|
||||
|
||||
We present FoundationPose, a unified foundation model for 6D object pose estimation and tracking, supporting both model-based and model-free setups. Our approach can be instantly applied at test-time to a novel object without fine-tuning, as long as its CAD model is given, or a small number of reference images are captured. We bridge the gap between these two setups with a neural implicit representation that allows for effective novel view synthesis, keeping the downstream pose estimation modules invariant under the same unified framework. Strong generalizability is achieved via large-scale synthetic training, aided by a large language model (LLM), a novel transformer-based architecture, and contrastive learning formulation. Extensive evaluation on multiple public datasets involving challenging scenarios and objects indicate our unified approach outperforms existing methods specialized for each task by a large margin. In addition, it even achieves comparable results to instance-level methods despite the reduced assumptions.
|
||||
|
||||
|
||||
<img src="assets/intro.jpg" width="70%">
|
||||
|
||||
**🤖 For ROS version, please check [Isaac ROS Pose Estimation](https://github.com/NVIDIA-ISAAC-ROS/isaac_ros_pose_estimation), which enjoys TRT fast inference and C++ speed up.**
|
||||
|
||||
\
|
||||
**🥇 No. 1 on the world-wide [BOP leaderboard](https://bop.felk.cvut.cz/leaderboards/pose-estimation-unseen-bop23/core-datasets/) (as of 2024/03) for model-based novel object pose estimation.**
|
||||
<img src="assets/bop.jpg" width="80%">
|
||||
|
||||
## Demos
|
||||
|
||||
Robotic Applications:
|
||||
|
||||
https://github.com/NVlabs/FoundationPose/assets/23078192/aa341004-5a15-4293-b3da-000471fd74ed
|
||||
|
||||
|
||||
AR Applications:
|
||||
|
||||
https://github.com/NVlabs/FoundationPose/assets/23078192/80e96855-a73c-4bee-bcef-7cba92df55ca
|
||||
|
||||
|
||||
Results on YCB-Video dataset:
|
||||
|
||||
https://github.com/NVlabs/FoundationPose/assets/23078192/9b5bedde-755b-44ed-a973-45ec85a10bbe
|
||||
|
||||
|
||||
|
||||
# Bibtex
|
||||
```bibtex
|
||||
@InProceedings{foundationposewen2024,
|
||||
author = {Bowen Wen, Wei Yang, Jan Kautz, Stan Birchfield},
|
||||
title = {{FoundationPose}: Unified 6D Pose Estimation and Tracking of Novel Objects},
|
||||
booktitle = {CVPR},
|
||||
year = {2024},
|
||||
}
|
||||
```
|
||||
|
||||
If you find the model-free setup useful, please also consider cite:
|
||||
|
||||
```bibtex
|
||||
@InProceedings{bundlesdfwen2023,
|
||||
author = {Bowen Wen and Jonathan Tremblay and Valts Blukis and Stephen Tyree and Thomas M\"{u}ller and Alex Evans and Dieter Fox and Jan Kautz and Stan Birchfield},
|
||||
title = {{BundleSDF}: {N}eural 6-{DoF} Tracking and {3D} Reconstruction of Unknown Objects},
|
||||
booktitle = {CVPR},
|
||||
year = {2023},
|
||||
}
|
||||
```
|
||||
|
||||
# Data prepare
|
||||
|
||||
|
||||
1) Download all network weights from [here](https://drive.google.com/drive/folders/1DFezOAD0oD1BblsXVxqDsl8fj0qzB82i?usp=sharing) and put them under the folder `weights/`. For the refiner, you will need `2023-10-28-18-33-37`. For scorer, you will need `2024-01-11-20-02-45`.
|
||||
|
||||
1) [Download demo data](https://drive.google.com/drive/folders/1pRyFmxYXmAnpku7nGRioZaKrVJtIsroP?usp=sharing) and extract them under the folder `demo_data/`
|
||||
|
||||
1) [Optional] Download our large-scale training data: ["FoundationPose Dataset"](https://drive.google.com/drive/folders/1s4pB6p4ApfWMiMjmTXOFco8dHbNXikp-?usp=sharing)
|
||||
|
||||
1) [Optional] Download our preprocessed reference views [here](https://drive.google.com/drive/folders/1PXXCOJqHXwQTbwPwPbGDN9_vLVe0XpFS?usp=sharing) in order to run model-free few-shot version.
|
||||
|
||||
# Env setup option 1: docker (recommended)
|
||||
```
|
||||
cd docker/
|
||||
docker pull wenbowen123/foundationpose && docker tag wenbowen123/foundationpose foundationpose # Or to build from scratch: docker build --network host -t foundationpose .
|
||||
bash docker/run_container.sh
|
||||
```
|
||||
|
||||
|
||||
If it's the first time you launch the container, you need to build extensions. Run this command *inside* the Docker container.
|
||||
```
|
||||
bash build_all.sh
|
||||
```
|
||||
|
||||
Later you can execute into the container without re-build.
|
||||
```
|
||||
docker exec -it foundationpose bash
|
||||
```
|
||||
|
||||
For more recent GPU such as 4090, refer to [this](https://github.com/NVlabs/FoundationPose/issues/27).
|
||||
In short, do the following:
|
||||
```
|
||||
docker pull shingarey/foundationpose_custom_cuda121:latest
|
||||
```
|
||||
Then modify the bash script to use this image instead of `foundationpose:latest`.
|
||||
|
||||
|
||||
# Env setup option 2: conda (local)
|
||||
|
||||
1) **Create the environment** (C++ build deps + Python; all on `conda-forge`):
|
||||
|
||||
```bash
|
||||
conda env create -f environment.yml
|
||||
conda activate foundationpose
|
||||
```
|
||||
|
||||
2) **Install PyTorch** with a CUDA build that matches your machine. The [PyTorch “Get Started”](https://pytorch.org/get-started/locally/) page lists the right `--index-url` (for example `cu124` works on most current NVIDIA drivers). Example:
|
||||
|
||||
```bash
|
||||
python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
|
||||
```
|
||||
|
||||
3) **Install PyTorch3D and NVDiffRast** (compile from source; needs the CUDA toolkit for `nvcc`. Point `CUDA_HOME` at your install, e.g. `/usr/local/cuda-12.8` or `/usr/local/cuda`):
|
||||
|
||||
```bash
|
||||
export CUDA_HOME=/usr/local/cuda # or e.g. /usr/local/cuda-12.8
|
||||
export PATH="$CUDA_HOME/bin:$PATH"
|
||||
python -m pip install --no-build-isolation "git+https://github.com/facebookresearch/pytorch3d.git"
|
||||
python -m pip install --no-build-isolation "git+https://github.com/NVlabs/nvdiffrast.git"
|
||||
```
|
||||
|
||||
`--no-build-isolation` is required so the build can see the `torch` you already installed.
|
||||
|
||||
4) **Install remaining Python dependencies** and **build the `mycpp` extension** (BundleSDF’s `mycuda` step is optional; skip it unless you need the model-free / NeRF path):
|
||||
|
||||
```bash
|
||||
python -m pip install -r requirements.txt
|
||||
bash build_all_conda.sh
|
||||
```
|
||||
|
||||
5) **Optional — Kaolin** (only for the model-free setup; version must match your PyTorch/CUDA—see the [Kaolin install docs](https://kaolin.readthedocs.io/en/latest/notes/installation.html)).
|
||||
|
||||
|
||||
# Run model-based demo
|
||||
The paths have been set in argparse by default. If you need to change the scene, you can pass the args accordingly. By running on the demo data, you should be able to see the robot manipulating the mustard bottle. Pose estimation is conducted on the first frame, then it automatically switches to tracking mode for the rest of the video. The resulting visualizations will be saved to the `debug_dir` specified in the argparse. (Note the first time running could be slower due to online compilation)
|
||||
```
|
||||
python run_demo.py
|
||||
```
|
||||
|
||||
|
||||
<img src="assets/demo.jpg" width="50%">
|
||||
|
||||
|
||||
Feel free to try on other objects (**no need to retrain**) such as driller, by changing the paths in argparse.
|
||||
|
||||
<img src="assets/demo_driller.jpg" width="50%">
|
||||
|
||||
|
||||
# Run on public datasets (LINEMOD, YCB-Video)
|
||||
|
||||
For this you first need to download LINEMOD dataset and YCB-Video dataset.
|
||||
|
||||
To run model-based version on these two datasets respectively, set the paths based on where you download. The results will be saved to `debug` folder
|
||||
```
|
||||
python run_linemod.py --linemod_dir /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/LINEMOD --use_reconstructed_mesh 0
|
||||
|
||||
python run_ycb_video.py --ycbv_dir /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video --use_reconstructed_mesh 0
|
||||
```
|
||||
|
||||
To run model-free few-shot version. You first need to train Neural Object Field. `ref_view_dir` is based on where you download in the above "Data prepare" section. Set the `dataset` flag to your interested dataset.
|
||||
```
|
||||
python bundlesdf/run_nerf.py --ref_view_dir /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video/bowen_addon/ref_views_16 --dataset ycbv
|
||||
```
|
||||
|
||||
Then run the similar command as the model-based version with some small modifications. Here we are using YCB-Video as example:
|
||||
```
|
||||
python run_ycb_video.py --ycbv_dir /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video --use_reconstructed_mesh 1 --ref_view_dir /mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video/bowen_addon/ref_views_16
|
||||
```
|
||||
|
||||
# Troubleshooting
|
||||
|
||||
|
||||
- For more recent GPU such as 4090, refer to [this](https://github.com/NVlabs/FoundationPose/issues/27).
|
||||
|
||||
- For setting up on Windows, refer to [this](https://github.com/NVlabs/FoundationPose/issues/148).
|
||||
|
||||
- If you are getting unreasonable results, check [this](https://github.com/NVlabs/FoundationPose/issues/44#issuecomment-2048141043) and [this](https://github.com/030422Lee/FoundationPose_manual)
|
||||
|
||||
# Training data download
|
||||
Our training data include scenes using 3D assets from GSO and Objaverse, rendered with high quality photo-realism and large domain randomization. Each data point includes **RGB, depth, object pose, camera pose, instance segmentation, 2D bounding box**. [[Google Drive]](https://drive.google.com/drive/folders/1s4pB6p4ApfWMiMjmTXOFco8dHbNXikp-?usp=sharing).
|
||||
|
||||
<img src="assets/train_data_vis.png" width="80%">
|
||||
|
||||
- To parse the camera params including extrinsics and intrinsics
|
||||
```
|
||||
glcam_in_cvcam = np.array([[1,0,0,0],
|
||||
[0,-1,0,0],
|
||||
[0,0,-1,0],
|
||||
[0,0,0,1]]).astype(float)
|
||||
W, H = camera_params["renderProductResolution"]
|
||||
with open(f'{base_dir}/camera_params/camera_params_000000.json','r') as ff:
|
||||
camera_params = json.load(ff)
|
||||
world_in_glcam = np.array(camera_params['cameraViewTransform']).reshape(4,4).T
|
||||
cam_in_world = np.linalg.inv(world_in_glcam)@glcam_in_cvcam
|
||||
world_in_cam = np.linalg.inv(cam_in_world)
|
||||
focal_length = camera_params["cameraFocalLength"]
|
||||
horiz_aperture = camera_params["cameraAperture"][0]
|
||||
vert_aperture = H / W * horiz_aperture
|
||||
focal_y = H * focal_length / vert_aperture
|
||||
focal_x = W * focal_length / horiz_aperture
|
||||
center_y = H * 0.5
|
||||
center_x = W * 0.5
|
||||
|
||||
fx, fy, cx, cy = focal_x, focal_y, center_x, center_y
|
||||
K = np.eye(3)
|
||||
K[0,0] = fx
|
||||
K[1,1] = fy
|
||||
K[0,2] = cx
|
||||
K[1,2] = cy
|
||||
```
|
||||
|
||||
|
||||
|
||||
# Notes
|
||||
Due to the legal restrictions of Stable-Diffusion that is trained on LAION dataset, we are not able to release the diffusion-based texture augmented data, nor the pretrained weights using it. We thus release the version without training on diffusion-augmented data. Slight performance degradation is expected.
|
||||
|
||||
# Acknowledgement
|
||||
|
||||
We would like to thank Jeff Smith for helping with the code release; NVIDIA Isaac Sim and Omniverse team for the support on synthetic data generation; Tianshi Cao for the valuable discussions. Finally, we are also grateful for the positive feebacks and constructive suggestions brought up by reviewers and AC at CVPR.
|
||||
|
||||
<img src="assets/cvpr_review.png" width="100%">
|
||||
|
||||
|
||||
# License
|
||||
The code and data are released under the NVIDIA Source Code License. Copyright © 2024, NVIDIA Corporation. All rights reserved.
|
||||
|
||||
|
||||
# Contact
|
||||
For questions, please contact [Bowen Wen](https://wenbowen123.github.io/).
|
||||
+43
@@ -0,0 +1,43 @@
|
||||
# Install PyTorch first (CUDA build must match your driver; cu124 works on most recent NVIDIA GPUs):
|
||||
# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
|
||||
#
|
||||
# Then GPU extensions (need nvcc; set CUDA_HOME to your toolkit, e.g. /usr/local/cuda):
|
||||
# pip install --no-build-isolation git+https://github.com/facebookresearch/pytorch3d.git
|
||||
# pip install --no-build-isolation git+https://github.com/NVlabs/nvdiffrast.git
|
||||
#
|
||||
# Then this file:
|
||||
# pip install -r requirements.txt
|
||||
|
||||
# Core numerics & IO
|
||||
numpy>=2
|
||||
scipy
|
||||
scikit-learn
|
||||
h5py
|
||||
joblib
|
||||
|
||||
# Config / serialization
|
||||
PyYAML>=6
|
||||
ruamel.yaml>=0.18
|
||||
|
||||
# Vision & 3D
|
||||
opencv-python
|
||||
imageio
|
||||
open3d
|
||||
trimesh
|
||||
transformations
|
||||
matplotlib
|
||||
pandas
|
||||
Pillow
|
||||
pyrender
|
||||
pyOpenGL>=3.1
|
||||
pyOpenGL_accelerate>=3.1
|
||||
|
||||
# Inference stack (torch installed separately)
|
||||
kornia>=0.7
|
||||
omegaconf>=2.3
|
||||
psutil
|
||||
tqdm
|
||||
warp-lang
|
||||
|
||||
# Optional: faster / utility (import guarded in code)
|
||||
scikit-image
|
||||
+79
@@ -0,0 +1,79 @@
|
||||
# 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.
|
||||
|
||||
|
||||
from estimater import *
|
||||
from datareader import *
|
||||
import argparse
|
||||
|
||||
|
||||
if __name__=='__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
parser.add_argument('--mesh_file', type=str, default=f'{code_dir}/demo_data/mustard0/mesh/textured_simple.obj')
|
||||
parser.add_argument('--test_scene_dir', type=str, default=f'{code_dir}/demo_data/mustard0')
|
||||
parser.add_argument('--est_refine_iter', type=int, default=5)
|
||||
parser.add_argument('--track_refine_iter', type=int, default=2)
|
||||
parser.add_argument('--debug', type=int, default=1)
|
||||
parser.add_argument('--debug_dir', type=str, default=f'{code_dir}/debug')
|
||||
args = parser.parse_args()
|
||||
|
||||
set_logging_format()
|
||||
set_seed(0)
|
||||
|
||||
mesh = trimesh.load(args.mesh_file)
|
||||
|
||||
debug = args.debug
|
||||
debug_dir = args.debug_dir
|
||||
os.system(f'rm -rf {debug_dir}/* && mkdir -p {debug_dir}/track_vis {debug_dir}/ob_in_cam')
|
||||
|
||||
to_origin, extents = trimesh.bounds.oriented_bounds(mesh)
|
||||
bbox = np.stack([-extents/2, extents/2], axis=0).reshape(2,3)
|
||||
|
||||
scorer = ScorePredictor()
|
||||
refiner = PoseRefinePredictor()
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
est = FoundationPose(model_pts=mesh.vertices, model_normals=mesh.vertex_normals, mesh=mesh, scorer=scorer, refiner=refiner, debug_dir=debug_dir, debug=debug, glctx=glctx)
|
||||
logging.info("estimator initialization done")
|
||||
|
||||
reader = YcbineoatReader(video_dir=args.test_scene_dir, shorter_side=None, zfar=np.inf)
|
||||
|
||||
for i in range(len(reader.color_files)):
|
||||
logging.info(f'i:{i}')
|
||||
color = reader.get_color(i)
|
||||
depth = reader.get_depth(i)
|
||||
if i==0:
|
||||
mask = reader.get_mask(0).astype(bool)
|
||||
pose = est.register(K=reader.K, rgb=color, depth=depth, ob_mask=mask, iteration=args.est_refine_iter)
|
||||
|
||||
if debug>=3:
|
||||
m = mesh.copy()
|
||||
m.apply_transform(pose)
|
||||
m.export(f'{debug_dir}/model_tf.obj')
|
||||
xyz_map = depth2xyzmap(depth, reader.K)
|
||||
valid = depth>=0.001
|
||||
pcd = toOpen3dCloud(xyz_map[valid], color[valid])
|
||||
o3d.io.write_point_cloud(f'{debug_dir}/scene_complete.ply', pcd)
|
||||
else:
|
||||
pose = est.track_one(rgb=color, depth=depth, K=reader.K, iteration=args.track_refine_iter)
|
||||
|
||||
os.makedirs(f'{debug_dir}/ob_in_cam', exist_ok=True)
|
||||
np.savetxt(f'{debug_dir}/ob_in_cam/{reader.id_strs[i]}.txt', pose.reshape(4,4))
|
||||
|
||||
if debug>=1:
|
||||
center_pose = pose@np.linalg.inv(to_origin)
|
||||
vis = draw_posed_3d_box(reader.K, img=color, ob_in_cam=center_pose, bbox=bbox)
|
||||
vis = draw_xyz_axis(color, ob_in_cam=center_pose, scale=0.1, K=reader.K, thickness=3, transparency=0, is_input_rgb=True)
|
||||
cv2.imshow('1', vis[...,::-1])
|
||||
cv2.waitKey(1)
|
||||
|
||||
|
||||
if debug>=2:
|
||||
os.makedirs(f'{debug_dir}/track_vis', exist_ok=True)
|
||||
imageio.imwrite(f'{debug_dir}/track_vis/{reader.id_strs[i]}.png', vis)
|
||||
|
||||
+149
@@ -0,0 +1,149 @@
|
||||
# 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.
|
||||
|
||||
|
||||
from Utils import *
|
||||
import json,uuid,joblib,os,sys
|
||||
import scipy.spatial as spatial
|
||||
from multiprocessing import Pool
|
||||
import multiprocessing
|
||||
from functools import partial
|
||||
from itertools import repeat
|
||||
import itertools
|
||||
from datareader import *
|
||||
from estimater import *
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(f'{code_dir}/mycpp/build')
|
||||
import yaml
|
||||
|
||||
|
||||
|
||||
def get_mask(reader, i_frame, ob_id, detect_type):
|
||||
if detect_type=='box':
|
||||
mask = reader.get_mask(i_frame, ob_id)
|
||||
H,W = mask.shape[:2]
|
||||
vs,us = np.where(mask>0)
|
||||
umin = us.min()
|
||||
umax = us.max()
|
||||
vmin = vs.min()
|
||||
vmax = vs.max()
|
||||
valid = np.zeros((H,W), dtype=bool)
|
||||
valid[vmin:vmax,umin:umax] = 1
|
||||
elif detect_type=='mask':
|
||||
mask = reader.get_mask(i_frame, ob_id)
|
||||
if mask is None:
|
||||
return None
|
||||
valid = mask>0
|
||||
elif detect_type=='detected':
|
||||
mask = cv2.imread(reader.color_files[i_frame].replace('rgb','mask_cosypose'), -1)
|
||||
valid = mask==ob_id
|
||||
else:
|
||||
raise RuntimeError
|
||||
return valid
|
||||
|
||||
|
||||
|
||||
def run_pose_estimation_worker(reader, i_frames, est:FoundationPose=None, debug=0, ob_id=None, device='cuda:0'):
|
||||
torch.cuda.set_device(device)
|
||||
est.to_device(device)
|
||||
est.glctx = dr.RasterizeCudaContext(device=device)
|
||||
|
||||
result = NestDict()
|
||||
|
||||
for i, i_frame in enumerate(i_frames):
|
||||
logging.info(f"{i}/{len(i_frames)}, i_frame:{i_frame}, ob_id:{ob_id}")
|
||||
video_id = reader.get_video_id()
|
||||
color = reader.get_color(i_frame)
|
||||
depth = reader.get_depth(i_frame)
|
||||
id_str = reader.id_strs[i_frame]
|
||||
H,W = color.shape[:2]
|
||||
|
||||
debug_dir =est.debug_dir
|
||||
|
||||
ob_mask = get_mask(reader, i_frame, ob_id, detect_type=detect_type)
|
||||
if ob_mask is None:
|
||||
logging.info("ob_mask not found, skip")
|
||||
result[video_id][id_str][ob_id] = np.eye(4)
|
||||
return result
|
||||
|
||||
est.gt_pose = reader.get_gt_pose(i_frame, ob_id)
|
||||
|
||||
pose = est.register(K=reader.K, rgb=color, depth=depth, ob_mask=ob_mask, ob_id=ob_id)
|
||||
logging.info(f"pose:\n{pose}")
|
||||
|
||||
if debug>=3:
|
||||
m = est.mesh_ori.copy()
|
||||
tmp = m.copy()
|
||||
tmp.apply_transform(pose)
|
||||
tmp.export(f'{debug_dir}/model_tf.obj')
|
||||
|
||||
result[video_id][id_str][ob_id] = pose
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def run_pose_estimation():
|
||||
wp.force_load(device='cuda')
|
||||
reader_tmp = LinemodReader(f'{opt.linemod_dir}/lm_test_all/test/000002', split=None)
|
||||
|
||||
debug = opt.debug
|
||||
use_reconstructed_mesh = opt.use_reconstructed_mesh
|
||||
debug_dir = opt.debug_dir
|
||||
|
||||
res = NestDict()
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
mesh_tmp = trimesh.primitives.Box(extents=np.ones((3)), transform=np.eye(4)).to_mesh()
|
||||
est = FoundationPose(model_pts=mesh_tmp.vertices.copy(), model_normals=mesh_tmp.vertex_normals.copy(), symmetry_tfs=None, mesh=mesh_tmp, scorer=None, refiner=None, glctx=glctx, debug_dir=debug_dir, debug=debug)
|
||||
|
||||
for ob_id in reader_tmp.ob_ids:
|
||||
ob_id = int(ob_id)
|
||||
if use_reconstructed_mesh:
|
||||
mesh = reader_tmp.get_reconstructed_mesh(ob_id, ref_view_dir=opt.ref_view_dir)
|
||||
else:
|
||||
mesh = reader_tmp.get_gt_mesh(ob_id)
|
||||
symmetry_tfs = reader_tmp.symmetry_tfs[ob_id]
|
||||
|
||||
args = []
|
||||
|
||||
video_dir = f'{opt.linemod_dir}/lm_test_all/test/{ob_id:06d}'
|
||||
reader = LinemodReader(video_dir, split=None)
|
||||
video_id = reader.get_video_id()
|
||||
est.reset_object(model_pts=mesh.vertices.copy(), model_normals=mesh.vertex_normals.copy(), symmetry_tfs=symmetry_tfs, mesh=mesh)
|
||||
|
||||
for i in range(len(reader.color_files)):
|
||||
args.append((reader, [i], est, debug, ob_id, "cuda:0"))
|
||||
|
||||
outs = []
|
||||
for arg in args:
|
||||
out = run_pose_estimation_worker(*arg)
|
||||
outs.append(out)
|
||||
|
||||
for out in outs:
|
||||
for video_id in out:
|
||||
for id_str in out[video_id]:
|
||||
for ob_id in out[video_id][id_str]:
|
||||
res[video_id][id_str][ob_id] = out[video_id][id_str][ob_id]
|
||||
|
||||
with open(f'{opt.debug_dir}/linemod_res.yml','w') as ff:
|
||||
yaml.safe_dump(make_yaml_dumpable(res), ff)
|
||||
|
||||
|
||||
if __name__=='__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
parser.add_argument('--linemod_dir', type=str, default="/mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/LINEMOD", help="linemod root dir")
|
||||
parser.add_argument('--use_reconstructed_mesh', type=int, default=0)
|
||||
parser.add_argument('--ref_view_dir', type=str, default="/mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video/bowen_addon/ref_views_16")
|
||||
parser.add_argument('--debug', type=int, default=0)
|
||||
parser.add_argument('--debug_dir', type=str, default=f'{code_dir}/debug')
|
||||
opt = parser.parse_args()
|
||||
set_seed(0)
|
||||
|
||||
detect_type = 'mask' # mask / box / detected
|
||||
|
||||
run_pose_estimation()
|
||||
+149
@@ -0,0 +1,149 @@
|
||||
# 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.
|
||||
|
||||
|
||||
from Utils import *
|
||||
import json,uuid,joblib,os,sys,argparse
|
||||
from datareader import *
|
||||
from estimater import *
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.append(f'{code_dir}/mycpp/build')
|
||||
import yaml
|
||||
|
||||
|
||||
def get_mask(reader, i_frame, ob_id, detect_type):
|
||||
if detect_type=='box':
|
||||
mask = reader.get_mask(i_frame, ob_id)
|
||||
H,W = mask.shape[:2]
|
||||
vs,us = np.where(mask>0)
|
||||
umin = us.min()
|
||||
umax = us.max()
|
||||
vmin = vs.min()
|
||||
vmax = vs.max()
|
||||
valid = np.zeros((H,W), dtype=bool)
|
||||
valid[vmin:vmax,umin:umax] = 1
|
||||
elif detect_type=='mask':
|
||||
mask = reader.get_mask(i_frame, ob_id, type='mask_visib')
|
||||
valid = mask>0
|
||||
elif detect_type=='cnos': #https://github.com/nv-nguyen/cnos
|
||||
mask = cv2.imread(reader.color_files[i_frame].replace('rgb','mask_cnos'), -1)
|
||||
valid = mask==ob_id
|
||||
else:
|
||||
raise RuntimeError
|
||||
|
||||
return valid
|
||||
|
||||
|
||||
|
||||
def run_pose_estimation_worker(reader, i_frames, est:FoundationPose, debug=False, ob_id=None, device:int=0):
|
||||
result = NestDict()
|
||||
torch.cuda.set_device(device)
|
||||
est.to_device(f'cuda:{device}')
|
||||
est.glctx = dr.RasterizeCudaContext(device)
|
||||
debug_dir = est.debug_dir
|
||||
|
||||
for i in range(len(i_frames)):
|
||||
i_frame = i_frames[i]
|
||||
id_str = reader.id_strs[i_frame]
|
||||
logging.info(f"{i}/{len(i_frames)}, video:{reader.get_video_id()}, id_str:{id_str}")
|
||||
color = reader.get_color(i_frame)
|
||||
depth = reader.get_depth(i_frame)
|
||||
|
||||
H,W = color.shape[:2]
|
||||
scene_ob_ids = reader.get_instance_ids_in_image(i_frame)
|
||||
video_id = reader.get_video_id()
|
||||
|
||||
logging.info(f"video:{reader.get_video_id()}, id_str:{id_str}, ob_id:{ob_id}")
|
||||
if ob_id not in scene_ob_ids:
|
||||
logging.info(f'skip {ob_id} as it does not exist in this scene')
|
||||
continue
|
||||
ob_mask = get_mask(reader, i_frame, ob_id, detect_type=detect_type)
|
||||
|
||||
est.gt_pose = reader.get_gt_pose(i_frame, ob_id)
|
||||
pose = est.register(K=reader.K, rgb=color, depth=depth, ob_mask=ob_mask, ob_id=ob_id, iteration=5)
|
||||
logging.info(f"pose:\n{pose}")
|
||||
|
||||
|
||||
if debug>=3:
|
||||
tmp = est.mesh_ori.copy()
|
||||
tmp.apply_transform(pose)
|
||||
tmp.export(f'{debug_dir}/model_tf.obj')
|
||||
|
||||
result[video_id][id_str][ob_id] = pose
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def run_pose_estimation():
|
||||
wp.force_load(device='cuda')
|
||||
video_dirs = sorted(glob.glob(f'{opt.ycbv_dir}/test/*'))
|
||||
res = NestDict()
|
||||
|
||||
debug = opt.debug
|
||||
use_reconstructed_mesh = opt.use_reconstructed_mesh
|
||||
debug_dir = opt.debug_dir
|
||||
|
||||
reader_tmp = YcbVideoReader(video_dirs[0])
|
||||
glctx = dr.RasterizeCudaContext()
|
||||
mesh_tmp = trimesh.primitives.Box(extents=np.ones((3)), transform=np.eye(4))
|
||||
est = FoundationPose(model_pts=mesh_tmp.vertices.copy(), model_normals=mesh_tmp.vertex_normals.copy(), symmetry_tfs=None, mesh=mesh_tmp, scorer=None, refiner=None, glctx=glctx, debug_dir=debug_dir, debug=debug)
|
||||
|
||||
ob_ids = reader_tmp.ob_ids
|
||||
|
||||
for ob_id in ob_ids:
|
||||
if use_reconstructed_mesh:
|
||||
mesh = reader_tmp.get_reconstructed_mesh(ob_id, ref_view_dir=opt.ref_view_dir)
|
||||
else:
|
||||
mesh = reader_tmp.get_gt_mesh(ob_id)
|
||||
symmetry_tfs = reader_tmp.symmetry_tfs[ob_id]
|
||||
|
||||
args = []
|
||||
for video_dir in video_dirs:
|
||||
reader = YcbVideoReader(video_dir, zfar=1.5)
|
||||
scene_ob_ids = reader.get_instance_ids_in_image(0)
|
||||
if ob_id not in scene_ob_ids:
|
||||
continue
|
||||
video_id = reader.get_video_id()
|
||||
|
||||
for i in range(len(reader.color_files)):
|
||||
if not reader.is_keyframe(i):
|
||||
continue
|
||||
args.append((reader, [i], est, debug, ob_id, 0))
|
||||
|
||||
est.reset_object(model_pts=mesh.vertices.copy(), model_normals=mesh.vertex_normals.copy(), symmetry_tfs=symmetry_tfs, mesh=mesh)
|
||||
outs = []
|
||||
for arg in args:
|
||||
out = run_pose_estimation_worker(*arg)
|
||||
outs.append(out)
|
||||
|
||||
for out in outs:
|
||||
for video_id in out:
|
||||
for id_str in out[video_id]:
|
||||
res[video_id][id_str][ob_id] = out[video_id][id_str][ob_id]
|
||||
|
||||
with open(f'{opt.debug_dir}/ycbv_res.yml','w') as ff:
|
||||
yaml.safe_dump(make_yaml_dumpable(res), ff)
|
||||
|
||||
|
||||
|
||||
if __name__=='__main__':
|
||||
parser = argparse.ArgumentParser()
|
||||
code_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
parser.add_argument('--ycbv_dir', type=str, default="/mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video", help="data dir")
|
||||
parser.add_argument('--use_reconstructed_mesh', type=int, default=0)
|
||||
parser.add_argument('--ref_view_dir', type=str, default="/mnt/9a72c439-d0a7-45e8-8d20-d7a235d02763/DATASET/YCB_Video/bowen_addon/ref_views_16")
|
||||
parser.add_argument('--debug', type=int, default=0)
|
||||
parser.add_argument('--debug_dir', type=str, default=f'{code_dir}/debug')
|
||||
opt = parser.parse_args()
|
||||
os.environ["YCB_VIDEO_DIR"] = opt.ycbv_dir
|
||||
|
||||
set_seed(0)
|
||||
|
||||
detect_type = 'mask' # mask / box / detected
|
||||
|
||||
run_pose_estimation()
|
||||
@@ -0,0 +1,39 @@
|
||||
lr: 0.0001
|
||||
c_in: 6
|
||||
zfar: 'Infinity'
|
||||
debug: null
|
||||
w_rot: 0.1
|
||||
n_view: 1
|
||||
run_id: null
|
||||
use_BN: true
|
||||
rot_rep: axis_angle
|
||||
ckpt_dir: null
|
||||
exp_name: 2023-10-28-18-33-37
|
||||
save_dir: /tmp/2023-10-28-18-33-37/
|
||||
loss_type: l2
|
||||
optimizer: adam
|
||||
trans_rep: tracknet
|
||||
batch_size: 64
|
||||
crop_ratio: 1.2
|
||||
use_normal: false
|
||||
BN_momentum: 0.1
|
||||
max_num_key: null
|
||||
warmup_step: -1
|
||||
input_resize:
|
||||
- 160
|
||||
- 160
|
||||
max_step_val: 1000
|
||||
normal_uint8: false
|
||||
vis_interval: 1000
|
||||
weight_decay: 0
|
||||
n_max_objects: null
|
||||
normalize_xyz: true
|
||||
clip_grad_norm: 'Infinity'
|
||||
rot_normalizer: 0.3490658503988659
|
||||
trans_normalizer:
|
||||
- 0.019999999552965164
|
||||
- 0.019999999552965164
|
||||
- 0.05000000074505806
|
||||
max_step_per_epoch: 25000
|
||||
val_epoch_interval: 10
|
||||
n_dataloader_workers: 60
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,41 @@
|
||||
lr: 0.0001
|
||||
c_in: 6
|
||||
zfar: 'Infinity'
|
||||
debug: null
|
||||
n_view: 1
|
||||
run_id: 3wy8qqex
|
||||
use_BN: true
|
||||
exp_name: 2024-01-11-20-02-45
|
||||
n_epochs: 62
|
||||
save_dir: /home/bowenw/debug/2024-01-11-20-02-45/
|
||||
use_mask: false
|
||||
loss_type: pairwise_valid
|
||||
optimizer: adam
|
||||
batch_size: 64
|
||||
crop_ratio: 1.1
|
||||
enable_amp: true
|
||||
use_normal: false
|
||||
max_num_key: null
|
||||
warmup_step: -1
|
||||
input_resize:
|
||||
- 160
|
||||
- 160
|
||||
max_step_val: 1000
|
||||
vis_interval: 1000
|
||||
weight_decay: 0
|
||||
normalize_xyz: true
|
||||
resume_run_id: null
|
||||
clip_grad_norm: 'Infinity'
|
||||
lr_epoch_decay: 500
|
||||
render_backend: nvdiffrast
|
||||
train_num_pair: 5
|
||||
lr_decay_epochs:
|
||||
- 50
|
||||
n_epochs_warmup: 1
|
||||
make_pair_online: false
|
||||
gradient_max_norm: 'Infinity'
|
||||
max_step_per_epoch: 10000
|
||||
n_rendering_workers: 1
|
||||
save_epoch_interval: 100
|
||||
n_dataloader_workers: 100
|
||||
split_objects_across_gpus: true
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,30 @@
|
||||
{
|
||||
"source": "community mirror gpue/foundationpose-weights; official Google Drive quota exceeded",
|
||||
"verification": "size and mirror LFS SHA256; not independently compared to official files",
|
||||
"files": [
|
||||
{
|
||||
"path": "2023-10-28-18-33-37/config.yml",
|
||||
"size": 708,
|
||||
"sha256": "28a6ba94a33230ee5fc3c51939486281578b0972542bd9e38ca6123e75605686",
|
||||
"url": "https://hf-mirror.com/gpue/foundationpose-weights/resolve/main/2023-10-28-18-33-37/config.yml"
|
||||
},
|
||||
{
|
||||
"path": "2023-10-28-18-33-37/model_best.pth",
|
||||
"size": 68220109,
|
||||
"sha256": "774700586ddc435d408fc01c9809c43e151232936369dfbea0f0f964ba471d60",
|
||||
"url": "https://hf-mirror.com/gpue/foundationpose-weights/resolve/main/2023-10-28-18-33-37/model_best.pth"
|
||||
},
|
||||
{
|
||||
"path": "2024-01-11-20-02-45/config.yml",
|
||||
"size": 778,
|
||||
"sha256": "a79db4de3b95885dd5ae86833b37b8698a75dad81e87d1086cd50b2fcd8dda3f",
|
||||
"url": "https://hf-mirror.com/gpue/foundationpose-weights/resolve/main/2024-01-11-20-02-45/config.yml"
|
||||
},
|
||||
{
|
||||
"path": "2024-01-11-20-02-45/model_best.pth",
|
||||
"size": 190229389,
|
||||
"sha256": "81924d384bf5c26c646ee4783104982ae3d1e049c181c36641b6a7aeae494c26",
|
||||
"url": "https://hf-mirror.com/gpue/foundationpose-weights/resolve/main/2024-01-11-20-02-45/model_best.pth"
|
||||
}
|
||||
]
|
||||
}
|
||||
+16
@@ -0,0 +1,16 @@
|
||||
# 按尺寸重新建立的理想化瓶子
|
||||
|
||||
高度 **157.856 mm**，最大直径 **54.110 mm**。这两个尺寸沿用用户提供的瓶子资产；不再使用其异常收腰的半径轮廓。
|
||||
|
||||
圆柱瓶身高 103 mm，锥形瓶肩至 139 mm，瓶颈直径 28 mm，顶端总高 157.856 mm。瓶肩和瓶颈比例是设计选择，没有宣称还原实物细节。顶底封闭，无内壁、瓶底凹槽、螺纹或独立瓶盖。
|
||||
|
||||
- `bottle.xml`：可直接加载的 MuJoCo 自由瓶子；视觉网格与三段碰撞几何使用同一轮廓。
|
||||
- `assets/bottle.obj`、`bottle.stl`、`bottle.glb`：网格格式，均以米为单位。
|
||||
- `preview.png`：预览；`validation.json`：尺寸、封闭性及加载检查。
|
||||
- `profile.npz`：新设计的高度/半径控制点。
|
||||
|
||||
碰撞为圆柱瓶身、凸锥台瓶肩、圆柱瓶颈，去掉旧版十段忽大忽小的圆柱。质量暂沿用旧 XML 的 0.2308 kg，惯量按均匀实心模型重算；质量、材料及质量分布未实测。原始三目资产未修改。
|
||||
|
||||
重建命令：`env PYTHONPATH= LD_LIBRARY_PATH= .dex/bin/python scripts/build_parametric_bottle.py`。
|
||||
|
||||
新资产已接入 `output/l20_2047635068/bottle_grasp_v2/contact_scene.xml` 并通过加载检查。该场景是后续抓取优化入口，尚未运行新版物理抓取验证。
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
BIN
Binary file not shown.
+17
@@ -0,0 +1,17 @@
|
||||
<mujoco model="dimension_based_bottle">
|
||||
<option gravity="0 0 -9.81" />
|
||||
<asset>
|
||||
<mesh name="bottle_visual" file="assets/bottle.obj" />
|
||||
<mesh name="bottle_shoulder" file="assets/shoulder.obj" />
|
||||
</asset>
|
||||
<worldbody>
|
||||
<body name="bottle">
|
||||
<freejoint name="bottle_freejoint" />
|
||||
<inertial mass=".2308" pos="1.10250386217449e-19 9.190620430899951e-19 0.06616881962256083" diaginertia="0.0004053028305440948 0.0004053028305440948 7.742359133762652e-05" />
|
||||
<geom type="mesh" mesh="bottle_visual" group="1" contype="0" conaffinity="0" rgba=".68 .74 .8 1" />
|
||||
<geom type="cylinder" size="0.027055 0.0515" pos="0 0 0.0515" group="3" friction=".6 .005 .001" condim="6" />
|
||||
<geom type="mesh" mesh="bottle_shoulder" group="3" friction=".6 .005 .001" condim="6" />
|
||||
<geom type="cylinder" size="0.014 0.009427999999999992" pos="0 0 0.148428" group="3" friction=".6 .005 .001" condim="6" />
|
||||
</body>
|
||||
</worldbody>
|
||||
</mujoco>
|
||||
BIN
Binary file not shown.
BIN
Binary file not shown.
@@ -0,0 +1,17 @@
|
||||
{
|
||||
"height_mm": 157.856,
|
||||
"max_diameter_mm": 54.11,
|
||||
"neck_diameter_mm": 28.0,
|
||||
"watertight": true,
|
||||
"winding_consistent": true,
|
||||
"volume_m3": 0.000297608913283581,
|
||||
"mujoco_load_passed": true,
|
||||
"collision_parts": 3,
|
||||
"extent_m": [
|
||||
0.05411,
|
||||
0.05411,
|
||||
0.157856
|
||||
],
|
||||
"shape_source": "Idealized cylinder body, tapered shoulder, cylinder neck; only height and maximum diameter inherited",
|
||||
"mass_source": "Previous XML mass 0.2308 kg retained; uniform solid inertia recomputed, not measured"
|
||||
}
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user