Update app.py

app.py

@@ -1,24 +1,24 @@
-import os
-from os.path import basename, splitext, join
+import sys
+from subprocess import check_call
 import tempfile
-import gradio as gr
+
+from os.path import basename, splitext, join
+from io import BytesIO
+
 import numpy as np
+from scipy.spatial import KDTree
 from PIL import Image
+
 import torch
-import cv2
+import torch.nn.functional as F
 from torchvision.transforms.functional import to_tensor, to_pil_image
-from torch import Tensor
-from genstereo import GenStereo, AdaptiveFusionLayer
-import ssl
+from einops import rearrange
+import gradio as gr
 from huggingface_hub import hf_hub_download
-import spaces
 
-from extern.DAM2.depth_anything_v2.dpt import DepthAnythingV2
-ssl._create_default_https_context = ssl._create_unverified_context
+from extern.ZoeDepth.zoedepth.utils.misc import colorize
 
-IMAGE_SIZE = 512
-DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
-CHECKPOINT_NAME = 'genstereo'
+from gradio_model3dgscamera import Model3DGSCamera
 
 def download_models():
     models = [
@@ -37,17 +37,10 @@ def download_models():
             'token': None
         },
         {
-            'repo': 'FQiao/GenStereo',
-            'sub': None,
-            'dst': 'checkpoints/genstereo',
-            'files': ['config.json', 'denoising_unet.pth', 'fusion_layer.pth', 'pose_guider.pth', 'reference_unet.pth'],
-            'token': None
-        },
-        {
-            'repo': 'depth-anything/Depth-Anything-V2-Large',
-            'sub': None,
+            'repo': 'Sony/genwarp',
+            'sub': 'multi1',
             'dst': 'checkpoints',
-            'files': [f'depth_anything_v2_vitl.pth'],
+            'files': ['config.json', 'denoising_unet.pth', 'pose_guider.pth', 'reference_unet.pth'],
             'token': None
         }
     ]
@@ -65,43 +58,53 @@ def download_models():
 # Setup.
 download_models()
 
-# DepthAnythingV2
-def get_dam2_model():
-    model_configs = {
-        'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]},
-        'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
-        'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
-    }
-
-    encoder = 'vitl'
-    encoder_size_map = {'vits': 'Small', 'vitb': 'Base', 'vitl': 'Large'}
-
-    if encoder not in encoder_size_map:
-        raise ValueError(f"Unsupported encoder: {encoder}. Supported: {list(encoder_size_map.keys())}")
-
-    dam2 = DepthAnythingV2(**model_configs[encoder])
-    dam2_checkpoint = f'checkpoints/depth_anything_v2_{encoder}.pth'
-    dam2.load_state_dict(torch.load(dam2_checkpoint, map_location='cpu'))
-    dam2 = dam2.to(DEVICE).eval()
-    return dam2
-
-# GenStereo
-def get_genstereo_model():
-    genwarp_cfg = dict(
-        pretrained_model_path='checkpoints',
-        checkpoint_name=CHECKPOINT_NAME,
-        half_precision_weights=True
-    )
-    genstereo = GenStereo(cfg=genwarp_cfg, device=DEVICE)
-    return genstereo
-
-# Adaptive Fusion
-def get_fusion_model():
-    fusion_model = AdaptiveFusionLayer()
-    fusion_checkpoint = join('checkpoints', CHECKPOINT_NAME, 'fusion_layer.pth')
-    fusion_model.load_state_dict(torch.load(fusion_checkpoint, map_location='cpu'))
-    fusion_model = fusion_model.to(DEVICE).eval()
-    return fusion_model
+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:
@@ -112,64 +115,211 @@ 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))
+    img = img.crop((left, top, right, bottom))
+    img = img.resize((IMAGE_SIZE, IMAGE_SIZE))
+    return img
+
+def save_as_splat(
+    filepath: str,
+    xyz: np.ndarray,
+    rgb: np.ndarray
+):
+    # To gaussian splat
+    inv_sigmoid = lambda x: np.log(x / (1 - x))
+    dist2 = np.clip(calc_dist2(xyz), a_min=0.0000001, a_max=None)
+    scales = np.repeat(np.log(np.sqrt(dist2))[..., np.newaxis], 3, axis=1)
+    rots = np.zeros((xyz.shape[0], 4))
+    rots[:, 0] = 1
+    opacities = inv_sigmoid(0.1 * np.ones((xyz.shape[0], 1)))
+
+    sorted_indices = np.argsort((
+        -np.exp(np.sum(scales, axis=-1, keepdims=True))
+        / (1 + np.exp(-opacities))
+    ).squeeze())
+
+    buffer = BytesIO()
+    for idx in sorted_indices:
+        position = xyz[idx]
+        scale = np.exp(scales[idx]).astype(np.float32)
+        rot = rots[idx].astype(np.float32)
+        color = np.concatenate(
+            (rgb[idx], 1 / (1 + np.exp(-opacities[idx]))),
+            axis=-1
+        )
+        buffer.write(position.tobytes())
+        buffer.write(scale.tobytes())
+        buffer.write((color * 255).clip(0, 255).astype(np.uint8).tobytes())
+        buffer.write(
+            ((rot / np.linalg.norm(rot)) * 128 + 128)
+            .clip(0, 255)
+            .astype(np.uint8)
+            .tobytes()
+        )
+
+    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)
+
+    cx = np.cos(euler[0])
+    sx = np.sin(euler[0])
+    cy = np.cos(euler[1])
+    sy = np.sin(euler[1])
+    cz = np.cos(euler[2])
+    sz = np.sin(euler[2])
+    R = np.array([
+        cy * cz + sy * sx * sz,
+        -cy * sz + sy * sx * cz,
+        sy * cx,
+        cx * sz,
+        cx * cz,
+        -sx,
+        -sy * cz + cy * sx * sz,
+        sy * sz + cy * sx * cz,
+        cy * cx
+    ])
+    view_mtx = np.array([
+        [R[0], R[1], R[2], 0],
+        [R[3], R[4], R[5], 0],
+        [R[6], R[7], R[8], 0],
+        [
+            -t[0] * R[0] - t[1] * R[3] - t[2] * R[6],
+            -t[0] * R[1] - t[1] * R[4] - t[2] * R[7],
+            -t[0] * R[2] - t[1] * R[5] - t[2] * R[8],
+            1
+        ]
+    ]).T
+
+    B = np.array([
+        [1, 0, 0, 0],
+        [0, -1, 0, 0],
+        [0, 0, -1, 0],
+        [0, 0, 0, 1]
+    ])
+    return B @ view_mtx
+
 
-# Gradio app
 with tempfile.TemporaryDirectory() as tmpdir:
     with gr.Blocks(
-        title='StereoGen Demo',
+        title='GenWarp Demo',
         css='img {display: inline;}'
     ) as demo:
         # Internal states.
-        src_image = gr.State()
-        src_depth = gr.State()
-
-        def normalize_disp(disp):
-            return (disp - disp.min()) / (disp.max() - disp.min())
+        image = gr.State()
+        depth = gr.State()
 
         # Callbacks
-        @spaces.GPU()        
+        @spaces.GPU()
         def cb_mde(image_file: str):
-            if not image_file:
-                # Return None if no image is provided (e.g., when file is cleared).
-                return None, None, None, None
+            # 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
 
-            image = crop(Image.open(image_file).convert('RGB'))  # Load image using PIL
-            image = image.resize((IMAGE_SIZE, IMAGE_SIZE))
-
-            image_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
-
-            dam2 = get_dam2_model()
-            depth_dam2 = dam2.infer_image(image_bgr)
-            depth = torch.tensor(depth_dam2).unsqueeze(0).unsqueeze(0).float()
-
-            depth_image = cv2.applyColorMap((normalize_disp(depth_dam2) * 255).astype(np.uint8), cv2.COLORMAP_JET)
-
-            return image, depth_image, 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(image, depth: Tensor, scale_factor):
-            norm_disp = normalize_disp(depth.cuda())
-            disp = norm_disp * scale_factor / 100 * IMAGE_SIZE
+        def cb_generate(viewer, image, depth):
+            if depth is None:
+                gr.Error('Image and Depth are not set. Try again.')
+                return None, None
 
-            genstereo = get_genstereo_model()
-            fusion_model = get_fusion_model()
+            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()
 
-            renders = genstereo(
-                src_image=image,
-                src_disparity=disp,
-                ratio=None,
+            # 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 = (renders['warped'] + 1) / 2
-            
-            synthesized = renders['synthesized']
-            mask = renders['mask']
-            fusion_image = fusion_model(synthesized.float(), warped.float(), mask.float())
+            warped_pil = to_pil_image(renders['warped'].cpu()[0])
+            synthesized_pil = to_pil_image(renders['synthesized'].cpu()[0])
 
-            warped_pil = to_pil_image(warped[0])
-            fusion_pil = to_pil_image(fusion_image[0])
+            return warped_pil, synthesized_pil
 
-            return warped_pil, fusion_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(
@@ -180,55 +330,86 @@ with tempfile.TemporaryDirectory() as tmpdir:
             [![Github](https://img.shields.io/badge/Github-Repo-orange?logo=github)](https://github.com/sony/genwarp/)  
             [![Models](https://img.shields.io/badge/Models-checkpoints-blue?logo=huggingface)](https://huggingface.co/Sony/genwarp)  
             [![arXiv](https://img.shields.io/badge/arXiv-2405.17251-red?logo=arxiv)](https://arxiv.org/abs/2405.17251)
-        
+
             ## Introduction
             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"
+            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='Left', file_types=['image'])
-        examples = gr.Examples(
-            examples=['./assets/COCO_val2017_000000070229.jpg',
-                    './assets/COCO_val2017_000000092839.jpg',
-                    './assets/KITTI2015_000003_10.png',
-                    './assets/KITTI2015_000147_10.png'],
-            inputs=file
-        )
+        file = gr.File(label='Reference Input', file_types=['image'])
         with gr.Row():
             image_widget = gr.Image(
-                label='Depth', type='filepath',
+                label='Reference View', type='filepath',
                 interactive=False
             )
             depth_widget = gr.Image(label='Estimated Depth', type='pil')
-
-        # Add scale factor slider
-        scale_slider = gr.Slider(
-            label='Scale Factor',
-            minimum=1.0,
-            maximum=30.0,
-            value=15.0,
-            step=0.1,
-        )
-
-        button = gr.Button('Generate a right image', size='lg', variant='primary')
+            viewer = Model3DGSCamera(
+                label = 'Unprojected 3DGS',
+                width=IMAGE_SIZE,
+                height=IMAGE_SIZE,
+                camera_width=IMAGE_SIZE,
+                camera_height=IMAGE_SIZE,
+                camera_fx=IMAGE_SIZE / (np.tan(FOVY / 2.)) / 2.,
+                camera_fy=IMAGE_SIZE / (np.tan(FOVY / 2.)) / 2.,
+                camera_near=NEAR,
+                camera_far=FAR
+            )
+        button = gr.Button('Generate a novel view', size='lg', variant='primary')
         with gr.Row():
             warped_widget = gr.Image(
                 label='Warped Image', type='pil', interactive=False
             )
             gen_widget = gr.Image(
-                label='Generated Right', type='pil', interactive=False
+                label='Generated View', type='pil', interactive=False
             )
+        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]
+            outputs=[image_widget, depth_widget, image, depth]
+        ).success(
+            fn=cb_3d,
+            inputs=[image_widget, image, depth],
+            outputs=viewer
         )
         button.click(
             fn=cb_generate,
-            inputs=[src_image, src_depth, scale_slider],
+            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()