implicitron_volumes.ipynb 33.8 KB
Newer Older
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
1
2
3
4
{
  "cells": [
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
5
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
6
7
      "metadata": {
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
8
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
9
10
        "customOutput": null,
        "executionStartTime": 1659619824914,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
11
12
13
14
        "executionStopTime": 1659619825485,
        "originalKey": "d38652e8-200a-413c-a36a-f4d349b78a9d",
        "requestMsgId": "641de8aa-0e42-4446-9304-c160a2d226bf",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
15
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
16
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
17
18
      "source": [
        "# Copyright (c) Meta Platforms, Inc. and affiliates. All rights reserved."
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
19
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
20
21
    },
    {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
22
      "attachments": {},
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
23
24
      "cell_type": "markdown",
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
25
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
26
        "originalKey": "a48a9dcf-e80f-474b-a0c4-2c9a765b15c5",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
27
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
28
29
30
31
32
33
34
      },
      "source": [
        "# A simple model using Implicitron\n",
        "\n",
        "In this demo, we use the VolumeRenderer from PyTorch3D as a custom implicit function in Implicitron. We will see\n",
        "* some of the main objects in Implicitron\n",
        "* how to plug in a custom part of a model"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
35
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
36
37
    },
    {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
38
      "attachments": {},
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
39
40
      "cell_type": "markdown",
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
41
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
42
        "originalKey": "51337c0e-ad27-4b75-ad6a-737dca5d7b95",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
43
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
44
45
46
47
      },
      "source": [
        "## 0. Install and import modules\n",
        "\n",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
48
49
        "Ensure `torch` and `torchvision` are installed. If `pytorch3d` is not installed, install it using the following cell:\n"
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
50
51
52
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
53
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
54
55
      "metadata": {
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
56
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
57
58
        "customOutput": null,
        "executionStartTime": 1659619898147,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
59
60
61
62
        "executionStopTime": 1659619898274,
        "originalKey": "76f1ecd4-6b73-4214-81b0-118ef8d86872",
        "requestMsgId": "deb6a860-6923-4227-abef-d31388b5142d",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
63
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
64
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
      "source": [
        "import os\n",
        "import sys\n",
        "import torch\n",
        "need_pytorch3d=False\n",
        "try:\n",
        "    import pytorch3d\n",
        "except ModuleNotFoundError:\n",
        "    need_pytorch3d=True\n",
        "if need_pytorch3d:\n",
        "    if torch.__version__.startswith(\"1.12.\") and sys.platform.startswith(\"linux\"):\n",
        "        # We try to install PyTorch3D via a released wheel.\n",
        "        pyt_version_str=torch.__version__.split(\"+\")[0].replace(\".\", \"\")\n",
        "        version_str=\"\".join([\n",
        "            f\"py3{sys.version_info.minor}_cu\",\n",
        "            torch.version.cuda.replace(\".\",\"\"),\n",
        "            f\"_pyt{pyt_version_str}\"\n",
        "        ])\n",
        "        !pip install fvcore iopath\n",
        "        !pip install --no-index --no-cache-dir pytorch3d -f https://dl.fbaipublicfiles.com/pytorch3d/packaging/wheels/{version_str}/download.html\n",
        "    else:\n",
        "        # We try to install PyTorch3D from source.\n",
        "        !curl -LO https://github.com/NVIDIA/cub/archive/1.10.0.tar.gz\n",
        "        !tar xzf 1.10.0.tar.gz\n",
        "        os.environ[\"CUB_HOME\"] = os.getcwd() + \"/cub-1.10.0\"\n",
        "        !pip install 'git+https://github.com/facebookresearch/pytorch3d.git@stable'"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
91
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
92
93
    },
    {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
94
      "attachments": {},
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
95
96
      "cell_type": "markdown",
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
97
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
98
        "originalKey": "2c1020e6-eb4a-4644-9719-9147500d8e4f",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
99
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
100
101
      },
      "source": [
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
102
103
        "Ensure omegaconf and visdom are installed. If not, run this cell. (It should not be necessary to restart the runtime.)"
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
104
105
106
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
107
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
108
109
      "metadata": {
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
110
111
112
        "customOutput": null,
        "originalKey": "9e751931-a38d-44c9-9ff1-ac2f7d3a3f99",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
113
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
114
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
115
      "source": [
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
116
117
        "!pip install omegaconf visdom"
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
118
119
120
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
121
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
122
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
123
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
124
125
126
127
        "collapsed": false,
        "customOutput": null,
        "executionStartTime": 1659612480556,
        "executionStopTime": 1659612480644,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
128
129
130
        "hidden_ranges": [],
        "originalKey": "86807e4a-1675-4520-a033-c7af85b233ec",
        "requestMsgId": "880a7e20-4a90-4b37-a5eb-bccc0b23cac6"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
131
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
132
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
      "source": [
        "import logging\n",
        "from typing import Tuple\n",
        "\n",
        "import matplotlib.animation as animation\n",
        "import matplotlib.pyplot as plt\n",
        "import numpy as np\n",
        "import torch\n",
        "import tqdm\n",
        "from IPython.display import HTML\n",
        "from omegaconf import OmegaConf\n",
        "from PIL import Image\n",
        "from pytorch3d.implicitron.dataset.dataset_base import FrameData\n",
        "from pytorch3d.implicitron.dataset.rendered_mesh_dataset_map_provider import RenderedMeshDatasetMapProvider\n",
        "from pytorch3d.implicitron.models.generic_model import GenericModel\n",
        "from pytorch3d.implicitron.models.implicit_function.base import ImplicitFunctionBase\n",
        "from pytorch3d.implicitron.models.renderer.base import EvaluationMode\n",
        "from pytorch3d.implicitron.tools.config import expand_args_fields, get_default_args, registry, remove_unused_components\n",
        "from pytorch3d.renderer import RayBundle\n",
        "from pytorch3d.renderer.implicit.renderer import VolumeSampler\n",
        "from pytorch3d.structures import Volumes\n",
        "from pytorch3d.vis.plotly_vis import plot_batch_individually, plot_scene"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
155
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
156
157
158
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
159
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
160
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
161
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
162
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
163
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
164
165
166
        "customOutput": null,
        "executionStartTime": 1659610929375,
        "executionStopTime": 1659610929383,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
167
168
169
170
        "hidden_ranges": [],
        "originalKey": "b2d9f5bd-a9d4-4f78-b21e-92f2658e0fe9",
        "requestMsgId": "7e43e623-4030-438b-af4e-b96170c9a052",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
171
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
172
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
173
174
      "source": [
        "output_resolution = 80"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
175
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
176
177
178
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
179
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
180
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
181
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
182
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
183
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
184
185
186
        "customOutput": null,
        "executionStartTime": 1659610930042,
        "executionStopTime": 1659610930050,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
187
188
189
190
        "hidden_ranges": [],
        "originalKey": "0b0c2087-4c86-4c57-b0ee-6f48a70a9c78",
        "requestMsgId": "46883aad-f00b-4fd4-ac17-eec0b2ac272a",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
191
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
192
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
193
194
      "source": [
        "torch.set_printoptions(sci_mode=False)"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
195
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
196
197
    },
    {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
198
      "attachments": {},
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
199
200
      "cell_type": "markdown",
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
201
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
202
        "originalKey": "37809d0d-b02e-42df-85b6-cdd038373653",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
203
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
204
205
206
207
208
209
210
      },
      "source": [
        "## 1. Load renders of a mesh (the cow mesh) as a dataset\n",
        "\n",
        "A dataset's train, val and test parts in Implicitron are represented as a `dataset_map`, and provided by an implementation of `DatasetMapProvider`. \n",
        "`RenderedMeshDatasetMapProvider` is one which generates a single-scene dataset with only a train component by taking a mesh and rendering it.\n",
        "We use it with the cow mesh."
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
211
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
212
213
    },
    {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
214
      "attachments": {},
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
215
216
217
      "cell_type": "markdown",
      "metadata": {
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
218
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
219
220
        "customOutput": null,
        "executionStartTime": 1659620739780,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
221
222
223
224
        "executionStopTime": 1659620739914,
        "originalKey": "cc68cb9c-b8bf-4e9e-bef1-2cfafdf6caa2",
        "requestMsgId": "398cfcae-5d43-4b6f-9c75-db3d297364d4",
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
225
226
227
228
      },
      "source": [
        "If running this notebook using **Google Colab**, run the following cell to fetch the mesh obj and texture files and save it at the path data/cow_mesh.\n",
        "If running locally, the data is already available at the correct path."
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
229
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
230
231
232
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
233
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
234
235
      "metadata": {
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
236
237
238
        "customOutput": null,
        "originalKey": "2c55e002-a885-4169-8fdc-af9078b05968",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
239
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
240
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
241
242
243
244
245
      "source": [
        "!mkdir -p data/cow_mesh\n",
        "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.obj\n",
        "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow.mtl\n",
        "!wget -P data/cow_mesh https://dl.fbaipublicfiles.com/pytorch3d/data/cow_mesh/cow_texture.png"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
246
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
247
248
249
250
    },
    {
      "cell_type": "markdown",
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
251
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
252
        "originalKey": "2a976be8-01bf-4a1c-a6e7-61d5d08c3dbd",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
253
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
254
255
256
257
258
259
260
      },
      "source": [
        "If we want to instantiate one of Implicitron's configurable objects, such as `RenderedMeshDatasetMapProvider`, without using the OmegaConf initialisation (get_default_args), we need to call `expand_args_fields` on the class first."
      ]
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
261
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
262
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
263
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
264
265
266
267
        "collapsed": false,
        "customOutput": null,
        "executionStartTime": 1659621652237,
        "executionStopTime": 1659621652903,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
268
269
270
271
        "hidden_ranges": [],
        "originalKey": "eb77aaec-048c-40bd-bd69-0e66b6ab60b1",
        "requestMsgId": "09b9975c-ff86-41c9-b4a9-975d23afc562",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
272
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
273
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
274
275
276
277
278
279
280
      "source": [
        "expand_args_fields(RenderedMeshDatasetMapProvider)\n",
        "cow_provider = RenderedMeshDatasetMapProvider(\n",
        "    data_file=\"data/cow_mesh/cow.obj\",\n",
        "    use_point_light=False,\n",
        "    resolution=output_resolution,\n",
        ")"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
281
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
282
283
284
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
285
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
286
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
287
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
288
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
289
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
290
291
292
        "customOutput": null,
        "executionStartTime": 1659610966145,
        "executionStopTime": 1659610966255,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
293
294
295
296
        "hidden_ranges": [],
        "originalKey": "8210e15b-da48-4306-a49a-41c4e7e7d42f",
        "requestMsgId": "c243edd2-a106-4fba-8471-dfa4f99a2088",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
297
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
298
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
299
300
301
      "source": [
        "dataset_map = cow_provider.get_dataset_map()\n",
        "tr_cameras = [training_frame.camera for training_frame in dataset_map.train]"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
302
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
303
304
305
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
306
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
307
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
308
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
309
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
310
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
311
312
313
        "customOutput": null,
        "executionStartTime": 1659610967703,
        "executionStopTime": 1659610967848,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
314
315
316
317
        "hidden_ranges": [],
        "originalKey": "458d72ad-d9a7-4f13-b5b7-90d2aec61c16",
        "requestMsgId": "7f9431f3-8717-4d89-a7fe-1420dd0e00c4",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
318
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
319
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
320
321
322
323
324
      "source": [
        "# The cameras are all in the XZ plane, in a circle about 2.7 from the origin\n",
        "centers = torch.cat([i.get_camera_center() for i in tr_cameras])\n",
        "print(centers.min(0).values)\n",
        "print(centers.max(0).values)"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
325
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
326
327
328
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
329
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
330
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
331
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
332
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
333
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
334
335
336
        "customOutput": null,
        "executionStartTime": 1659552920194,
        "executionStopTime": 1659552923122,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
337
338
339
340
        "hidden_ranges": [],
        "originalKey": "931e712b-b141-437a-97fb-dc2a07ce3458",
        "requestMsgId": "931e712b-b141-437a-97fb-dc2a07ce3458",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
341
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
342
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
343
344
345
346
347
      "source": [
        "# visualization of the cameras\n",
        "plot = plot_scene({\"k\": {i: camera for i, camera in enumerate(tr_cameras)}}, camera_scale=0.25)\n",
        "plot.layout.scene.aspectmode = \"data\"\n",
        "plot"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
348
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
349
350
    },
    {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
351
      "attachments": {},
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
352
353
      "cell_type": "markdown",
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
354
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
355
        "originalKey": "afa9c02d-f76b-4f68-83e9-9733c615406b",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
356
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
357
358
359
360
361
362
363
364
365
366
367
368
369
      },
      "source": [
        "## 2. Custom implicit function 🧊\n",
        "\n",
        "At the core of neural rendering methods are functions of spatial coordinates called implicit functions, which are used in some kind of rendering process.\n",
        "(Often those functions can additionally take other data as well, such as view direction.)\n",
        "A common rendering process is ray marching over densities and colors provided by an implicit function.\n",
        "In our case, taking samples from a 3D volume grid is a very simple function of spatial coordinates. \n",
        "\n",
        "Here we define our own implicit function, which uses PyTorch3D's existing functionality for sampling from a volume grid.\n",
        "We do this by subclassing `ImplicitFunctionBase`.\n",
        "We need to register our subclass with a special decorator.\n",
        "We use Python's dataclass annotations for configuring the module."
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
370
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
371
372
373
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
374
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
375
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
376
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
377
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
378
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
379
380
381
        "customOutput": null,
        "executionStartTime": 1659613575850,
        "executionStopTime": 1659613575940,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
382
383
384
385
        "hidden_ranges": [],
        "originalKey": "61b55043-dc52-4de7-992e-e2195edd2123",
        "requestMsgId": "dfaace3c-098c-4ffe-9240-6a7ae0ff271e",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
386
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
387
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
      "source": [
        "@registry.register\n",
        "class MyVolumes(ImplicitFunctionBase, torch.nn.Module):\n",
        "    grid_resolution: int = 50  # common HWD of volumes, the number of voxels in each direction\n",
        "    extent: float = 1.0  # In world coordinates, the volume occupies is [-extent, extent] along each axis\n",
        "\n",
        "    def __post_init__(self):\n",
        "        # We have to call this explicitly if there are other base classes like Module\n",
        "        super().__init__()\n",
        "\n",
        "        # We define parameters like other torch.nn.Module objects.\n",
        "        # In this case, both our parameter tensors are trainable; they govern the contents of the volume grid.\n",
        "        density = torch.full((self.grid_resolution, self.grid_resolution, self.grid_resolution), -2.0)\n",
        "        self.density = torch.nn.Parameter(density)\n",
        "        color = torch.full((3, self.grid_resolution, self.grid_resolution, self.grid_resolution), 0.0)\n",
        "        self.color = torch.nn.Parameter(color)\n",
        "        self.density_activation = torch.nn.Softplus()\n",
        "\n",
        "    def forward(\n",
        "        self,\n",
        "        ray_bundle: RayBundle,\n",
        "        fun_viewpool=None,\n",
        "        global_code=None,\n",
        "    ):\n",
        "        densities = self.density_activation(self.density[None, None])\n",
        "        voxel_size = 2.0 * float(self.extent) / self.grid_resolution\n",
        "        features = self.color.sigmoid()[None]\n",
        "\n",
        "        # Like other PyTorch3D structures, the actual Volumes object should only exist as long\n",
        "        # as one iteration of training. It is local to this function.\n",
        "\n",
        "        volume = Volumes(densities=densities, features=features, voxel_size=voxel_size)\n",
        "        sampler = VolumeSampler(volumes=volume)\n",
        "        densities, features = sampler(ray_bundle)\n",
        "\n",
        "        # When an implicit function is used for raymarching, i.e. for MultiPassEmissionAbsorptionRenderer,\n",
        "        # it must return (densities, features, an auxiliary tuple)\n",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
425
426
        "        return densities, features, {}\n"
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
427
428
    },
    {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
429
      "attachments": {},
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
430
431
      "cell_type": "markdown",
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
432
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
433
        "originalKey": "abaf2cd6-1b68-400e-a142-8fb9f49953f3",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
434
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
435
436
437
438
439
440
      },
      "source": [
        "## 3. Construct the model object.\n",
        "\n",
        "The main model object in PyTorch3D is `GenericModel`, which has pluggable components for the major steps, including the renderer and the implicit function(s).\n",
        "There are two ways to construct it which are equivalent here."
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
441
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
442
443
444
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
445
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
446
447
      "metadata": {
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
448
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
449
450
        "customOutput": null,
        "executionStartTime": 1659621267561,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
451
452
453
454
        "executionStopTime": 1659621267938,
        "originalKey": "f26c3dce-fbae-4592-bd0e-e4a8abc57c2c",
        "requestMsgId": "9213687e-1caf-46a8-a4e5-a9c531530092",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
455
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
456
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
      "source": [
        "CONSTRUCT_MODEL_FROM_CONFIG = True\n",
        "if CONSTRUCT_MODEL_FROM_CONFIG:\n",
        "    # Via a DictConfig - this is how our training loop with hydra works\n",
        "    cfg = get_default_args(GenericModel)\n",
        "    cfg.implicit_function_class_type = \"MyVolumes\"\n",
        "    cfg.render_image_height=output_resolution\n",
        "    cfg.render_image_width=output_resolution\n",
        "    cfg.loss_weights={\"loss_rgb_huber\": 1.0}\n",
        "    cfg.tqdm_trigger_threshold=19000\n",
        "    cfg.raysampler_AdaptiveRaySampler_args.scene_extent= 4.0\n",
        "    gm = GenericModel(**cfg)\n",
        "else:\n",
        "    # constructing GenericModel directly\n",
        "    expand_args_fields(GenericModel)\n",
        "    gm = GenericModel(\n",
        "        implicit_function_class_type=\"MyVolumes\",\n",
        "        render_image_height=output_resolution,\n",
        "        render_image_width=output_resolution,\n",
        "        loss_weights={\"loss_rgb_huber\": 1.0},\n",
        "        tqdm_trigger_threshold=19000,\n",
        "        raysampler_AdaptiveRaySampler_args = {\"scene_extent\": 4.0}\n",
        "    )\n",
        "\n",
        "    # In this case we can get the equivalent DictConfig cfg object to the way gm is configured as follows\n",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
482
483
        "    cfg = OmegaConf.structured(gm)\n"
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
484
485
    },
    {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
486
      "attachments": {},
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
487
488
      "cell_type": "markdown",
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
489
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
490
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
491
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
492
493
494
        "customOutput": null,
        "executionStartTime": 1659611214689,
        "executionStopTime": 1659611214748,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
495
496
497
498
        "hidden_ranges": [],
        "originalKey": "4e659f7d-ce66-4999-83de-005eb09d7705",
        "requestMsgId": "7b815b2b-cf19-44d0-ae89-76fde6df35ec",
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
499
500
501
      },
      "source": [
        " The default renderer is an emission-absorbtion raymarcher. We keep that default."
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
502
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
503
504
505
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
506
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
507
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
508
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
509
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
510
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
511
512
513
        "customOutput": null,
        "executionStartTime": 1659621268007,
        "executionStopTime": 1659621268190,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
514
515
516
517
        "hidden_ranges": [],
        "originalKey": "d37ae488-c57c-44d3-9def-825dc1a6495b",
        "requestMsgId": "71143ec1-730f-4876-8a14-e46eea9d6dd1",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
518
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
519
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
520
521
522
523
524
      "source": [
        "# We can display the configuration in use as follows.\n",
        "remove_unused_components(cfg)\n",
        "yaml = OmegaConf.to_yaml(cfg, sort_keys=False)\n",
        "%page -r yaml"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
525
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
526
527
528
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
529
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
530
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
531
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
532
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
533
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
534
535
536
        "customOutput": null,
        "executionStartTime": 1659621268727,
        "executionStopTime": 1659621268776,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
537
538
539
540
        "hidden_ranges": [],
        "originalKey": "52e53179-3c6e-4c1f-a38a-3a6d803687bb",
        "requestMsgId": "05de9bc3-3f74-4a6f-851c-9ec919b59506",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
541
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
542
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
543
544
545
546
      "source": [
        "device = torch.device(\"cuda:0\")\n",
        "gm.to(device)\n",
        "assert next(gm.parameters()).is_cuda"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
547
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
548
549
550
551
    },
    {
      "cell_type": "markdown",
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
552
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
553
        "originalKey": "528a7d53-c645-49c2-9021-09adbb18cd23",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
554
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
555
556
557
558
559
560
561
      },
      "source": [
        "## 4. train the model "
      ]
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
562
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
563
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
564
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
565
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
566
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
567
568
569
        "customOutput": null,
        "executionStartTime": 1659621270236,
        "executionStopTime": 1659621270446,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
570
571
572
573
        "hidden_ranges": [],
        "originalKey": "953280bd-3161-42ba-8dcb-0c8ef2d5cc25",
        "requestMsgId": "9bba424b-7bfd-4e5a-9d79-ae316e20bab0",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
574
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
575
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
576
577
      "source": [
        "train_data_collated = [FrameData.collate([frame.to(device)]) for frame in dataset_map.train]"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
578
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
579
580
581
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
582
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
583
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
584
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
585
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
586
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
587
588
589
        "customOutput": null,
        "executionStartTime": 1659621270815,
        "executionStopTime": 1659621270948,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
590
591
592
593
        "hidden_ranges": [],
        "originalKey": "2fcf07f0-0c28-49c7-8c76-1c9a9d810167",
        "requestMsgId": "821deb43-6084-4ece-83c3-dee214562c47",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
594
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
595
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
596
597
598
      "source": [
        "gm.train()\n",
        "optimizer = torch.optim.Adam(gm.parameters(), lr=0.1)"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
599
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
600
601
602
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
603
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
604
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
605
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
606
607
608
609
        "collapsed": false,
        "customOutput": null,
        "executionStartTime": 1659621271875,
        "executionStopTime": 1659621298146,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
610
611
612
613
        "hidden_ranges": [],
        "originalKey": "105099f7-ed0c-4e7f-a976-61a93fd0a8fe",
        "requestMsgId": "0c87c108-83e3-4129-ad02-85e0140f1368",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
614
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
615
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
616
617
618
619
620
621
622
623
624
625
626
      "source": [
        "iterator = tqdm.tqdm(range(2000))\n",
        "for n_batch in iterator:\n",
        "    optimizer.zero_grad()\n",
        "\n",
        "    frame = train_data_collated[n_batch % len(dataset_map.train)]\n",
        "    out = gm(**frame, evaluation_mode=EvaluationMode.TRAINING)\n",
        "    out[\"objective\"].backward()\n",
        "    if n_batch % 100 == 0:\n",
        "        iterator.set_postfix_str(f\"loss: {float(out['objective']):.5f}\")\n",
        "    optimizer.step()"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
627
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
628
629
    },
    {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
630
      "attachments": {},
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
631
632
633
      "cell_type": "markdown",
      "metadata": {
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
634
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
635
636
        "customOutput": null,
        "executionStartTime": 1659535024768,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
637
638
639
640
        "executionStopTime": 1659535024906,
        "originalKey": "e3cd494a-536b-48bc-8290-c048118c82eb",
        "requestMsgId": "e3cd494a-536b-48bc-8290-c048118c82eb",
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
641
642
643
644
645
      },
      "source": [
        "## 5. Evaluate the module\n",
        "\n",
        "We generate complete images from all the viewpoints to see how they look."
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
646
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
647
648
649
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
650
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
651
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
652
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
653
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
654
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
655
656
657
        "customOutput": null,
        "executionStartTime": 1659621299859,
        "executionStopTime": 1659621311133,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
658
659
660
661
        "hidden_ranges": [],
        "originalKey": "fbe1b2ea-cc24-4b20-a2d7-0249185e34a5",
        "requestMsgId": "771ef1f8-5eee-4932-9e81-33604bf0512a",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
662
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
663
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
      "source": [
        "def to_numpy_image(image):\n",
        "    # Takes an image of shape (C, H, W) in [0,1], where C=3 or 1\n",
        "    # to a numpy uint image of shape (H, W, 3)\n",
        "    return (image * 255).to(torch.uint8).permute(1, 2, 0).detach().cpu().expand(-1, -1, 3).numpy()\n",
        "def resize_image(image):\n",
        "    # Takes images of shape (B, C, H, W) to (B, C, output_resolution, output_resolution)\n",
        "    return torch.nn.functional.interpolate(image, size=(output_resolution, output_resolution))\n",
        "\n",
        "gm.eval()\n",
        "images = []\n",
        "expected = []\n",
        "masks = []\n",
        "masks_expected = []\n",
        "for frame in tqdm.tqdm(train_data_collated):\n",
        "    with torch.no_grad():\n",
        "        out = gm(**frame, evaluation_mode=EvaluationMode.EVALUATION)\n",
        "\n",
        "    image_rgb = to_numpy_image(out[\"images_render\"][0])\n",
        "    mask = to_numpy_image(out[\"masks_render\"][0])\n",
        "    expd = to_numpy_image(resize_image(frame.image_rgb)[0])\n",
        "    mask_expected = to_numpy_image(resize_image(frame.fg_probability)[0])\n",
        "\n",
        "    images.append(image_rgb)\n",
        "    masks.append(mask)\n",
        "    expected.append(expd)\n",
        "    masks_expected.append(mask_expected)"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
691
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
692
693
    },
    {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
694
      "attachments": {},
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
695
696
697
      "cell_type": "markdown",
      "metadata": {
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
698
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
699
700
        "customOutput": null,
        "executionStartTime": 1659614622542,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
701
702
703
704
        "executionStopTime": 1659614622757,
        "originalKey": "24953039-9780-40fd-bd81-5d63e9f40069",
        "requestMsgId": "7af895a3-dfe4-4c28-ac3b-4ff0fbb40c7f",
        "showInput": false
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
      },
      "source": [
        "We draw a grid showing predicted image and expected image, followed by predicted mask and expected mask, from each viewpoint. \n",
        "This is a grid of four rows of images, wrapped in to several large rows, i.e..\n",
        "<small><center>\n",
        "```\n",
        "┌────────┬────────┐           ┌────────┐\n",
        "│pred    │pred    │           │pred    │\n",
        "│image   │image   │           │image   │\n",
        "│1       │2       │           │n       │\n",
        "├────────┼────────┤           ├────────┤\n",
        "│expected│expected│           │expected│\n",
        "│image   │image   │  ...      │image   │\n",
        "│1       │2       │           │n       │\n",
        "├────────┼────────┤           ├────────┤\n",
        "│pred    │pred    │           │pred    │\n",
        "│mask    │mask    │           │mask    │\n",
        "│1       │2       │           │n       │\n",
        "├────────┼────────┤           ├────────┤\n",
        "│expected│expected│           │expected│\n",
        "│mask    │mask    │           │mask    │\n",
        "│1       │2       │           │n       │\n",
        "├────────┼────────┤           ├────────┤\n",
        "│pred    │pred    │           │pred    │\n",
        "│image   │image   │           │image   │\n",
        "│n+1     │n+1     │           │2n      │\n",
        "├────────┼────────┤           ├────────┤\n",
        "│expected│expected│           │expected│\n",
        "│image   │image   │  ...      │image   │\n",
        "│n+1     │n+2     │           │2n      │\n",
        "├────────┼────────┤           ├────────┤\n",
        "│pred    │pred    │           │pred    │\n",
        "│mask    │mask    │           │mask    │\n",
        "│n+1     │n+2     │           │2n      │\n",
        "├────────┼────────┤           ├────────┤\n",
        "│expected│expected│           │expected│\n",
        "│mask    │mask    │           │mask    │\n",
        "│n+1     │n+2     │           │2n      │\n",
        "└────────┴────────┘           └────────┘\n",
        "           ...\n",
        "```\n",
        "</center></small>"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
747
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
748
749
750
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
751
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
752
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
753
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
754
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
755
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
756
757
758
        "customOutput": null,
        "executionStartTime": 1659621313894,
        "executionStopTime": 1659621314042,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
759
760
761
762
        "hidden_ranges": [],
        "originalKey": "c488a34a-e46d-4649-93fb-4b1bb5a0e439",
        "requestMsgId": "4221e632-fca1-4fe5-b2e3-f92c37aa40e4",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
763
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
764
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
      "source": [
        "images_to_display = [images.copy(), expected.copy(), masks.copy(), masks_expected.copy()]\n",
        "n_rows = 4\n",
        "n_images = len(images)\n",
        "blank_image = images[0] * 0\n",
        "n_per_row = 1+(n_images-1)//n_rows\n",
        "for _ in range(n_per_row*n_rows - n_images):\n",
        "    for group in images_to_display:\n",
        "        group.append(blank_image)\n",
        "\n",
        "images_to_display_listed = [[[i] for i in j] for j in images_to_display]\n",
        "split = []\n",
        "for row in range(n_rows):\n",
        "    for group in images_to_display_listed:\n",
        "        split.append(group[row*n_per_row:(row+1)*n_per_row])  \n",
        "\n",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
781
782
        "Image.fromarray(np.block(split))\n"
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
783
784
785
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
786
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
787
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
788
        "code_folding": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
789
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
790
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
791
792
793
        "customOutput": null,
        "executionStartTime": 1659621323795,
        "executionStopTime": 1659621323820,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
794
795
796
797
        "hidden_ranges": [],
        "originalKey": "49eab9e1-4fe2-4fbe-b4f3-7b6953340170",
        "requestMsgId": "85b402ad-f903-431f-a13e-c2d697e869bb",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
798
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
799
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
800
801
802
      "source": [
        "# Print the maximum channel intensity in the first image.\n",
        "print(images[1].max()/255)"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
803
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
804
805
806
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
807
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
808
      "metadata": {
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
809
810
        "code_folding": [],
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
811
812
        "customInput": null,
        "customOutput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
813
814
        "executionStartTime": 1659621408642,
        "executionStopTime": 1659621409559,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
815
        "hidden_ranges": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
816
        "originalKey": "137d2c43-d39d-4266-ac5e-2b714da5e0ee",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
817
        "requestMsgId": "8e27ec57-c2d6-4ae0-be69-b63b6af929ff",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
818
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
819
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
820
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
821
822
823
824
825
826
827
      "source": [
        "plt.ioff()\n",
        "fig, ax = plt.subplots(figsize=(3,3))\n",
        "\n",
        "ax.grid(None)\n",
        "ims = [[ax.imshow(im, animated=True)] for im in images]\n",
        "ani = animation.ArtistAnimation(fig, ims, interval=80, blit=True)\n",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
828
829
        "ani_html = ani.to_jshtml()\n"
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
830
831
832
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
833
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
834
835
      "metadata": {
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
836
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
837
838
        "customOutput": null,
        "executionStartTime": 1659621409620,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
839
840
841
842
        "executionStopTime": 1659621409725,
        "originalKey": "783e70d6-7cf1-4d76-a126-ba11ffc2f5be",
        "requestMsgId": "b6843506-c5fa-4508-80fc-8ecae51a934a",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
843
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
844
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
845
846
      "source": [
        "HTML(ani_html)"
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
847
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
848
849
850
    },
    {
      "cell_type": "code",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
851
      "execution_count": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
852
853
      "metadata": {
        "collapsed": false,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
854
        "customInput": null,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
855
856
        "customOutput": null,
        "executionStartTime": 1659614670081,
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
857
858
859
860
        "executionStopTime": 1659614670168,
        "originalKey": "0286c350-2362-4f47-8181-2fc2ba51cfcf",
        "requestMsgId": "976f4db9-d4c7-466c-bcfd-218234400226",
        "showInput": true
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
861
      },
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
862
      "outputs": [],
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
863
864
865
866
      "source": [
        "# If you want to see the output of the model with the volume forced to opaque white, run this and re-evaluate\n",
        "# with torch.no_grad():\n",
        "#      gm._implicit_functions[0]._fn.density.fill_(9.0)\n",
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
867
868
        "#      gm._implicit_functions[0]._fn.color.fill_(9.0)\n"
      ]
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
869
    }
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
  ],
  "metadata": {
    "bento_stylesheets": {
      "bento/extensions/flow/main.css": true,
      "bento/extensions/kernel_selector/main.css": true,
      "bento/extensions/kernel_ui/main.css": true,
      "bento/extensions/new_kernel/main.css": true,
      "bento/extensions/system_usage/main.css": true,
      "bento/extensions/theme/main.css": true
    },
    "captumWidgetMessage": {},
    "dataExplorerConfig": {},
    "kernelspec": {
      "display_name": "pytorch3d",
      "language": "python",
      "metadata": {
        "cinder_runtime": false,
        "fbpkg_supported": true,
        "is_prebuilt": true,
        "kernel_name": "bento_kernel_pytorch3d",
        "nightly_builds": true
      },
      "name": "bento_kernel_pytorch3d"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3"
    },
    "last_base_url": "https://9177.od.fbinfra.net:443/",
    "last_kernel_id": "bb33cd83-7924-489a-8bd8-2d9d62eb0126",
    "last_msg_id": "99f7088e-d22b355b859660479ef0574e_5743",
    "last_server_session_id": "2944b203-9ea8-4c0e-9634-645dfea5f26b",
    "outputWidgetContext": {}
  },
  "nbformat": 4,
  "nbformat_minor": 2
Jeremy Reizenstein's avatar
Jeremy Reizenstein committed
913
}