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

0
examples/tutorial/__init__.py Normal file
View File

BIN
examples/tutorial/assets/ndc_camera_space.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 210 KiB

BIN
examples/tutorial/assets/ndc_image_space.png Normal file
View File

Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 34 KiB

1628
examples/tutorial/bbox_tutorial.ipynb Normal file
View File

File diff suppressed because one or more lines are too long

1373
examples/tutorial/camera_and_rasterization.ipynb Normal file
View File

File diff suppressed because it is too large Load Diff

492
examples/tutorial/dibr_tutorial.ipynb Normal file
View File

File diff suppressed because one or more lines are too long

602
examples/tutorial/diffuse_lighting.ipynb Normal file
View File

File diff suppressed because one or more lines are too long

88
examples/tutorial/dmtet_network.py Normal file
View File

@@ -0,0 +1,88 @@
import torch
from tqdm import tqdm
# MLP + Positional Encoding
class Decoder(torch.nn.Module):
def __init__(self, input_dims = 3, internal_dims = 128, output_dims = 4, hidden = 5, multires = 2):
super().__init__()
self.embed_fn = None
if multires > 0:
embed_fn, input_ch = get_embedder(multires)
self.embed_fn = embed_fn
input_dims = input_ch
net = (torch.nn.Linear(input_dims, internal_dims, bias=False), torch.nn.ReLU())
for i in range(hidden-1):
net = net + (torch.nn.Linear(internal_dims, internal_dims, bias=False), torch.nn.ReLU())
net = net + (torch.nn.Linear(internal_dims, output_dims, bias=False),)
self.net = torch.nn.Sequential(*net)
def forward(self, p):
if self.embed_fn is not None:
p = self.embed_fn(p)
out = self.net(p)
return out
def pre_train_sphere(self, iter):
print ("Initialize SDF to sphere")
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.Adam(list(self.parameters()), lr=1e-4)
for i in tqdm(range(iter)):
p = torch.rand((1024,3), device='cuda') - 0.5
ref_value = torch.sqrt((p**2).sum(-1)) - 0.3
output = self(p)
loss = loss_fn(output[...,0], ref_value)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("Pre-trained MLP", loss.item())
# Positional Encoding from https://github.com/yenchenlin/nerf-pytorch/blob/1f064835d2cca26e4df2d7d130daa39a8cee1795/run_nerf_helpers.py
class Embedder:
def __init__(self, **kwargs):
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 embed(self, inputs):
return torch.cat([fn(inputs) for fn in self.embed_fns], -1)
def get_embedder(multires):
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],
}
embedder_obj = Embedder(**embed_kwargs)
embed = lambda x, eo=embedder_obj : eo.embed(x)
return embed, embedder_obj.out_dim

369
examples/tutorial/dmtet_tutorial.ipynb Normal file
View File

@@ -0,0 +1,369 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "ee093335",
"metadata": {
"tags": []
},
"source": [
"# Reconstructing Point Cloud with DMTet\n",
"\n",
"Deep Marching Tetrahedra (DMTet) is a hybrid 3D representation that combines both implicit and explicit 3D surface representations. It represents a shape with a discrete SDF defined on vertices of a deformable tetrahedral grid. The SDF is converted to triangular mesh using a differentiable marching tetrahedra layer (MT), allowing explicit supervision on the extracted surface to be back-propagated to SDF and change mesh topology. In this tutorial, we demonstrate this by optimizing DMTet to reconstruct point cloud by minimizing the Chamfer Distance. The key functions used in this tutorial are in `kaolin.ops.conversions.trianglemesh`. See detailed [API documentation](https://kaolin.readthedocs.io/en/latest/modules/kaolin.ops.conversions.html#kaolin-ops-conversions).\n",
"\n",
"In addition, we demonstrate the use of [Kaolin's 3D checkpoints and training visualization](https://kaolin.readthedocs.io/en/latest/modules/kaolin.visualize.html) with the [Omniverse Kaolin App](https://docs.omniverse.nvidia.com/app_kaolin/app_kaolin/user_manual.html)."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "31d9198f",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import torch\n",
"import kaolin\n",
"import numpy as np\n",
"from dmtet_network import Decoder\n",
"\n",
"# path to the point cloud to be reconstructed\n",
"pcd_path = \"../samples/bear_pointcloud.usd\"\n",
"# path to the output logs (readable with the training visualizer in the omniverse app)\n",
"logs_path = './logs/'\n",
"\n",
"# We initialize the timelapse that will store USD for the visualization apps\n",
"timelapse = kaolin.visualize.Timelapse(logs_path)"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "58c9c196",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# arguments and hyperparameters\n",
"device = 'cuda'\n",
"lr = 1e-3\n",
"laplacian_weight = 0.1\n",
"iterations = 5000\n",
"save_every = 100\n",
"multires = 2\n",
"grid_res = 128"
]
},
{
"cell_type": "markdown",
"id": "16a2899d",
"metadata": {},
"source": [
"# Loading Point Cloud\n",
"\n",
"In this example, we use the point cloud generated by [Omniverse Kaolin App](https://docs.omniverse.nvidia.com/app_kaolin/app_kaolin/user_manual.html#data-generator). We load the pre-generated point cloud in `examples/samples/` and normalize it to the range of the tetrahedral grid. The normalized point cloud is saved to the checkpoint which can be visualized using [the Omniverse app](https://docs.omniverse.nvidia.com/app_kaolin/app_kaolin).\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "5674d9a2",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"True\n",
"torch.Size([89164, 3])\n"
]
}
],
"source": [
"points = kaolin.io.usd.import_pointclouds(pcd_path)[0].points.to(device)\n",
"if points.shape[0] > 100000:\n",
" idx = list(range(points.shape[0]))\n",
" np.random.shuffle(idx)\n",
" idx = torch.tensor(idx[:100000], device=points.device, dtype=torch.long) \n",
" points = points[idx]\n",
"\n",
"# The reconstructed object needs to be slightly smaller than the grid to get watertight surface after MT.\n",
"points = kaolin.ops.pointcloud.center_points(points.unsqueeze(0), normalize=True).squeeze(0) * 0.9\n",
"timelapse.add_pointcloud_batch(category='input',\n",
" pointcloud_list=[points.cpu()], points_type = \"usd_geom_points\")"
]
},
{
"cell_type": "markdown",
"id": "8b39b36c",
"metadata": {},
"source": [
"# Loading the Tetrahedral Grid\n",
"\n",
"DMTet starts from a uniform tetrahedral grid of predefined resolution, and uses a network to predict the SDF value as well as deviation vector at each grid vertex. \n",
"\n",
"Here we load the pre-generated tetrahedral grid using [Quartet](https://github.com/crawforddoran/quartet) at resolution 128, which has roughly the same number of vertices as a voxel grid of resolution 65. We use a simple MLP + positional encoding to predict the SDF and deviation vectors in DMTet, and initialize the encoded SDF to represent a sphere. "
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "33ab4b6f",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"torch.Size([277410, 3]) torch.Size([1524684, 4])\n",
"Initialize SDF to sphere\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"100%|██████████| 1000/1000 [00:03<00:00, 279.25it/s]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Pre-trained MLP 5.480436811922118e-06\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\n"
]
}
],
"source": [
"tet_verts = torch.tensor(np.load('../samples/{}_verts.npz'.format(grid_res))['data'], dtype=torch.float, device=device)\n",
"tets = torch.tensor(([np.load('../samples/{}_tets_{}.npz'.format(grid_res, i))['data'] for i in range(4)]), dtype=torch.long, device=device).permute(1,0)\n",
"print (tet_verts.shape, tets.shape)\n",
"\n",
"# Initialize model and create optimizer\n",
"model = Decoder(multires=multires).to(device)\n",
"model.pre_train_sphere(1000)\n"
]
},
{
"cell_type": "markdown",
"id": "73fe95a7",
"metadata": {},
"source": [
"# Preparing the Losses and Regularizer\n",
"\n",
"During training we will use two losses defined on the surface mesh:\n",
"- We use Chamfer Distance as the reconstruction loss. At each step, we randomly sample points from the surface mesh and compute the point-to-point distance to the GT point cloud.\n",
"- DMTet can employ direct regularization on the surface mesh to impose useful geometric constraints. We demonstrate this with a Laplacian loss which encourages the surface to be smooth.\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "78ad11ef",
"metadata": {},
"outputs": [],
"source": [
"# Laplacian regularization using umbrella operator (Fujiwara / Desbrun).\n",
"# https://mgarland.org/class/geom04/material/smoothing.pdf\n",
"def laplace_regularizer_const(mesh_verts, mesh_faces):\n",
" term = torch.zeros_like(mesh_verts)\n",
" norm = torch.zeros_like(mesh_verts[..., 0:1])\n",
"\n",
" v0 = mesh_verts[mesh_faces[:, 0], :]\n",
" v1 = mesh_verts[mesh_faces[:, 1], :]\n",
" v2 = mesh_verts[mesh_faces[:, 2], :]\n",
"\n",
" term.scatter_add_(0, mesh_faces[:, 0:1].repeat(1,3), (v1 - v0) + (v2 - v0))\n",
" term.scatter_add_(0, mesh_faces[:, 1:2].repeat(1,3), (v0 - v1) + (v2 - v1))\n",
" term.scatter_add_(0, mesh_faces[:, 2:3].repeat(1,3), (v0 - v2) + (v1 - v2))\n",
"\n",
" two = torch.ones_like(v0) * 2.0\n",
" norm.scatter_add_(0, mesh_faces[:, 0:1], two)\n",
" norm.scatter_add_(0, mesh_faces[:, 1:2], two)\n",
" norm.scatter_add_(0, mesh_faces[:, 2:3], two)\n",
"\n",
" term = term / torch.clamp(norm, min=1.0)\n",
"\n",
" return torch.mean(term**2)\n",
"\n",
"def loss_f(mesh_verts, mesh_faces, points, it):\n",
" pred_points = kaolin.ops.mesh.sample_points(mesh_verts.unsqueeze(0), mesh_faces, 50000)[0][0]\n",
" chamfer = kaolin.metrics.pointcloud.chamfer_distance(pred_points.unsqueeze(0), points.unsqueeze(0)).mean()\n",
" if it > iterations//2:\n",
" lap = laplace_regularizer_const(mesh_verts, mesh_faces)\n",
" return chamfer + lap * laplacian_weight\n",
" return chamfer\n"
]
},
{
"cell_type": "markdown",
"id": "0f96974c",
"metadata": {},
"source": [
"# Setting up Optimizer"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "a5d4a42f",
"metadata": {},
"outputs": [],
"source": [
"vars = [p for _, p in model.named_parameters()]\n",
"optimizer = torch.optim.Adam(vars, lr=lr)\n",
"scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: max(0.0, 10**(-x*0.0002))) # LR decay over time"
]
},
{
"cell_type": "markdown",
"id": "e7917ee1",
"metadata": {},
"source": [
"# Training\n",
"\n",
"At every iteration, we first predict SDF and deviation vector at each vertex with the network. Next, we extract the triangular mesh by running Marching Tetrahedra on the grid. We then compute loss functions on the extracted mesh and backpropagate gradient to the network weights. Notice that the topology of the mesh is changing during training, as shown in the output message. The training takes ~5 minutes on a TITAN RTX GPU."
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "583bec8b",
"metadata": {
"scrolled": true
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Iteration 0 - loss: 0.02473130077123642, # of mesh vertices: 18110, # of mesh faces: 36216\n",
"Iteration 100 - loss: 0.002605137648060918, # of mesh vertices: 24234, # of mesh faces: 48464\n",
"Iteration 200 - loss: 0.0003765518486034125, # of mesh vertices: 26862, # of mesh faces: 53720\n",
"Iteration 300 - loss: 0.0010241996496915817, # of mesh vertices: 31508, # of mesh faces: 63012\n",
"Iteration 400 - loss: 0.0001085952389985323, # of mesh vertices: 28300, # of mesh faces: 56596\n",
"Iteration 500 - loss: 7.9919038398657e-05, # of mesh vertices: 28710, # of mesh faces: 57416\n",
"Iteration 600 - loss: 0.00010018410830525681, # of mesh vertices: 27400, # of mesh faces: 54796\n",
"Iteration 700 - loss: 6.0749654949177057e-05, # of mesh vertices: 28494, # of mesh faces: 56984\n",
"Iteration 800 - loss: 0.0002924088039435446, # of mesh vertices: 27660, # of mesh faces: 55316\n",
"Iteration 900 - loss: 9.263768151868135e-05, # of mesh vertices: 28512, # of mesh faces: 57020\n",
"Iteration 1000 - loss: 7.250437192851678e-05, # of mesh vertices: 28598, # of mesh faces: 57192\n",
"Iteration 1100 - loss: 6.00546263740398e-05, # of mesh vertices: 28352, # of mesh faces: 56700\n",
"Iteration 1200 - loss: 4.965237167198211e-05, # of mesh vertices: 28606, # of mesh faces: 57208\n",
"Iteration 1300 - loss: 4.5047825551591814e-05, # of mesh vertices: 28934, # of mesh faces: 57864\n",
"Iteration 1400 - loss: 4.2731968278530985e-05, # of mesh vertices: 28878, # of mesh faces: 57752\n",
"Iteration 1500 - loss: 8.582305599702522e-05, # of mesh vertices: 28790, # of mesh faces: 57576\n",
"Iteration 1600 - loss: 4.140706005273387e-05, # of mesh vertices: 28924, # of mesh faces: 57844\n",
"Iteration 1700 - loss: 3.995447332272306e-05, # of mesh vertices: 28850, # of mesh faces: 57696\n",
"Iteration 1800 - loss: 3.944659692933783e-05, # of mesh vertices: 29064, # of mesh faces: 58128\n",
"Iteration 1900 - loss: 3.890909647452645e-05, # of mesh vertices: 28994, # of mesh faces: 57984\n",
"Iteration 2000 - loss: 3.9877151721157134e-05, # of mesh vertices: 28832, # of mesh faces: 57660\n",
"Iteration 2100 - loss: 3.8087084249127656e-05, # of mesh vertices: 28942, # of mesh faces: 57880\n",
"Iteration 2200 - loss: 3.8198602851480246e-05, # of mesh vertices: 29116, # of mesh faces: 58228\n",
"Iteration 2300 - loss: 3.789698894252069e-05, # of mesh vertices: 29188, # of mesh faces: 58372\n",
"Iteration 2400 - loss: 3.733349876711145e-05, # of mesh vertices: 28986, # of mesh faces: 57968\n",
"Iteration 2500 - loss: 3.886773993144743e-05, # of mesh vertices: 28728, # of mesh faces: 57452\n",
"Iteration 2600 - loss: 3.7754220102215186e-05, # of mesh vertices: 29132, # of mesh faces: 58260\n",
"Iteration 2700 - loss: 3.751121403183788e-05, # of mesh vertices: 28962, # of mesh faces: 57920\n",
"Iteration 2800 - loss: 3.733678022399545e-05, # of mesh vertices: 28942, # of mesh faces: 57880\n",
"Iteration 2900 - loss: 3.712274701683782e-05, # of mesh vertices: 28970, # of mesh faces: 57936\n",
"Iteration 3000 - loss: 3.738816667464562e-05, # of mesh vertices: 29154, # of mesh faces: 58304\n",
"Iteration 3100 - loss: 3.6861980333924294e-05, # of mesh vertices: 29090, # of mesh faces: 58176\n",
"Iteration 3200 - loss: 3.7955178413540125e-05, # of mesh vertices: 29228, # of mesh faces: 58452\n",
"Iteration 3300 - loss: 3.692376412800513e-05, # of mesh vertices: 28990, # of mesh faces: 57976\n",
"Iteration 3400 - loss: 3.6803434340981767e-05, # of mesh vertices: 29032, # of mesh faces: 58060\n",
"Iteration 3500 - loss: 3.666708289529197e-05, # of mesh vertices: 29006, # of mesh faces: 58008\n",
"Iteration 3600 - loss: 3.6867546441499144e-05, # of mesh vertices: 28916, # of mesh faces: 57828\n",
"Iteration 3700 - loss: 3.673196624731645e-05, # of mesh vertices: 28876, # of mesh faces: 57748\n",
"Iteration 3800 - loss: 3.683008617372252e-05, # of mesh vertices: 28868, # of mesh faces: 57732\n",
"Iteration 3900 - loss: 3.696472413139418e-05, # of mesh vertices: 28932, # of mesh faces: 57860\n",
"Iteration 4000 - loss: 3.699162698467262e-05, # of mesh vertices: 29188, # of mesh faces: 58372\n",
"Iteration 4100 - loss: 3.622782969614491e-05, # of mesh vertices: 28980, # of mesh faces: 57956\n",
"Iteration 4200 - loss: 3.6102632293477654e-05, # of mesh vertices: 28990, # of mesh faces: 57976\n",
"Iteration 4300 - loss: 3.6840694519924e-05, # of mesh vertices: 28888, # of mesh faces: 57772\n",
"Iteration 4400 - loss: 3.603967707022093e-05, # of mesh vertices: 28992, # of mesh faces: 57980\n",
"Iteration 4500 - loss: 3.609260966186412e-05, # of mesh vertices: 29044, # of mesh faces: 58084\n",
"Iteration 4600 - loss: 3.623321754275821e-05, # of mesh vertices: 29112, # of mesh faces: 58220\n",
"Iteration 4700 - loss: 3.591994391172193e-05, # of mesh vertices: 29116, # of mesh faces: 58228\n",
"Iteration 4800 - loss: 3.641782677732408e-05, # of mesh vertices: 29148, # of mesh faces: 58292\n",
"Iteration 4900 - loss: 3.601510252337903e-05, # of mesh vertices: 29078, # of mesh faces: 58152\n",
"Iteration 4999 - loss: 3.580914199119434e-05, # of mesh vertices: 29056, # of mesh faces: 58108\n"
]
}
],
"source": [
"for it in range(iterations):\n",
" pred = model(tet_verts) # predict SDF and per-vertex deformation\n",
" sdf, deform = pred[:,0], pred[:,1:]\n",
" verts_deformed = tet_verts + torch.tanh(deform) / grid_res # constraint deformation to avoid flipping tets\n",
" mesh_verts, mesh_faces = kaolin.ops.conversions.marching_tetrahedra(verts_deformed.unsqueeze(0), tets, sdf.unsqueeze(0)) # running MT (batched) to extract surface mesh\n",
" mesh_verts, mesh_faces = mesh_verts[0], mesh_faces[0]\n",
"\n",
" loss = loss_f(mesh_verts, mesh_faces, points, it)\n",
" optimizer.zero_grad()\n",
" loss.backward()\n",
" optimizer.step()\n",
" scheduler.step()\n",
" if (it) % save_every == 0 or it == (iterations - 1): \n",
" print ('Iteration {} - loss: {}, # of mesh vertices: {}, # of mesh faces: {}'.format(it, loss, mesh_verts.shape[0], mesh_faces.shape[0]))\n",
" # save reconstructed mesh\n",
" timelapse.add_mesh_batch(\n",
" iteration=it+1,\n",
" category='extracted_mesh',\n",
" vertices_list=[mesh_verts.cpu()],\n",
" faces_list=[mesh_faces.cpu()]\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "e49e8f67",
"metadata": {},
"source": [
"# Visualize Training\n",
"\n",
"You can now use [the Omniverse app](https://docs.omniverse.nvidia.com/app_kaolin/app_kaolin) to visualize the mesh optimization over training by using the training visualizer on \"./logs/\", where we stored the checkpoints.\n",
"\n",
"Alternatively, you can use [kaolin-dash3d](https://kaolin.readthedocs.io/en/latest/notes/checkpoints.html?highlight=usd#visualizing-with-kaolin-dash3d) to visualize the checkpoint by running <code>kaolin-dash3d --logdir=$logs_path --port=8080</code>. This command will launch a web server that will stream geometry to web clients. You can view the input point cloud and the reconstructed mesh at [localhost:8080](localhost:8080) as shown below. You can change the *global iteration* on the left to see how the mesh evolves during training. \n",
"\n",
"![alt text](../samples/dash3d_mesh.png \"Title\")\n",
"![alt text](../samples/dash3d_pcd.png \"Title\")"
]
}
],
"metadata": {
"interpreter": {
"hash": "4040fd28a16387d31474220157706b1752bd7f86ecfd14350c5c940438c26826"
},
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}

508
examples/tutorial/gltf_viz.ipynb Normal file
View File

File diff suppressed because one or more lines are too long

1709
examples/tutorial/interactive_visualizer.ipynb Normal file
View File

File diff suppressed because one or more lines are too long

450
examples/tutorial/sg_specular_lighting.ipynb Normal file
View File

File diff suppressed because one or more lines are too long

59
examples/tutorial/spc_formatting.py Normal file
View File

@@ -0,0 +1,59 @@
from collections import deque
from termcolor import colored
def color_by_level(level):
_colormap = ["red", "blue", "green", "yellow", "magenta", "cyan", "grey"]
return _colormap[level % len(_colormap)]
def push_pop_octree(q, oct_item):
prefix = q.popleft()
bit_idx = 0
parsed_bits = oct_item.item()
while parsed_bits:
bit_idx += 1
if parsed_bits & 1:
if len(prefix) > 0:
q.append(prefix + f'-{bit_idx}')
else:
q.append(prefix + f'{bit_idx}')
parsed_bits >>= 1
return prefix
def format_octree_str(octree_byte, octree_path, level_idx, max_level):
text = []
level_color = color_by_level(level_idx - 1)
text += ['Level ' + colored(f'#{level_idx}, ', level_color)]
colored_path = []
for i in range(len(octree_path)):
level_color = color_by_level(i // 2)
if i % 2 == 0:
colored_path += [colored(octree_path[i], level_color)]
else:
colored_path += [octree_path[i]]
colored_path = ''.join(colored_path)
text += [f'Path{colored_path}, ']
text += [' ' for _ in range((max_level - level_idx) * 2)]
text += ['{0:08b}'.format(octree_byte)]
return ''.join(text)
def describe_octree(octree, level, limit_levels=None):
bit_counter = lambda x: bin(x).count('1')
level_idx, curr_level_remaining_cells, next_level_cells = 1, 1, 0
octree_paths = deque('*')
for oct_idx, octree_byte in enumerate(octree):
octree_path = push_pop_octree(octree_paths, octree_byte)
if limit_levels is None or level_idx in limit_levels:
print(format_octree_str(octree_byte, octree_path, level_idx, level))
curr_level_remaining_cells -= 1
next_level_cells += bit_counter(octree_byte)
if not curr_level_remaining_cells:
level_idx += 1
curr_level_remaining_cells = next_level_cells
next_level_cells = 0

35
examples/tutorial/tutorial_common.py Normal file
View File

@@ -0,0 +1,35 @@
# Copyright (c) 2019,20-22 NVIDIA CORPORATION & AFFILIATES.
# All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
import os
import torch
try:
import matplotlib.pyplot as plt
except Exception as e:
pass
FILE_DIR = os.path.abspath(os.path.dirname(os.path.abspath(__file__)))
# Folder with common data samples (<kaolin_root>/sample_data)
COMMON_DATA_DIR = os.path.realpath(os.path.join(FILE_DIR, os.pardir, os.pardir, 'sample_data'))
# Folder with data specific to examples (<kaolin_root>/examples/samples)
EXAMPLES_DATA_DIR = os.path.realpath(os.path.join(FILE_DIR, os.pardir, 'samples'))

3463
examples/tutorial/understanding_spcs_tutorial.ipynb Normal file
View File

File diff suppressed because one or more lines are too long

73
examples/tutorial/usd_kitchenset.py Normal file
View File

@@ -0,0 +1,73 @@
#!/usr/bin/env python
# Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import argparse
from kaolin.io import usd
from kaolin.io.utils import mesh_handler_naive_triangulate
def import_kitchen_set(kitchen_set_usd):
# The Kitchen Set example organizes assets in a particular way. Since we want to import complete objects and not
# not each separate part of an object, we'll find all the paths that are named :code:`Geom`:
scene_paths = usd.get_scene_paths(kitchen_set_usd, r'.*/Geom$')
# The meshes in this dataset have a heterogeneous topology, meaning the number of vertices
# for each polygon varies. To deal with those, we'll pass in a handler function that will
# homogenize those meshes to homogenous triangle meshes.
usd_meshes = usd.import_meshes(
kitchen_set_usd,
scene_paths=scene_paths,
heterogeneous_mesh_handler=mesh_handler_naive_triangulate
)
return usd_meshes
def save_kitchen_set_dataset(meshes, out_dir):
for i, m in enumerate(meshes):
out_path = os.path.join(out_dir, f'mesh_{i}.usd')
usd.export_mesh(
file_path=out_path,
vertices=m.vertices[..., [0, 2, 1]], # flipping Y and Z to make models Y-up
faces=m.faces
)
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description=r'Convert Pixar\'s Kitchen Set scene (http://graphics.pixar.com/usd/downloads.html) '
'into a dataset of Pytorch Tensors, ready to be use to train our next awesome model.')
parser.add_argument('--kitchen_set_dir', type=str, required=True,
help='Location of the kitchen_set data.')
parser.add_argument('--output_dir', type=str, required=True,
help='Output directory to export the dataset to; must exist.')
args = parser.parse_args()
# We will be importing Pixar's Kitchen Set scene (http://graphics.pixar.com/usd/downloads.html) as a
# dataset of Pytorch Tensors, ready to be used to train our next awesome model.
kitchen_set_usd = os.path.join(args.kitchen_set_dir, 'Kitchen_set.usd')
meshes = import_kitchen_set(kitchen_set_usd)
print(len(meshes)) # 426
# And just like that, we have a dataset of 426 diverse objects for our use!
# Now let's save our dataset so we can use it again later.
save_kitchen_set_dataset(meshes, args.output_dir)
# We can now fire up Omniverse Kaolin and use the Dataset Visualizer extension to
# see what this dataset looks like and start using it in our next project!

137
examples/tutorial/visualize_main.py Normal file
View File

@@ -0,0 +1,137 @@
#!/usr/bin/env python
# Copyright (c) 2019-2020, NVIDIA CORPORATION. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import argparse
import logging
import os
import random
import torch
import sys
import kaolin
logger = logging.getLogger(__name__)
def __normalize_vertices(vertices):
"""
Normalizes vertices to fit an [-1...1] bounding box,
common during training, but not necessary for visualization.
"""
return kaolin.ops.pointcloud.center_points(res.vertices.unsqueeze(0), normalize=True).squeeze(0) * 2
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description='Example exporting 3D data during training as timed USDs; '
'also demonstrates OBJ import and mesh to pointcloud conversions.')
parser.add_argument('--test_objs', type=str, required=True,
help='Comma separated list of several example obj files.')
parser.add_argument('--output_dir', type=str, required=True,
help='Output directory to write checkpoints to; must exist.')
parser.add_argument('--iterations', type=int, default=101,
help='How many training iterations to emulate.')
parser.add_argument('--checkpoint_interval', type=int, default=10,
help='Frequency with which to write out checkpoints.')
parser.add_argument('--skip_normalization', action='store_true',
help='If not set, will normalize bounding box of each input '
'to be within -1..1 cube.')
args = parser.parse_args()
logging.basicConfig(level=logging.INFO, stream=sys.stdout)
if not os.path.isdir(args.output_dir):
raise RuntimeError(
'Output directory does not exist: --output_dir={}'.format(
args.output_dir))
# Read test 3D models & setup fake training ----------------------------------
obj_files = args.test_objs.split(',')
logger.info('Parsing {} OBJ files: '.format(len(obj_files)))
face_list = []
gt_vert_list = []
input_pt_clouds = []
delta_list = []
delta_pt_list = []
# TODO: add textured example
for f in obj_files:
res = kaolin.io.obj.import_mesh(f)
vertices = res.vertices if args.skip_normalization else __normalize_vertices(vertices)
num_samples = random.randint(1000, 1500) # Vary to ensure robustness
pts = kaolin.ops.mesh.sample_points(
vertices.unsqueeze(0), res.faces, num_samples)[0].squeeze(0)
# Randomly displace vertices to emulate training
delta = (2.0 * torch.rand(vertices.shape, dtype=vertices.dtype) - 1.0) * 0.25
delta_pts = (2.0 * torch.rand(pts.shape, dtype=pts.dtype) - 1.0) * 0.25
face_list.append(res.faces)
gt_vert_list.append(vertices)
delta_list.append(delta)
input_pt_clouds.append(pts)
delta_pt_list.append(delta_pts)
# Emulate visualizing during training -------------------------------------
logger.info('Emulating training run for {} iterations'.format(args.iterations))
# Create a Timelapse instance
timelapse = kaolin.visualize.Timelapse(args.output_dir)
# Save static objects such as ground truth or inputs that do not change with iterations
# just once.
timelapse.add_mesh_batch(
category='ground_truth',
faces_list=face_list,
vertices_list=gt_vert_list)
timelapse.add_pointcloud_batch(
category='input',
pointcloud_list=input_pt_clouds)
for iteration in range(args.iterations):
if iteration % args.checkpoint_interval == 0:
# Emulate a training update
out_pt_clouds = []
out_vert_list = []
out_voxels = []
for i in range(len(gt_vert_list)):
delta_weight = 1.0 - iteration / (args.iterations - 1)
out_vert_list.append(gt_vert_list[i] * (1.0 + delta_list[i] * delta_weight))
out_pt_clouds.append(input_pt_clouds[i] * (1.0 + delta_pt_list[i] * delta_weight))
vg = kaolin.ops.conversions.trianglemeshes_to_voxelgrids(
out_vert_list[-1].unsqueeze(0), face_list[i], 30)
out_voxels.append(vg.squeeze(0).bool())
# Save model predictions to track training progress over time
timelapse.add_mesh_batch(
iteration=iteration,
category='output',
faces_list=face_list,
vertices_list=out_vert_list)
timelapse.add_pointcloud_batch(
iteration=iteration,
category='output',
pointcloud_list=out_pt_clouds)
timelapse.add_voxelgrid_batch(
iteration=iteration,
category='output',
voxelgrid_list=out_voxels)
logger.info('Emulated training complete!\n'
'You can now view created USD files by running:\n\n'
f'kaolin-dash3d --logdir={args.output_dir}\n\n'
'And then navigating to localhost:8080\n')
# TODO(mshugrina): once dash3d is also integrated, write an integration test
# to ensure timelapse output is properly parsed by the visualizer

846
examples/tutorial/working_with_meshes.ipynb Normal file
View File

@@ -0,0 +1,846 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "e504f9ed",
"metadata": {},
"source": [
"# Working with Meshes\n",
"\n",
"This tutorial shows how to expedite working with kaolin operations using the `SurfaceMesh` container class. We will cover import/export, batching strategies, managing mesh data, rendering and visualization. \n",
"\n",
"Note that material support of `SurfaceMesh` is currently limited and is on the roadmap."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "4d8db10b",
"metadata": {},
"outputs": [],
"source": [
"import copy\n",
"import logging\n",
"import numpy as np\n",
"import os\n",
"import sys\n",
"import torch\n",
"\n",
"import kaolin as kal\n",
"from kaolin.rep import SurfaceMesh\n",
"\n",
"from tutorial_common import COMMON_DATA_DIR\n",
"\n",
"def sample_mesh_path(fname):\n",
" return os.path.join(COMMON_DATA_DIR, 'meshes', fname)\n",
"\n",
"def print_tensor(t, **kwargs):\n",
" print(kal.utils.testing.tensor_info(t, **kwargs))"
]
},
{
"cell_type": "markdown",
"id": "16c66902",
"metadata": {},
"source": [
"## Understanding the SurfaceMesh Container\n",
"\n",
"`SurfaceMesh` can store information about a single mesh and a batch of meshes, following three batching\n",
"strategies:\n",
" * `NONE` - a single mesh, not batched\n",
" * `FIXED` - a batch of meshes with fixed topology (faces are fixed)\n",
" * `LIST` - a list of variable topology meshes\n",
"\n",
"Automatically converting between these batching strategies allows quickly connecting to various Kaolin operations. "
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "7ea7171c",
"metadata": {
"scrolled": true
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Expected SurfaceMesh contents for batching strategy FIXED\n",
" vertices: (torch.FloatTensor) of shape ['B', 'V', 3]\n",
" faces: (torch.IntTensor) of shape ['F', 'FSz']\n",
" face_vertices: (torch.FloatTensor) of shape ['B', 'F', 'FSz', 3]\n",
" normals: (torch.FloatTensor) of shape ['B', 'VN', 3]\n",
" face_normals_idx: (torch.IntTensor) of shape ['B', 'F', 'FSz']\n",
" face_normals: (torch.FloatTensor) of shape ['B', 'F', 'FSz', 3]\n",
" uvs: (torch.FloatTensor) of shape ['B', 'U', 2]\n",
" face_uvs_idx: (torch.IntTensor) of shape ['B', 'F', 'FSz']\n",
" face_uvs: (torch.FloatTensor) of shape ['B', 'F', 'FSz', 2]\n",
" vertex_normals: (torch.FloatTensor) of shape ['B', 'V', 3]\n",
" vertex_tangents: (torch.FloatTensor) of shape ['B', 'V', 3]\n",
"material_assignments: (torch.IntTensor) of shape ['B', 'F']\n",
" materials: non-tensor attribute\n",
"\n",
"Key: B - batch size, F - number of faces, FSz - face size, V - number of vertices,\n",
" VN - number of normals, U - number of UVs\n"
]
}
],
"source": [
"# To get a sense for what the mesh can contain for different batching strategies, run:\n",
"\n",
"print(SurfaceMesh.attribute_info_string(SurfaceMesh.Batching.FIXED))\n",
"print('\\nKey: B - batch size, F - number of faces, FSz - face size, V - number of vertices,'\n",
" '\\n VN - number of normals, U - number of UVs')"
]
},
{
"cell_type": "markdown",
"id": "bf6648b8",
"metadata": {},
"source": [
"## Constructor and Auto-computable Attributes\n",
"\n",
"A `SurfaceMesh` can be constructed from torch tensors with names, types and sizes as described above. Only `faces` and `vertices` are required, both of which are allowed to contain zero elements, and **many attributes can be computed automatically**. \n",
"\n",
"Important settings passed to the constructor:\n",
"* `unset_attributes_return_none` (default: `True`) - set this to `False` to raise an error when accessing mesh attributes that are missing\n",
"* `allow_auto_compute` (default: `True`) - set this to `False` to disable computation of attributes such as `face_uvs` and `vertex_normals`\n",
"* `strict_checks` (default: `True`) - set this to `False` to allow setting attributes to unexpected shapes\n",
"\n",
"You can also set `mesh.unset_attributes_return_none` or `mesh.allow_auto_compute` later to change mesh behavior."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "8de1ea19",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"vertices: [10, 3] (torch.float32)[cpu] \n",
"faces: [5, 3] (torch.int64)[cpu] \n",
"face_vertices: None\n",
"face_vertices (auto-computed): [5, 3, 3] (torch.float32)[cpu] \n"
]
}
],
"source": [
"# Let's construct a simple unbatched mesh\n",
"V, F, Fsz = 10, 5, 3\n",
"mesh = kal.rep.SurfaceMesh(vertices=torch.rand((V, 3)).float(),\n",
" faces=torch.randint(0, V, (F, Fsz)).long(),\n",
" allow_auto_compute=False) # disable auto-compute for now\n",
"print_tensor(mesh.vertices, name='vertices')\n",
"print_tensor(mesh.faces, name='faces')\n",
"print_tensor(mesh.face_vertices, name='face_vertices')\n",
"\n",
"# Now let's enable auto-compute\n",
"mesh.allow_auto_compute=True\n",
"print_tensor(mesh.face_vertices, name='face_vertices (auto-computed)') "
]
},
{
"cell_type": "markdown",
"id": "44055abc",
"metadata": {},
"source": [
"Batched meshes can also be instantiated by passing batched inputs to the constructor, for example:"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "59a0b773",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Instantiated mesh with batching FIXED and length 3\n",
"Instantiated mesh with batching LIST and length 2\n"
]
}
],
"source": [
"# FIXED: inputs are batched tensors with fixed faces across batches\n",
"B, VN = 3, 20\n",
"mesh_fixed = kal.rep.SurfaceMesh(vertices=torch.rand((B, V, 3)).float(),\n",
" faces=torch.randint(0, V, (F, Fsz)).long(),\n",
" normals=torch.rand((B, VN, 3)).float(),\n",
" face_normals_idx=torch.randint(0, VN, (B, F, Fsz)))\n",
"print(f'Instantiated mesh with batching {mesh_fixed.batching} and length {len(mesh_fixed)}')\n",
"\n",
"# LIST: all inputs are lists of equal length\n",
"V2, F2 = 12, 20\n",
"mesh_list = kal.rep.SurfaceMesh(\n",
" vertices=[torch.rand((V, 3)).float(), torch.rand((V2, 3)).float()],\n",
" faces=[torch.randint(0, V, (F, Fsz)).long(), torch.randint(0, V2, (F2, Fsz)).long()])\n",
"print(f'Instantiated mesh with batching {mesh_list.batching} and length {len(mesh_list)}')"
]
},
{
"cell_type": "markdown",
"id": "96d7d2be",
"metadata": {},
"source": [
"## Inspecting SurfaceMesh Objects\n",
"\n",
"Working with many batched mesh attributes can be confusing, and details really matter. `SurfaceMesh` provides multiple ways to inspect its contents. These print statements also make it clear, which attributes can be auto-computed and how."
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "fb432078",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Mesh with batching NONE and length 1\n",
"\n",
"Attributes ['vertices', 'faces', 'face_vertices']\n",
"\n",
"Are face_normals set? False\n",
"\n",
"Are face_normals auto-computable? True\n",
"\n",
"Attributes (after accessing face_normals) ['vertices', 'faces', 'face_vertices', 'face_normals']\n",
"\n",
"Face normals face_normals: [5, 3, 3] (torch.float32)[cpu] \n",
"\n",
"\n",
"Does the mesh have expected shapes? True\n",
"SurfaceMesh object with batching strategy NONE\n",
" vertices: [10, 3] (torch.float32)[cpu] \n",
" faces: [5, 3] (torch.int64)[cpu] \n",
" face_vertices: [5, 3, 3] (torch.float32)[cpu] \n",
" face_normals: [5, 3, 3] (torch.float32)[cpu] \n",
" face_uvs: if possible, computed on access from: (uvs, face_uvs_idx)\n",
" vertex_normals: if possible, computed on access from: (faces, face_normals)\n",
" vertex_tangents: if possible, computed on access from: (faces, vertices, face_uvs)\n"
]
}
],
"source": [
"# Get batching strategy and batch size (length)\n",
"print(f'\\nMesh with batching {mesh.batching} and length {len(mesh)}')\n",
"\n",
"# Get currently set attributes\n",
"print(f'\\nAttributes {mesh.get_attributes(only_tensors=True)}')\n",
"\n",
"# Check if an attribute is set without causing the mesh to auto-compute it\n",
"print(f'\\nAre face_normals set? {mesh.has_attribute(\"face_normals\")}')\n",
"\n",
"# Check if the attribute likely can be auto-computed without actually trying to\n",
"print(f'\\nAre face_normals auto-computable? {mesh.probably_can_compute_attribute(\"face_normals\")}')\n",
"\n",
"# Let's access face_normals and cause them to be computed\n",
"mesh.face_normals\n",
"print(f'\\nAttributes (after accessing face_normals) {mesh.get_attributes(only_tensors=True)}')\n",
"\n",
"# Check that face_normals are now set\n",
"print(f'\\nFace normals{mesh.describe_attribute(\"face_normals\")}\\n')\n",
"\n",
"# Check if mesh tensor sizes follow expected conventions\n",
"print(f'\\nDoes the mesh have expected shapes? {mesh.check_sanity()}')\n",
"\n",
"# Print mesh contents (and computable attributes)\n",
"print(mesh)\n",
"\n",
"# We can also convert mesh to string with more details and tensor stats\n",
"# print(f'\\nDetailed string of {mesh.to_string(detailed=True, print_stats=True)}')"
]
},
{
"cell_type": "markdown",
"id": "10b7b49f",
"metadata": {},
"source": [
"## Explicit API\n",
"\n",
"In addition to default `SurfaceMesh` API that allows compute on access and automatic caching, this class also supports alternative more verbose API that makes these actions explicit."
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "d81f1dfe",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Deleting face_vertices\n",
"Deleting face_normals\n",
"\n",
"Mesh attributes after deletion: ['vertices', 'faces']\n",
"\n",
"Get face_normals without computing: None\n",
"\n",
"Computed face_normals shape is torch.Size([5, 3, 3])\n",
"\n",
"Did mesh cache computed face_normals (and face_vertices required to compute them)?\n",
"False, False\n",
"\n",
"Did mesh cache computed face_normals (and face_vertices required to compute them)?\n",
"True, True\n"
]
}
],
"source": [
"# Let's delete attributes we just computed\n",
"mesh.face_vertices = None\n",
"mesh.face_normals = None\n",
"\n",
"# Check attributes were removed\n",
"print(f'\\nMesh attributes after deletion: {mesh.get_attributes(only_tensors=True)}')\n",
"\n",
"# Get attribute without any auto-compute magic\n",
"print(f'\\nGet face_normals without computing: {mesh.get_attribute(\"face_normals\")}')\n",
"\n",
"# Compute attribute, but don't cache\n",
"face_normals = mesh.get_or_compute_attribute('face_normals', should_cache=False)\n",
"print(f'\\nComputed face_normals shape is {face_normals.shape}')\n",
"\n",
"# Verify attributes were not cached\n",
"print('\\nDid mesh cache computed face_normals (and face_vertices required to compute them)?')\n",
"print(f'{mesh.has_attribute(\"face_normals\")}, {mesh.has_attribute(\"face_vertices\")}')\n",
"\n",
"# Compute and cache\n",
"face_normals = mesh.get_or_compute_attribute('face_normals', should_cache=True)\n",
"\n",
"print('\\nDid mesh cache computed face_normals (and face_vertices required to compute them)?')\n",
"print(f'{mesh.has_attribute(\"face_normals\")}, {mesh.has_attribute(\"face_vertices\")}')"
]
},
{
"cell_type": "markdown",
"id": "787ebbdc",
"metadata": {
"pycharm": {
"name": "#%% md\n"
}
},
"source": [
"## Importing Data\n",
"\n",
"Since version 0.14.0, kaolin `obj` and `usd` importers return a `SurfaceMesh`, which is nearly backward-compatible with the previous `named_tuple` return type, while providing mutability and convenient data management. \n",
"\n",
"**Porting from earlier versions:** If porting from kaolin<=0.13.0, `obj` importer now correctly uses `face_normals_idx` (previously `face_normals`) to refer to the face-vertex indices into normals and `normals` (previously `vertex_normals`) to refer to the normals array that may or may not have the same number of elements as vertices. In addition, `materials` are now imported in name-sorted order and `material_order` has been replaced with `material_assignments` tensor of shape `(num_faces,)`, with integer value indicating the material index assigned to the corresponding face."
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "f316587e",
"metadata": {
"pycharm": {
"name": "#%%\n"
}
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Mesh imported from obj: SurfaceMesh object with batching strategy NONE\n",
" vertices: [42, 3] (torch.float32)[cpu] \n",
" faces: [80, 3] (torch.int64)[cpu] \n",
" normals: [80, 3] (torch.float32)[cpu] \n",
" face_normals_idx: [80, 3] (torch.int64)[cpu] \n",
" uvs: [63, 2] (torch.float32)[cpu] \n",
" face_uvs_idx: [80, 3] (torch.int64)[cpu] \n",
"material_assignments: [80] (torch.int16)[cpu] \n",
" materials: list of length 1\n",
" face_vertices: if possible, computed on access from: (faces, vertices)\n",
" face_normals: if possible, computed on access from: (normals, face_normals_idx) or (vertices, faces)\n",
" face_uvs: if possible, computed on access from: (uvs, face_uvs_idx)\n",
" vertex_normals: if possible, computed on access from: (faces, face_normals)\n",
" vertex_tangents: if possible, computed on access from: (faces, vertices, face_uvs)\n",
"\n",
"Mesh imported from usd: SurfaceMesh object with batching strategy NONE\n",
" vertices: [42, 3] (torch.float32)[cpu] \n",
" faces: [80, 3] (torch.int64)[cpu] \n",
" face_normals: [80, 3, 3] (torch.float32)[cpu] \n",
" uvs: [240, 2] (torch.float32)[cpu] \n",
" face_uvs_idx: [80, 3] (torch.int64)[cpu] \n",
"material_assignments: [80] (torch.int16)[cpu] \n",
" materials: list of length 1\n",
" face_vertices: if possible, computed on access from: (faces, vertices)\n",
" face_uvs: if possible, computed on access from: (uvs, face_uvs_idx)\n",
" vertex_normals: if possible, computed on access from: (faces, face_normals)\n",
" vertex_tangents: if possible, computed on access from: (faces, vertices, face_uvs)\n",
"dict_keys(['batching', 'allow_auto_compute', 'unset_attributes_return_none', 'materials', 'vertices', 'face_normals', 'uvs', 'faces', 'face_uvs_idx', 'material_assignments'])\n"
]
}
],
"source": [
"import_args = {'with_materials' : True, 'with_normals' : True}\n",
"\n",
"# Let's import a single mesh from OBJ\n",
"mesh_obj = kal.io.obj.import_mesh(sample_mesh_path('ico_flat.obj'), **import_args)\n",
"\n",
"# Let's import the same mesh from its USD version\n",
"mesh_usd = kal.io.usd.import_mesh(sample_mesh_path('ico_flat.usda'), **import_args)\n",
"\n",
"# Let's inspect contents of both meshes (notice consistent naming of attributes)\n",
"print(f'\\nMesh imported from obj: {mesh_obj}')\n",
"print(f'\\nMesh imported from usd: {mesh_usd}')\n",
"\n",
"# Note: if you prefer to work with raw values, SurfaceMesh is convertible to dict\n",
"mesh_dict = mesh_usd.as_dict()\n",
"print(mesh_dict.keys())"
]
},
{
"cell_type": "markdown",
"id": "fdfec845",
"metadata": {},
"source": [
"Although geometrically these objects are the same, you will notice that USD stroes UVs and normals differently from OBJ, resulting in different imported arrays. Despite these differences, actual UVs and normals auto-computed and assigned to faces are actually the same."
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "5733dc8d",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n",
"Are face_uvs same? True\n",
"\n",
"Are face_normals same? True\n"
]
}
],
"source": [
"print(f'\\nAre face_uvs same? {torch.allclose(mesh_obj.face_uvs, mesh_usd.face_uvs, atol=1e-4)}')\n",
"print(f'\\nAre face_normals same? {torch.allclose(mesh_obj.face_normals, mesh_usd.face_normals, atol=1e-4)}')"
]
},
{
"cell_type": "markdown",
"id": "5c53e057",
"metadata": {},
"source": [
"## Working with Batches\n",
"\n",
"`SurfaceMesh` objects can be converted between batching strategies, as long as it is possible (for example list of meshes of variable topology cannot be converted to `Batching.FIXED`). "
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "4f5476b2",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"SurfaceMesh object with batching strategy FIXED\n",
" vertices: [1, 42, 3] (torch.float32)[cpu] \n",
" faces: [80, 3] (torch.int64)[cpu] \n",
" normals: [1, 80, 3] (torch.float32)[cpu] \n",
" face_normals_idx: [1, 80, 3] (torch.int64)[cpu] \n",
" face_normals: [1, 80, 3, 3] (torch.float32)[cpu] \n",
" uvs: [1, 63, 2] (torch.float32)[cpu] \n",
" face_uvs_idx: [1, 80, 3] (torch.int64)[cpu] \n",
" face_uvs: [1, 80, 3, 2] (torch.float32)[cpu] \n",
"material_assignments: [1, 80] (torch.int16)[cpu] \n",
" materials: [\n",
" 0: list of length 1\n",
" ]\n",
" face_vertices: if possible, computed on access from: (faces, vertices)\n",
" vertex_normals: if possible, computed on access from: (faces, face_normals)\n",
" vertex_tangents: if possible, computed on access from: (faces, vertices, face_uvs)\n",
"\n",
"SurfaceMesh object with batching strategy LIST\n",
" vertices: [\n",
" 0: [42, 3] (torch.float32)[cpu] \n",
" ]\n",
" faces: [\n",
" 0: [80, 3] (torch.int64)[cpu] \n",
" ]\n",
" face_normals: [\n",
" 0: [80, 3, 3] (torch.float32)[cpu] \n",
" ]\n",
" uvs: [\n",
" 0: [240, 2] (torch.float32)[cpu] \n",
" ]\n",
" face_uvs_idx: [\n",
" 0: [80, 3] (torch.int64)[cpu] \n",
" ]\n",
" face_uvs: [\n",
" 0: [80, 3, 2] (torch.float32)[cpu] \n",
" ]\n",
"material_assignments: [\n",
" 0: [80] (torch.int16)[cpu] \n",
" ]\n",
" materials: [\n",
" 0: list of length 1\n",
" ]\n",
" face_vertices: if possible, computed on access from: (faces, vertices)\n",
" vertex_normals: if possible, computed on access from: (faces, face_normals)\n",
" vertex_tangents: if possible, computed on access from: (faces, vertices, face_uvs)\n"
]
}
],
"source": [
"# Convert unbatched mesh to most commonly used FIXED batching\n",
"mesh_obj.to_batched() # Shortcut for mesh_usd.set_batching(SurfaceMesh.Batching.FIXED)\n",
"print(mesh_obj)\n",
"\n",
"# Convert mesh to list batching\n",
"mesh_usd.set_batching(SurfaceMesh.Batching.LIST)\n",
"print(f'\\n{mesh_usd}')"
]
},
{
"cell_type": "markdown",
"id": "d6c2a73c",
"metadata": {},
"source": [
"We can also concatenate meshes of any batching strategy, with the output using `FIXED` (if `fixed_toplogy`) or `LIST` batching. Errors will be raised if concatentation is not possible for `vertices` or `faces`, and other attributes will be handled if possible. "
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "9cdad2dd",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Cannot cat uvs arrays of given shapes; trying to concatenate face_uvs instead, due to: stack expects each tensor to be equal size, but got [63, 2] at entry 0 and [240, 2] at entry 1\n"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"SurfaceMesh object with batching strategy FIXED\n",
" vertices: [2, 42, 3] (torch.float32)[cpu] \n",
" faces: [80, 3] (torch.int64)[cpu] \n",
" face_normals: [2, 80, 3, 3] (torch.float32)[cpu] \n",
" face_uvs: [2, 80, 3, 2] (torch.float32)[cpu] \n",
"material_assignments: [2, 80] (torch.int16)[cpu] \n",
" materials: [\n",
" 0: list of length 1\n",
" 1: list of length 1\n",
" ]\n",
" face_vertices: if possible, computed on access from: (faces, vertices)\n",
" vertex_normals: if possible, computed on access from: (faces, face_normals)\n",
" vertex_tangents: if possible, computed on access from: (faces, vertices, face_uvs)\n"
]
}
],
"source": [
"mesh = SurfaceMesh.cat([mesh_obj, mesh_usd], fixed_topology=True)\n",
"\n",
"# Notice that the concatenated mesh:\n",
"# 1. does not have uvs, as those could not be stacked, but face_uvs were computed and stacked instead.\n",
"# 2. does not have normals, as only first mesh had them, but face_normals were computed and stacked. \n",
"print(mesh)"
]
},
{
"cell_type": "markdown",
"id": "0091b207",
"metadata": {},
"source": [
"With `fixed_topology=False`, it is also possible to concatenate meshes of variable topology into a list representation. For example:"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "541462ee",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"SurfaceMesh object with batching strategy LIST\n",
" vertices: [\n",
" 0: [42, 3] (torch.float32)[cpu] \n",
" 1: [42, 3] (torch.float32)[cpu] \n",
" 2: [482, 3] (torch.float32)[cpu] \n",
" ]\n",
" faces: [\n",
" 0: [80, 3] (torch.int64)[cpu] \n",
" 1: [80, 3] (torch.int64)[cpu] \n",
" 2: [960, 3] (torch.int64)[cpu] \n",
" ]\n",
" face_normals: [\n",
" 0: [80, 3, 3] (torch.float32)[cpu] \n",
" 1: [80, 3, 3] (torch.float32)[cpu] \n",
" 2: [960, 3, 3] (torch.float32)[cpu] \n",
" ]\n",
"material_assignments: [\n",
" 0: [80] (torch.int16)[cpu] \n",
" 1: [80] (torch.int16)[cpu] \n",
" 2: [960] (torch.int16)[cpu] \n",
" ]\n",
" materials: [\n",
" 0: list of length 1\n",
" 1: list of length 1\n",
" 2: list of length 2\n",
" ]\n",
" face_vertices: if possible, computed on access from: (faces, vertices)\n",
" face_uvs: if possible, computed on access from: (uvs, face_uvs_idx)\n",
" vertex_normals: if possible, computed on access from: (faces, face_normals)\n",
" vertex_tangents: if possible, computed on access from: (faces, vertices, face_uvs)\n",
"\n",
"Note that auto-compute is still supported, e.g. after access:\n",
" vertex_normals: [\n",
" 0: [42, 3] (torch.float32)[cpu] \n",
" 1: [42, 3] (torch.float32)[cpu] \n",
" 2: [482, 3] (torch.float32)[cpu] \n",
" ]\n"
]
}
],
"source": [
"tmp = SurfaceMesh.cat([mesh, kal.io.usd.import_mesh(sample_mesh_path('pizza.usda'), **import_args)],\n",
" fixed_topology=False)\n",
"print(tmp)\n",
"tmp.vertex_normals\n",
"print(f'\\nNote that auto-compute is still supported, e.g. after access:')\n",
"print(f'{tmp.describe_attribute(\"vertex_normals\")}')\n",
"\n",
"del tmp # not needed later"
]
},
{
"cell_type": "markdown",
"id": "4b7855ae",
"metadata": {},
"source": [
"## Convenience Methods and Mutability\n",
"\n",
"Now let's see a few useful capabilities of `SurfaceMesh`. "
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "640359ed",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"SurfaceMesh object with batching strategy FIXED\n",
" vertices: [2, 42, 3] (torch.float32)[cuda:0] \n",
" faces: [80, 3] (torch.int64)[cuda:0] \n",
" face_normals: [2, 80, 3, 3] (torch.float32)[cuda:0] \n",
" face_uvs: [2, 80, 3, 2] (torch.float32)[cuda:0] \n",
" vertex_normals: [2, 42, 3] (torch.float32)[cpu] \n",
"material_assignments: [2, 80] (torch.int16)[cuda:0] \n",
" materials: [\n",
" 0: list of length 1\n",
" 1: list of length 1\n",
" ]\n",
" face_vertices: if possible, computed on access from: (faces, vertices)\n",
" vertex_tangents: if possible, computed on access from: (faces, vertices, face_uvs)\n",
" vertices: [2, 42, 3] (torch.float32)[cuda:0] - [min -1.0000, max 1.0000, mean -0.0000] \n",
" vertices: [2, 42, 3] (torch.float32)[cuda:0] - [min -0.5000, max 0.5000, mean -0.0000] \n",
"dict_keys(['vertices', 'faces', 'face_normals', 'face_uvs', 'vertex_normals', 'material_assignments'])\n"
]
}
],
"source": [
"# Recall that mesh contains two fixed topology meshes\n",
"\n",
"# Let's move it to cuda (you can also specify particular attributes to move)\n",
"mesh = mesh.cuda()\n",
"\n",
"# Let's say we actually don't need vertex_normals on the GPU\n",
"mesh = mesh.cpu(attributes=['vertex_normals'])\n",
"print(mesh)\n",
"\n",
"# We can also directly set mesh attributes, for example:\n",
"print(mesh.describe_attribute('vertices', print_stats=True))\n",
"mesh.vertices = kal.ops.pointcloud.center_points(mesh.vertices, normalize=True)\n",
"print(mesh.describe_attribute('vertices', print_stats=True))\n",
"\n",
"# Mesh also supports copy and deepcopy\n",
"mesh_copy = copy.copy(mesh)\n",
"mesh_copy = copy.deepcopy(mesh)\n",
"\n",
"# Finally, mesh can be converted to a simple dict\n",
"mesh_dict = mesh.as_dict(only_tensors=True)\n",
"print(mesh_dict.keys())"
]
},
{
"cell_type": "markdown",
"id": "bdd71c5b",
"metadata": {},
"source": [
"## Optimization and Gradients\n",
"\n",
"It is also possible to optimize mesh attributes by going through auto-computed attributes. However, take care to set `requires_grad` before auto-computed attribute is cached. This causes auto-computed attributes to be computed every time, allowing gradients to flow."
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "9b053b52",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Has face_vertices? False\n",
"computed face_vertices: [2, 80, 3, 3] (torch.float32)[cuda:0] \n",
"Were face_vertices cached? False\n",
"\n",
"Sample loss 1.7655433416366577\n"
]
}
],
"source": [
"# Let's try to optimize vertices\n",
"mesh.vertices.requires_grad = True\n",
"\n",
"# Check that mesh does not cache face_vertices\n",
"print(f'Has face_vertices? {mesh.has_attribute(\"face_vertices\")}')\n",
"\n",
"# Check that we can actually compute them\n",
"face_vertices = mesh.face_vertices\n",
"print_tensor(face_vertices, name='computed face_vertices')\n",
"\n",
"# However, because mesh.vertices.requires_grad, this value is not cached\n",
"print(f'Were face_vertices cached? {mesh.has_attribute(\"face_vertices\")}')\n",
"\n",
"# Now we can use mesh.face_vertices in a loss function, while optimizing mesh.vertices, e.g.:\n",
"sample_pt_cloud = torch.randn((2, 100, 3), dtype=mesh.vertices.dtype, device=mesh.vertices.device)\n",
"sample_loss = kal.metrics.trianglemesh.point_to_mesh_distance(sample_pt_cloud, mesh.face_vertices)[0].mean()\n",
"print(f'\\nSample loss {sample_loss}')"
]
},
{
"cell_type": "markdown",
"id": "67097c5d",
"metadata": {},
"source": [
"## Exporting\n",
"\n",
"Automatic conversion to `LIST` batching also makes it easy to export a batch of USD meshes to file."
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "bf0591c9",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"Exporting to USD: 100%|█████████████████████████████████████████████| 2/2 [00:00<00:00, 190.03mesh/s]"
]
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"SurfaceMesh object with batching strategy LIST\n",
" vertices: [\n",
" 0: [42, 3] (torch.float32)[cpu] \n",
" 1: [42, 3] (torch.float32)[cpu] \n",
" ]\n",
" faces: [\n",
" 0: [80, 3] (torch.int64)[cpu] \n",
" 1: [80, 3] (torch.int64)[cpu] \n",
" ]\n",
" face_normals: [\n",
" 0: [80, 3, 3] (torch.float32)[cpu] \n",
" 1: [80, 3, 3] (torch.float32)[cpu] \n",
" ]\n",
" face_vertices: if possible, computed on access from: (faces, vertices)\n",
" face_uvs: if possible, computed on access from: (uvs, face_uvs_idx)\n",
" vertex_normals: if possible, computed on access from: (faces, face_normals)\n",
" vertex_tangents: if possible, computed on access from: (faces, vertices, face_uvs)\n",
"True\n",
"True\n",
"True\n"
]
},
{
"name": "stderr",
"output_type": "stream",
"text": [
"\n",
"/kaolin/kaolin/io/usd/mesh.py:371: UserWarning: Some child prims for /World/Meshes/mesh_0 are missing uvs; skipping importing uvs.\n",
" warnings.warn(f'Some child prims for {scene_path} are missing {k}; skipping importing {k}.', UserWarning)\n",
"/kaolin/kaolin/io/usd/mesh.py:371: UserWarning: Some child prims for /World/Meshes/mesh_0 are missing face_uvs_idx; skipping importing face_uvs_idx.\n",
" warnings.warn(f'Some child prims for {scene_path} are missing {k}; skipping importing {k}.', UserWarning)\n",
"/kaolin/kaolin/io/usd/mesh.py:371: UserWarning: Some child prims for /World/Meshes/mesh_1 are missing uvs; skipping importing uvs.\n",
" warnings.warn(f'Some child prims for {scene_path} are missing {k}; skipping importing {k}.', UserWarning)\n",
"/kaolin/kaolin/io/usd/mesh.py:371: UserWarning: Some child prims for /World/Meshes/mesh_1 are missing face_uvs_idx; skipping importing face_uvs_idx.\n",
" warnings.warn(f'Some child prims for {scene_path} are missing {k}; skipping importing {k}.', UserWarning)\n"
]
}
],
"source": [
"mesh = mesh.set_batching(SurfaceMesh.Batching.LIST)\n",
"\n",
"# Note: you can only run this once due to USD caching; restart Kernel to rerun cell without errors\n",
"kal.io.usd.export_meshes('/tmp/out.usd', vertices=mesh.vertices, faces=mesh.faces, face_normals=mesh.face_normals)\n",
"\n",
"# Verify we can read back the same meshes we exported\n",
"imported_meshes = SurfaceMesh.cat(\n",
" kal.io.usd.import_meshes('/tmp/out.usd', with_normals=True), fixed_topology=False)\n",
"mesh = mesh.cpu()\n",
"print(imported_meshes)\n",
"print(kal.utils.testing.contained_torch_equal(mesh.vertices, imported_meshes.vertices, approximate=True))\n",
"print(kal.utils.testing.contained_torch_equal(mesh.faces, imported_meshes.faces))\n",
"print(kal.utils.testing.contained_torch_equal(mesh.face_normals, imported_meshes.face_normals, approximate=True))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.16"
}
},
"nbformat": 4,
"nbformat_minor": 5
}