Upload app.py

app.py

@@ -41,25 +41,22 @@ pipe = DAIPipeline(
 ).to(device)
 
 @spaces.GPU
-def process_image(input_image):
+def process_image(input_image, resolution_choice):
     # 将 Gradio 输入转换为 PIL 图像
     input_image = Image.fromarray(input_image)
 
-    # 根据输入图像的大小设置分辨率
-    max_side = max(input_image.size)
-    if max_side < 768:
-        resolution = 768
-    elif max_side > 2560:
-        resolution = 2560
+    # 根据用户选择设置处理分辨率
+    if resolution_choice == "768":
+        processing_resolution = None
     else:
-        resolution = 0
+        processing_resolution = 0  # 使用原始分辨率
 
     # 处理图像
     pipe_out = pipe(
         image=input_image,
         prompt="remove glass reflection",
         vae_2=vae_2,
-        processing_resolution=resolution,
+        processing_resolution=processing_resolution,
     )
 
     # 将输出转换为图像
@@ -73,19 +70,26 @@ def process_image(input_image):
 def create_gradio_interface():
     # 示例图像
     example_images = [
-        os.path.join("files", "image", f"{i}.png") for i in range(1, 14)
+        [os.path.join("files", "image", f"{i}.png"), "768"] for i in range(1, 14)
     ]
     title = "# Dereflection Any Image"
     description = """Official demo for **Dereflection Any Image**.
     Please refer to our [paper](), [project page](https://abuuu122.github.io/DAI.github.io/), and [github](https://github.com/Abuuu122/Dereflection-Any-Image) for more details."""
 
-
     with gr.Blocks() as demo:
         gr.Markdown(title)
         gr.Markdown(description)
         with gr.Row():
             with gr.Column():
                 input_image = gr.Image(label="Input Image", type="numpy")
+                resolution_choice = gr.Radio(
+                    choices=["768", "Original Resolution"],
+                    label="Processing Resolution",
+                    value="768",  # 默认选择原始分辨率
+                )
+                gr.Markdown(
+                    "Select the resolution for processing the image. Higher resolution may take longer to process."
+                )
                 submit_btn = gr.Button("Remove Reflection", variant="primary")
             with gr.Column():
                 output_image = gr.Image(label="Processed Image")
@@ -93,7 +97,7 @@ def create_gradio_interface():
         # 添加示例
         gr.Examples(
             examples=example_images,
-            inputs=input_image,
+            inputs=[input_image, resolution_choice],  # 输入组件列表
             outputs=output_image,
             fn=process_image,
             cache_examples=False,  # 缓存结果以加快加载速度
@@ -103,7 +107,7 @@ def create_gradio_interface():
         # 绑定按钮点击事件
         submit_btn.click(
             fn=process_image,
-            inputs=input_image,
+            inputs=[input_image, resolution_choice],  # 输入组件列表
             outputs=output_image,
         )