Skip to content

GitLab

Commit 1bc644a1 authored by Nianchen Deng's avatar Nianchen Deng
Browse files

sync

parent 6294701e
Expand all Hide whitespace changes
Inline Side-by-side
.clang-format deleted 100644 → 0
View file @ 6294701e
---
Language: Cpp
# BasedOnStyle: LLVM
AccessModifierOffset: -4
AlignAfterOpenBracket: Align
AlignConsecutiveAssignments: false
AlignConsecutiveDeclarations: false
AlignEscapedNewlines: Right
AlignOperands: true
AlignTrailingComments: true
AllowAllParametersOfDeclarationOnNextLine: true
AllowShortBlocksOnASingleLine: false
AllowShortCaseLabelsOnASingleLine: false
AllowShortFunctionsOnASingleLine: All
AllowShortIfStatementsOnASingleLine: false
AllowShortLoopsOnASingleLine: false
AlwaysBreakAfterDefinitionReturnType: None
AlwaysBreakAfterReturnType: None
AlwaysBreakBeforeMultilineStrings: false
AlwaysBreakTemplateDeclarations: false
BinPackArguments: true
BinPackParameters: true
BraceWrapping:
AfterClass: true
AfterControlStatement: false
AfterEnum: true
AfterFunction: true
AfterNamespace: true
AfterObjCDeclaration: true
AfterStruct: true
AfterUnion: true
BeforeCatch: false
BeforeElse: false
IndentBraces: false
SplitEmptyFunction: true
SplitEmptyRecord: true
SplitEmptyNamespace: true
BreakBeforeBinaryOperators: None
BreakBeforeBraces: Attach
BreakBeforeInheritanceComma: false
BreakBeforeTernaryOperators: true
BreakConstructorInitializersBeforeComma: false
BreakConstructorInitializers: BeforeColon
BreakAfterJavaFieldAnnotations: false
BreakStringLiterals: true
ColumnLimit: 100
CommentPragmas: '^ IWYU pragma:'
CompactNamespaces: false
ConstructorInitializerAllOnOneLineOrOnePerLine: true
ConstructorInitializerIndentWidth: 4
ContinuationIndentWidth: 4
Cpp11BracedListStyle: true
DerivePointerAlignment: false
DisableFormat: false
ExperimentalAutoDetectBinPacking: false
FixNamespaceComments: true
ForEachMacros:
- foreach
- Q_FOREACH
- BOOST_FOREACH
IncludeCategories:
- Regex: '^"(llvm|llvm-c|clang|clang-c)/'
Priority: 2
- Regex: '^(<|"(gtest|gmock|isl|json)/)'
Priority: 3
- Regex: '.*'
Priority: 1
IncludeIsMainRegex: '(Test)?$'
IndentCaseLabels: false
IndentWidth: 4
IndentWrappedFunctionNames: false
JavaScriptQuotes: Leave
JavaScriptWrapImports: true
KeepEmptyLinesAtTheStartOfBlocks: true
MacroBlockBegin: ''
MacroBlockEnd: ''
MaxEmptyLinesToKeep: 1
NamespaceIndentation: None
ObjCBlockIndentWidth: 4
ObjCSpaceAfterProperty: false
ObjCSpaceBeforeProtocolList: true
PenaltyBreakAssignment: 2
PenaltyBreakBeforeFirstCallParameter: 19
PenaltyBreakComment: 300
PenaltyBreakFirstLessLess: 120
PenaltyBreakString: 1000
PenaltyExcessCharacter: 1000000
PenaltyReturnTypeOnItsOwnLine: 60
PointerAlignment: Right
ReflowComments: true
SortIncludes: false
SortUsingDeclarations: true
SpaceAfterCStyleCast: false
SpaceAfterTemplateKeyword: true
SpaceBeforeAssignmentOperators: true
SpaceBeforeParens: ControlStatements
SpaceInEmptyParentheses: false
SpacesBeforeTrailingComments: 1
SpacesInAngles: false
SpacesInContainerLiterals: true
SpacesInCStyleCastParentheses: false
SpacesInParentheses: false
SpacesInSquareBrackets: false
Standard: Cpp11
TabWidth: 4
UseTab: Never
...
.vscode/c_cpp_properties.json
View file @ 1bc644a1
......@@ -4,13 +4,17 @@
"name": "Linux",
"includePath": [
"${workspaceFolder}/cpp/**",
"/usr/local/cuda/include"
"${workspaceFolder}/clib/include",
"/usr/local/cuda/include",
"/home/dengnc/libtorch/include",
"/home/dengnc/libtorch/include/torch/csrc/api/include",
"/home/dengnc/miniconda3/include/**"
],
"defines": [],
"compilerPath": "/usr/bin/gcc",
"cStandard": "gnu17",
"cppStandard": "gnu++14",
"intelliSenseMode": "gcc-x64"
"cppStandard": "gnu++17",
"intelliSenseMode": "${default}"
}
],
"version": 4
......
.vscode/launch.json
View file @ 1bc644a1
......@@ -4,8 +4,17 @@
// 欲了解更多信息，请访问: https://go.microsoft.com/fwlink/?linkid=830387
"version": "0.2.0",
"configurations": [
{
"name": "convert_nerf_checkpoint",
"type": "python",
"request": "launch",
"program": "tools/convert_nerf_checkpoint.py",
"args": [
"/home/dengnc/Work/ref_code/nerf-pytorch/logs/dvs_gas_nearrange/200000.tar"
],
"console": "integratedTerminal",
"justMyCode": false
},
{
"name": "Debug/Voxel Sampler Export 3D",
"type": "python",
......@@ -24,7 +33,7 @@
"program": "train.py",
"args": [
"-c",
"_lr_snerf_voxels+ls",
"nerf_llff",
//"/home/dengnc/dvs/data/classroom/_nets/pano_t0.8/smnerf_voxels+ls+lbl/checkpoint_35.tar",
//"--prune",
//"1",
......@@ -34,9 +43,9 @@
//"100",
//"--views",
//"5",
"data/classroom/lr_view_t0.8_r360x80_train"
"/home/dengnc/Work/fov_nerf/data/__thesis/trex/train.json"
],
"justMyCode": false,
"justMyCode": true,
"console": "integratedTerminal"
},
{
......@@ -45,16 +54,48 @@
"request": "launch",
"program": "test.py",
"args": [
"-m",
"/home/dengnc/dvs/data/classroom/_nets/ms_train_t0.8/_cnerf/checkpoint_50.tar",
"-o",
"perf",
"color",
"--output-type",
"--media",
"image",
"/home/dengnc/dvs/data/classroom/lr_view_t0.8_r360x80_test.json",
"--views",
"1"
"3",
//"-r",
//"100x200",
"/home/dengnc/Work/fov_nerf/data/__thesis/barbershop_old/_nets/train_t0.3/eval@snerffast4-rgb_e6_fc256x8_d1.20-6.00_s64_~p/checkpoint_50.tar",
"/home/dengnc/Work/fov_nerf/data/__thesis/barbershop_old/test_t0.3.json",
//"--batch",
//"8192"
],
"justMyCode": false,
"console": "integratedTerminal"
},
{
"name": "Convert Colmap",
"type": "python",
"request": "launch",
"program": "tools/data/colmap2dataset.py",
"args": [
"data/__captured/sittingroom",
"--scale-down",
"4"
],
"console": "integratedTerminal"
},
{
"name": "Generate Video",
"type": "python",
"request": "launch",
"program": "tools/gen_video.py",
"args": [
"data/__captured/sittingroom",
"data/__demo/realvideo/sittingroom_1.json",
"-f",
"40",
"-s",
"-d",
"0.6"
],
"console": "integratedTerminal"
}
......
.vscode/settings.json deleted 100644 → 0
View file @ 6294701e
{
"files.watcherExclude": {
"**/data/**": true
},
"files.associations": {
"string": "cpp",
"__functional_03": "cpp",
"functional": "cpp",
"vector": "cpp",
"__config": "cpp",
"__nullptr": "cpp"
},
"python.pythonPath": "/home/dengnc/miniconda3/bin/python",
}
\ No newline at end of file
LICENSE
View file @ 1bc644a1
MIT License
Copyright (c) 2020 Jiannan Ye
Copyright (c) 2022 Nianchen Deng
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
......
README.md
View file @ 1bc644a1
# Configure environment
## 1. Install Conda packages:
* Pytorch 1.8.1 with CUDA
```
$ conda install pytorch torchvision torchaudio cudatoolkit=<your cuda version> -c pytorch -c nvidia
```
Or ref to https://pytorch.org/get-started/locally/ for install guide
* Pytorch with CUDA
```bash
$ conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
```
Or ref to https://pytorch.org/get-started/locally/ for install guide
* matplotlib
* tensorboard
* plyfile
```
$ conda install -c conda-forge plyfile
```
* tqdm
* configargparse
* (Optional) dash
```
$ conda install dash pandas
```
```bash
$ conda install dash pandas
```
## 2. Install Pip packages:
* pyGlm
* tensorboardX
* (optional) opencv-python
* torch_tb_profiler
* opencv-python
* ipympl
* lpips
* (optional) thop
* (optional) ConcurrentLogHandler
# Useful commands
## 1. Video generate:
```
$ ffmpeg -y -r 50 -i %04d.png -c:v libx264 -vframes 600 ../classroom_hmd_mono_hint.mp4
## 4. Build extension "clib._ext"
```bash
$ python setup.py build_ext
```
If build successed, a _ext.\*.so will be generated under build/lib.\*/clib directory. Move this file to clib/.
## 2. Convert onnx to tensorRT
```
$ trtexec --onnx=in.onnx --fp16 --saveEngine=out.trt --workspace=4096
```
## 4. (Optional) Install FFMpeg with Extra Codecs:
# Install FFMpeg with Extra Codecs:
```
```bash
sudo apt-get update -qq && sudo apt-get -y install \
autoconf \
automake \
......@@ -109,4 +97,25 @@ PATH="$HOME/bin:$PATH" PKG_CONFIG_PATH="$HOME/ffmpeg_build/lib/pkgconfig" ./conf
PATH="$HOME/bin:$PATH" make && \
make install && \
hash -r
```
# Useful commands
## 1. Video generate:
```bash
$ ffmpeg -y -r 50 -i %04d.png -c:v libx264 -vframes 600 ../classroom_hmd_mono_hint.mp4
```
## 2. Extract frames:
```bash
$ ffmpeg -i <video_path> -f image2 -q:v 2 -vf fps=<fps> <out_dir>/image%04d.png
```
## 3. Convert onnx to tensorRT
```bash
$ trtexec --onnx=in.onnx --fp16 --saveEngine=out.trt --workspace=4096
```
## 4. Generate dataset of specific path
```bash
$ python tools/data/gen_seq.py -s helix|look_around|scan_around -n <frames> --ref <train_dataset.json> <dataset_dir>
```
\ No newline at end of file
a.drawio 0 → 100644
View file @ 1bc644a1
This diff is collapsed.
args_test.py 0 → 100644
View file @ 1bc644a1
from configargparse import ArgumentParser
from utils.args import BaseArgs
class Args(BaseArgs):
a: int | None = 20
b: str = "hello"
c: list[int] = [4, 5]
flag: bool = False
parser = ArgumentParser()
args = Args()
args.parse(debug=True)
#args.setup_parser(parser, True)
#parser.add_argument('--data', nargs='+', type=int,default=[10])
#parser.parse_args("--flag --data 20 30", namespace=args)
#setattr(args, "c", [4, 5])
print(args)
\ No newline at end of file
batch_collect_video.sh
View file @ 1bc644a1
#/usr/bin/bash
#!/usr/bin/bash
curdir=$(pwd)
datadir="$curdir/data/__new/classroom_fovea_r360x80_t0.6"
videodir="$datadir/eval_video"
......@@ -6,7 +6,7 @@ epochs=50
if [ ! -d "$videodir" ]; then
echo "make directory for Video"
mkdir $videodir
mkdir "$videodir"
fi
# nets: 1, 2, 4, 8
......@@ -22,7 +22,7 @@ for n_nets in 1 2 4 8; do
dst_path="$videodir/$exportname.mp4"
if [ -f "$videodir/$src_path" ]; then
if [ ! -f "$dst_path" ]; then
ln -s $src_path $dst_path
ln -s $src_path "$dst_path"
fi
fi
done
......
batch_export_net.sh
View file @ 1bc644a1
#/usr/bin/bash
#!/usr/bin/bash
datadir='data/__new/classroom_fovea_r360x80_t0.6'
onnxdir="$datadir/eval_onnx"
......@@ -23,7 +23,7 @@ for n_nets in 1 2 4 8; do
for nf in 64 128 256 512 1024; do
for n_samples in 8 16 32 64 128; do
configid="eval@snerffast${n_nets}-rgb_e6_fc${nf}x${n_layers}_d1.00-7.00_s${n_samples}_~p"
if (( $n_samples == 64 )); then
if ((n_samples == 64)); then
exportname="eval_${n_nets}x${nf}x${n_layers}"
else
exportname="eval_${n_nets}x${nf}x${n_layers}_${n_samples}"
......@@ -45,4 +45,4 @@ for n_nets in 1 2 4 8; do
done
done
done
done
\ No newline at end of file
done
batch_infer.sh
View file @ 1bc644a1
#/usr/bin/bash
#!/usr/bin/bash
testcase=$1
datadir='data/__new/classroom_fovea_r360x80_t0.6'
......@@ -19,7 +19,7 @@ for nf in 64 128 256 512 1024; do
configid="eval@snerffast${n_nets}-rgb_e6_fc${nf}x${n_layers}_d1.00-7.00_s64_~p"
if [ ! -f "$datadir/$configid/model-epoch_$epochs.pth" ]; then
cont_epoch=0
for ((i=$epochs-1;i>0;i--)) do
for ((i=epochs-1;i>0;i--)) do
if [ -f "$datadir/$configid/model-epoch_$i.pth" ]; then
cont_epoch=$i
break
......
batch_test.sh
View file @ 1bc644a1
......@@ -3,11 +3,24 @@
test_dataset=$1
test_model_dir=$2
for i in "$test_model_dir"*
for misc_path in "$test_model_dir"/*/_misc/checkpoint_30.tar
do
python test.py -m "$(pwd)/$i/checkpoint_50.tar" -o perf color --output-type image "$test_dataset"
[ -f "$misc_path" ] || continue
misc_dir=$(dirname "$misc_path")
echo mv "$misc_path" "${misc_dir%/*}"
mv "$misc_path" "${misc_dir%/*}"
done
for model_path in "$test_model_dir"/*/checkpoint_30.tar
do
model_dir=$(dirname "$model_path")
#model_name=${model_dir#"$test_model_dir/"}
[ -d "$model_dir/output_30" ] || python test.py "$model_path" "$test_dataset" -o perf --media image
done
echo Test Finished
ls $test_model_dir/*/output_50/perf* | awk -F"/" '{print $6, "\t", $8}'
\ No newline at end of file
cd "$test_model_dir"
ls */output_30/perf* | sed -r "s/()\/output_30\/perf.+_([0-9\.]+)ms_([0-9\.e\-]+)\.csv/\1\t\2\t\3/"
blender/gen_utils.py
View file @ 1bc644a1
......@@ -9,9 +9,9 @@ from itertools import product
class Gen:
def __init__(self, root_dir: str, dataset_name: str, *,
res: Tuple[int, int],
res: tuple[int, int],
fov: float,
samples: List[int]) -> None:
samples: list[int]) -> None:
self.res = res
self.fov = fov
self.samples = samples
......@@ -47,11 +47,11 @@ class Gen:
with open(self.data_desc_file, 'w') as fp:
json.dump(self.desc, fp, indent=4)
def add_sample(self, i, x: List[float], render_only=False):
def add_sample(self, i, x: list[float], render_only=False):
self.cam_obj.location = x[:3]
if len(x) > 3:
self.cam_obj.rotation_euler = [math.radians(x[4]), math.radians(x[3]), 0]
self.scene.render.filepath = self.data_dir + self.desc['view_file_pattern'] % i
self.scene.render.filepath = self.data_dir + self.desc['color_file'] % i
bpy.ops.render.render(write_still=True)
if not render_only:
self.desc['view_centers'].append(x[:3])
......@@ -84,8 +84,8 @@ class Gen:
# Render missing views in data desc
for i in range(len(self.desc['view_centers'])):
if not os.path.exists(self.data_dir + self.desc['view_file_pattern'] % i):
x: List[float] = self.desc['view_centers'][i]
if not os.path.exists(self.data_dir + self.desc['color_file'] % i):
x: list[float] = self.desc['view_centers'][i]
if 'view_rots' in self.desc:
x += self.desc['view_rots'][i]
self.add_sample(i, x, render_only=True)
......@@ -99,15 +99,15 @@ class Gen:
class GenView(Gen):
def __init__(self, root_dir: str, dataset_name: str, *,
res: Tuple[int, int], fov: float, samples: List[int],
tbox: Tuple[float, float, float], rbox: Tuple[float, float]) -> None:
res: tuple[int, int], fov: float, samples: list[int],
tbox: tuple[float, float, float], rbox: tuple[float, float]) -> None:
super().__init__(root_dir, dataset_name, res=res, fov=fov, samples=samples)
self.tbox = tbox
self.rbox = rbox
def init_desc(self):
return {
'view_file_pattern': 'view_%04d.png',
'color_file': 'view_%04d.png',
"gl_coord": True,
'view_res': {
'x': self.res[0],
......@@ -143,8 +143,8 @@ class GenView(Gen):
class GenPano(Gen):
def __init__(self, root_dir: str, dataset_name: str, *,
samples: List[int], depth_range: Tuple[float, float],
tbox: Tuple[float, float, float] = None) -> None:
samples: list[int], depth_range: tuple[float, float],
tbox: tuple[float, float, float] = None) -> None:
self.depth_range = depth_range
self.tbox = tbox
super().__init__(root_dir, dataset_name, res=[4096, 2048], fov=-1, samples=samples)
......@@ -157,13 +157,15 @@ class GenPano(Gen):
}
} if self.tbox else {}
return {
"type": "pano",
'view_file_pattern': 'view_%04d.png',
'color_file': 'view_%04d.png',
"gl_coord": True,
'view_res': {
'x': self.res[0],
'y': self.res[1]
},
"cam_params": {
"type": "pano"
},
**range,
"depth_range": {
"min": self.depth_range[0],
......
clib/__init__.py
View file @ 1bc644a1
......@@ -2,6 +2,9 @@
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
#
# To install the _ext library, run the following command:
# > python setup.py build_ext --inplace
''' Modified based on: https://github.com/erikwijmans/Pointnet2_PyTorch '''
from __future__ import (
......@@ -11,32 +14,16 @@ from __future__ import (
print_function,
unicode_literals,
)
import os
import sys
from typing import Tuple
import torch
import torch.nn.functional as F
from torch.autograd import Function
import torch.nn as nn
import sys
import numpy as np
from torch.autograd import Function
from torch.autograd.function import FunctionCtx, once_differentiable
import clib._ext as _ext
from utils.geometry import discretize_points
from utils import math
try:
import builtins
except:
import __builtin__ as builtins
try:
import clib._ext as _ext
except ImportError:
raise ImportError(
"Could not import _ext module.\n"
"Please see the setup instructions in the README"
)
class BallRayIntersect(Function):
@staticmethod
......@@ -59,47 +46,8 @@ class BallRayIntersect(Function):
ball_ray_intersect = BallRayIntersect.apply
class AABBRayIntersect(Function):
@staticmethod
def forward(ctx, voxelsize, n_max, points, ray_start, ray_dir):
# HACK: speed-up ray-voxel intersection by batching...
G = min(2048, int(2 * 10 ** 9 / points.numel())) # HACK: avoid out-of-memory
S, N = ray_start.shape[:2]
K = int(np.ceil(N / G))
G, K = 1, N # HACK
H = K * G
if H > N:
ray_start = torch.cat([ray_start, ray_start[:, :H - N]], 1)
ray_dir = torch.cat([ray_dir, ray_dir[:, :H - N]], 1)
ray_start = ray_start.reshape(S * G, K, 3)
ray_dir = ray_dir.reshape(S * G, K, 3)
points = points[None].expand(S * G, *points.size()).contiguous()
inds, min_depth, max_depth = _ext.aabb_intersect(
ray_start.float(), ray_dir.float(), points.float(), voxelsize, n_max)
min_depth = min_depth.type_as(ray_start)
max_depth = max_depth.type_as(ray_start)
inds = inds.reshape(S, H, -1)
min_depth = min_depth.reshape(S, H, -1)
max_depth = max_depth.reshape(S, H, -1)
if H > N:
inds = inds[:, :N]
min_depth = min_depth[:, :N]
max_depth = max_depth[:, :N]
ctx.mark_non_differentiable(inds)
ctx.mark_non_differentiable(min_depth)
ctx.mark_non_differentiable(max_depth)
return inds, min_depth, max_depth
@staticmethod
def backward(ctx, a, b, c):
return None, None, None, None, None
def aabb_ray_intersect(voxelsize: float, n_max: int, points: torch.Tensor, ray_start: torch.Tensor,
ray_dir: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
ray_dir: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
AABB-Ray intersect test
......@@ -112,7 +60,33 @@ def aabb_ray_intersect(voxelsize: float, n_max: int, points: torch.Tensor, ray_s
:return `Tensor(S, N, n_max)`: min depths of every intersected voxels
:return `Tensor(S, N, n_max)`: max depths of every intersected voxels
"""
return AABBRayIntersect.apply(voxelsize, n_max, points, ray_start, ray_dir)
# HACK: speed-up ray-voxel intersection by batching...
G = min(2048, int(2e9 / points.numel())) # HACK: avoid out-of-memory
S, N = ray_start.shape[:2]
K = math.ceil(N / G)
G, K = 1, N # HACK
H = K * G
if H > N:
ray_start = torch.cat([ray_start, ray_start[:, :H - N]], 1)
ray_dir = torch.cat([ray_dir, ray_dir[:, :H - N]], 1)
ray_start = ray_start.reshape(S * G, K, 3)
ray_dir = ray_dir.reshape(S * G, K, 3)
points = points[None].expand(S * G, *points.size()).contiguous()
inds, min_depth, max_depth = _ext.aabb_intersect(
ray_start.float(), ray_dir.float(), points.float(), voxelsize, n_max)
min_depth = min_depth.type_as(ray_start)
max_depth = max_depth.type_as(ray_start)
inds = inds.reshape(S, H, -1)
min_depth = min_depth.reshape(S, H, -1)
max_depth = max_depth.reshape(S, H, -1)
if H > N:
inds = inds[:, :N]
min_depth = min_depth[:, :N]
max_depth = max_depth[:, :N]
return inds, min_depth, max_depth
class SparseVoxelOctreeRayIntersect(Function):
......@@ -122,7 +96,7 @@ class SparseVoxelOctreeRayIntersect(Function):
G = min(2048, int(2 * 10 ** 9 / (points.numel() + children.numel())))
S, N = ray_start.shape[:2]
K = int(np.ceil(N / G))
G, K = 1, N # HACK
G, K = 1, N # HACK
H = K * G
if H > N:
ray_start = torch.cat([ray_start, ray_start[:, :H - N]], 1)
......@@ -156,7 +130,7 @@ class SparseVoxelOctreeRayIntersect(Function):
def octree_ray_intersect(voxelsize: float, n_max: int, points: torch.Tensor, children: torch.Tensor,
ray_start: torch.Tensor, ray_dir: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
ray_start: torch.Tensor, ray_dir: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Octree-Ray intersect test
......@@ -267,7 +241,7 @@ class UniformRaySampling(Function):
def uniform_ray_sampling(pts_idx: torch.Tensor, min_depth: torch.Tensor, max_depth: torch.Tensor,
step_size: float, max_ray_length: float, deterministic: bool = False) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
step_size: float, max_ray_length: float, deterministic: bool = False) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Sample along rays uniformly
......@@ -357,7 +331,7 @@ class InverseCDFRaySampling(Function):
def inverse_cdf_sampling(pts_idx: torch.Tensor, min_depth: torch.Tensor, max_depth: torch.Tensor,
probs: torch.Tensor, steps: torch.Tensor, fixed_step_size: float = -1,
deterministic: bool = False) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
deterministic: bool = False) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Sample along rays by inverse CDF
......@@ -461,7 +435,7 @@ def parallel_ray_sampling(MARCH_SIZE, pts_idx, min_depth, max_depth, determinist
return sampled_idx, sampled_depth, sampled_dists
def build_easy_octree(points: torch.Tensor, half_voxel: float) -> Tuple[torch.Tensor, torch.Tensor]:
def build_easy_octree(points: torch.Tensor, half_voxel: float) -> tuple[torch.Tensor, torch.Tensor]:
"""
Build an octree.
......@@ -476,4 +450,69 @@ def build_easy_octree(points: torch.Tensor, half_voxel: float) -> Tuple[torch.Te
center = (coords.max(0)[0] + coords.min(0)[0]) / 2
centers, children = _ext.build_octree(center, coords, int(depths))
centers = centers.float() * half_voxel + residual # transform back to float
return centers, children
\ No newline at end of file
return centers, children
class MultiresHashEncode(Function):
@staticmethod
def forward(ctx: FunctionCtx, levels: int, coarse_levels: int, res_list: torch.Tensor,
hash_table_offsets: torch.Tensor, x: torch.Tensor, hash_table: torch.Tensor,
grad_enabled: bool) -> torch.Tensor:
"""
[summary]
:param ctx `FunctionCtx`: [description]
:param levels `int`: [description]
:param coarse_levels `int`: [description]
:param res_list `Tensor(L, D)`: [description]
:param hash_table_offsets `Tensor(L+1)`: [description]
:param x `Tensor(N, D)`: [description]
:param hash_table `Tensor(T, F)`: [description]
:return `Tensor(L, N, F)`: [description]
"""
x = x.contiguous()
res_list = res_list.int().contiguous()
hash_table_offsets = hash_table_offsets.int().contiguous()
if grad_enabled and hash_table.requires_grad:
encoded, weights, indices = _ext.multires_hash_encode_with_grad(
levels, coarse_levels, x, res_list, hash_table, hash_table_offsets)
ctx.save_for_backward(weights, indices.long())
ctx.hash_table_shape = hash_table.shape
return encoded
print(hash_table)
return _ext.multires_hash_encode(levels, coarse_levels, x, res_list, hash_table,
hash_table_offsets)
@staticmethod
@once_differentiable
def backward(ctx: FunctionCtx, grad_output: torch.Tensor):
"""
[summary]
:param ctx `FunctionCtx`: [description]
:param grad_output `Tensor(L, N, F)`: [description]
:return: [description]
"""
weights, indices = ctx.saved_tensors # (L, N, C)
t = grad_output[..., None, :] * weights[..., None] # (L, N, C, F)
grad_hash_table = grad_output.new_zeros(*ctx.hash_table_shape)
grad_hash_table.index_put_([indices], t, accumulate=True)
return None, None, None, None, None, grad_hash_table, None
def multires_hash_encode(levels: int, coarse_levels: int, res_list: torch.Tensor,
hash_table_offsets: torch.Tensor, x: torch.Tensor, hash_table: torch.Tensor) -> torch.Tensor:
"""
:param levels `int`: [description]
:param coarse_levels `int`: [description]
:param res_list `Tensor(L, D)`: [description]
:param hash_table_offsets `Tensor(L+1)`: [description]
:param x `Tensor(N, D)`: [description]
:param hash_table `Tensor(T, F)`: [description]
:return `Tensor(L, N, F)`: [description]
"""
return MultiresHashEncode.apply(levels, coarse_levels, res_list, hash_table_offsets, x,
hash_table, torch.is_grad_enabled())
clib/_ext.pyi 0 → 100644
View file @ 1bc644a1
import torch
# intersect.h
def ball_intersect(ray_start: torch.Tensor, ray_dir: torch.Tensor, points: torch.Tensor,
radius: float, n_max: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...
def aabb_intersect(ray_start: torch.Tensor, ray_dir: torch.Tensor, points: torch.Tensor,
voxelsize: float, n_max: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...
def svo_intersect(ray_start: torch.Tensor, ray_dir: torch.Tensor, points: torch.Tensor, children: torch.Tensor,
voxelsize: float, n_max: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...
def triangle_intersect(ray_start: torch.Tensor, ray_dir: torch.Tensor, face_points: torch.Tensor,
cagesize: float, blur: float, n_max: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...
# octree.h
def build_octree(center: torch.Tensor, points: torch.Tensor,
depth: int) -> tuple[torch.Tensor, torch.Tensor]: ...
# encode.h
def multires_hash_encode(levels: int, coarse_levels: int, x: torch.Tensor, res_list: torch.Tensor,
hash_table: torch.Tensor, hash_table_offsets: torch.Tensor) -> torch.Tensor: ...
def multires_hash_encode_with_grad(levels: int, coarse_levels: int, x: torch.Tensor, res_list: torch.Tensor,
hash_table: torch.Tensor, hash_table_offsets: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: ...
clib/include/cuda_utils.h
View file @ 1bc644a1
// Copyright (c) Facebook, Inc. and its affiliates.
//
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.
......@@ -9,38 +9,51 @@
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <cmath>
#include <vector>
#include <cuda.h>
#include <cuda_runtime.h>
#include <device_launch_parameters.h>
#include <vector>
#define TOTAL_THREADS 512
#define LOG2_TOTAL_THREADS 10
#define TOTAL_THREADS (2 << LOG2_TOTAL_THREADS)
inline int opt_n_threads(int work_size) {
const int pow_2 = std::log(static_cast<double>(work_size)) / std::log(2.0);
return max(min(1 << pow_2, TOTAL_THREADS), 1);
inline uint opt_n_threads(uint work_size) {
const uint pow_2 = std::log(work_size) / std::log(2.0);
return 1 << min(pow_2, LOG2_TOTAL_THREADS);
}
inline dim3 opt_block_config(int x, int y) {
const int x_threads = opt_n_threads(x);
const int y_threads =
max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1);
dim3 block_config(x_threads, y_threads, 1);
const int x_threads = opt_n_threads(x);
const int y_threads = max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1);
dim3 block_config(x_threads, y_threads, 1);
return block_config;
return block_config;
}
#define CUDA_CHECK_ERRORS() \
do { \
cudaError_t err = cudaGetLastError(); \
if (cudaSuccess != err) { \
fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \
cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \
__FILE__); \
exit(-1); \
} \
} while (0)
#define CUDA_CHECK_ERRORS() \
do { \
cudaError_t err = cudaGetLastError(); \
if (cudaSuccess != err) { \
fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \
cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, __FILE__); \
exit(-1); \
} \
} while (0)
#endif
template <typename T, uint N_ELEMS> struct vec {
__host__ __device__ T &operator[](uint idx) { return data[idx]; }
__host__ __device__ T operator[](uint idx) const { return data[idx]; }
T data[N_ELEMS];
static constexpr uint N = N_ELEMS;
};
template <uint N_FLOATS> using fvec = vec<float, N_FLOATS>;
template <uint N_HALFS> using hvec = vec<__half, N_HALFS>;
template <uint N_UINTS> using uvec = vec<uint, N_UINTS>;
clib/include/encode.h 0 → 100644
View file @ 1bc644a1
#pragma once
#include <torch/extension.h>
at::Tensor multires_hash_encode(const int levels, const int coarse_levels, at::Tensor x,
at::Tensor res_list, at::Tensor hash_table,
at::Tensor hash_table_offsets);
std::tuple<at::Tensor, at::Tensor, at::Tensor>
multires_hash_encode_with_grad(const int levels, const int coarse_levels, at::Tensor x,
at::Tensor res_list, at::Tensor hash_table,
at::Tensor hash_table_offsets);
\ No newline at end of file
clib/include/encode_debug.h 0 → 100644
View file @ 1bc644a1
#pragma once
#include <torch/extension.h>
std::tuple<at::Tensor, at::Tensor, at::Tensor>
multires_hash_encode_debug(const int levels, const int coarse_levels, at::Tensor x,
at::Tensor res_list, at::Tensor hash_table,
at::Tensor hash_table_offsets);
\ No newline at end of file
clib/include/utils.h
View file @ 1bc644a1
// Copyright (c) Facebook, Inc. and its affiliates.
//
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.
......@@ -7,24 +7,27 @@
#include <ATen/cuda/CUDAContext.h>
#include <torch/extension.h>
#define CHECK_CUDA(x) \
do { \
TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor"); \
} while (0)
#define CHECK_CUDA(x) \
do { \
TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor"); \
} while (0)
#define CHECK_CONTIGUOUS(x) \
do { \
TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \
} while (0)
#define CHECK_CONTIGUOUS(x) \
do { \
TORCH_CHECK(x.is_contiguous(), #x " must be a contiguous tensor"); \
} while (0)
#define CHECK_IS_INT(x) \
do { \
TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, \
#x " must be an int tensor"); \
} while (0)
#define CHECK_IS_INT(x) \
do { \
TORCH_CHECK(x.scalar_type() == at::ScalarType::Int, #x " must be an int tensor"); \
} while (0)
#define CHECK_IS_FLOAT(x) \
do { \
TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, \
#x " must be a float tensor"); \
} while (0)
#define CHECK_IS_FLOAT(x) \
do { \
TORCH_CHECK(x.scalar_type() == at::ScalarType::Float, #x " must be a float tensor"); \
} while (0)
#define CHECK_CUDA_CONT_TENSOR(__TYPE__, __VAR__) \
CHECK_CONTIGUOUS(__VAR__); \
CHECK_IS_##__TYPE__(__VAR__); \
CHECK_CUDA(__VAR__);
clib/src/binding.cpp
View file @ 1bc644a1
// Copyright (c) Facebook, Inc. and its affiliates.
//
//
// This source code is licensed under the MIT license found in the
// LICENSE file in the root directory of this source tree.
#include "intersect.h"
#include "octree.h"
#include "sample.h"
#include "encode.h"
#include "encode_debug.h"
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("ball_intersect", &ball_intersect);
m.def("aabb_intersect", &aabb_intersect);
m.def("svo_intersect", &svo_intersect);
m.def("triangle_intersect", &triangle_intersect);
m.def("ball_intersect", &ball_intersect);
m.def("aabb_intersect", &aabb_intersect);
m.def("svo_intersect", &svo_intersect);
m.def("triangle_intersect", &triangle_intersect);
m.def("uniform_ray_sampling", &uniform_ray_sampling);
m.def("inverse_cdf_sampling", &inverse_cdf_sampling);
m.def("uniform_ray_sampling", &uniform_ray_sampling);
m.def("inverse_cdf_sampling", &inverse_cdf_sampling);
m.def("build_octree", &build_octree);
m.def("build_octree", &build_octree);
m.def("multires_hash_encode", &multires_hash_encode);
m.def("multires_hash_encode_with_grad", &multires_hash_encode_with_grad);
m.def("multires_hash_encode_debug", &multires_hash_encode_debug);
}
\ No newline at end of file
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment