Upload app.py

app.py

@@ -1,9 +1,10 @@
-import spaces  # 必须放在最前面
+import spaces
 import os
 import numpy as np
 import torch
 from PIL import Image
 import gradio as gr
+from gradio_imageslider import ImageSlider  # 导入 ImageSlider 组件
 
 # 延迟 CUDA 初始化
 weight_dtype = torch.float32
@@ -58,7 +59,8 @@ def process_image(input_image):
     processed_frame = (processed_frame[0] * 255).astype(np.uint8)
     processed_frame = Image.fromarray(processed_frame)
 
-    return processed_frame
+    # 返回输入图像和处理后的图像
+    return input_image, processed_frame
 
 # 创建 Gradio 界面
 def create_gradio_interface():
@@ -74,13 +76,18 @@ def create_gradio_interface():
                 input_image = gr.Image(label="Input Image", type="numpy")
                 submit_btn = gr.Button("Remove Reflection", variant="primary")
             with gr.Column():
-                output_image = gr.Image(label="Processed Image")
+                # 使用 ImageSlider 显示前后对比
+                output_slider = ImageSlider(
+                    label="Before & After",
+                    show_download_button=True,
+                    show_share_button=True,
+                )
 
         # 添加示例
         gr.Examples(
             examples=example_images,
             inputs=input_image,
-            outputs=output_image,
+            outputs=output_slider,
             fn=process_image,
             cache_examples=False,  # 缓存结果以加快加载速度
             label="Example Images",
@@ -90,7 +97,7 @@ def create_gradio_interface():
         submit_btn.click(
             fn=process_image,
             inputs=input_image,
-            outputs=output_image,
+            outputs=output_slider,
         )
 
     return demo