{ "cells": [ { "cell_type": "code", "execution_count": null, "id": "69114969", "metadata": {}, "outputs": [], "source": [ "!pip install matplotlib --quiet\n", "import glob\n", "import math\n", "from matplotlib import pyplot as plt\n", "import nvdiffrast\n", "import os\n", "import torch\n", "from IPython.display import display, clear_output\n", "\n", "from tutorial_common import COMMON_DATA_DIR\n", "import kaolin as kal\n", "\n", "nvctx = nvdiffrast.torch.RasterizeCudaContext(device='cuda')" ] }, { "cell_type": "markdown", "id": "dd8f9d9e", "metadata": {}, "source": [ "## Load Mesh information" ] }, { "cell_type": "code", "execution_count": null, "id": "4963d2ce", "metadata": {}, "outputs": [], "source": [ "# Set KAOLIN_TEST_SHAPENETV2_PATH env variable, or replace by your shapenet path\n", "SHAPENETV2_PATH = os.getenv('KAOLIN_TEST_SHAPENETV2_PATH')\n", "\n", "if SHAPENETV2_PATH is not None:\n", " ds = kal.io.shapenet.ShapeNetV2(root=SHAPENETV2_PATH,\n", " categories=['car'],\n", " train=True, split=1.,\n", " with_materials=True,\n", " output_dict=True)\n", " mesh = ds[0]['mesh']\n", "else:\n", " # Load a specific obj instead\n", " OBJ_PATH = os.path.join(COMMON_DATA_DIR, 'meshes', 'fox.obj')\n", " mesh = kal.io.obj.import_mesh(OBJ_PATH, with_materials=True, triangulate=True)\n", "\n", "# Batch, move to GPU and center and normalize vertices in the range [-0.5, 0.5]\n", "mesh = mesh.to_batched().cuda()\n", "mesh.vertices = kal.ops.pointcloud.center_points(mesh.vertices, normalize=True)\n", "print(mesh.to_string(print_stats=True))\n", "\n", "# Preprocess OBJ materials into a format we will use for rendering \n", "materials = [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", "\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))\n", "mesh.face_uvs_idx[mesh.face_uvs_idx == -1] = mesh.uvs.shape[1] - 1" ] }, { "cell_type": "markdown", "id": "03e898f4", "metadata": {}, "source": [ "## Instantiate a camera\n", "\n", "With the general constructor `Camera.from_args()` the underlying constructors are `CameraExtrinsics.from_lookat()` and `PinholeIntrinsics.from_fov`" ] }, { "cell_type": "code", "execution_count": null, "id": "c6eee7ab", "metadata": {}, "outputs": [], "source": [ "cam = kal.render.camera.Camera.from_args(eye=torch.tensor([2., 0., 0.]),\n", " at=torch.tensor([0., 0., 0.]),\n", " up=torch.tensor([0., 1., 0.]),\n", " fov=math.pi * 45 / 180,\n", " width=512, height=512, device='cuda')" ] }, { "cell_type": "markdown", "id": "4fff8eb1", "metadata": {}, "source": [ "## Rendering a mesh\n", "\n", "Here we are rendering the loaded mesh with [nvdiffrast](https://github.com/NVlabs/nvdiffrast) using the camera object created above" ] }, { "cell_type": "code", "execution_count": null, "id": "5e4b8a49", "metadata": {}, "outputs": [], "source": [ "def render(mesh, cam):\n", " vertices_camera = cam.extrinsics.transform(mesh.vertices)\n", " face_vertices_camera = kal.ops.mesh.index_vertices_by_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", " \n", " # Create a fake W (See nvdiffrast documentation)\n", " proj = cam.projection_matrix()[None]\n", " homogeneous_vecs = kal.render.camera.up_to_homogeneous(\n", " vertices_camera\n", " )[..., None]\n", "\n", " vertices_clip = (proj @ homogeneous_vecs).squeeze(-1)\n", "\n", " rast = nvdiffrast.torch.rasterize(\n", " nvctx, vertices_clip, mesh.faces.int(),\n", " (cam.height, cam.width), grad_db=False\n", " )\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, rast0, mesh.face_uvs_idx[0, ...].int()\n", " )[0] % 1.\n", "\n", " img = torch.zeros((1, 512, 512, 3), dtype=torch.float, device='cuda')\n", " # Obj meshes can be composed of multiple materials\n", " # so at rendering we need to interpolate from corresponding materials\n", " im_material_idx = mesh.material_assignments[0, ...][face_idx] # Assume single assignments map\n", " im_material_idx[face_idx == -1] = -1\n", " for i, material in enumerate(materials):\n", " mask = im_material_idx == i\n", " mask_idx = torch.nonzero(mask, as_tuple=False)\n", " _texcoords = uv_map[mask]\n", " if _texcoords.shape[0] > 0:\n", " pixel_val = nvdiffrast.torch.texture(\n", " materials[i].contiguous(),\n", " _texcoords.reshape(1, 1, -1, 2).contiguous(),\n", " filter_mode='linear'\n", " )\n", " img[mask] = pixel_val[0, 0]\n", "\n", " # Need to flip the image because opengl\n", " return torch.clamp(img * hard_mask, 0., 1.)[0]" ] }, { "cell_type": "markdown", "id": "4978003c", "metadata": {}, "source": [ "# Moving the camera\n", "\n", "Once the camera is created you can move it using `cam.move_up()`, `cam.move_right` and `cam.move_forward()`.\n", "\n", "To be noted that in OpenGL `forward` in the camera space is actually toward the viewer (so it actually move away from the object looked at)\n", "\n", "\"Markdown\n", "\"Markdown\n", "\n", "Below is a simple interactive rendering, where buttons are linked to camera methods for moving it." ] }, { "cell_type": "code", "execution_count": null, "id": "83dd84cf", "metadata": { "tags": [] }, "outputs": [], "source": [ "%matplotlib inline\n", "from ipywidgets import Button, Output, HBox, VBox\n", "\n", "output = Output()\n", "\n", "def update():\n", " \"\"\"Update the image buffer\"\"\"\n", " with output:\n", " clear_output(wait=True)\n", " fig, ax = plt.subplots(figsize=(6, 6))\n", " ax.axis('off')\n", " image_data = render(mesh, cam).cpu()\n", " ax.imshow(image_data)\n", " plt.show()\n", "\n", "def on_button_up_clicked(b):\n", " \"\"\"Callback on Up\"\"\"\n", " cam.move_up(0.1)\n", " update()\n", "\n", "def on_button_down_clicked(b):\n", " \"\"\"Callback on Down\"\"\"\n", " cam.move_up(-0.1)\n", " update()\n", "\n", "def on_button_left_clicked(b):\n", " \"\"\"Callback on Left\"\"\"\n", " cam.move_right(-0.1)\n", " update()\n", "\n", "def on_button_right_clicked(b):\n", " \"\"\"Callback on Right\"\"\"\n", " cam.move_right(0.1)\n", " update()\n", "\n", "def on_button_forward_clicked(b):\n", " \"\"\"Callback on Forward\n", " \n", " Note: Forward is actually on the back of the camera\n", " \"\"\"\n", " cam.move_forward(0.1)\n", " update()\n", " \n", "def on_button_backward_clicked(b):\n", " \"\"\"Callback on Backward\n", " \n", " Note: Forward is actually on the back of the camera\n", " \"\"\"\n", " cam.move_forward(-0.1)\n", " update()\n", "\n", "def on_button_increase_pitch_clicked(b):\n", " \"\"\"Callback on increasing Pitch\"\"\"\n", " cam.rotate(pitch=0.1)\n", " update()\n", " \n", "def on_button_decrease_pitch_clicked(b):\n", " \"\"\"Callback on decreasing Pitch\"\"\"\n", " cam.rotate(pitch=-0.1)\n", " update()\n", " \n", "def on_button_increase_yaw_clicked(b):\n", " \"\"\"Callback on increasing Yaw\"\"\"\n", " cam.rotate(yaw=0.1)\n", " update()\n", " \n", "def on_button_decrease_yaw_clicked(b):\n", " \"\"\"Callback on decreasing Yaw\"\"\"\n", " cam.rotate(yaw=-0.1)\n", " update()\n", " \n", "def on_button_increase_roll_clicked(b):\n", " \"\"\"Callback on increasing Roll\"\"\"\n", " cam.rotate(roll=0.1)\n", " update()\n", " \n", "def on_button_decrease_roll_clicked(b):\n", " \"\"\"Callback on decreasing Roll\"\"\"\n", " cam.rotate(roll=-0.1)\n", " update()\n", "\n", "btn_up = Button(description=\"Up\")\n", "btn_down = Button(description=\"Bottom\")\n", "btn_left = Button(description=\"Left\")\n", "btn_right = Button(description=\"Right\")\n", "btn_forward = Button(description=\"Forward\")\n", "btn_backward = Button(description=\"Backward\")\n", "btn_increase_pitch = Button(description=\"Increase\\nPitch\")\n", "btn_decrease_pitch = Button(description=\"Decrease\\nPitch\")\n", "btn_increase_yaw = Button(description=\"Increase\\nYaw\")\n", "btn_decrease_yaw = Button(description=\"Decrease\\nYaw\")\n", "btn_increase_roll = Button(description=\"Increase\\nRoll\")\n", "btn_decrease_roll = Button(description=\"Decrease\\nRoll\")\n", "\n", "btn_up.on_click(on_button_up_clicked)\n", "btn_down.on_click(on_button_down_clicked)\n", "btn_left.on_click(on_button_left_clicked)\n", "btn_right.on_click(on_button_right_clicked)\n", "btn_forward.on_click(on_button_forward_clicked)\n", "btn_backward.on_click(on_button_backward_clicked)\n", "btn_increase_pitch.on_click(on_button_increase_pitch_clicked)\n", "btn_decrease_pitch.on_click(on_button_decrease_pitch_clicked)\n", "btn_increase_yaw.on_click(on_button_increase_yaw_clicked)\n", "btn_decrease_yaw.on_click(on_button_decrease_yaw_clicked)\n", "btn_increase_roll.on_click(on_button_increase_roll_clicked)\n", "btn_decrease_roll.on_click(on_button_decrease_roll_clicked)\n", "\n", "row1 = HBox([btn_up, btn_left, btn_down, btn_right, btn_forward, btn_backward])\n", "row2 = HBox([btn_increase_pitch, btn_decrease_pitch, btn_increase_yaw, btn_decrease_yaw, btn_increase_roll, btn_decrease_roll])\n", "controls = VBox([row1, row2])\n", "\n", "display(controls, output)\n", "update()" ] } ], "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.12.3" } }, "nbformat": 4, "nbformat_minor": 5 }