Upload app.py

app.py

@@ -4,6 +4,7 @@ import numpy as np
 import torch
 from PIL import Image
 import gradio as gr
+from gradio_imageslider import ImageSlider
 
 # 延迟 CUDA 初始化
 weight_dtype = torch.float32
@@ -45,11 +46,22 @@ def process_image(input_image, resolution_choice):
     # 将 Gradio 输入转换为 PIL 图像
     input_image = Image.fromarray(input_image)
 
-    # 根据用户选择设置处理分辨率
+    # 如果 resolution_choice 为 '768'，将 input_image resize 到最大边 768
     if resolution_choice == "768":
-        processing_resolution = None
-    else:
-        processing_resolution = 0  # 使用原始分辨率
+        max_size = 768
+        width, height = input_image.size
+        if max(width, height) > max_size:
+            scaling_factor = max_size / max(width, height)
+            new_width = int(width * scaling_factor)
+            new_height = int(height * scaling_factor)
+            input_image = input_image.resize((new_width, new_height), Image.LANCZOS)
+
+    # 根据用户选择设置处理分辨率
+    # if resolution_choice == "768":
+    #     processing_resolution = None
+    # else:
+    #     processing_resolution = 0  # 使用原始分辨率
+    processing_resolution = 0  # 使用原始分辨率
 
     # 处理图像
     pipe_out = pipe(
@@ -64,7 +76,7 @@ def process_image(input_image, resolution_choice):
     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():
@@ -88,17 +100,18 @@ def create_gradio_interface():
                     value="768",  # 默认选择原始分辨率
                 )
                 gr.Markdown(
-                    "Select the resolution for processing the image. Higher resolution may take longer to process."
+                    "Select the resolution for processing the image. Higher resolution may take longer to process. 768 is recommended for faster processing and stable performance."
                 )
                 submit_btn = gr.Button("Remove Reflection", variant="primary")
             with gr.Column():
-                output_image = gr.Image(label="Processed Image")
+                # output_image = gr.Image(label="Processed Image")
+                output_slider = ImageSlider(label="Processed image", type="pil")
 
         # 添加示例
         gr.Examples(
             examples=example_images,
             inputs=[input_image, resolution_choice],  # 输入组件列表
-            outputs=output_image,
+            outputs=output_slider,
             fn=process_image,
             cache_examples=False,  # 缓存结果以加快加载速度
             label="Example Images",
@@ -108,7 +121,7 @@ def create_gradio_interface():
         submit_btn.click(
             fn=process_image,
             inputs=[input_image, resolution_choice],  # 输入组件列表
-            outputs=output_image,
+            outputs=output_slider,
         )
 
     return demo