Update to 2026-09-17 pipeline snapshot; add weights, L20 assets and recording via Git LFS
Source: RGB-D -> Dyn-HaMR -> L20 retargeting -> FoundationPose -> reference repair -> SPIDER, documented in docs/PIPELINE_LATEST.md and docs/SETUP_AND_WEIGHTS.md. Adds FoundationPose and nvdiffrast upstream snapshots, requirements/pipeline_venv.txt and the FoundationPose weight manifest/downloader. Assets (Git LFS): weights/ (WiLoR detector, HandFlow denoiser, UniDepth-L), FoundationPose checkpoints, HaMeR checkpoint, Dyn-HaMR HMP model and BMC constraints, L20 URDF/meshes, the 20260915_171525 D405 recording and the two box CADs. MANO models are not redistributed (third_party/hamer/_DATA/data/mano/README.txt). Environments, caches and run outputs excluded. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
liyang
@@ -0,0 +1,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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user