Update app.py

app.py

@@ -1,24 +1,24 @@
-import sys
-from subprocess import check_call
-import tempfile
-
+import os
 from os.path import basename, splitext, join
-from io import BytesIO
-
+import tempfile
+import gradio as gr
 import numpy as np
-from scipy.spatial import KDTree
 from PIL import Image
-
 import torch
-import torch.nn.functional as F
+import cv2
 from torchvision.transforms.functional import to_tensor, to_pil_image
-from einops import rearrange
-import gradio as gr
+from torch import Tensor
+from genstereo import GenStereo, AdaptiveFusionLayer
+import ssl
 from huggingface_hub import hf_hub_download
+import spaces
 
-from extern.ZoeDepth.zoedepth.utils.misc import colorize
+from extern.DAM2.depth_anything_v2.dpt import DepthAnythingV2
+ssl._create_default_https_context = ssl._create_unverified_context
 
-from gradio_model3dgscamera import Model3DGSCamera
+IMAGE_SIZE = 512
+DEVICE = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
+CHECKPOINT_NAME = 'genstereo'
 
 def download_models():
     models = [
@@ -37,10 +37,17 @@ def download_models():
             'token': None
         },
         {
-            'repo': 'Sony/genwarp',
-            'sub': 'multi1',
+            '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,
             'dst': 'checkpoints',
-            'files': ['config.json', 'denoising_unet.pth', 'pose_guider.pth', 'reference_unet.pth'],
+            'files': [f'depth_anything_v2_vitl.pth'],
             'token': None
         }
     ]
@@ -58,53 +65,43 @@ def download_models():
 # 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
-)
+# 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
 
 # Crop the image to the shorter side.
 def crop(img: Image) -> Image:
@@ -115,301 +112,130 @@ def crop(img: Image) -> Image:
     else:
         left, right = np.ceil((W - H) / 2.), np.floor((W - H) / 2.) + H
         top, bottom = 0, H
-    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
+    return img.crop((left, top, right, bottom))
 
+def normalize_disp(disp):
+    return (disp - disp.min()) / (disp.max() - disp.min())
 
+# Gradio app
 with tempfile.TemporaryDirectory() as tmpdir:
     with gr.Blocks(
-        title='GenWarp Demo',
+        title='StereoGen Demo',
         css='img {display: inline;}'
     ) as demo:
         # Internal states.
-        image = gr.State()
-        depth = gr.State()
+        src_image = gr.State()
+        src_depth = gr.State()
 
         # Callbacks
-        @spaces.GPU()
+        @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
+            if not image_file:
+                # Return None if no image is provided (e.g., when file is cleared).
+                return None, None, None, None
 
-        @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)
+            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_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.
+        def cb_generate(image, depth: Tensor, scale_factor):
+            norm_disp = normalize_disp(depth.cuda())
+            disp = norm_disp * scale_factor / 100 * IMAGE_SIZE
+
+            genstereo = get_genstereo_model()
+            fusion_model = get_fusion_model()
+
+            renders = genstereo(
+                src_image=image,
+                src_disparity=disp,
+                ratio=None,
             )
+            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(warped[0])
+            fusion_pil = to_pil_image(fusion_image[0])
+
+            return warped_pil, fusion_pil
 
         # Blocks.
         gr.Markdown(
             """
-            # GenWarp: Single Image to Novel Views with Semantic-Preserving Generative Warping
-            [![Project Site](https://img.shields.io/badge/Project-Web-green)](https://genwarp-nvs.github.io/)  
-            [![Spaces](https://img.shields.io/badge/Spaces-Demo-yellow?logo=huggingface)](https://huggingface.co/spaces/Sony/GenWarp)  
-            [![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)  
+            # StereoGen: Towards Open-World Generation of Stereo Images and Unsupervised Matching
+            [![Project Site](https://img.shields.io/badge/Project-Web-green)](https://qjizhi.github.io/genstereo)  
+            [![Spaces](https://img.shields.io/badge/Spaces-Demo-yellow?logo=huggingface)](https://huggingface.co/spaces/FQiao/GenStereo)  
+            [![Github](https://img.shields.io/badge/Github-Repo-orange?logo=github)](https://github.com/Qjizhi/GenStereo)  
+            [![Models](https://img.shields.io/badge/Models-checkpoints-blue?logo=huggingface)](https://huggingface.co/FQiao/GenStereo/tree/main)  
             [![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).
-
+        
+            ## Introduction 
+            This is an official demo for the paper "[Towards Open-World Generation of Stereo Images and Unsupervised Matching](https://qjizhi.github.io/genstereo)". Given an arbitrary reference image, GenStereo can generate the corresponding right-view image.
+        
             ## 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.
+            1. Upload a reference image to "Left Image"
+                - You can also select an image from "Examples"
+            3. Hit "Generate a right image" button and check the result
 
             """
         )
-        file = gr.File(label='Reference Input', file_types=['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
+        )
         with gr.Row():
             image_widget = gr.Image(
-                label='Reference View', type='filepath',
+                label='Depth', type='filepath',
                 interactive=False
             )
             depth_widget = gr.Image(label='Estimated Depth', type='pil')
-            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')
+
+        # 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')
         with gr.Row():
             warped_widget = gr.Image(
                 label='Warped Image', type='pil', interactive=False
             )
             gen_widget = gr.Image(
-                label='Generated View', type='pil', interactive=False
+                label='Generated Right', 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.upload(
+        file.change(
             fn=cb_mde,
             inputs=file,
-            outputs=[image_widget, depth_widget, image, depth]
-        ).success(
-            fn=cb_3d,
-            inputs=[image_widget, image, depth],
-            outputs=viewer
+            outputs=[image_widget, depth_widget, src_image, src_depth]
         )
         button.click(
             fn=cb_generate,
-            inputs=[viewer, image, depth],
+            inputs=[src_image, src_depth, scale_slider],
             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()