caching-version + variable baseline control

app.py

@@ -11,6 +11,8 @@ import sys
 import tempfile
 from PIL import Image
 from gradio_patches.radio import Radio
+import torch
+from omegaconf import OmegaConf
 
 REPO_URL = "https://github.com/prs-eth/stereospace.git"
 COMMIT_SHA = "d7bbae6"
@@ -80,9 +82,35 @@ def find_all_output_files(output_dir, base_name):
             pass
     return outputs
 
+class Args:
+    def __init__(
+        self,
+        input,
+        output,
+        baseline=0.15,
+        batch_size=1,
+        src_intrinsics=None,
+        tgt_intrinsics=None,
+    ):
+        self.input = input
+        self.output = output
+        self.baseline = baseline
+        self.batch_size = batch_size
+        self.src_intrinsics = src_intrinsics
+        self.tgt_intrinsics = tgt_intrinsics
+
+stereo_nvs = None
+config = OmegaConf.load(os.path.join(REPO_DIR, "configs/stereospace.yaml"))
 
 @spaces.GPU
-def process_all_modes(input_image):
+def process_all_modes(input_image, baseline=0.15):
+    global stereo_nvs
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    if stereo_nvs is None:
+        stereo_nvs = StereoSpace(config, device=device)
+
     if input_image is None:
         raise gr.Error("Please upload an image or select an example.")
     
@@ -96,27 +124,16 @@ def process_all_modes(input_image):
     try:
         with tempfile.TemporaryDirectory() as tmp_output:
             base_name = os.path.splitext(os.path.basename(input_path))[0]
-            
-            inference_script = "inference.py"
-            cmd = [
-                "python", inference_script,
-                "--input", input_path,
-                "--output", tmp_output
-            ]
-            
-            print(f"Running command: {' '.join(cmd)}")
-            print(f"Working directory: {REPO_DIR}")
-            result = subprocess.run(
-                cmd,
-                cwd=REPO_DIR,
-                capture_output=True,
-                text=True
+        
+            args = Args(
+                input=input_path,
+                output=tmp_output,
+                baseline=float(baseline),
+                batch_size=1,
             )
+
+            generate_novel_view(args=args, config=config, stereo_nvs=stereo_nvs)
             
-            if result.returncode != 0:
-                error_msg = f"Inference failed:\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}"
-                print(error_msg)
-                raise gr.Error(error_msg)
             
             try:
                 output_files = find_all_output_files(tmp_output, base_name)
@@ -198,12 +215,24 @@ with gr.Blocks(
         )
     
     with gr.Row():
-        image = gr.Image(
+        image_in = gr.Image(
+            type="filepath",
+            height=480,
+            visible=False,
+        )
+        image_out = gr.Image(
             type="filepath",
             label="Input/Output Image",
             elem_classes="result-image",
             height=480,
         )
+
+    with gr.Row():
+        baseline_slider = gr.Slider(
+            minimum=0.025, maximum=0.4, value=0.15, step=0.005,
+            label="Baseline length (meters)"
+        )
+        run_btn = gr.Button("Run")
     
     outputs_gallery = gr.Gallery(
         visible=False,
@@ -243,12 +272,12 @@ with gr.Blocks(
         
         return result
     
-    def process_new_image(img, current_mode):
+    def process_new_image(img, current_mode, baseline=0.15):
         if img is None:
             placeholder = create_placeholder_image()
             return [], placeholder, gr.update()
         
-        all_outputs = process_all_modes(img)        
+        all_outputs = process_all_modes(img, baseline=baseline)        
         gallery_data = []
         available_modes = []
         for mode_name, mode_image in all_outputs.items():
@@ -259,12 +288,12 @@ with gr.Blocks(
         radio_value = current_mode if current_mode in available_modes else (available_modes[0] if available_modes else "Anaglyph")
         return gallery_data, placeholder, gr.update(choices=available_modes, value=radio_value)
     
-    def process_example_simple(img):
+    def process_example_simple(img, baseline=0.15):
         if img is None:
             placeholder = create_placeholder_image()
             return [], placeholder, gr.update()
         
-        all_outputs = process_all_modes(img)
+        all_outputs = process_all_modes(img, baseline=baseline)
         
         gallery_data = []
         available_modes = []
@@ -277,45 +306,43 @@ with gr.Blocks(
         return gallery_data, placeholder, gr.update(choices=available_modes, value=radio_value)
     
     def clear_image():
-        return None, []
+        return None, None, []
     
     examples_list = get_example_images()
     if examples_list:
-        def process_example_wrapper(img):
-            gallery_data, placeholder_image, radio_update = process_example_simple(img)
-            return gallery_data, placeholder_image, radio_update
-        
         examples_component = gr.Examples(
             examples=examples_list,
-            inputs=[image],
-            outputs=[outputs_gallery, image, output_mode],
-            fn=process_example_wrapper,
-            cache_examples=True,
-            cache_mode="lazy",
+            inputs=[image_in],
             label="Example Images",
             elem_id="example-images-gallery",
         )
     
-    def process_upload_wrapper(img, current_mode):
-        gallery_data, blocked_image, radio_update = process_new_image(img, current_mode)
+    image_in.change(
+        fn=lambda img: (img, []),
+        inputs=[image_in],
+        outputs=[image_out, outputs_gallery],
+    )
+    
+    def run_wrapper(img, current_mode, baseline=0.15):
+        gallery_data, blocked_image, radio_update = process_new_image(img, current_mode, baseline=baseline)
         return gallery_data, blocked_image, radio_update
     
-    upload_event = image.upload(
-        fn=process_upload_wrapper,
-        inputs=[image, output_mode],
-        outputs=[outputs_gallery, image, output_mode]
+    run_btn.click(
+        fn=run_wrapper,
+        inputs=[image_in, output_mode, baseline_slider],
+        outputs=[outputs_gallery, image_out, output_mode],
     )
     
     def on_gallery_change(gallery_data, current_mode, current_image):
         if not gallery_data or len(gallery_data) == 0:
-            return current_image, gr.update(interactive=True)
+            return current_image
         updated_image = update_image_from_gallery(gallery_data, current_mode, current_image)
-        return updated_image, gr.update(interactive=True)
+        return updated_image
     
     gallery_change_event = outputs_gallery.change(
         fn=on_gallery_change,
-        inputs=[outputs_gallery, output_mode, image],
-        outputs=[image, image]
+        inputs=[outputs_gallery, output_mode, image_in],
+        outputs=image_out
     )
     
     def switch_mode_handler(current_mode, gallery_data, current_image):
@@ -324,13 +351,19 @@ with gr.Blocks(
     
     output_mode.change(
         fn=switch_mode_handler,
-        inputs=[output_mode, outputs_gallery, image],
-        outputs=image
+        inputs=[output_mode, outputs_gallery, image_out],
+        outputs=image_out
+    )
+
+    image_out.upload(
+        fn=lambda img: [img, img],
+        inputs=[image_out],
+        outputs=[image_in, image_out]
     )
     
-    image.clear(
+    image_out.clear(
         fn=clear_image,
-        outputs=[image, outputs_gallery]
+        outputs=[image_in, image_out, outputs_gallery]
     )
 
 if __name__ == "__main__":