Skip to content

Commit

Permalink
add query feature and fix a UV bug in examples using nvdiffrast (#742)
Browse files Browse the repository at this point in the history 
Signed-off-by: Clement Fuji Tsang <cfujitsang@nvidia.com>
  • Loading branch information
@Caenorst
Caenorst committed Jul 13, 2023
1 parent 31f4e1e commit 96f19cf
Show file tree
Hide file tree
Showing 4 changed files with 220 additions and 113 deletions.
101 changes: 67 additions & 34 deletions examples/tutorial/camera_and_rasterization.ipynb
View file

Large diffs are not rendered by default.

135 changes: 75 additions & 60 deletions examples/tutorial/interactive_visualizer.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,10 +36,7 @@
"import kaolin as kal\n",
"\n",
"import nvdiffrast\n",
"glctx = nvdiffrast.torch.RasterizeGLContext(False, device='cuda')\n",
"\n",
"def print_tensor(t, **kwargs):\n",
" print(kal.utils.testing.tensor_info(t, **kwargs))"
"glctx = nvdiffrast.torch.RasterizeGLContext(False, device='cuda')"
]
},
{
Expand All @@ -63,17 +60,20 @@
"output_type": "stream",
"text": [
"SurfaceMesh object with batching strategy FIXED\n",
" vertices: [1, 52081, 3] (torch.float32)[cuda:0] \n",
" uvs: [1, 633, 2] (torch.float32)[cuda:0] \n",
" faces: [200018, 3] (torch.int64)[cuda:0] \n",
" face_uvs_idx: [1, 200018, 3] (torch.int64)[cuda:0] \n",
"material_assignments: [1, 200018] (torch.int16)[cuda:0] \n",
" materials: list of length 44\n",
" vertex_normals: will be computed on access from (if present): (faces, face_normals)\n",
" face_normals: will be computed on access from (if present): (normals, face_normals_idx) or (vertices, faces)\n",
" face_uvs: will be computed on access from (if present): (uvs, face_uvs_idx)\n",
" face_vertices: will be computed on access from (if present): (faces, vertices)\n",
"['batching', 'allow_auto_compute', 'unset_attributes_return_none', 'materials', 'vertices', 'uvs', 'faces', 'face_uvs_idx', 'material_assignments']\n"
" vertices: [1, 5002, 3] (torch.float32)[cuda:0] \n",
" normals: [1, 5002, 3] (torch.float32)[cuda:0] \n",
" uvs: [1, 5505, 2] (torch.float32)[cuda:0] \n",
" faces: [10000, 3] (torch.int64)[cuda:0] \n",
" face_normals_idx: [1, 10000, 3] (torch.int64)[cuda:0] \n",
" face_uvs_idx: [1, 10000, 3] (torch.int64)[cuda:0] \n",
"material_assignments: [1, 10000] (torch.int16)[cuda:0] \n",
" materials: [\n",
" 0: list of length 1\n",
" ]\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",
" vertex_normals: if possible, computed on access from: (faces, face_normals)\n",
" face_uvs: if possible, computed on access from: (uvs, face_uvs_idx)\n"
]
}
],
Expand All @@ -98,15 +98,15 @@
"mesh.vertices = kal.ops.pointcloud.center_points(mesh.vertices, normalize=True)\n",
"print(mesh)\n",
"\n",
"diffuse_maps = [m['map_Kd'].permute(2, 0, 1).unsqueeze(0).cuda().float() / 255. if 'map_Kd' in m else\n",
" m['Kd'].reshape(1, 3, 1, 1).cuda()\n",
"diffuse_maps = [m['map_Kd'].unsqueeze(0).cuda().float() / 255. if 'map_Kd' in m else\n",
" m['Kd'].reshape(1, 1, 1, 3).cuda()\n",
" for m in mesh.materials[0]]\n",
"specular_maps = [m['map_Ks'].permute(2, 0, 1).unsqueeze(0).cuda().float() / 255. if 'map_Ks' in m else\n",
" m['Ks'].reshape(1, 3, 1, 1).cuda()\n",
"specular_maps = [m['map_Ks'].unsqueeze(0).cuda().float() / 255. if 'map_Ks' in m else\n",
" m['Ks'].reshape(1, 1, 1, 3).cuda()\n",
" for m in mesh.materials[0]]\n",
"\n",
"# Use a single diffuse color as backup when map doesn't exist (and face_uvs_idx == -1)\n",
"mesh.uvs = torch.nn.functional.pad(mesh.uvs, (0, 0, 0, 1)) % 1.\n",
"mesh.uvs = torch.nn.functional.pad(mesh.uvs, (0, 0, 0, 1))\n",
"mesh.face_uvs_idx[mesh.face_uvs_idx == -1] = mesh.uvs.shape[1] - 1"
]
},
Expand Down Expand Up @@ -204,31 +204,41 @@
"\n",
"def base_render(mesh, camera, height, width):\n",
" \"\"\"Base function for rendering using separate height and width, assuming batch_size=1\"\"\"\n",
" transformed_vertices = camera.transform(mesh.vertices)\n",
" vertices_camera = camera.extrinsics.transform(mesh.vertices)\n",
" face_vertices_camera = kal.ops.mesh.index_vertices_by_faces(\n",
" transformed_vertices, mesh.faces)\n",
" vertices_camera, mesh.faces)\n",
" face_normals_z = kal.ops.mesh.face_normals(\n",
" face_vertices_camera,\n",
" unit=True\n",
" )[..., -1:].contiguous()\n",
" # Create a fake W (See nvdiffrast documentation)\n",
" pos = torch.nn.functional.pad(\n",
" transformed_vertices, (0, 1), mode='constant', value=1.\n",
" ).contiguous()\n",
"\n",
" # Projection: nvdiffrast take clip coordinates as input to apply barycentric perspective correction.\n",
" # Using `camera.intrinsics.transform(vertices_camera) would return the normalized device coordinates.\n",
" proj = camera.projection_matrix()[None]\n",
" homogeneous_vecs = kal.render.camera.up_to_homogeneous(\n",
" vertices_camera\n",
" )[..., None]\n",
" vertices_clip = (proj @ homogeneous_vecs).squeeze(-1)\n",
"\n",
" rast = nvdiffrast.torch.rasterize(\n",
" glctx, pos, mesh.faces.int(), (height, width), grad_db=False)\n",
" hard_mask = rast[0][:, :, :, -1:] != 0\n",
" face_idx = (rast[0][..., -1].long() - 1).contiguous()\n",
" glctx, vertices_clip, mesh.faces.int(),\n",
" (height, width), grad_db=False\n",
" )\n",
" # nvdiffrast rasteriztion output is y-up, we need to flip as our display is y-down\n",
" rast0 = torch.flip(rast[0], dims=(1,))\n",
" hard_mask = rast0[:, :, :, -1:] != 0\n",
" face_idx = (rast0[..., -1].long() - 1).contiguous()\n",
"\n",
" uv_map = nvdiffrast.torch.interpolate(\n",
" mesh.uvs, rast[0], mesh.face_uvs_idx[0, ...].int())[0]\n",
" mesh.uvs, rast0, mesh.face_uvs_idx[0, ...].int()\n",
" )[0] % 1.\n",
" \n",
" if mesh.has_attribute('normals') and mesh.has_attribute('face_normals_idx'):\n",
" im_world_normals = nvdiffrast.torch.interpolate(\n",
" mesh.normals, rast[0], mesh.face_normals_idx[0, ...].int())[0]\n",
" mesh.normals, rast0, mesh.face_normals_idx[0, ...].int())[0]\n",
" else:\n",
" im_world_normals = nvdiffrast.torch.interpolate(\n",
" mesh.face_normals.reshape(len(mesh), -1, 3), rast[0],\n",
" mesh.face_normals.reshape(len(mesh), -1, 3), rast0,\n",
" torch.arange(mesh.faces.shape[0] * 3, device='cuda', dtype=torch.int).reshape(-1, 3)\n",
" )[0]\n",
" \n",
Expand All @@ -253,18 +263,21 @@
" for i, material in enumerate(diffuse_maps):\n",
" mask = im_material_idx == i\n",
" mask_idx = torch.nonzero(mask, as_tuple=False)\n",
" _texcoords = uv_map[mask] * 2. - 1.\n",
" _texcoords = uv_map[mask]\n",
" _texcoords[:, 1] = -_texcoords[:, 1]\n",
" pixel_val = torch.nn.functional.grid_sample(\n",
" diffuse_maps[i], _texcoords.reshape(1, 1, -1, 2),\n",
" mode='bilinear', align_corners=False,\n",
" padding_mode='border')\n",
" albedo[mask] = pixel_val[0, :, 0].permute(1, 0)\n",
" pixel_val = torch.nn.functional.grid_sample(\n",
" specular_maps[i], _texcoords.reshape(1, 1, -1, 2),\n",
" mode='bilinear', align_corners=False,\n",
" padding_mode='border')\n",
" spec_albedo[mask] = pixel_val[0, :, 0].permute(1, 0)\n",
" if _texcoords.shape[0] > 0:\n",
" pixel_val = nvdiffrast.torch.texture(\n",
" diffuse_maps[i].contiguous(),\n",
" _texcoords.reshape(1, 1, -1, 2).contiguous(),\n",
" filter_mode='linear'\n",
" )\n",
" albedo[mask] = pixel_val[0, 0]\n",
" pixel_val = nvdiffrast.torch.texture(\n",
" specular_maps[i].contiguous(),\n",
" _texcoords.reshape(1, 1, -1, 2).contiguous(),\n",
" filter_mode='linear'\n",
" )\n",
" spec_albedo[mask] = pixel_val[0, 0] #.permute(1, 0)\n",
" img = torch.zeros((1, height, width, 3),\n",
" dtype=torch.float, device='cuda')\n",
" sg_x, sg_y, sg_z = kal.ops.coords.spherical2cartesian(azimuth, elevation)\n",
Expand All @@ -291,10 +304,11 @@
" )\n",
" img[hard_mask.squeeze(-1)] += specular_effect\n",
"\n",
" # Need to flip the image because opengl\n",
" return (torch.flip(torch.clamp(\n",
" img * hard_mask, 0., 1.\n",
" )[0], dims=(0,)) * 255.).to(torch.uint8)\n",
" # 'img' is the displayed image, while the other value `face_idx` is only printed on query\n",
" return {\n",
" 'img': (torch.clamp(img * hard_mask, 0., 1.)[0] * 255.).to(torch.uint8),\n",
" 'face_idx': face_idx[0]\n",
" }\n",
"\n",
"def render(camera):\n",
" \"\"\"Render using camera dimension.\n",
Expand Down Expand Up @@ -328,13 +342,14 @@
"execution_count": 5,
"id": "f10f6ab7-662c-4cf6-8af4-5f28ae79d795",
"metadata": {
"scrolled": false,
"tags": []
},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "675199e8811d4315be9106a4b9839459",
"model_id": "68ea5e664b064a2eb97ef8ccdca004b9",
"version_major": 2,
"version_minor": 0
},
Expand All @@ -348,7 +363,7 @@
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "e5d1372ce3cd4a29be14e2738c80d94f",
"model_id": "5083ae6b7e594e3fad555f5c03d75ea3",
"version_major": 2,
"version_minor": 0
},
Expand Down Expand Up @@ -395,7 +410,7 @@
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "dd7e3bc6947241e98e97211a9611027a",
"model_id": "33ecbe9c9d92447691b7c93ac72c1c79",
"version_major": 2,
"version_minor": 0
},
Expand All @@ -409,7 +424,7 @@
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "b152d54454194a1a9d5377b0ca4625c4",
"model_id": "ddd1612cd0224eaeb6b11ee7d08d9d21",
"version_major": 2,
"version_minor": 0
},
Expand Down Expand Up @@ -456,12 +471,12 @@
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "84164d50ea1343bcba76b9625c93962b",
"model_id": "9bcd10fb4d7542ebbe094fd29ee73643",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"HBox(children=(Canvas(height=512, width=512), interactive(children=(FloatSlider(value=0.0, description='Elevat…"
"HBox(children=(Canvas(height=512, width=512), interactive(children=(FloatSlider(value=1.0471975803375244, desc…"
]
},
"metadata": {},
Expand All @@ -470,7 +485,7 @@
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "d288140ef7c34e1d8d7a8e2d2dfc2174",
"model_id": "82638dc05bfa4970a09aee4d197f95d0",
"version_major": 2,
"version_minor": 0
},
Expand Down Expand Up @@ -514,7 +529,7 @@
" visualizer.render_update()\n",
" \n",
"elevation_slider = FloatSlider(\n",
" value=0.,\n",
" value=elevation.item(),\n",
" min=-math.pi / 2.,\n",
" max=math.pi / 2.,\n",
" step=0.1,\n",
Expand All @@ -525,7 +540,7 @@
")\n",
"\n",
"azimuth_slider = FloatSlider(\n",
" value=0.,\n",
" value=azimuth.item(),\n",
" min=-math.pi,\n",
" max=math.pi,\n",
" step=0.1,\n",
Expand All @@ -536,7 +551,7 @@
")\n",
"\n",
"amplitude_slider = FloatSlider(\n",
" value=5.,\n",
" value=amplitude[0,0].item(),\n",
" min=0.1,\n",
" max=20.,\n",
" step=0.1,\n",
Expand All @@ -547,7 +562,7 @@
")\n",
"\n",
"sharpness_slider = FloatSlider(\n",
" value=5.,\n",
" value=sharpness.item(),\n",
" min=0.1,\n",
" max=20.,\n",
" step=0.1,\n",
Expand All @@ -573,7 +588,7 @@
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
Expand All @@ -587,7 +602,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.13"
"version": "3.10.10"
}
},
"nbformat": 4,
Expand Down

0 comments on commit 96f19cf

Please sign in to comment.