update examples to use cache

app.py

@@ -20,10 +20,6 @@ from extern.ZoeDepth.zoedepth.utils.misc import colorize
 
 from gradio_model3dgscamera import Model3DGSCamera
 
-IMAGE_SIZE = 512
-NEAR, FAR = 0.01, 100
-FOVY = np.deg2rad(55)
-
 def download_models():
     models = [
         {
@@ -59,6 +55,57 @@ def download_models():
                 token=model['token']
             )
 
+# Setup.
+download_models()
+
+mde = torch.hub.load(
+    './extern/ZoeDepth',
+    'ZoeD_N',
+    source='local',
+    pretrained=True,
+    trust_repo=True
+)
+
+import spaces
+
+check_call([
+    sys.executable, '-m', 'pip', 'install',
+    'extern/splatting-0.0.1-py3-none-any.whl'
+])
+
+from genwarp import GenWarp
+from genwarp.ops import (
+    camera_lookat, get_projection_matrix, get_viewport_matrix
+)
+
+# GenWarp
+genwarp_cfg = dict(
+    pretrained_model_path='checkpoints',
+    checkpoint_name='multi1',
+    half_precision_weights=True
+)
+genwarp_nvs = GenWarp(cfg=genwarp_cfg, device='cpu')
+
+# Fixed parameters.
+IMAGE_SIZE = 512
+NEAR, FAR = 0.01, 100
+FOVY = np.deg2rad(55)
+PROJ_MTX = get_projection_matrix(
+    fovy=torch.ones(1) * FOVY,
+    aspect_wh=1.,
+    near=NEAR,
+    far=FAR
+)
+VIEW_MTX = camera_lookat(
+    torch.tensor([[0., 0., 0.]]),
+    torch.tensor([[0., 0., 1.]]),
+    torch.tensor([[0., -1., 0.]])
+)
+VIEWPORT_MTX = get_viewport_matrix(
+    IMAGE_SIZE, IMAGE_SIZE,
+    batch_size=1
+)
+
 # Crop the image to the shorter side.
 def crop(img: Image) -> Image:
     W, H = img.size
@@ -68,68 +115,9 @@ def crop(img: Image) -> Image:
     else:
         left, right = np.ceil((W - H) / 2.), np.floor((W - H) / 2.) + H
         top, bottom = 0, H
-    return img.crop((left, top, right, bottom))
-
-def unproject(depth):
-    fovy_deg = 55
-    H, W = depth.shape[2:4]
-
-    mean_depth = depth.mean(dim=(2, 3)).squeeze().item()
-
-    viewport_mtx = get_viewport_matrix(
-        IMAGE_SIZE, IMAGE_SIZE,
-        batch_size=1
-    ).to(depth)
-
-    # Projection matrix.
-    fovy = torch.ones(1) * FOVY
-    proj_mtx = get_projection_matrix(
-        fovy=fovy,
-        aspect_wh=1.,
-        near=NEAR,
-        far=FAR
-    ).to(depth)
-
-    view_mtx = camera_lookat(
-        torch.tensor([[0., 0., 0.]]),
-        torch.tensor([[0., 0., 1.]]),
-        torch.tensor([[0., -1., 0.]])
-    ).to(depth)
-
-    scr_mtx = (viewport_mtx @ proj_mtx).to(depth)
-
-    grid = torch.stack(torch.meshgrid(
-        torch.arange(W), torch.arange(H), indexing='xy'), dim=-1
-    ).to(depth)[None]  # BHW2
-
-    screen = F.pad(grid, (0, 1), 'constant', 0)
-    screen = F.pad(screen, (0, 1), 'constant', 1)
-    screen_flat = rearrange(screen, 'b h w c -> b (h w) c')
-
-    eye = screen_flat @ torch.linalg.inv_ex(
-        scr_mtx.float()
-    )[0].mT.to(depth)
-    eye = eye * rearrange(depth, 'b c h w -> b (h w) c')
-    eye[..., 3] = 1
-
-    points = eye @ torch.linalg.inv_ex(view_mtx.float())[0].mT.to(depth)
-    points = points[0, :, :3]
-
-    # Translate to the origin.
-    points[..., 2] -= mean_depth
-    camera_pos = (0, 0, -mean_depth)
-    view_mtx = camera_lookat(
-        torch.tensor([[0., 0., -mean_depth]]),
-        torch.tensor([[0., 0., 0.]]),
-        torch.tensor([[0., -1., 0.]])
-    ).to(depth)
-
-    return points, camera_pos, view_mtx, proj_mtx
-
-def calc_dist2(points: np.ndarray):
-    dists, _ = KDTree(points).query(points, k=4)
-    mean_dists = (dists[:, 1:] ** 2).mean(1)
-    return mean_dists
+    img = img.crop((left, top, right, bottom))
+    img = img.resize((IMAGE_SIZE, IMAGE_SIZE))
+    return img
 
 def save_as_splat(
     filepath: str,
@@ -171,6 +159,44 @@ def save_as_splat(
     with open(filepath, "wb") as f:
         f.write(buffer.getvalue())
 
+def calc_dist2(points: np.ndarray):
+    dists, _ = KDTree(points).query(points, k=4)
+    mean_dists = (dists[:, 1:] ** 2).mean(1)
+    return mean_dists
+
+def unproject(depth):
+    H, W = depth.shape[2:4]
+    mean_depth = depth.mean(dim=(2, 3)).squeeze().item()
+
+    # Matrices.
+    viewport_mtx = VIEWPORT_MTX.to(depth)
+    proj_mtx = PROJ_MTX.to(depth)
+    view_mtx = VIEW_MTX.to(depth)
+    scr_mtx = (viewport_mtx @ proj_mtx).to(depth)
+
+    grid = torch.stack(torch.meshgrid(
+        torch.arange(W), torch.arange(H), indexing='xy'), dim=-1
+    ).to(depth)[None]  # BHW2
+
+    screen = F.pad(grid, (0, 1), 'constant', 0)
+    screen = F.pad(screen, (0, 1), 'constant', 1)
+    screen_flat = rearrange(screen, 'b h w c -> b (h w) c')
+
+    eye = screen_flat @ torch.linalg.inv_ex(
+        scr_mtx.float()
+    )[0].mT.to(depth)
+    eye = eye * rearrange(depth, 'b c h w -> b (h w) c')
+    eye[..., 3] = 1
+
+    points = eye @ torch.linalg.inv_ex(view_mtx.float())[0].mT.to(depth)
+    points = points[0, :, :3]
+
+    # Translate to the origin.
+    points[..., 2] -= mean_depth
+    camera_pos = (0, 0, -mean_depth)
+
+    return points, camera_pos
+
 def view_from_rt(position, rotation):
     t = np.array(position)
     euler = np.array(rotation)
@@ -213,48 +239,87 @@ def view_from_rt(position, rotation):
     return B @ view_mtx
 
 
-# Setup.
-download_models()
-
-mde = torch.hub.load(
-    './extern/ZoeDepth',
-    'ZoeD_N',
-    source='local',
-    pretrained=True,
-    trust_repo=True
-)
-
-import spaces
-
-check_call([
-    sys.executable, '-m', 'pip', 'install',
-    'extern/splatting-0.0.1-py3-none-any.whl'
-])
-
-from genwarp import GenWarp
-from genwarp.ops import (
-    camera_lookat, get_projection_matrix, get_viewport_matrix
-)
-
-# GenWarp
-genwarp_cfg = dict(
-    pretrained_model_path='checkpoints',
-    checkpoint_name='multi1',
-    half_precision_weights=True
-)
-genwarp_nvs = GenWarp(cfg=genwarp_cfg, device='cpu')
-
-
 with tempfile.TemporaryDirectory() as tmpdir:
     with gr.Blocks(
         title='GenWarp Demo',
         css='img {display: inline;}'
     ) as demo:
         # Internal states.
-        src_image = gr.State()
-        src_depth = gr.State()
-        proj_mtx = gr.State()
-        src_view_mtx = gr.State()
+        image = gr.State()
+        depth = gr.State()
+
+        # Callbacks
+        @spaces.GPU()
+        def cb_mde(image_file: str):
+            # Load an image.
+            image_pil = crop(Image.open(image_file).convert('RGB'))
+            image = to_tensor(image_pil)[None].detach()
+            # Get depth.
+            depth = mde.cuda().infer(image.cuda()).cpu().detach()
+            depth_pil = to_pil_image(colorize(depth[0]))
+            return image_pil, depth_pil, image, depth
+
+        @spaces.GPU()
+        def cb_3d(image_file, image, depth):
+            # Unproject.
+            xyz, camera_pos = unproject(depth.cuda())
+            xyz = xyz.cpu().detach().numpy()
+            # Save as a splat.
+            ## Output filename.
+            splat_file = join(
+                tmpdir, f'./{splitext(basename(image_file))[0]}.splat')
+            rgb = rearrange(image, 'b c h w -> b (h w) c')[0].numpy()
+            save_as_splat(splat_file, xyz, rgb)
+            return splat_file, camera_pos, (0, 0, 0)
+
+        @spaces.GPU()
+        def cb_generate(viewer, image, depth):
+            if depth is None:
+                gr.Error('Image and Depth are not set. Try again.')
+                return None, None
+
+            mean_depth = depth.mean(dim=(2, 3)).squeeze().item()
+            src_view_mtx = camera_lookat(
+                torch.tensor([[0., 0., -mean_depth]]),
+                torch.tensor([[0., 0., 0.]]),
+                torch.tensor([[0., -1., 0.]])
+            ).to(depth)
+            tar_camera_pos, tar_camera_rot = viewer[1:3]
+            tar_view_mtx = torch.from_numpy(view_from_rt(
+                tar_camera_pos, tar_camera_rot
+            ))
+            rel_view_mtx = (
+                tar_view_mtx @ torch.linalg.inv(src_view_mtx.double())
+            ).half().cuda()
+            proj_mtx = PROJ_MTX.half().cuda()
+
+            # GenWarp.
+            renders = genwarp_nvs.to('cuda')(
+                src_image=image.half().cuda(),
+                src_depth=depth.half().cuda(),
+                rel_view_mtx=rel_view_mtx,
+                src_proj_mtx=proj_mtx,
+                tar_proj_mtx=proj_mtx
+            )
+            warped_pil = to_pil_image(renders['warped'].cpu()[0])
+            synthesized_pil = to_pil_image(renders['synthesized'].cpu()[0])
+
+            return warped_pil, synthesized_pil
+
+        def process_example(image_file):
+            gr.Error('')
+            image_pil, depth_pil, image, depth = cb_mde(image_file)
+            viewer = cb_3d(image_file, image, depth)
+            # Fixed angle for examples.
+            viewer = (viewer[0], (-2.020, -0.727, -5.236), (-0.132, 0.378, 0.0))
+            warped_pil, synthsized_pil = cb_generate(
+                viewer, image, depth
+            )
+            return (
+                image_pil, depth_pil, viewer,
+                warped_pil, synthsized_pil,
+                None, None  # Clear internal states.
+            )
 
         # Blocks.
         gr.Markdown(
@@ -270,23 +335,23 @@ with tempfile.TemporaryDirectory() as tmpdir:
             This is an official demo for the paper "[GenWarp: Single Image to Novel Views with Semantic-Preserving Generative Warping](https://genwarp-nvs.github.io/)". Genwarp can generate novel view images from a single input conditioned on camera poses. In this demo, we offer a basic use of inference of the model. For detailed information, please refer to the [paper](https://arxiv.org/abs/2405.17251).
 
             ## How to Use
+
+            ### Try examples
+            - Examples are in the bottom section of the page
+
+            ### Upload your own images
             1. Upload a reference image to "Reference Input"
-                - You can also select a image from "Examples"
             2. Move the camera to your desired view in "Unprojected 3DGS" 3D viewer
             3. Hit "Generate a novel view" button and check the result
 
             ## Tips
+            - This model is mainly trained for indoor/outdoor scenery. It might not work well for object-centric inputs. For details on training the model, please check our [paper](https://arxiv.org/abs/2405.17251).
             - Extremely large camera movement from the input view might cause low performance results due to the unexpected deviation from the training distribution, which is not the scope of this model. Instead, you can feed the generation result for the small camera movement repeatedly and progressively move towards a desired view.
             - 3D viewer might take some time to update especially when trying different images back to back. Wait until it fully updates to the new image.
 
             """
         )
         file = gr.File(label='Reference Input', file_types=['image'])
-        examples = gr.Examples(
-            examples=['./assets/pexels-heyho-5998120_19mm.jpg',
-                    './assets/pexels-itsterrymag-12639296_24mm.jpg'],
-            inputs=file
-        )
         with gr.Row():
             image_widget = gr.Image(
                 label='Reference View', type='filepath',
@@ -312,68 +377,39 @@ with tempfile.TemporaryDirectory() as tmpdir:
             gen_widget = gr.Image(
                 label='Generated View', type='pil', interactive=False
             )
-
-        # Callbacks
-        @spaces.GPU
-        def cb_mde(image_file: str):
-            image = to_tensor(crop(Image.open(
-                image_file
-            ).convert('RGB')).resize((IMAGE_SIZE, IMAGE_SIZE)))[None].cuda()
-            depth = mde.cuda().infer(image)
-            depth_image = to_pil_image(colorize(depth[0]))
-            return to_pil_image(image[0]), depth_image, image.cpu().detach(), depth.cpu().detach()
-
-        @spaces.GPU
-        def cb_3d(image, depth, image_file):
-            xyz, camera_pos, view_mtx, proj_mtx = unproject(depth.cuda())
-            rgb = rearrange(image, 'b c h w -> b (h w) c')[0]
-            splat_file = join(tmpdir, f'./{splitext(basename(image_file))[0]}.splat')
-            save_as_splat(splat_file, xyz.cpu().detach().numpy(), rgb.cpu().detach().numpy())
-            return (splat_file, camera_pos, None), view_mtx.cpu().detach(), proj_mtx.cpu().detach()
-
-        @spaces.GPU
-        def cb_generate(viewer, image, depth, src_view_mtx, proj_mtx):
-            image = image.cuda()
-            depth = depth.cuda()
-            src_view_mtx = src_view_mtx.cuda()
-            proj_mtx = proj_mtx.cuda()
-            src_camera_pos = viewer[1]
-            src_camera_rot = viewer[2]
-            tar_view_mtx = view_from_rt(src_camera_pos, src_camera_rot)
-            tar_view_mtx = torch.from_numpy(tar_view_mtx).to(image)
-            rel_view_mtx = (
-                tar_view_mtx @ torch.linalg.inv(src_view_mtx.to(image))
-            ).to(image)
-
-            # GenWarp.
-            renders = genwarp_nvs.to('cuda')(
-                src_image=image.half(),
-                src_depth=depth.half(),
-                rel_view_mtx=rel_view_mtx.half(),
-                src_proj_mtx=proj_mtx.half(),
-                tar_proj_mtx=proj_mtx.half()
-            )
-
-            warped = renders['warped']
-            synthesized = renders['synthesized']
-            warped_pil = to_pil_image(warped[0])
-            synthesized_pil = to_pil_image(synthesized[0])
-
-            return warped_pil, synthesized_pil
+        examples = gr.Examples(
+            examples=[
+                './assets/pexels-heyho-5998120_19mm.jpg',
+                './assets/pexels-itsterrymag-12639296_24mm.jpg'
+            ],
+            fn=process_example,
+            inputs=file,
+            outputs=[image_widget, depth_widget, viewer,
+                     warped_widget, gen_widget,
+                     image, depth]
+        )
 
         # Events
-        file.change(
+        file.upload(
             fn=cb_mde,
             inputs=file,
-            outputs=[image_widget, depth_widget, src_image, src_depth]
-        ).then(
+            outputs=[image_widget, depth_widget, image, depth]
+        ).success(
             fn=cb_3d,
-            inputs=[src_image, src_depth, image_widget],
-            outputs=[viewer, src_view_mtx, proj_mtx])
+            inputs=[image_widget, image, depth],
+            outputs=viewer
+        )
         button.click(
             fn=cb_generate,
-            inputs=[viewer, src_image, src_depth, src_view_mtx, proj_mtx],
-            outputs=[warped_widget, gen_widget])
+            inputs=[viewer, image, depth],
+            outputs=[warped_widget, gen_widget]
+        )
+        # To re-calculate the uncached depth for examples in background.
+        examples.load_input_event.success(
+            fn=lambda x: cb_mde(x)[2:4],
+            inputs=file,
+            outputs=[image, depth]
+        )
 
     if __name__ == '__main__':
         demo.launch()