Update app.py

app.py

@@ -12,7 +12,6 @@ from torch import Tensor
 from genstereo import GenStereo, AdaptiveFusionLayer
 import ssl
 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
@@ -146,39 +145,31 @@ with tempfile.TemporaryDirectory() as tmpdir:
                 IMAGE_SIZE = 768
                 CHECKPOINT_NAME = 'genstereo-v2.1'
                 print(f"Switched to GenStereo {sd_version_choice}. IMAGE_SIZE: {IMAGE_SIZE}, CHECKPOINT: {CHECKPOINT_NAME}")
-            return None, None, None, None, None, None
+            return None, None, None, None, None, None, None
 
         @spaces.GPU()
-        def cb_generate(image, depth, scale_factor, sd_version):
-            depth_tensor = torch.tensor(depth).unsqueeze(0).unsqueeze(0).float()
-            norm_disp = normalize_disp(depth_tensor.cuda())
-            disp = norm_disp * scale_factor / 100 * IMAGE_SIZE
+        def cb_mde(image_file: str, sd_version):
+            if not image_file:
+                return None, None, None, None
 
-            genstereo = get_genstereo_model(sd_version)
-            fusion_model = get_fusion_model()
+            image = crop(Image.open(image_file).convert('RGB'))
+            if sd_version == "v1.5":
+                image = image.resize((IMAGE_SIZE, IMAGE_SIZE))
+            elif sd_version == "v2.1":
+                image = image.resize((IMAGE_SIZE, IMAGE_SIZE))
+            else:
+                gr.Warning(f"Unknown SD version: {sd_version}. Defaulting to {IMAGE_SIZE}.")
+                image = image.resize((IMAGE_SIZE, IMAGE_SIZE))
+            gr.Info(f"Generating with GenStereo {sd_version} at {IMAGE_SIZE}px resolution.")
 
-            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())
+            image_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
 
-            warped_pil = to_pil_image(warped[0])
-            fusion_pil = to_pil_image(fusion_image[0])
+            dam2 = get_dam2_model()
+            depth_dam2 = dam2.infer_image(image_bgr)
 
-            # Create full SBS for Quest 2
-            left_resized = image.resize((1832, 1920))
-            right_resized = fusion_pil.resize((1832, 1920))
-            sbs = Image.new('RGB', (3664, 1920))
-            sbs.paste(left_resized, (0, 0))
-            sbs.paste(right_resized, (1832, 0))
+            depth_image = cv2.applyColorMap((normalize_disp(depth_dam2) * 255).astype(np.uint8), cv2.COLORMAP_JET)
 
-            return warped_pil, fusion_pil, sbs
+            return image, depth_image, image, depth_dam2
 
         @spaces.GPU()
         def cb_generate(image, depth, scale_factor, sd_version):
@@ -195,7 +186,7 @@ with tempfile.TemporaryDirectory() as tmpdir:
                 ratio=None,
             )
             warped = (renders['warped'] + 1) / 2
-            
+
             synthesized = renders['synthesized']
             mask = renders['mask']
             fusion_image = fusion_model(synthesized.float(), warped.float(), mask.float())
@@ -209,21 +200,20 @@ with tempfile.TemporaryDirectory() as tmpdir:
             sbs = Image.new('RGB', (3664, 1920))
             sbs.paste(left_resized, (0, 0))
             sbs.paste(right_resized, (1832, 0))
-            sbs.save('/home/user/app/sbs_quest2.jpg', quality=95)
 
-            return warped_pil, fusion_pil
+            return warped_pil, fusion_pil, sbs
 
         # Blocks.
         gr.Markdown(
             """
             # [ICCV 2025] Towards Open-World Generation of Stereo Images and Unsupervised Matching
             [Project Site](https://qjizhi.github.io/genstereo) | [Spaces](https://huggingface.co/spaces/FQiao/GenStereo) | [Github](https://github.com/Qjizhi/GenStereo) | [Models](https://huggingface.co/FQiao/GenStereo-sd2.1/tree/main) | [arXiv](https://arxiv.org/abs/2503.12720)
-        
-            ## Introduction 
+
+            ## 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
-            
+
             1. Select the GenStereo version
                 - v1.5: 512px, faster.
                 - v2.1: 768px, better performance, high resolution, takes more time.
@@ -241,60 +231,8 @@ with tempfile.TemporaryDirectory() as tmpdir:
         )
 
         with gr.Row():
-
             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='Left Image', 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')
-        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
-            )
-
-        # Events
-        sd_version_radio.change(
-            fn=cb_update_sd_version,
-            inputs=sd_version_radio,
-            outputs=[
-                image_widget, depth_widget, # Clear image displays
-                src_image, src_depth,         # Clear internal states
-                warped_widget, gen_widget     # Clear generation outputs
-            ]
-        )        
-        file.change(
-            fn=cb_mde,
-            inputs=[file, sd_version_radio],
-            outputs=[image_widget, depth_widget, src_image, src_depth]
-        )
-        button.click(
-            fn=cb_generate,
-            inputs=[src_image, src_depth, scale_slider, sd_version_radio],
-            outputs=[warped_widget, gen_widget]
-        )
-
-    if __name__ == '__main__':
-        demo.launch()
+                        './asset