google-research

1
{
2
 "cells": [
3
  {
4
   "cell_type": "markdown",
5
   "id": "804a617d-b335-49d2-8883-e4c3992695ee",
6
   "metadata": {},
7
   "source": [
8
    "Copyright 2023 Google LLC. SPDX-License-Identifier: Apache-2.0\n",
9
    "\n",
10
    "Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at\n",
11
    "\n",
12
    "https://www.apache.org/licenses/LICENSE-2.0\n",
13
    "\n",
14
    "Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License."
15
   ]
16
  },
17
  {
18
   "cell_type": "code",
19
   "execution_count": null,
20
   "id": "50ff8ba0",
21
   "metadata": {},
22
   "outputs": [],
23
   "source": [
24
    "%env CUDA_DEVICE_ORDER=PCI_BUS_ID\n",
25
    "%env CUDA_VISIBLE_DEVICES=0"
26
   ]
27
  },
28
  {
29
   "cell_type": "code",
30
   "execution_count": null,
31
   "id": "2f34a8c8",
32
   "metadata": {},
33
   "outputs": [],
34
   "source": [
35
    "import torch\n",
36
    "from PIL import Image\n",
37
    "import pickle\n",
38
    "import numpy as np\n",
39
    "from IPython.display import display\n",
40
    "from ipywidgets import HTML\n",
41
    "from ipyevents import Event\n",
42
    "from baukit import renormalize\n",
43
    "\n",
44
    "from models.layout import model_full\n",
45
    "from utils import sky_util, soat_util, camera_util, render_settings, show, filters\n"
46
   ]
47
  },
48
  {
49
   "cell_type": "code",
50
   "execution_count": null,
51
   "id": "91216085",
52
   "metadata": {},
53
   "outputs": [],
54
   "source": [
55
    "torch.set_grad_enabled(False)\n",
56
    "device = 'cuda'"
57
   ]
58
  },
59
  {
60
   "cell_type": "markdown",
61
   "id": "c4cc1e6e",
62
   "metadata": {},
63
   "source": [
64
    "# load models"
65
   ]
66
  },
67
  {
68
   "cell_type": "code",
69
   "execution_count": null,
70
   "id": "530a4d16",
71
   "metadata": {},
72
   "outputs": [],
73
   "source": [
74
    "full_model = model_full.ModelFull('pretrained/model_terrain.pkl', 'pretrained/model_sky_360.pkl').to(device).eval()"
75
   ]
76
  },
77
  {
78
   "cell_type": "code",
79
   "execution_count": null,
80
   "id": "b3aeafb5",
81
   "metadata": {},
82
   "outputs": [],
83
   "source": [
84
    "full_model.set_nerf_params(**render_settings.nerf_render_interactive)\n",
85
    "G_layout = full_model.terrain_model.layout_model\n",
86
    "G_soat = soat_util.init_soat_model(G_layout)\n",
87
    "G_sky = full_model.sky_model\n",
88
    "grid = sky_util.make_grid(G_sky.G)\n",
89
    "input_layer = G_sky.G.synthesis.input"
90
   ]
91
  },
92
  {
93
   "cell_type": "markdown",
94
   "id": "59e45118",
95
   "metadata": {},
96
   "source": [
97
    "# generate initial layout and skydome env map"
98
   ]
99
  },
100
  {
101
   "cell_type": "code",
102
   "execution_count": null,
103
   "id": "65820c01",
104
   "metadata": {},
105
   "outputs": [],
106
   "source": [
107
    "seed = 944 # np.random.randint(0, 1000)\n",
108
    "truncation = 0.8\n",
109
    "grid_size = 5\n",
110
    "\n",
111
    "layout = soat_util.generate_layout(seed, grid_h=grid_size, grid_w=grid_size, device=device, truncation_psi=truncation)\n",
112
    "z = torch.randn(1, G_layout.layout_generator.z_dim, device=device)  \n",
113
    "c = None\n",
114
    "noise_input = torch.randn_like(layout)[:, :1]\n",
115
    "\n",
116
    "sampled_Rt = G_layout.trajectory_sampler.sample_trajectories(G_layout.layout_decoder, layout)\n",
117
    "initial_camera = camera_util.camera_from_pose(sampled_Rt.squeeze())\n",
118
    "\n",
119
    "print(initial_camera)"
120
   ]
121
  },
122
  {
123
   "cell_type": "code",
124
   "execution_count": null,
125
   "id": "ddf0303b",
126
   "metadata": {},
127
   "outputs": [],
128
   "source": [
129
    "# sample initial frame\n",
130
    "Rt = camera_util.pose_from_camera(initial_camera)[None].to(device)\n",
131
    "camera_params = camera_util.get_full_image_parameters(\n",
132
    "    G_layout, G_layout.layout_decoder_kwargs.nerf_out_res, batch_size=1, device=device, Rt=Rt)\n",
133
    "outputs = full_model(z, c, camera_params, truncation=truncation, \n",
134
    "                     nerf_kwargs=dict(cached_layout=layout,\n",
135
    "                                      extras=[],\n",
136
    "                                      noise_input=noise_input,\n",
137
    "                                     )\n",
138
    "                    )\n",
139
    "\n",
140
    "# generate sky texture based on initial frame\n",
141
    "sky_encoder_ws = G_sky.encode(outputs['rgb_up'] * outputs['sky_mask'])\n",
142
    "sky_z = z[:, :G_sky.G.z_dim]\n",
143
    "start_grid = sky_util.generate_start_grid(seed, input_layer, grid)\n",
144
    "sky_pano = sky_util.generate_pano_transform(G_sky.G, sky_z, sky_encoder_ws, start_grid)\n",
145
    "sky_texture = sky_pano[None]\n",
146
    "# show(renormalize.as_image((outputs['rgb_up'] * outputs['sky_mask'])[0]))\n",
147
    "show(renormalize.as_image(sky_texture[0]))"
148
   ]
149
  },
150
  {
151
   "cell_type": "markdown",
152
   "id": "9ecbb7a2",
153
   "metadata": {},
154
   "source": [
155
    "# interactive widget"
156
   ]
157
  },
158
  {
159
   "cell_type": "code",
160
   "execution_count": null,
161
   "id": "c76602f2",
162
   "metadata": {},
163
   "outputs": [],
164
   "source": [
165
    "l = HTML(\"\")\n",
166
    "h = HTML(\"\")\n",
167
    "d = Event(source=l, watched_events=['keydown'])\n",
168
    "display_size = (256, 256)\n",
169
    "\n",
170
    "camera = initial_camera\n",
171
    "camera_util.INITIAL_CAMERA = initial_camera\n",
172
    "\n",
173
    "# How fast we adjust. Too large and it will overshoot.\n",
174
    "# Too small and it will not react in time to avoid mountains.\n",
175
    "tilt_velocity_scale = .3    # Keep this small, otherwise you'll get motion sickness.\n",
176
    "offset_velocity_scale = .5\n",
177
    "\n",
178
    "# How far up the image should the horizon be, ideally.\n",
179
    "# Suggested range: 0.5 to 0.7.\n",
180
    "horizon_target = 0.65\n",
181
    "\n",
182
    "# What proportion of the depth map should be \"near\" the camera, ideally.\n",
183
    "# The smaller the number, the higher up the camera will fly.\n",
184
    "# Suggested range: 0.05 to 0.2\n",
185
    "near_target = 0.2\n",
186
    "\n",
187
    "offset = 0\n",
188
    "tilt = 0\n",
189
    "initial_stabilize_frames = 10\n",
190
    "\n",
191
    "def generate_frame_from_camera(camera):\n",
192
    "    Rt = camera_util.pose_from_camera(camera)[None].to(device)\n",
193
    "    camera_params = camera_util.get_full_image_parameters(\n",
194
    "        G_layout, G_layout.layout_decoder_kwargs.nerf_out_res, batch_size=1, device=device, Rt=Rt)\n",
195
    "    outputs = full_model(z, c, camera_params, truncation=truncation, \n",
196
    "                         nerf_kwargs=dict(extras = ['camera_points'],\n",
197
    "                                          cached_layout=layout,\n",
198
    "                                          noise_input=noise_input, \n",
199
    "                                         ),\n",
200
    "                         sky_texture=sky_texture\n",
201
    "                        )\n",
202
    "    return outputs\n",
203
    "    \n",
204
    "def update_display(outputs, camera):\n",
205
    "    composite_rgb_url = renormalize.as_url(outputs['rgb_overlay_upsample'][0], size=display_size)\n",
206
    "\n",
207
    "    vis_rays =  camera_util.visualize_rays(G_layout, outputs['extras']['Rt'], outputs['extras']['camera_points'],\n",
208
    "                                           outputs['extras']['layout'], display_size[0])\n",
209
    "    cam_img = renormalize.as_image(vis_rays)    \n",
210
    "    cam_url = renormalize.as_url(cam_img, size=display_size)\n",
211
    "    img_html = ('<div class=\"row\"> <img src=\"%s\"/> <img src=\"%s\"/> </div>' % (composite_rgb_url, cam_url))\n",
212
    "    l.value = img_html\n",
213
    "    h.value = str(camera)\n",
214
    "    \n",
215
    "\n",
216
    "def handle_event(event):\n",
217
    "    global camera, offset, tilt\n",
218
    "    camera = camera_util.update_camera(camera, event['key'], auto_adjust_height_and_tilt=True)\n",
219
    "    c = camera_util.adjust_camera_vertically(camera, offset, tilt)\n",
220
    "    outputs = generate_frame_from_camera(c)\n",
221
    "    outputs = filters.smooth_mask(outputs) # optional mask smoothing\n",
222
    "    update_display(outputs, c)\n",
223
    "    tilt, offset = camera_util.update_tilt_and_offset(outputs, tilt, offset,\n",
224
    "                                                      horizon_target=horizon_target,\n",
225
    "                                                      near_target=near_target, \n",
226
    "                                                      tilt_velocity_scale=tilt_velocity_scale,\n",
227
    "                                                      offset_velocity_scale=offset_velocity_scale)                                                    \n",
228
    "\n",
229
    "    \n",
230
    "for x in range(initial_stabilize_frames):\n",
231
    "    outputs = generate_frame_from_camera(camera_util.adjust_camera_vertically(camera, offset, tilt))\n",
232
    "    tilt, offset = camera_util.update_tilt_and_offset(outputs, tilt, offset, \n",
233
    "                                                      horizon_target=horizon_target,\n",
234
    "                                                      near_target=near_target, \n",
235
    "                                                      tilt_velocity_scale=tilt_velocity_scale,\n",
236
    "                                                      offset_velocity_scale=offset_velocity_scale)\n",
237
    "\n",
238
    "d.on_dom_event(handle_event)\n",
239
    "display(h, l)\n",
240
    "handle_event({'key': 'x'})"
241
   ]
242
  },
243
  {
244
   "cell_type": "code",
245
   "execution_count": null,
246
   "id": "8c958a58",
247
   "metadata": {},
248
   "outputs": [],
249
   "source": []
250
  }
251
 ],
252
 "metadata": {
253
  "kernelspec": {
254
   "display_name": "persistentnature",
255
   "language": "python",
256
   "name": "persistentnature"
257
  },
258
  "language_info": {
259
   "codemirror_mode": {
260
    "name": "ipython",
261
    "version": 3
262
   },
263
   "file_extension": ".py",
264
   "mimetype": "text/x-python",
265
   "name": "python",
266
   "nbconvert_exporter": "python",
267
   "pygments_lexer": "ipython3",
268
   "version": "3.8.16"
269
  }
270
 },
271
 "nbformat": 4,
272
 "nbformat_minor": 5
273
}
274