Skip to content

Commit

Permalink
Simulating INRIA splats using Simplicits in the notebook (#819)
Browse files Browse the repository at this point in the history 
Small updates to cleanup all simplicits notebooks

User-friendly changes and explanations; disable simulate click when simulation is running

Small usability / wording changes to interactive notebook

Minor updates to notebooks. Commented out splat install. Added new vars for visualizers. Added original sim_threading back for interactive turntable

Swapped splat repos to inria main

Fixed splatting memory issue by adding torch.no_grad() in run_sim()

Signed-off-by: Vismay <vismaym@nvidia.com>
Co-authored-by: Vismay <vismaym@nvidia.com>
  • Loading branch information
@itsvismay
itsvismay and Vismay authored Jul 27, 2024
1 parent e986d42 commit c7593f0
Show file tree
Hide file tree
Showing 6 changed files with 846 additions and 695 deletions.
0 examples/tutorial/physics_overview.md → ...ples/tutorial/physics/physics_overview.md
View file
File renamed without changes.
203 changes: 111 additions & 92 deletions examples/tutorial/physics/simplicits_easy_api.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,14 @@
"metadata": {},
"source": [
"# Using Simplicit's Easy API To Simulate Avocado Mesh\n",
"A simple way to use the simplicit's code base. We can create a simple object, train it, simulate it and visualize all in a very few lines of code."
"[Simplicits](https://research.nvidia.com/labs/toronto-ai/simplicits/) is a mesh-free, representation-agnostic way to simulation elastic deformations. \n",
"\n",
"Here's a simple way to use the simplicit's code base. We can create a simple object, train it, simulate it and visualize all in a very few lines of code via our `easy_api`."
]
},
{
"cell_type": "code",
"execution_count": 7,
"execution_count": null,
"id": "d6dcc14f-a4a3-4c9c-87f8-5a28d6f5a2d4",
"metadata": {},
"outputs": [],
Expand All @@ -27,8 +29,9 @@
"from IPython.display import display\n",
"from ipywidgets import Button, HBox, VBox\n",
"\n",
"logging.basicConfig(level=logging.DEBUG, stream=sys.stdout)\n",
"logging.basicConfig(level=logging.INFO, stream=sys.stdout)\n",
"logger = logging.getLogger(__name__)\n",
"logging.getLogger('kaolin.physics').setLevel(logging.DEBUG)\n",
"\n",
"sys.path.append(str(Path(\"..\")))\n",
"from tutorial_common import COMMON_DATA_DIR\n",
Expand All @@ -46,50 +49,20 @@
"metadata": {},
"source": [
"## Loading Geometry\n",
"Simplicies works with any geometry. Meshes, pointclouds, SDFs, Gaussian splats, and more. For this example we "
"Simplicits works with any geometry: meshes, pointclouds, SDFs, Gaussian splats, and more. For this example, we will use a mesh:"
]
},
{
"cell_type": "code",
"execution_count": 8,
"execution_count": null,
"id": "40a49d68-c73e-4a06-9708-695e2f201318",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"WARNING:root:Cannot convert Ka from obj mtl to PBR spec; use raw_materials=True to import raw values\n",
"SurfaceMesh object with batching strategy NONE\n",
" vertices: [406, 3] (torch.float32)[cuda:0] \n",
" faces: [682, 3] (torch.int64)[cuda:0] \n",
" normals: [377, 3] (torch.float32)[cuda:0] \n",
" face_normals_idx: [682, 3] (torch.int64)[cuda:0] \n",
" uvs: [406, 2] (torch.float32)[cuda:0] \n",
" face_uvs_idx: [682, 3] (torch.int64)[cuda:0] \n",
"material_assignments: [682] (torch.int16)[cuda:0] \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 (vertex_normals, faces) 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, face_vertices, face_uvs, vertex_normals)\n",
" vertex_colors: if possible, computed on access from: (faces, face_colors)\n",
" vertex_features: if possible, computed on access from: (faces, face_features)\n",
" face_tangents: if possible, computed on access from: (faces, vertex_tangents)\n",
" face_colors: if possible, computed on access from: (faces, vertex_colors)\n",
" face_features: if possible, computed on access from: (faces, vertex_features)\n"
]
}
],
"outputs": [],
"source": [
"# Import and triangulate to enable rasterization; move to GPU\n",
"mesh = kal.io.import_mesh(sample_mesh_path('avocado.obj'), triangulate=True).cuda()\n",
"# Normalize so it is easy to set up default camera\n",
"mesh.vertices = kal.ops.pointcloud.center_points(mesh.vertices.unsqueeze(0), normalize=True).squeeze(0) \n",
"orig_vertices = mesh.vertices.clone() # Also save original undeformed vertices\n",
"\n",
"# Inspect\n",
"print(mesh)"
]
},
Expand All @@ -98,13 +71,13 @@
"id": "c6e54cb8-e2d2-4921-814b-62ef9a7b431e",
"metadata": {},
"source": [
"## Setup\n",
"Lets set up the vertices and material parameters of our object"
"## Sample Geometry\n",
"To enable simulation we need point samples (vertices in this case), and physical material parameters per point. Lets set this up."
]
},
{
"cell_type": "code",
"execution_count": 4,
"execution_count": null,
"id": "9c11c603-4b6f-4229-9303-94c40c8d06b1",
"metadata": {
"scrolled": true
Expand All @@ -129,37 +102,37 @@
"id": "66146660-e708-4cf7-815c-a1fd5d337053",
"metadata": {},
"source": [
"## Create a SimplicitsObject\n",
"We make a \"simplicits object\" using our api. Then we train it for num_steps."
"## Create and Train a SimplicitsObject\n",
"We encapsulate everything Simplicits method needs to know about the simulated object in a `SimplicitsObject` instance. Once the object is created, we need to run training to learn reduced degrees of freedom our simulator can use. \n",
"\n",
"**This will take a couple of minutes.** Please be patient."
]
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"id": "d333391d-618a-4e25-8866-acaec1e3fac4",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"DEBUG:kaolin.physics.simplicits.easy_api:Training step: 0, le: 428.7832946777344, lo: 32387982.0\n",
"DEBUG:kaolin.physics.simplicits.easy_api:Training step: 100, le: 23.801427841186523, lo: 154803.328125\n",
"DEBUG:kaolin.physics.simplicits.easy_api:Training step: 200, le: 112.9085464477539, lo: 116969.1796875\n",
"DEBUG:kaolin.physics.simplicits.easy_api:Training step: 300, le: 264.17999267578125, lo: 80777.2734375\n",
"DEBUG:kaolin.physics.simplicits.easy_api:Training step: 400, le: 308.72802734375, lo: 80147.8515625\n",
"DEBUG:kaolin.physics.simplicits.easy_api:Training step: 500, le: 221.305908203125, lo: 80178.9375\n",
"DEBUG:kaolin.physics.simplicits.easy_api:Training step: 600, le: 271.9401550292969, lo: 79970.1328125\n",
"DEBUG:kaolin.physics.simplicits.easy_api:Training step: 700, le: 165.41661071777344, lo: 77258.0703125\n",
"DEBUG:kaolin.physics.simplicits.easy_api:Training step: 800, le: 857.9734497070312, lo: 40415.76953125\n",
"DEBUG:kaolin.physics.simplicits.easy_api:Training step: 900, le: 632.1121215820312, lo: 42898.03125\n"
]
}
],
"outputs": [],
"source": [
"# Initialize and train a Simpicits object to enable simulation\n",
"sim_obj = kal.physics.simplicits.SimplicitsObject(pts, yms, prs, rhos, torch.tensor([approx_volume], dtype=torch.float32, device=\"cuda\"), num_handles=5)\n",
"sim_obj.train(num_steps=10000)"
"print('Training simplicits object. This will take 2-3min. ')\n",
"sim_obj.train(num_steps=10000) # TODO: with next patch add log_interval=1000\n",
"print('Object ready to simulate.')\n",
"\n",
"# sim_obj.load_model('/tmp/avocado_simplicit.pt') # if you saved previously trained object, you can load it instead"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "861e7e85-e160-4c12-8ef5-e1b771901dd2",
"metadata": {},
"outputs": [],
"source": [
"# Optionally, you can save/load this pre-trained object\n",
"# sim_obj.save_model('/tmp/avocado_simplicit.pt')"
]
},
{
Expand All @@ -168,7 +141,7 @@
"metadata": {},
"source": [
"## Create a Scene\n",
"Now that we're trained the object. Lets simulate it by creating a scene"
"Now we are ready to set up all the forces in the scene to simulated as well as simulation settings. For example, here we will add gravity and a floor plane."
]
},
{
Expand All @@ -178,9 +151,17 @@
"metadata": {},
"outputs": [],
"source": [
"# Create a default scene\n",
"scene = kal.physics.simplicits.SimplicitsScene() # default empty scene\n",
"scene.max_newton_steps = 3 #Convergence might not be guaranteed, but runs very fast"
"#Convergence might not be guaranteed with few newton iterations, but runs very fast\n",
"scene.max_newton_steps = 3 "
]
},
{
"cell_type": "markdown",
"id": "09cc0c0d-1e78-42c0-845d-9c6665667c61",
"metadata": {},
"source": [
" The same `SimplicitsObject` can be added to multiple scene. Let's add it to our scene. Not we can reference it within the scene using `obj_idx`."
]
},
{
Expand All @@ -190,82 +171,114 @@
"metadata": {},
"outputs": [],
"source": [
"# Add our object to the scene. The scene copies it into an internal SimulatableObject utility class\n",
"obj_idx = scene.add_object(sim_obj)"
]
},
{
"cell_type": "markdown",
"id": "6935147a-1e0e-444b-a0b9-a50232e4bf73",
"metadata": {},
"source": [
"Lets set gravity and floor forces on the scene"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a7c5640e-1008-44d1-a9ae-caa1aad58576",
"metadata": {},
"outputs": [],
"source": [
"# Add gravity to the scene\n",
"scene.set_scene_gravity(acc_gravity=torch.tensor([0, 9.8, 0]))\n",
"# Add floor to the scene\n",
"scene.set_scene_floor(floor_height=-0.8, floor_axis=1, floor_penalty=1000)"
]
},
{
"cell_type": "markdown",
"id": "63defc99-fffa-4757-872b-34cc4ab27b6d",
"metadata": {},
"source": [
"We can play around with the material parameters of the object, indicated via `object_idx`"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1c51bb5c-f0ff-4350-aa01-ae244d93f9c3",
"metadata": {},
"outputs": [],
"source": [
"# Make the object softer by updating the material parameter\n",
"scene.set_object_materials(0, yms=torch.tensor(5e4, device='cuda', dtype=torch.float))"
"scene.set_object_materials(obj_idx, yms=torch.tensor(5e4, device='cuda', dtype=torch.float))"
]
},
{
"cell_type": "markdown",
"id": "60ed3fec-cc3a-4b0a-b51f-3df7f8333b3c",
"id": "36b45eb6-6d0f-4443-90f5-bf805a4c4e53",
"metadata": {},
"source": [
"## Thats it! Onto display\n",
"All we need to do now is display the object using Kaolin's in-notebook visualizer."
"## Set Up Rendering\n",
"Let's set up rendering of our mesh so we can view it in a notebook."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "30963bf0-06be-4186-b4e0-ab7d377b1409",
"id": "d06cff49-da6d-4806-bebb-588df336e6ac",
"metadata": {},
"outputs": [],
"source": [
"resolution = 512\n",
"camera = kal.render.easy_render.default_camera(resolution).cuda()\n",
"\n",
"azimuth = torch.zeros((1,), device='cuda')\n",
"elevation = torch.full((1,), math.pi / 3., device='cuda')\n",
"amplitude = torch.full((1, 3), 3., device='cuda')\n",
"sharpness = torch.full((1,), 5., device='cuda')\n",
"\n",
"def current_lighting():\n",
" direction = kal.render.lighting.sg_direction_from_azimuth_elevation(azimuth, elevation)\n",
" return kal.render.lighting.SgLightingParameters(\n",
" amplitude=amplitude, sharpness=sharpness, direction=direction)\n",
"# light that works well for Avocado\n",
"light_direction = kal.render.lighting.sg_direction_from_azimuth_elevation(1., 1.)\n",
"lighting = kal.render.lighting.SgLightingParameters(amplitude=3., sharpness=5., direction=light_direction).cuda()\n",
" \n",
"def render(in_cam, **kwargs):\n",
"def render(in_cam):\n",
" # render\n",
" active_pass=kal.render.easy_render.RenderPass.render\n",
" render_res = kal.render.easy_render.render_mesh(in_cam, mesh, lighting=current_lighting(), **kwargs)\n",
" \n",
" render_res = kal.render.easy_render.render_mesh(in_cam, mesh, lighting=lighting)\n",
"\n",
" # create white background\n",
" img = render_res[active_pass]\n",
" background_mask = (render_res[kal.render.easy_render.RenderPass.face_idx] < 0).bool()\n",
" img2 = torch.clamp(img, 0, 1)[0]\n",
" img2[background_mask[0]] = 1\n",
" final = (img2 * 255.).to(torch.uint8)\n",
" return {\"img\":final, \"face_idx\": render_res[kal.render.easy_render.RenderPass.face_idx].squeeze(0).unsqueeze(-1)}\n",
"\n",
"# faster low-res render during mouse motion\n",
"def fast_render(in_cam, factor=8):\n",
" lowres_cam = copy.deepcopy(in_cam)\n",
" lowres_cam.width = in_cam.width // factor\n",
" lowres_cam.height = in_cam.height // factor\n",
" return render(lowres_cam)\n",
" return render(lowres_cam)"
]
},
{
"cell_type": "markdown",
"id": "ef421ba7-e245-4780-88f5-8fe10c2bb43d",
"metadata": {},
"source": [
"## That's it! Let's Run and View or Physics Simulation\n",
"All we need to do now is run simulation and display the object using Kaolin's in-notebook visualizer."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "30963bf0-06be-4186-b4e0-ab7d377b1409",
"metadata": {},
"outputs": [],
"source": [
"# Reset mesh to its rest state\n",
"mesh.vertices = orig_vertices\n",
"\n",
"global sim_thread_open, sim_thread\n",
"sim_thread_open = False\n",
"sim_thread = None\n",
"\n",
"def run_sim():\n",
" scene.reset_object(0)\n",
" scene.reset() # reset internal simultion state\n",
"\n",
" for s in range(100):\n",
" with visualizer.out:\n",
Expand All @@ -275,16 +288,22 @@
" visualizer.render_update()\n",
"\n",
"def start_simulation(b):\n",
" # run_sim()\n",
" t = threading.Thread(target=run_sim, daemon=True)\n",
" t.start()\n",
" global sim_thread_open, sim_thread\n",
" with visualizer.out:\n",
" if(sim_thread_open):\n",
" sim_thread.join()\n",
" sim_thread_open = False\n",
" sim_thread_open = True\n",
" sim_thread = threading.Thread(target=run_sim, daemon=True)\n",
" sim_thread.start()\n",
"\n",
"scene.reset_object(0)\n",
"scene.reset_object(obj_idx)\n",
"button = Button(description='Run Sim')\n",
"button.on_click(start_simulation)\n",
"visualizer = kal.visualize.IpyTurntableVisualizer(\n",
" resolution, resolution, copy.deepcopy(camera), render, fast_render=fast_render,\n",
" max_fps=24, world_up_axis=1)\n",
"visualizer.render_update() # render first frame\n",
"display(HBox([visualizer.canvas, button]), visualizer.out)"
]
},
Expand Down
Loading

0 comments on commit c7593f0

Please sign in to comment.