Update space

app.py

@@ -1,18 +1,70 @@
 import gradio as gr
+import cv2
+import numpy as np
+import os
+from PIL import Image
+import spaces
+import torch
+from gradio_imageslider import ImageSlider
 
-def add_numbers(num1, num2):
-    # 这个函数接收两个参数并返回它们的和。
-    return num1 + num2
-
-# 创建Gradio的界面
-demo = gr.Interface(
-    fn=add_numbers,  # 指定用于计算的函数
-    inputs=[gr.Number(label="第一个数"), gr.Number(label="第二个数")],  # 输入组件，这里是两个数字输入框
-    outputs=gr.Number(label="两数之和"),  # 输出组件，这里是一个数字输出框
-    title="简单加法器",  # 界面标题
-    description="输入两个数字，点击提交按钮以查看它们的和。"  # 界面描述
-)
-
-# 启动应用
-if __name__ == "__main__":
-    demo.launch()
+css = """
+#img-display-container {
+    max-height: 100vh;
+    }
+#img-display-input {
+    max-height: 80vh;
+    }
+#img-display-output {
+    max-height: 80vh;
+    }
+"""
+DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
+title = "# Stereo Anything"
+description = """Official demo for **Stereo Anything: Unifying Stereo Matching with Large-Scale Mixed Data**.
+Please refer to our [paper](https://arxiv.org/abs/2411.14053), [github](https://github.com/XiandaGuo/OpenStereo/) for more details."""
+
+@spaces.GPU
+@torch.no_grad()
+def predict_depth(model, image):
+    return model(image)
+
+
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(title)
+    gr.Markdown(description)
+    gr.Markdown("### Depth Prediction demo")
+    gr.Markdown("You can slide the output to compare the depth prediction with input image")
+
+    with gr.Row():
+        left_image = gr.Image(label="Left Image", type='numpy', elem_id='img-display-input')
+        right_image = gr.Image(label="Right Image", type='numpy', elem_id='img-display-input')
+        depth_image_slider = ImageSlider(label="Depth Map with Slider View", elem_id='img-display-output', position=0.5,)
+    # raw_file = gr.File(label="16-bit raw depth (can be considered as disparity)")
+    submit = gr.Button("Submit")
+
+    def on_submit(left_image,right_image):
+        sample = {
+            'left': left_image,
+            'right': right_image,
+        }
+        sample['left'] = sample['left'].unsqueeze(0)
+        sample['right'] = sample['right'].unsqueeze(0)
+
+        model.eval()
+        for k, v in sample.items():
+            sample[k] = v.to(0) if torch.is_tensor(v) else v
+
+        model_pred = model(sample)
+
+        return [model_pred]
+
+    submit.click(on_submit, inputs=[left_image,right_image], outputs=[depth_image_slider])
+
+    example_files = os.listdir('examples')
+    example_files.sort()
+    example_files = [os.path.join('examples', filename) for filename in example_files]
+    examples = gr.Examples(examples=example_files, inputs=[left_image,right_image], outputs=[depth_image_slider], fn=on_submit, cache_examples=True)
+
+
+if __name__ == '__main__':
+    demo.queue().launch()