huayinshuke/kaolin
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

提交测试

Browse Source
This commit is contained in:
moyin
2024-01-16 17:22:21 +08:00
parent 92862c0372
commit 73635fda01
654 changed files with 178015 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
examples/recipes/README.md Normal file
View File

@@ -0,0 +1,18 @@
# Quick Start: Kaolin Recipes
<hr>
For a quick start with Kaolin, see the example snippets included below. <br>
In depth guides are available in the [tutorials](https://kaolin.readthedocs.io/en/latest/notes/tutorial_index.html) section.
## Data
### Converting Data
<hr>
* [Point cloud to SPC]("https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/recipes/dataload/spc_from_pointcloud.py")
## 3D Formats
### SPC / Octree based Ops
<hr>
* [SPC: Basic Usage]("https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/recipes/spc/spc_basics.py")

25
examples/recipes/camera/camera_coordinate_systems.py Normal file
View File

@@ -0,0 +1,25 @@
# ==============================================================================================================
# The following snippet demonstrates how to change the coordinate system of the camera.
# ==============================================================================================================
import math
import torch
import numpy as np
from kaolin.render.camera import Camera, blender_coords
device = 'cuda'
camera = Camera.from_args(
eye=torch.tensor([4.0, 4.0, 4.0]),
at=torch.tensor([0.0, 0.0, 0.0]),
up=torch.tensor([0.0, 1.0, 0.0]),
fov=30 * np.pi / 180, # In radians
width=800, height=800,
device=device
)
print(camera.basis_change_matrix)
camera.change_coordinate_system(blender_coords())
print(camera.basis_change_matrix)
camera.reset_coordinate_system()
print(camera.basis_change_matrix)

88
examples/recipes/camera/camera_init_explicit.py Normal file
View File

@@ -0,0 +1,88 @@
# ==============================================================================================================
# The following snippet demonstrates how to initialize instances of kaolin's pinhole / ortho cameras
# explicitly.
# Also review `camera_init_simple` which greatly simplifies the construction methods shown here.
# ==============================================================================================================
import math
import torch
from kaolin.render.camera import Camera, CameraExtrinsics, PinholeIntrinsics, OrthographicIntrinsics
#################################################################
# Camera 1: from eye, at, up and focal length (Perspective) #
#################################################################
# Build the camera extrinsics object from lookat
eye = torch.tensor([0.0, 0.0, -1.0], device='cuda') # Camera positioned here in world coords
at = torch.tensor([0.0, 0.0, 0.0], device='cuda') # Camera observing this world point
up = torch.tensor([0.0, 1.0, 0.0], device='cuda') # Camera up direction vector
extrinsics = CameraExtrinsics.from_lookat(eye, at, up)
# Build a pinhole camera's intrinsics. This time we use focal length (other useful args: focal_y, x0, y0)
intrinsics = PinholeIntrinsics.from_focal(width=800, height=600, focal_x=1.0, device='cuda')
# Combine extrinsics and intrinsics to obtain the full camera object
camera_1 = Camera(extrinsics=extrinsics, intrinsics=intrinsics)
print('--- Camera 1 ---')
print(camera_1)
########################################################################
# Camera 2: from camera position, orientation and fov (Perspective) #
########################################################################
# Build the camera extrinsics object from lookat
cam_pos = torch.tensor([0.0, 0.0, -1.0], device='cuda')
cam_dir = torch.tensor([[1.0, 0.0, 0.0],
[0.0, 1.0, 0.0],
[0.0, 0.0, 1.0]], device='cuda') # 3x3 orientation within the world
extrinsics = CameraExtrinsics.from_camera_pose(cam_pos=cam_pos, cam_dir=cam_dir)
# Use pinhole camera intrinsics, construct using field-of-view (other useful args: camera_fov_direction, x0, y0)
intrinsics = PinholeIntrinsics.from_fov(width=800, height=600, fov=math.radians(45.0), device='cuda')
camera_2 = Camera(extrinsics=extrinsics, intrinsics=intrinsics)
print('--- Camera 2 ---')
print(camera_2)
####################################################################
# Camera 3: camera view matrix, (Orthographic) #
####################################################################
# Build the camera extrinsics object from lookat
world2cam = torch.tensor([[1.0, 0.0, 0.0, 0.5],
[0.0, 1.0, 0.0, 0.5],
[0.0, 0.0, 1.0, 0.5],
[0.0, 0.0, 0.0, 1.0]], device='cuda') # 3x3 orientation within the world
extrinsics = CameraExtrinsics.from_view_matrix(view_matrix=world2cam)
# Use pinhole camera intrinsics, construct using field-of-view (other useful args: camera_fov_direction, x0, y0)
intrinsics = OrthographicIntrinsics.from_frustum(width=800, height=600, near=-800, far=800,
fov_distance=1.0, device='cuda')
camera_3 = Camera(extrinsics=extrinsics, intrinsics=intrinsics)
print('--- Camera 3 ---')
print(camera_3)
##########################################################
# Camera 4: Combining cameras #
##########################################################
# Must be of the same intrinsics type, and non params fields such as width, height, near, far
# (currently we don't perform validation)
camera_4 = Camera.cat((camera_1, camera_2))
print('--- Camera 4 ---')
print(camera_4)
##########################################################
# Camera 5: constructing a batch of cameras together #
##########################################################
# Extrinsics are created using batched tensors. The intrinsics are automatically broadcasted.
camera_5 = Camera.from_args(
eye=torch.tensor([[4.0, 4.0, 4.0], [4.0, 4.0, 4.0]]),
at=torch.tensor([[0.0, 0.0, 0.0], [4.0, 4.0, 4.0]]),
up=torch.tensor([[0.0, 1.0, 0.0], [4.0, 4.0, 4.0]]),
width=800, height=600, focal_x=300.0
)
print('--- Camera 5 ---')
print(camera_5)

65
examples/recipes/camera/camera_init_simple.py Normal file
View File

@@ -0,0 +1,65 @@
# ==============================================================================================================
# The following snippet demonstrates how to initialize instances of kaolin's pinhole / ortho cameras.
# ==============================================================================================================
import math
import torch
import numpy as np
from kaolin.render.camera import Camera
device = 'cuda'
perspective_camera_1 = Camera.from_args(
eye=torch.tensor([4.0, 4.0, 4.0]),
at=torch.tensor([0.0, 0.0, 0.0]),
up=torch.tensor([0.0, 1.0, 0.0]),
fov=30 * np.pi / 180, # In radians
x0=0.0, y0=0.0,
width=800, height=800,
near=1e-2, far=1e2,
dtype=torch.float64,
device=device
)
print('--- Perspective Camera 1 ---')
print(perspective_camera_1)
perspective_camera_2 = Camera.from_args(
eye=torch.tensor([4.0, 4.0, 4.0]),
at=torch.tensor([0.0, 0.0, 0.0]),
up=torch.tensor([0.0, 1.0, 0.0]),
fov=30 * np.pi / 180, # In radians
width=800, height=800,
device=device
)
print('--- Perspective Camera 2 ---')
print(perspective_camera_2)
ortho_camera_1 = Camera.from_args(
eye=torch.tensor([4.0, 4.0, 4.0]),
at=torch.tensor([0.0, 0.0, 0.0]),
up=torch.tensor([0.0, 1.0, 0.0]),
width=800, height=800,
near=-800, far=800,
fov_distance=1.0,
dtype=torch.float64,
device=device
)
print('--- Orthographic Camera 1 ---')
print(ortho_camera_1)
ortho_camera_2 = Camera.from_args(
view_matrix=torch.tensor([[1.0, 0.0, 0.0, 0.5],
[0.0, 1.0, 0.0, 0.5],
[0.0, 0.0, 1.0, 0.5],
[0.0, 0.0, 0.0, 1.0]]),
width=800, height=800,
dtype=torch.float64,
device=device
)
print('--- Orthographic Camera 2 ---')
print(ortho_camera_2)

27
examples/recipes/camera/camera_movement.py Normal file
View File

@@ -0,0 +1,27 @@
# ==============================================================================================================
# The following snippet demonstrates how to manipulate kaolin's camera.
# ==============================================================================================================
import torch
from kaolin.render.camera import Camera
camera = Camera.from_args(
eye=torch.tensor([0.0, 0.0, -1.0]),
at=torch.tensor([0.0, 0.0, 0.0]),
up=torch.tensor([0.0, 1.0, 0.0]),
width=800, height=600,
fov=1.0,
device='cuda'
)
# Extrinsic rigid transformations managed by CameraExtrinsics
camera.move_forward(amount=10.0) # Translate forward in world coordinates (this is wisp's mouse zoom)
camera.move_right(amount=-5.0) # Translate left in world coordinates
camera.move_up(amount=5.0) # Translate up in world coordinates
camera.rotate(yaw=0.1, pitch=0.02, roll=1.0) # Rotate the camera
# Intrinsic lens transformations managed by CameraIntrinsics
# Zoom in to decrease field of view - for Orthographic projection the internal implementation differs
# as there is no acual fov or depth concept (hence we use a "made up" fov distance parameter, see the projection matrix)
camera.zoom(amount=0.5)

57
examples/recipes/camera/camera_opengl_shaders.py Normal file
View File

@@ -0,0 +1,57 @@
# ==============================================================================================================
# The following snippet demonstrates how to use the camera for generating a view-projection matrix
# as used in opengl shaders.
# ==============================================================================================================
import torch
import numpy as np
from kaolin.render.camera import Camera
# !!! This example is not runnable -- it is minimal to contain integration between the opengl shader and !!!
# !!! the camera matrix !!!
try:
from glumpy import gloo
except:
class DummyGloo(object):
def Program(self, vertex, fragment):
# see: https://glumpy.readthedocs.io/en/latest/api/gloo-shader.html#glumpy.gloo.Program
return dict([])
gloo = DummyGloo()
device = 'cuda'
camera = Camera.from_args(
eye=torch.tensor([4.0, 4.0, 4.0]),
at=torch.tensor([0.0, 0.0, 0.0]),
up=torch.tensor([0.0, 1.0, 0.0]),
fov=30 * np.pi / 180, # In radians
x0=0.0, y0=0.0,
width=800, height=800,
near=1e-2, far=1e2,
dtype=torch.float64,
device=device
)
vertex = """
uniform mat4 u_viewprojection;
attribute vec3 position;
attribute vec4 color;
varying vec4 v_color;
void main()
{
v_color = color;
gl_Position = u_viewprojection * vec4(position, 1.0f);
} """
fragment = """
varying vec4 v_color;
void main()
{
gl_FragColor = v_color;
} """
# Compile GL program
gl_program = gloo.Program(vertex, fragment)
gl_program["u_viewprojection"] = camera.view_projection_matrix()[0].cpu().numpy().T

47
examples/recipes/camera/camera_properties.py Normal file
View File

@@ -0,0 +1,47 @@
# ==============================================================================================================
# The following snippet demonstrates various camera properties
# ==============================================================================================================
import math
import torch
import numpy as np
from kaolin.render.camera import Camera
device = 'cuda'
camera = Camera.from_args(
eye=torch.tensor([4.0, 4.0, 4.0]),
at=torch.tensor([0.0, 0.0, 0.0]),
up=torch.tensor([0.0, 1.0, 0.0]),
fov=30 * np.pi / 180, # In radians
width=800, height=800,
dtype=torch.float32,
device=device
)
print(camera.width)
print(camera.height)
print(camera.lens_type)
print(camera.device)
camera = camera.cpu()
print(camera.device)
# Create a batched camera and view single element
camera = Camera.cat((camera, camera))
print(camera)
camera = camera[0]
print(camera)
print(camera.dtype)
camera = camera.half()
print(camera.dtype)
camera = camera.double()
print(camera.dtype)
camera = camera.float()
print(camera.dtype)
print(camera.extrinsics.requires_grad)
print(camera.intrinsics.requires_grad)
print(camera.to('cuda', torch.float64))

71
examples/recipes/camera/camera_ray_tracing.py Normal file
View File

@@ -0,0 +1,71 @@
# ==============================================================================================================
# The following snippet demonstrates how to use the camera for implementing a ray-generation function
# for ray based applications.
# ==============================================================================================================
import torch
import numpy as np
from typing import Tuple
from kaolin.render.camera import Camera, CameraFOV
def generate_pixel_grid(res_x=None, res_y=None, device='cuda'):
h_coords = torch.arange(res_x, device=device)
w_coords = torch.arange(res_y, device=device)
pixel_y, pixel_x = torch.meshgrid(h_coords, w_coords)
pixel_x = pixel_x + 0.5
pixel_y = pixel_y + 0.5
return pixel_y, pixel_x
def generate_perspective_rays(camera: Camera, pixel_grid: Tuple[torch.Tensor, torch.Tensor]):
# coords_grid should remain immutable (a new tensor is implicitly created here)
pixel_y, pixel_x = pixel_grid
pixel_x = pixel_x.to(camera.device, camera.dtype)
pixel_y = pixel_y.to(camera.device, camera.dtype)
# Account for principal point offset from canvas center
pixel_x = pixel_x - camera.x0
pixel_y = pixel_y + camera.y0
# pixel values are now in range [-1, 1], both tensors are of shape res_y x res_x
pixel_x = 2 * (pixel_x / camera.width) - 1.0
pixel_y = 2 * (pixel_y / camera.height) - 1.0
ray_dir = torch.stack((pixel_x * camera.tan_half_fov(CameraFOV.HORIZONTAL),
-pixel_y * camera.tan_half_fov(CameraFOV.VERTICAL),
-torch.ones_like(pixel_x)), dim=-1)
ray_dir = ray_dir.reshape(-1, 3) # Flatten grid rays to 1D array
ray_orig = torch.zeros_like(ray_dir)
# Transform from camera to world coordinates
ray_orig, ray_dir = camera.extrinsics.inv_transform_rays(ray_orig, ray_dir)
ray_dir /= torch.linalg.norm(ray_dir, dim=-1, keepdim=True)
ray_orig, ray_dir = ray_orig[0], ray_dir[0] # Assume a single camera
return ray_orig, ray_dir, camera.near, camera.far
camera = Camera.from_args(
eye=torch.tensor([4.0, 4.0, 4.0]),
at=torch.tensor([0.0, 0.0, 0.0]),
up=torch.tensor([0.0, 1.0, 0.0]),
fov=30 * np.pi / 180, # In radians
x0=0.0, y0=0.0,
width=800, height=800,
near=1e-2, far=1e2,
dtype=torch.float64,
device='cuda'
)
pixel_grid = generate_pixel_grid(200, 200)
ray_orig, ray_dir, near, far = generate_perspective_rays(camera, pixel_grid)
print('Ray origins:')
print(ray_orig)
print('Ray directions:')
print(ray_dir)
print('Near clipping plane:')
print(near)
print('Far clipping plane:')
print(far)

59
examples/recipes/camera/camera_transforms.py Normal file
View File

@@ -0,0 +1,59 @@
# ==============================================================================================================
# The following snippet demonstrates how to use the camera transform directly on vectors
# ==============================================================================================================
import math
import torch
import numpy as np
from kaolin.render.camera import Camera
device = 'cuda'
camera = Camera.from_args(
eye=torch.tensor([4.0, 4.0, 4.0]),
at=torch.tensor([0.0, 0.0, 0.0]),
up=torch.tensor([0.0, 1.0, 0.0]),
fov=30 * np.pi / 180, # In radians
width=800, height=800,
dtype=torch.float32,
device=device
)
print('View projection matrix')
print(camera.view_projection_matrix())
print('View matrix: world2cam')
print(camera.view_matrix())
print('Inv View matrix: cam2world')
print(camera.inv_view_matrix())
print('Projection matrix')
print(camera.projection_matrix())
vectors = torch.randn(10, 3).to(camera.device, camera.dtype) # Create a batch of points
# For ortho and perspective: this is equivalent to multiplying camera.projection_matrix() @ vectors
# and then dividing by the w coordinate (perspective division)
print(camera.transform(vectors))
# For ray tracing we have camera.inv_transform_rays which performs multiplication with inv_view_matrix()
# (just for the extrinsics part)
# Can also access properties directly:
# --
# View matrix components (camera space)
print(camera.R)
print(camera.t)
# Camera axes and position in world coordinates
print(camera.cam_pos())
print(camera.cam_right())
print(camera.cam_pos())
print(camera.cam_forward())
print(camera.focal_x)
print(camera.focal_y)
print(camera.x0)
print(camera.y0)

65
examples/recipes/camera/cameras_differentiable.py Normal file
View File

@@ -0,0 +1,65 @@
# ====================================================================================================================
# The following snippet demonstrates how cameras can be used for optimizing specific extrinsic / intrinsic parameters
# ====================================================================================================================
import torch
import torch.optim as optim
from kaolin.render.camera import Camera
# Create simple perspective camera
cam = Camera.from_args(
eye=torch.tensor([4.0, 4.0, 4.0]),
at=torch.tensor([0.0, 0.0, 0.0]),
up=torch.tensor([0.0, 1.0, 0.0]),
width=800, height=600, focal_x=300.0
)
# When requires_grad is on, the camera will automatically switch to differentiation friendly backend
# (implicitly calling cam.switch_backend('matrix_6dof_rotation') )
cam.requires_grad_(True)
# Constraint camera to optimize only fov and camera position (cannot rotate)
ext_mask, int_mask = cam.gradient_mask('t', 'focal_x', 'focal_y')
ext_params, int_params = cam.parameters()
ext_params.register_hook(lambda grad: grad * ext_mask.float())
grad_scale = 1e5 # Used to move the projection matrix elements faster
int_params.register_hook(lambda grad: grad * int_mask.float() * grad_scale)
# Make the camera a bit noisy
# Currently can't copy the camera here after requires_grad is true because we're still missing a camera.detach() op
target = Camera.from_args(
eye=torch.tensor([4.0, 4.0, 4.0]),
at=torch.tensor([0.0, 0.0, 0.0]),
up=torch.tensor([0.0, 1.0, 0.0]),
width=800, height=600, focal_x=300.0
)
target.t = target.t + torch.randn_like(target.t)
target.focal_x = target.focal_x + torch.randn_like(target.focal_x)
target.focal_y = target.focal_y + torch.randn_like(target.focal_y)
target_mat = target.view_projection_matrix()
# Save for later so we have some comparison of what changed
initial_view = cam.view_matrix().detach().clone()
initial_proj = cam.projection_matrix().detach().clone()
# Train a few steps
optimizer = optim.SGD(cam.parameters(), lr=0.1, momentum=0.9)
for idx in range(10):
view_proj = cam.view_projection_matrix()
optimizer.zero_grad()
loss = torch.nn.functional.mse_loss(target_mat, view_proj)
loss.backward()
optimizer.step()
print(f'Iteration {idx}:')
print(f'Loss: {loss.item()}')
print(f'Extrinsics: {cam.extrinsics.parameters()}')
print(f'Intrinsics: {cam.intrinsics.parameters()}')
# Projection matrix grads are much smaller as they're scaled by the view-frustum dimensions..
print(f'View matrix before: {initial_view}')
print(f'View matrix after: {cam.view_matrix()}')
print(f'Projection matrix before: {initial_proj}')
print(f'Projection matrix after: {cam.projection_matrix()}')
print('Did the camera change?')
print(not torch.allclose(cam, target))

52
examples/recipes/dataload/spc_from_pointcloud.py Normal file
View File

@@ -0,0 +1,52 @@
# ==============================================================================================================
# The following snippet demonstrates how to build kaolin's compressed octree,
# "Structured Point Cloud (SPC)", from raw point cloud data.
# ==============================================================================================================
# See also:
#
# - Tutorial: Understanding Structured Point Clouds (SPCs)
# https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/understanding_spcs_tutorial.ipynb
#
# - Documentation: Structured Point Clouds
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.spc.html?highlight=spc#kaolin-ops-spc
# ==============================================================================================================
import torch
import kaolin
# Create some point data with features
# Point coordinates are expected to be normalized to the range [-1, 1].
points = torch.tensor([
[-1.0, -1.0, -1.0],
[-0.9, -0.95, -1.0],
[1.0, 0.0, 0.0],
[0.0, -0.1, 0.3],
[1.0, 1.0, 1.0]
], device='cuda')
features = torch.tensor([
[0.1, 1.1, 2.1],
[0.2, 1.2, 2.2],
[0.3, 1.3, 2.3],
[0.4, 1.4, 2.4],
[0.5, 1.5, 2.5],
], device='cuda')
# Structured Point Cloud will be using 3 levels of detail
level = 3
# In kaolin, operations are batched by default
# Here, in contrast, we use a single point cloud and therefore invoke an unbatched conversion function.
# For more information about batched operations, see:
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.batch.html#kaolin-ops-batch
spc = kaolin.ops.conversions.pointcloud.unbatched_pointcloud_to_spc(pointcloud=points,
level=level,
features=features)
# SPC is an object which keep tracks of the various octree component
print(spc)
print(f'SPC keeps track of the following cells in {level} levels of detail (parents + leaves):\n'
f' {spc.point_hierarchies}\n')
# Note that the point cloud coordinates are quantized to integer coordinates.
# During conversion, when points fall within the same cell, their features are averaged
print(f'Features for leaf cells:\n {spc.features}')

140
examples/recipes/preprocess/fast_mesh_sampling.py Normal file
View File

@@ -0,0 +1,140 @@
# ==============================================================================================================
# The following snippet shows how to use kaolin to preprocess a shapenet dataset
# To quickly sample point clouds from the mesh at runtime
# ==============================================================================================================
# See also:
# - Documentation: ShapeNet dataset
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.io.shapenet.html#kaolin.io.shapenet.ShapeNetV2
# - Documentation: CachedDataset
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.io.dataset.html#kaolin.io.dataset.CachedDataset
# - Documentation: Mesh Ops:
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.mesh.html
# - Documentation: Obj loading:
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.io.obj.html
# ==============================================================================================================
import argparse
import os
import torch
import kaolin as kal
parser = argparse.ArgumentParser(description='')
parser.add_argument('--shapenet-dir', type=str, default=os.getenv('KAOLIN_TEST_SHAPENETV2_PATH'),
help='Path to shapenet (v2)')
parser.add_argument('--cache-dir', type=str, default='/tmp/dir',
help='Path where output of the dataset is cached')
parser.add_argument('--num-samples', type=int, default=10,
help='Number of points to sample on the mesh')
parser.add_argument('--cache-at-runtime', action='store_true',
help='run the preprocessing lazily')
parser.add_argument('--num-workers', type=int, default=0,
help='Number of workers during preprocessing (not used with --cache-at-runtime)')
args = parser.parse_args()
def preprocessing_transform(inputs):
"""This the transform used in shapenet dataset __getitem__.
Three tasks are done:
1) Get the areas of each faces, so it can be used to sample points
2) Get a proper list of RGB diffuse map
3) Get the material associated to each face
"""
mesh = inputs['mesh']
vertices = mesh.vertices.unsqueeze(0)
faces = mesh.faces
# Some materials don't contain an RGB texture map, so we are considering the single value
# to be a single pixel texture map (1, 3, 1, 1)
# we apply a modulo 1 on the UVs because ShapeNet follows GL_REPEAT behavior (see: https://open.gl/textures)
uvs = torch.nn.functional.pad(mesh.uvs.unsqueeze(0) % 1, (0, 0, 0, 1)) * 2. - 1.
uvs[:, :, 1] = -uvs[:, :, 1]
face_uvs_idx = mesh.face_uvs_idx
face_material_idx = mesh.material_assignments
materials = [m['map_Kd'].permute(2, 0, 1).unsqueeze(0).float() / 255. if 'map_Kd' in m else
m['Kd'].reshape(1, 3, 1, 1)
for m in mesh.materials]
mask = face_uvs_idx == -1
face_uvs_idx[mask] = 0
face_uvs = kal.ops.mesh.index_vertices_by_faces(
uvs, face_uvs_idx
)
face_uvs[:, mask] = 0.
outputs = {
'vertices': vertices,
'faces': faces,
'face_areas': kal.ops.mesh.face_areas(vertices, faces),
'face_uvs': face_uvs,
'materials': materials,
'face_material_idx': face_material_idx,
'name': inputs['name']
}
return outputs
class SamplePointsTransform(object):
def __init__(self, num_samples):
self.num_samples = num_samples
def __call__(self, inputs):
coords, face_idx, feature_uvs = kal.ops.mesh.sample_points(
inputs['vertices'],
inputs['faces'],
num_samples=self.num_samples,
areas=inputs['face_areas'],
face_features=inputs['face_uvs']
)
coords = coords.squeeze(0)
face_idx = face_idx.squeeze(0)
feature_uvs = feature_uvs.squeeze(0)
# Interpolate the RGB values from the texture map
point_materials_idx = inputs['face_material_idx'][face_idx]
all_point_colors = torch.zeros((self.num_samples, 3))
for i, material in enumerate(inputs['materials']):
mask = point_materials_idx == i
point_color = torch.nn.functional.grid_sample(
material,
feature_uvs[mask].reshape(1, 1, -1, 2),
mode='bilinear',
align_corners=False,
padding_mode='border')
all_point_colors[mask] = point_color[0, :, 0, :].permute(1, 0)
outputs = {
'coords': coords,
'face_idx': face_idx,
'colors': all_point_colors,
'name': inputs['name']
}
return outputs
# Make ShapeNet dataset with preprocessing transform
ds = kal.io.shapenet.ShapeNetV2(root=args.shapenet_dir,
categories=['dishwasher'],
train=True,
split=0.1,
with_materials=True,
output_dict=True,
transform=preprocessing_transform)
# Cache the result of the preprocessing transform
# and apply the sampling at runtime
pc_ds = kal.io.dataset.CachedDataset(ds,
cache_dir=args.cache_dir,
save_on_disk=True,
num_workers=args.num_workers,
transform=SamplePointsTransform(args.num_samples),
cache_at_runtime=args.cache_at_runtime,
force_overwrite=True)
for data in pc_ds:
print("coords:\n", data['coords'])
print("face_idx:\n", data['face_idx'])
print("colors:\n", data['colors'])
print("name:\n", data['name'])

48
examples/recipes/preprocess/occupancy_sampling.py Normal file
View File

@@ -0,0 +1,48 @@
# ==============================================================================================================
# The following snippet shows how to use kaolin to test sampled values of an occupancy function
# against a watertight mesh.
# ==============================================================================================================
# See also:
# - Documentation: Triangular meshes
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.mesh.html#triangular-meshes
# ==============================================================================================================
import os
import torch
import kaolin
FILE_DIR = os.path.abspath(os.path.dirname(os.path.abspath(__file__)))
mesh_path = os.path.join(FILE_DIR, os.pardir, os.pardir, "samples", "sphere.obj") # Path to some .obj file with textures
num_samples = 100000 # Number of sample points
# 1. Load a watertight mesh from obj file
mesh = kaolin.io.obj.import_mesh(mesh_path)
print(f'Loaded mesh with {len(mesh.vertices)} vertices and {len(mesh.faces)} faces.')
# 2. Preprocess mesh:
# Move tensors to CUDA device
vertices = mesh.vertices.cuda()
faces = mesh.faces.cuda()
# Kaolin assumes an exact batch format, we make sure to convert from: (V, 3) to (1, V, 3), where 1 is the batch size
vertices = vertices.unsqueeze(0)
# 3. Sample random points uniformly in space, from the bounding box of the mesh + 10% margin
min_bound, _ = vertices.min(dim=1)
max_bound, _ = vertices.max(dim=1)
margin = (max_bound - min_bound) * 0.1
max_bound += margin
min_bound -= margin
occupancy_coords = (max_bound - min_bound) * torch.rand(1, num_samples, 3, device='cuda') + min_bound
# 4. Calculate occupancy value
occupancy_value = kaolin.ops.mesh.check_sign(vertices, faces, occupancy_coords)
# Unbatch to obtain a torch.Tensor of (V, 3) and (V, 1)
occupancy_coords = occupancy_coords.squeeze(0)
occupancy_value = occupancy_value.squeeze(0)
percent_in_mesh = torch.count_nonzero(occupancy_value) / len(occupancy_value)
print(f'Sampled a tensor of points uniformly in space '
f'with {occupancy_coords.shape[0]} points of {occupancy_coords.shape[1]}D coordinates.')
print(f'{"{:.3f}".format(percent_in_mesh)}% of the sampled points are inside the mesh volume.')

54
examples/recipes/spc/spc_basics.py Normal file
View File

@@ -0,0 +1,54 @@
# ==============================================================================================================
# The following snippet demonstrates the basic usage of kaolin's compressed octree,
# termed "Structured Point Cloud (SPC)".
# Note this is a low level structure: practitioners are encouraged to visit the references below.
# ==============================================================================================================
# See also:
#
# - Code: kaolin.ops.spc.SPC
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.rep.html?highlight=SPC#kaolin.rep.Spc
#
# - Tutorial: Understanding Structured Point Clouds (SPCs)
# https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/understanding_spcs_tutorial.ipynb
#
# - Documentation: Structured Point Clouds
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.spc.html?highlight=spc#kaolin-ops-spc
# ==============================================================================================================
import torch
import kaolin
# Construct SPC from some points data. Point coordinates are expected to be normalized to the range [-1, 1].
points = torch.tensor([[-1.0, -1.0, -1.0], [-0.9, -0.95, -1.0], [1.0, 1.0, 1.0]], device='cuda')
# In kaolin, operations are batched by default
# Here, in contrast, we use a single point cloud and therefore invoke an unbatched conversion function.
# The Structured Point Cloud will be using 3 levels of detail
spc = kaolin.ops.conversions.pointcloud.unbatched_pointcloud_to_spc(pointcloud=points, level=3)
# SPC is a batched object, and most of its fields are packed.
# (see: https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.batch.html#kaolin-ops-batch )
# spc.length defines the boundaries between different batched SPC instances the same object holds.
# Here we keep track of a single entry batch, which has 8 octree non-leaf cells.
print(f'spc.batch_size: {spc.batch_size}')
print(f'spc.lengths (cells per batch entry): {spc.lengths}')
# SPC is hierarchical and keeps information for every level of detail from 0 to 3.
# spc.point_hierarchies keeps the sparse, zero indexed coordinates of each occupied cell, per level.
print(f'SPC keeps track of total of {spc.point_hierarchies.shape[0]} parent + leaf cells:')
# To separate the boundaries, the spc.pyramids field is used.
# This field is not-packed, unlike the other SPC fields.
pyramid_of_first_entry_in_batch = spc.pyramids[0]
cells_per_level = pyramid_of_first_entry_in_batch[0]
cumulative_cells_per_level = pyramid_of_first_entry_in_batch[1]
for i, lvl_cells in enumerate(cells_per_level[:-1]):
print(f'LOD #{i} has {lvl_cells} cells.')
# The spc.octrees field keeps track of the fundamental occupancy information of each cell in the octree.
print('The occupancy of each octant parent cell, in Morton / Z-curve order is:')
print(['{0:08b}'.format(octree_byte) for octree_byte in spc.octrees])
# Since SPCs are low level objects, they require bookkeeping of multiple fields.
# For ease of use, these fields are collected and tracked within a single class: kaolin.ops.spc.SPC
# See references at the header for elaborate information on how to use this object.

113
examples/recipes/spc/spc_conv3d_example.py Normal file
View File

@@ -0,0 +1,113 @@
# ==============================================================================================================
# The following code demonstrates the usage of kaolin's "Structured Point Cloud (SPC)" 3d convolution
# functionality. Note that this sample does NOT demonstrate how to use Kaolin's Pytorch 3d convolution layers.
# Rather, 3d convolutions are used to 'filter' color data useful for level-of-detail management during
# rendering. This can be thought of as the 3d analog of generating a 2d mipmap.
#
# Note this is a low level interface: practitioners are encouraged to visit the references below.
# ==============================================================================================================
# See also:
#
# - Code: kaolin.ops.spc.SPC
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.rep.html?highlight=SPC#kaolin.rep.Spc
#
# - Tutorial: Understanding Structured Point Clouds (SPCs)
# https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/understanding_spcs_tutorial.ipynb
#
# - Documentation: Structured Point Clouds
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.spc.html?highlight=spc#kaolin-ops-spc
# ==============================================================================================================
import torch
import kaolin
# The following function applies a series of SPC convolutions to encode the entire hierarchy into a single tensor.
# Each step applies a convolution on the "highest" level of the SPC with some averaging kernel.
# Therefore, each step locally averages the "colored point hierarchy", where each "colored point"
# corresponds to a point in the SPC point hierarchy.
# For a description of inputs 'octree', 'point_hierachy', 'level', 'pyramids', and 'exsum', as well a
# detailed description of the mathematics of SPC convolutions, see:
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.spc.html?highlight=SPC#kaolin.ops.spc.Conv3d
# The input 'color' is Pytorch tensor containing color features corresponding to some 'level' of the hierarchy.
def encode(colors, octree, point_hierachy, pyramids, exsum, level):
# SPC convolutions are characterized by a set of 'kernel vectors' and corresponding 'weights'.
# kernel_vectors is the "kernel support" -
# a listing of 3D coordinates where the weights of the convolution are non-null,
# in this case a it's a simple dense 2x2x2 grid.
kernel_vectors = torch.tensor([[0,0,0],[0,0,1],[0,1,0],[0,1,1],
[1,0,0],[1,0,1],[1,1,0],[1,1,1]],
dtype=torch.short, device='cuda')
# The weights specify how the input colors 'under' the kernel are mapped to an output color,
# in this case a simple average.
weights = torch.diag(torch.tensor([0.125, 0.125, 0.125, 0.125],
dtype=torch.float32, device='cuda')) # Tensor of (4, 4)
weights = weights.repeat(8,1,1).contiguous() # Tensor of (8, 4, 4)
# Storage for the output color hierarchy is allocated. This includes points at the bottom of the hierarchy,
# as well as intermediate SPC levels (which may store different features)
color_hierarchy = torch.empty((pyramids[0,1,level+1],4), dtype=torch.float32, device='cuda')
# Copy the input colors into the highest level of color_hierarchy. pyramids is used here to select all leaf
# points at the bottom of the hierarchy and set them to some pre-sampled random color. Points at intermediate
# levels are left empty.
color_hierarchy[pyramids[0,1,level]:pyramids[0,1,level+1]] = colors[:]
# Performs the 3d convolutions in a bottom up fashion to 'filter' colors from the previous level
for l in range(level,0,-1):
# Apply the 3d convolution. Note that jump=1 means the inputs and outputs differ by 1 level
# This is analogous to to a stride=2 in grid based convolutions
colors, ll = kaolin.ops.spc.conv3d(octree,
point_hierachy,
l,
pyramids,
exsum,
colors,
weights,
kernel_vectors,
jump=1)
# Copy the output colors into the color hierarchy
color_hierarchy[pyramids[0,1,ll]:pyramids[0,1,l]] = colors[:]
print(f"At level {l}, output feature shape is:\n{colors.shape}")
# Normalize the colors.
color_hierarchy /= color_hierarchy[:,3:]
# Normalization is needed here due to the sparse nature of SPCs. When a point under a kernel is not
# present in the point hierarchy, the corresponding data is treated as zeros. Normalization is equivalent
# to having the filter weights sum to one. This may not always be desirable, e.g. alpha blending.
return color_hierarchy
# Highest level of SPC
level = 3
# Construct a fully occupied Structured Point Cloud with N levels of detail
# See https://kaolin.readthedocs.io/en/latest/modules/kaolin.rep.html?highlight=SPC#kaolin.rep.Spc
spc = kaolin.rep.Spc.make_dense(level, device='cuda')
# In kaolin, operations are batched by default, the spc object above contains a single item batch, hence [0]
num_points_last_lod = spc.num_points(level)[0]
# Create tensor of random colors for all points in the highest level of detail
colors = torch.rand((num_points_last_lod, 4), dtype=torch.float32, device='cuda')
# Set 4th color channel to one for subsequent color normalization
colors[:,3] = 1
print(f'Input SPC features: {colors.shape}')
# Encode color hierarchy by invoking a series of convolutions, until we end up with a single tensor.
color_hierarchy = encode(colors=colors,
octree=spc.octrees,
point_hierachy=spc.point_hierarchies,
pyramids=spc.pyramids,
exsum=spc.exsum,
level=level)
# Print root node color
print(f'Final encoded value (average of averages):')
print(color_hierarchy[0])
# This will be the average of averages, over the entire spc hierarchy. Since the initial random colors
# came from a uniform distribution, this should approach [0.5, 0.5, 0.5, 1.0] as 'level' increases

73
examples/recipes/spc/spc_dual_octree.py Normal file
View File

@@ -0,0 +1,73 @@
# ==============================================================================================================
# The following snippet demonstrates the basic usage of kaolin's dual octree, an octree which keeps features
# at the 8 corners of each cell (the primary octree keeps a single feature at each cell center).
# The implementation is realized through kaolin's "Structured Point Cloud (SPC)".
# Note this is a low level structure: practitioners are encouraged to visit the references below.
# ==============================================================================================================
# See also:
#
# - Code: kaolin.ops.spc.SPC
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.rep.html?highlight=SPC#kaolin.rep.Spc
#
# - Tutorial: Understanding Structured Point Clouds (SPCs)
# https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/understanding_spcs_tutorial.ipynb
#
# - Documentation: Structured Point Clouds
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.spc.html?highlight=spc#kaolin-ops-spc
# ==============================================================================================================
import torch
import kaolin
# Construct SPC from some points data. Point coordinates are expected to be normalized to the range [-1, 1].
# To keep the example readable, by default we set the SPC level to 1: root + 8 cells
# (note that with a single LOD, only 2 cells should be occupied due to quantization)
level = 1
points = torch.tensor([[-1.0, -1.0, -1.0], [-0.9, -0.95, -1.0], [1.0, 1.0, 1.0]], device='cuda')
spc = kaolin.ops.conversions.pointcloud.unbatched_pointcloud_to_spc(pointcloud=points, level=level)
# Construct the dual octree with an unbatched operation, each cell is now converted to 8 corners
# More info about batched / packed tensors at:
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.batch.html#kaolin-ops-batch
pyramid = spc.pyramids[0] # The pyramids field is batched, we select the singleton entry, #0
point_hierarchy = spc.point_hierarchies # point_hierarchies is a packed tensor, so no need to unbatch
point_hierarchy_dual, pyramid_dual = kaolin.ops.spc.unbatched_make_dual(point_hierarchy=point_hierarchy,
pyramid=pyramid)
# Let's compare the primary and dual octrees.
# The function 'unbatched_get_level_points' yields a tensor which lists all points / sparse cell coordinates occupied
# at a certain level.
# [Primary octree] [Dual octree]
# . . . . . . . . X . . .X. . . X
# | . X . X | . | . . | .
# | . . . . . . . . ===> | X . . X . . . X
# | | . X . | X . X | . . | .
# | | . . . . . . . . | | X . . .X. . . X
# | | | | | | X | | |
# . .|. . | . . . | ===> X .|. . X . . X |
# .| X |. X . | .| |. . X
# . . | . . . . . | X . | . X . . X |
# . | X . X . | . | . . |
# . . . . . . . . X . . X . . . X
#
primary_lod0 = kaolin.ops.spc.unbatched_get_level_points(point_hierarchy, pyramid, level=0)
primary_lod1 = kaolin.ops.spc.unbatched_get_level_points(point_hierarchy, pyramid, level=1)
dual_lod0 = kaolin.ops.spc.unbatched_get_level_points(point_hierarchy_dual, pyramid_dual, level=0)
dual_lod1 = kaolin.ops.spc.unbatched_get_level_points(point_hierarchy_dual, pyramid_dual, level=1)
print(f'Primary octree: Level 0 (root cells): \n{primary_lod0}')
print(f'Dual octree: Level 0 (root corners): \n{dual_lod0}')
print(f'Primary octree: Level 1 (cells): \n{primary_lod1}')
print(f'Dual octree: Level 1 (corners): \n{dual_lod1}')
# kaolin allows for interchangeable usage of the primary and dual octrees.
# First we have to create a mapping between them:
trinkets, _ = kaolin.ops.spc.unbatched_make_trinkets(point_hierarchy, pyramid, point_hierarchy_dual, pyramid_dual)
# trinkets are indirection pointers (in practice, indices) from the nodes of the primary octree
# to the nodes of the dual octree. The nodes of the dual octree represent the corners of the voxels
# defined by the primary octree.
print(f'point_hierarchy is of shape {point_hierarchy.shape}')
print(f'point_hierarchy_dual is of shape {point_hierarchy_dual.shape}')
print(f'trinkets is of shape {trinkets.shape}')
print(f'Trinket indices are multilevel: {trinkets}')
# See also spc_trilinear_interp.py for a practical application which uses the dual octree & trinkets

69
examples/recipes/spc/spc_trilinear_interp.py Normal file
View File

@@ -0,0 +1,69 @@
# ==============================================================================================================
# The following snippet demonstrates the basic usage of kaolin's dual octree, an octree which keeps features
# at the 8 corners of each cell (the primary octree keeps a single feature at each cell center).
# In this example we sample an interpolated value according to the 8 corners of a cell.
# The implementation is realized through kaolin's "Structured Point Cloud (SPC)".
# Note this is a low level structure: practitioners are encouraged to visit the references below.
# ==============================================================================================================
# See also:
#
# - Code: kaolin.ops.spc.SPC
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.rep.html?highlight=SPC#kaolin.rep.Spc
#
# - Tutorial: Understanding Structured Point Clouds (SPCs)
# https://github.com/NVIDIAGameWorks/kaolin/blob/master/examples/tutorial/understanding_spcs_tutorial.ipynb
#
# - Documentation: Structured Point Clouds
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.spc.html?highlight=spc#kaolin-ops-spc
# ==============================================================================================================
import torch
import kaolin
# Construct SPC from some points data. Point coordinates are expected to be normalized to the range [-1, 1].
# To keep the example readable, by default we set the SPC level to 1: root + 8 cells
# (note that with a single LOD, only 2 cells should be occupied due to quantization)
level = 1
points = torch.tensor([[-1.0, -1.0, -1.0], [-0.9, -0.95, -1.0], [1.0, 1.0, 1.0]], device='cuda')
spc = kaolin.ops.conversions.pointcloud.unbatched_pointcloud_to_spc(pointcloud=points, level=level)
# Construct the dual octree with an unbatched operation, each cell is now converted to 8 corners
# More info about batched / packed tensors at:
# https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.batch.html#kaolin-ops-batch
pyramid = spc.pyramids[0] # The pyramids field is batched, we select the singleton entry, #0
point_hierarchy = spc.point_hierarchies # point_hierarchies is a packed tensor, so no need to unbatch
point_hierarchy_dual, pyramid_dual = kaolin.ops.spc.unbatched_make_dual(point_hierarchy=point_hierarchy,
pyramid=pyramid)
# kaolin allows for interchangeable usage of the primary and dual octrees via the "trinkets" mapping
# trinkets are indirection pointers (in practice, indices) from the nodes of the primary octree
# to the nodes of the dual octree. The nodes of the dual octree represent the corners of the voxels
# defined by the primary octree.
trinkets, _ = kaolin.ops.spc.unbatched_make_trinkets(point_hierarchy, pyramid, point_hierarchy_dual, pyramid_dual)
# We'll now apply the dual octree and trinkets to perform trilinaer interpolation.
# First we'll generate some features for the corners.
# The first dimension of pyramid / pyramid_dual specifies how many unique points exist per level.
# For the pyramid_dual, this means how many "unique corners" are in place (as neighboring cells may share corners!)
num_of_corners_at_last_lod = pyramid_dual[0, level]
feature_dims = 32
feats = torch.rand([num_of_corners_at_last_lod, feature_dims], device='cuda')
# Create some query coordinate with normalized values in the range [-1, 1], here we pick (0.5, 0.5, 0.5).
# We'll also modify the dimensions of the query tensor to match the interpolation function api:
# batch dimension refers to the unique number of spc cells we're querying.
# samples_count refers to the number of interpolations we perform per cell.
query_coord = points.new_tensor((0.5, 0.5, 0.5)).unsqueeze(0) # Tensor of (batch, 3), in this case batch=1
sampled_query_coords = query_coord.unsqueeze(1) # Tensor of (batch, samples_count, 3), in this case samples_count=1
# unbatched_query converts from normalized coordinates to the index of the cell containing this point.
# The query_index can be used to pick the point from point_hierarchy
query_index = kaolin.ops.spc.unbatched_query(spc.octrees, spc.exsum, query_coord, level, with_parents=False)
# The unbatched_interpolate_trilinear function uses the query coordinates to perform trilinear interpolation.
# Here, unbatched specifies this function supports only a single SPC at a time.
# Per single SPC, we may interpolate a batch of coordinates and samples
interpolated = kaolin.ops.spc.unbatched_interpolate_trilinear(coords=sampled_query_coords,
pidx=query_index.int(),
point_hierarchy=point_hierarchy,
trinkets=trinkets, feats=feats, level=level)
print(f'Interpolated a tensor of shape {interpolated.shape} with values: {interpolated}')