{
"cells": [
{
"cell_type": "markdown",
"id": "2db540bc",
"metadata": {},
"source": [
"# Fitting a 3D Bounding Box using a Differentiable Renderer\n",
"\n",
"Differentiable rendering enables optimization of 3D object properties like the geometry of a mesh. Unlike traditional rendering, differentiable rendering can backpropagate gradients from image space to 3D geometry.\n",
"\n",
"In this tutorial we fit a 3D bounding box around an object given ground truth images of the object. We assume images come with a segmentation mask of the object of interest and include camera extrinsics. During fitting we optimize a 3D bounding box to fit tightly around the object, despite never seeing any 3D ground truth data.\n",
"\n",
"We leverage kaolin's rendering capabilities with meshes, particularly the DIB-R rasterizer in `kaolin.render.mesh.dibr_rasterization` ([API documentation](https://kaolin.readthedocs.io/en/latest/modules/kaolin.render.mesh.html)). Please note this notebook is intended to illustrate key kaolin functionality and is not intended for use in a production system. A more sophisticated use of differentiable rendering can be found in the [Optimizing a mesh using a Differentiable Renderer notebook](examples/tutorial/dibr_tutorial.ipynb). Additional DIB-R examples can be found in [this repository](https://github.com/nv-tlabs/DIB-R-Single-Image-3D-Reconstruction).\n",
"\n",
"Before starting the tutorial uncompress the training data found in [examples/samples/rendered_clock.zip](examples/samples/rendered_clock.zip)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "98082a4a-21da-4774-a5f5-02c17f8e102a",
"metadata": {},
"outputs": [],
"source": [
"!pip install -q matplotlib"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "736eb3ef",
"metadata": {},
"outputs": [],
"source": [
"import glob\n",
"import os\n",
"from pathlib import Path\n",
"from typing import List, Tuple\n",
"\n",
"import matplotlib.animation as animation\n",
"import numpy as np\n",
"import torch\n",
"from torch import Tensor\n",
"from matplotlib import pyplot as plt\n",
"from IPython.display import display, clear_output, HTML\n",
"\n",
"from torch.utils.data import DataLoader\n",
"\n",
"import kaolin"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "27a9c0fc-0135-422b-ae7a-5f78dc44488d",
"metadata": {},
"outputs": [],
"source": [
"import zipfile\n",
"\n",
"# Unpack assets required by notebook, if n/a\n",
"cwd = Path.cwd().resolve()\n",
"\n",
"for p in [cwd, *cwd.parents]:\n",
" samples_dir = p / \"examples\" / \"samples\"\n",
" zip_path = samples_dir / \"rendered_clock.zip\"\n",
" rendered_path = samples_dir / \"rendered_clock\"\n",
" if zip_path.exists() or rendered_path.exists():\n",
" break\n",
"else:\n",
" raise RuntimeError(f\"Could not find examples/samples from {cwd}\")\n",
"\n",
"if not rendered_path.exists() or not list(rendered_path.glob(\"*_rgb.png\")):\n",
" if not zip_path.exists():\n",
" raise FileNotFoundError(f\"Missing sample archive: {zip_path}\")\n",
" print(f\"Unzipping {zip_path} -> {samples_dir}\")\n",
" with zipfile.ZipFile(zip_path, \"r\") as zf:\n",
" zf.extractall(samples_dir)\n",
"\n",
"num_rgb = len(list(rendered_path.glob(\"*_rgb.png\")))\n",
"if num_rgb == 0:\n",
" raise RuntimeError(f\"No *_rgb.png files found in {rendered_path}\")\n",
"\n",
"print(f\"OK: {num_rgb} rendered_clock RGB views available\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ce0ee9a1",
"metadata": {},
"outputs": [],
"source": [
"# set hyperparameters\n",
"batch_size_hyper = 2\n",
"mask_weight_hyper = 1.0\n",
"mask_occupancy_hyper = 0.05\n",
"mask_overlap_hyper = 1.0\n",
"lr = 5e-2\n",
"scheduler_step_size = 5\n",
"scheduler_gamma = 0.5\n",
"num_epoch = 30"
]
},
{
"cell_type": "markdown",
"id": "06e8b4ee",
"metadata": {},
"source": [
"# Preparing training data\n",
"\n",
"First we prepare the training 2D image data and differentiable 3D bounding box mesh.\n",
"\n",
"## Generating training data\n",
"\n",
"Our training data consists of segmentation masks of objects with known camera properties. One way to generate this data is to use the Data Generator in the [Omniverse Kaolin App](https://docs.omniverse.nvidia.com/app_kaolin/app_kaolin/user_manual.html#data-generator). We provide sample output from the app in `examples/samples/`."
]
},
{
"cell_type": "markdown",
"id": "c69d6434",
"metadata": {},
"source": [
"## Parsing synthetic data\n",
"\n",
"We first need to parse the synthetic data generated by the Omniverse app.\n",
"The app produces one output file for each category of data - among depth map, RGB image, or segmentation map - along with an additional metadata JSON file. The JSON file includes the camera_properties, with data related to the camera settings including \"clipping_range\", \"horizontal_aperture\", \"focal_length\", \"tf_mat\".\n",
"\n",
"Below we parse this training data using the `kaolin.io.render.import_synthetic_view` function."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "fc9741fa",
"metadata": {},
"outputs": [],
"source": [
"# set dataset of renders and the 3D bounding box mesh\n",
"rendered_path = Path(\"../samples/rendered_clock/\")\n",
"\n",
"# prepare the 2D render dataset including camera extrinsics:\n",
"num_views = len(glob.glob(os.path.join(rendered_path, \"*_rgb.png\")))\n",
"train_data = []\n",
"for i in range(num_views):\n",
" data = kaolin.io.render.import_synthetic_view(rendered_path, i, rgb=True, semantic=True)\n",
" train_data.append(data)\n",
"dataloader = DataLoader(train_data, batch_size=batch_size_hyper, shuffle=True, pin_memory=True)"
]
},
{
"cell_type": "markdown",
"id": "a238c89e",
"metadata": {},
"source": [
"## Defining the 3D bounding box\n",
"\n",
"We will use a differentiable renderer to create 2D renders of a mesh, specifically for a 3D bounding box. In 3D bounding box estimation our goal is to fit the 3D box around an object while keeping the box as some cuboid. To enforce this constraint we will take a fixed 3D cuboid and apply transformations that rotate, translate, or scale the box. Note that this constrains learning by preventing and shearing or other transformations that would deform the cuboid shape.\n",
"\n",
"We create this model by loading a 3D cube mesh (provided for this tutorial in [examples/samples/bbox.obj](examples/samples/bbox.obj)) and defining a set of learnable parameters for the object rotation, translation, and scaling. When we access the vertices of the object we apply those transforms and return the modified cuboid."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "56f6020e",
"metadata": {},
"outputs": [],
"source": [
"class DifferentiableBBox:\n",
" \"\"\"Represents a differentiable 3D bounding box.\n",
"\n",
" Box is parametrized in terms of a fixed mesh and a learned (optimized)\n",
" transformation in terms of rotation, scaling, and translation.\n",
" \"\"\"\n",
"\n",
" def __init__(self):\n",
" \"\"\"Construct a 3D bounding from a mesh filepath.\n",
"\n",
" Args:\n",
" mesh_file (str): Filepath to load the mesh. OBJ format.\n",
" \"\"\"\n",
" super().__init__()\n",
" mesh_file = Path(\"../samples/bbox.obj\")\n",
" assert mesh_file.exists(), f\"File missing: {mesh_file.absolute()}.\"\n",
" \n",
"\n",
" self._mesh = kaolin.io.obj.import_mesh(mesh_file, with_materials=True)\n",
" self._vertices = None\n",
" self._faces = None\n",
" \n",
" # define the learnable parameters\n",
" self._centers = torch.zeros((3,), dtype=torch.float, device=\"cuda\", requires_grad=True)\n",
" self._scales = torch.ones((3,), dtype=torch.float, device=\"cuda\", requires_grad=True)\n",
" self._rotations = torch.tensor([0.0, 0.0, 0.0, 1.0], device=\"cuda\", requires_grad=True) # rotations as quaternion\n",
" \n",
" # prepare the vertices for learning\n",
" self._preprocess()\n",
"\n",
" def _preprocess(self):\n",
" # scale up vertices to roughly the image scale\n",
" self._vertices = self._mesh.vertices.cuda().unsqueeze(0) * 0.75\n",
" \n",
" # disable gradient propagation to the vertices to fix the baseline 3D bounding box\n",
" self._vertices.requires_grad = False\n",
" \n",
" self._faces = self._mesh.faces.cuda()\n",
"\n",
" @property\n",
" def vertices(self):\n",
" # convert the rotation quaternion to a 3x3 rotation matrix (see below)\n",
" rot = quaternion_to_matrix33(self._rotations)\n",
" # apply scaling, rotation, and translation to vertices\n",
" return (torch.matmul(self._vertices, rot) * self._scales) + self._centers\n",
"\n",
" @property\n",
" def faces(self):\n",
" return self._faces\n",
" \n",
" @property\n",
" def parameters(self):\n",
" return [self._centers, self._scales, self._rotations]"
]
},
{
"cell_type": "markdown",
"id": "20fc26be",
"metadata": {},
"source": [
"### Mesh Rotation\n",
"\n",
"The mesh parameters represent the rotation as a quaternion. The image data represents points a 3D vectors, so for convenience we define a helper function that converts the quaternion rotation representation into a 3D rotation matrix. This simplifies the mesh rotation operation to a matrix multiplication.\n",
"\n",
"(As an alternative you could also modify the `vertices` property of `DifferentiableBBox` to rotate the vertices using the quaternion directly.)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c3c198ad",
"metadata": {},
"outputs": [],
"source": [
"# helper functions to convert from quaternion to rotation matrix\n",
"def vector_normalize(vec: Tensor) -> Tensor:\n",
" \"\"\"Normalize a 1d vector using the L2 norm.\n",
"\n",
" Args:\n",
" vec (Tensor): A 1d vector of shape (b,1).\n",
"\n",
" Returns:\n",
" Tensor: A normalized version of the input vector of shape (b,1).\n",
" \"\"\" \n",
" return vec / vec.norm(p=2, dim=-1, keepdim=True)\n",
"\n",
"\n",
"def quaternion_to_matrix33(quat: Tensor) -> Tensor:\n",
" \"\"\"Convert a quaternion to a 3x3 rotation matrix.\n",
"\n",
" Args:\n",
" quat (Tensor): Rotation quaternion of shape (4).\n",
"\n",
" Returns:\n",
" Tensor: Rotation matrix of shape (3,3).\n",
" \"\"\" \n",
" # reference: http://www.euclideanspace.com/maths/geometry/rotations/conversions/quaternionToMatrix/index.htm\n",
" q = vector_normalize(quat)\n",
" \n",
" qx, qy, qz, qw = q[0], q[1], q[2], q[3]\n",
" sqw = qw ** 2\n",
" sqx = qx ** 2\n",
" sqy = qy ** 2\n",
" sqz = qz ** 2\n",
" qxy = qx * qy\n",
" qzw = qz * qw\n",
" qxz = qx * qz\n",
" qyw = qy * qw\n",
" qyz = qy * qz\n",
" qxw = qx * qw\n",
"\n",
" invs = 1 / (sqx + sqy + sqz + sqw)\n",
" m00 = (sqx - sqy - sqz + sqw) * invs\n",
" m11 = (-sqx + sqy - sqz + sqw) * invs\n",
" m22 = (-sqx - sqy + sqz + sqw) * invs\n",
" m10 = 2.0 * (qxy + qzw) * invs\n",
" m01 = 2.0 * (qxy - qzw) * invs\n",
" m20 = 2.0 * (qxz - qyw) * invs\n",
" m02 = 2.0 * (qxz + qyw) * invs\n",
" m21 = 2.0 * (qyz + qxw) * invs\n",
" m12 = 2.0 * (qyz - qxw) * invs\n",
" r0 = torch.stack([m00, m01, m02])\n",
" r1 = torch.stack([m10, m11, m12])\n",
" r2 = torch.stack([m20, m21, m22])\n",
" mat33 = torch.stack([r0, r1, r2]).T\n",
" return mat33"
]
},
{
"cell_type": "markdown",
"id": "cbcbb6ba",
"metadata": {},
"source": [
"# Setting up the fitting loop\n",
"\n",
"In this section we set up the optimization and losses for training.\n",
"\n",
"## Setting up the optimizer\n",
"\n",
"Below we create the 3D bounding box and prepare an optimizer to tune the box's parameters. We add a learning rate scheduler so that we can gradually decrease the step size of changes made to the box parameters during the fitting process."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "56085cce",
"metadata": {},
"outputs": [],
"source": [
"# set up model & optimization parameters\n",
"bbox = DifferentiableBBox()\n",
"optim = torch.optim.Adam(params=bbox.parameters, lr=lr)\n",
"scheduler = torch.optim.lr_scheduler.StepLR(optim, step_size=scheduler_step_size, gamma=scheduler_gamma)"
]
},
{
"cell_type": "markdown",
"id": "8979dd00",
"metadata": {},
"source": [
"## Differentiable rendering\n",
"\n",
"During training we want to optimize the 3D bounding box to overlap with the semantic segmentation mask of the object in the 2D renders. We first need to project the 3D mesh into the 2D image space for each image. The function below uses ground truth data on the camera properties to project the 3D mesh vertices onto a 2D image to produce a 2D silhouette of the mesh."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2b6bda5e",
"metadata": {},
"outputs": [],
"source": [
"def project_to_2d(\n",
" bbox: DifferentiableBBox,\n",
" batch_size: int,\n",
" image_shape: Tuple[int, int],\n",
" camera_transform: Tensor,\n",
" camera_projection: Tensor,\n",
") -> Tensor:\n",
" \"\"\"Render a mesh onto a 2D image given a viewing camera's transform and projection.\n",
"\n",
" Args:\n",
" bbox (DifferentiableBBox): Differentiable bounding box representing 3D mesh.\n",
" batch_size (int): Number of elements in the data batch.\n",
" image_shape (Tuple[int, int]): Tuple of image dimensions (height, width).\n",
" camera_transform (Tensor): Camera transform of shape (b, 4, 3).\n",
" camera_projection (Tensor): Camera projection of shape (b, 3, 1).\n",
"\n",
" Returns:\n",
" Tensor: Soft mask for mesh silhouette of shape (b, h, w). (h, w) given by `image_shape`.\n",
" \"\"\"\n",
" (face_vertices_camera, face_vertices_image, face_normals,) = kaolin.render.mesh.prepare_vertices(\n",
" bbox.vertices.repeat(batch_size, 1, 1),\n",
" bbox.faces,\n",
" camera_projection,\n",
" camera_transform=camera_transform,\n",
" )\n",
"\n",
" nb_faces = bbox.faces.shape[0]\n",
" face_attributes = [torch.ones((batch_size, nb_faces, 3, 1), device=\"cuda\")]\n",
"\n",
" image_features, soft_mask, face_idx = kaolin.render.mesh.dibr_rasterization(\n",
" image_shape[0],\n",
" image_shape[1],\n",
" face_vertices_camera[:, :, :, -1],\n",
" face_vertices_image,\n",
" face_attributes,\n",
" face_normals[:, :, -1],\n",
" )\n",
" return soft_mask # only aligning images by 2D silhouette\n"
]
},
{
"cell_type": "markdown",
"id": "ba4bfddb",
"metadata": {},
"source": [
"## Defining losses\n",
"\n",
"Assuming that we have a 2D mask of the mesh in the image and the ground truth segmentation mask for that same image we can compare them. We define two losses to encourage the mesh to fit around the target object tightly:\n",
"1. `overlap` is a loss that encourages the two masks to overlap as much as possible.\n",
"2. `occupany` is a loss that encourages the 2D projection mask to be as large as possible.\n",
"\n",
"With overlap alone the projection will not stay outside the shape: it's better to cut off a few corners of the object in exchange for slightly less penalty from the empty space that's saved. Adding occupany counter-balances this effect. The end result is the mesh will try to fit around the outside of the shape (convex hull)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "75bc78d2",
"metadata": {},
"outputs": [],
"source": [
"def overlap(lhs_mask: Tensor, rhs_mask: Tensor) -> Tensor:\n",
" \"\"\"Compute the overlap of two 2D masks as the intersection over union.\n",
"\n",
" Args:\n",
" lhs_mask (Tensor): 2D mask of shape (b, h, w).\n",
" rhs_mask (Tensor): 2D mask of shape (b, h, w).\n",
"\n",
" Returns:\n",
" Tensor: Fraction of overlap of the two masks of shape (1). Averaged over batch samples.\n",
" \"\"\" \n",
" batch_size, height, width = lhs_mask.shape\n",
" assert rhs_mask.shape == lhs_mask.shape\n",
" sil_mul = lhs_mask * rhs_mask\n",
" sil_area = torch.sum(sil_mul.reshape(batch_size, -1), dim=1)\n",
"\n",
" return 1 - torch.mean(sil_area / (height * width))\n",
"\n",
"\n",
"def occupancy(mask: Tensor) -> Tensor:\n",
" \"\"\"Compute what fraction of a total image is occupied by a 2D mask.\n",
"\n",
" Args:\n",
" mask (Tensor): 2D mask of shape (b, h, w).\n",
"\n",
" Returns:\n",
" Tensor: Fraction of the full image occupied by the mask of shape (1). Averaged over batch samples.\n",
" \"\"\" \n",
" batch_size, height, width = mask.shape\n",
" mask_area = torch.sum(mask.reshape(batch_size, -1), dim=1)\n",
" return torch.mean(mask_area / (height * width))"
]
},
{
"cell_type": "markdown",
"id": "8e48917f",
"metadata": {},
"source": [
"## Visualization utilities\n",
"\n",
"During training it's helpful to watch the progression of bounding box fitting. Below we define helper functions to plot images that show the bounding box silhouette compared to the ground truth silhouette."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "eb5748e6",
"metadata": {},
"outputs": [],
"source": [
"def draw_image(gt_mask: Tensor, pred_mask: Tensor) -> np.ndarray:\n",
" \"\"\"Compute an image array showing ground truth vs predicted masks.\n",
"\n",
" Args:\n",
" gt_mask (Tensor): Ground truth mask of shape (w, h).\n",
" pred_mask (Tensor): Predicted mask of shape (w, h).\n",
"\n",
" Returns:\n",
" np.ndarray: [0,1] normalized numpy array of shape (w, h, 3).\n",
" \"\"\" \n",
" # mask shape is [w,h]\n",
" canvas = np.zeros((gt_mask.shape[0], gt_mask.shape[1], 3))\n",
" canvas[..., 2] = pred_mask.cpu().detach().numpy()\n",
" canvas[..., 1] = gt_mask.cpu().detach().numpy()\n",
"\n",
" return np.clip(canvas, 0.0, 1.0)\n",
"\n",
"\n",
"def show_renders(bbox: DifferentiableBBox, dataset: Tensor) -> List[np.ndarray]:\n",
" \"\"\"Generate images comparing ground truth and predicted semantic segmentations for a given mesh.\n",
"\n",
" The mesh is projected to 2D to match the camera views of each ground truth 2D render.\n",
" Images plot the true silhouette of the object overlayed with the mesh silhouette.\n",
"\n",
" Args:\n",
" bbox (DifferentiableBBox): Differentiable bounding box representing 3D mesh.\n",
" dataset (Tensor): Batch of ground truth 2D renders with camera extrinsics.\n",
"\n",
" Returns:\n",
" List[np.ndarray]: [0,1] normalized images comparing ground truth and mesh silhouette.\n",
" \"\"\" \n",
" with torch.no_grad():\n",
" images = []\n",
" for sample in dataset:\n",
" gt_mask = sample[\"semantic\"].cuda()\n",
" camera_transform = sample[\"metadata\"][\"cam_transform\"].cuda()\n",
" camera_projection = sample[\"metadata\"][\"cam_proj\"].cuda()\n",
"\n",
" # project model mesh onto 2D image\n",
" img_shape = (gt_mask.shape[0], gt_mask.shape[1])\n",
" soft_mask = project_to_2d(\n",
" bbox,\n",
" batch_size=1,\n",
" image_shape=img_shape,\n",
" camera_transform=camera_transform,\n",
" camera_projection=camera_projection,\n",
" )\n",
" image = draw_image(gt_mask.squeeze(), soft_mask)\n",
" images.append(image)\n",
"\n",
" return images"
]
},
{
"cell_type": "markdown",
"id": "864bc6b5",
"metadata": {},
"source": [
"# Fitting\n",
"\n",
"The fitting loop below pulls everything together. Each iteration samples a batch of ground truth 2D renders and their associated camera data. The mesh is projected to the same 2D images based on the camera properties and the projected mesh silhouette is compared to the ground truth silhouette (segmentation mask). Losses are calculated by comparing the two silhouettes and a gradient step backpropagates the error through the mesh parameters (that is, the bounding box rotation, translation, and scaling from the initial mesh position)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "345eca8a",
"metadata": {},
"outputs": [],
"source": [
"# set up for plotting progress during training\n",
"test_batch_ids = [2, 5, 10] # pick canonical test render views\n",
"num_subplots = len(test_batch_ids)\n",
"fig, ax = plt.subplots(ncols=num_subplots)\n",
"plot_handle = display(fig, display_id=True)\n",
"# run fitting loop\n",
"image_list = []\n",
"for epoch in range(num_epoch):\n",
" for idx, data in enumerate(dataloader):\n",
" optim.zero_grad() # zero out gradients for new batch\n",
"\n",
" # get image mask and camera extrinsics from the true 2D renders\n",
" gt_mask = data[\"semantic\"].cuda()\n",
" camera_transform = data[\"metadata\"][\"cam_transform\"].cuda()\n",
" camera_projection = data[\"metadata\"][\"cam_proj\"].cuda()\n",
"\n",
" # project the 3D mesh (bbox) onto a canvas matching the ground truth shape\n",
" # and using the ground truth camera\n",
" img_shape = (gt_mask.shape[1], gt_mask.shape[2])\n",
" soft_mask = project_to_2d(\n",
" bbox,\n",
" batch_size=batch_size_hyper,\n",
" image_shape=img_shape,\n",
" camera_transform=camera_transform,\n",
" camera_projection=camera_projection,\n",
" )\n",
"\n",
" # compute loss by rewarding overlap with the 2D render silhouette\n",
" # and rewarding masks that are bigger\n",
" # mask_overlap = reward more overlap\n",
" mask_overlap = overlap(soft_mask, gt_mask.squeeze(-1))\n",
" # mask_occupancy = reward larger masks\n",
" mask_occupancy = occupancy(soft_mask)\n",
" \n",
" loss = mask_occupancy * mask_occupancy_hyper + mask_overlap * mask_overlap_hyper\n",
"\n",
" # propagate losses\n",
" loss.backward()\n",
" optim.step()\n",
"\n",
" # view training progress\n",
" if idx % 10 == 0: \n",
" test_viz = [train_data[idx] for idx in test_batch_ids]\n",
" # only keep 1 in 10 renders to reduce animation processing time\n",
" image_list.append(show_renders(bbox, test_viz))\n",
" for i in range(num_subplots):\n",
" ax[i].clear()\n",
" ax[i].imshow(image_list[-1][i])\n",
" fig.canvas.draw()\n",
" plot_handle.update(fig)\n",
"\n",
" # step the learning rate schedule each epoch\n",
" # reduces the size of changes to the bbox parameters to fine-tune in later stages of fitting\n",
" scheduler.step()\n",
" print(f\"loss on epoch {epoch:<}: {loss}\")"
]
},
{
"cell_type": "markdown",
"id": "a2d19daf",
"metadata": {},
"source": [
"## Animation\n",
"\n",
"The code below can be used to compile the renders generated into an animation to view later (ex: when debugging learning after the fact)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7e4fbc90",
"metadata": {},
"outputs": [],
"source": [
"# generate an animation of the training loop\n",
"num_subplots = len(test_batch_ids)\n",
"fig, ax = plt.subplots(ncols=num_subplots)\n",
"\n",
"ims = []\n",
"for i in range(len(image_list)):\n",
" sp_ims = []\n",
" for j in range(num_subplots):\n",
" im = ax[j].imshow(image_list[i][j], animated=True)\n",
" sp_ims.append(im)\n",
" # show first to not have blinking animation\n",
" if i == 0:\n",
" for j in range(num_subplots):\n",
" ax[j].imshow(image_list[i][j])\n",
" ims.append(sp_ims)\n",
"ani = animation.ArtistAnimation(fig, ims, interval=60, blit=True, repeat_delay=100)\n",
"plt.close(fig)\n",
"display(HTML(ani.to_jshtml()))"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0d3be94a",
"metadata": {},
"outputs": [],
"source": [
"# ani.save(\"animation.gif\", writer=animation.PillowWriter()) # optionally save the animation"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python (kaolin)",
"language": "python",
"name": "kaolin"
},
"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.10.20"
}
},
"nbformat": 4,
"nbformat_minor": 5
}