Upload app.py

app.py

@@ -41,19 +41,16 @@ pipe = DAIPipeline(
 ).to(device)
 
 @spaces.GPU
-def process_image(input_image, resolution_option):
+def process_image(input_image):
     # 将 Gradio 输入转换为 PIL 图像
     input_image = Image.fromarray(input_image)
 
-    # 根据用户选择的分辨率选项设置 processing_resolution
-    processing_resolution = 2560 if resolution_option == "2560" else None
-
     # 处理图像
     pipe_out = pipe(
         image=input_image,
         prompt="remove glass reflection",
         vae_2=vae_2,
-        processing_resolution=processing_resolution,
+        processing_resolution=None,
     )
 
     # 将输出转换为图像
@@ -67,67 +64,41 @@ def process_image(input_image, resolution_option):
 def create_gradio_interface():
     # 示例图像
     example_images = [
-        os.path.join("files", "image", f"{i}.png") for i in range(1, 10)
+        os.path.join("files", "image", filename)
+        for filename in os.listdir(os.path.join("files", "image"))
+        if filename.endswith((".png", ".jpg", ".jpeg"))
     ]
     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)
-
-        # 创建两个 Tab
-        with gr.Tabs():
-            with gr.Tab("Image"):
-                with gr.Row():
-                    with gr.Column():
-                        input_image_default = gr.Image(label="Input Image", type="numpy")
-                        submit_btn_default = gr.Button("Remove Reflection", variant="primary")
-                    with gr.Column():
-                        output_image_default = gr.Image(label="Processed Image")
-
-                # 添加示例
-                gr.Examples(
-                    examples=example_images,
-                    inputs=input_image_default,
-                    outputs=output_image_default,
-                    fn=lambda x: process_image(x, "None"),
-                    cache_examples=False,  # 缓存结果以加快加载速度
-                    label="Example Images",
-                )
-
-                # 绑定按钮点击事件
-                submit_btn_default.click(
-                    fn=lambda x: process_image(x, "None"),
-                    inputs=input_image_default,
-                    outputs=output_image_default,
-                )
-
-            with gr.Tab("2K (2560)"):
-                with gr.Row():
-                    with gr.Column():
-                        input_image_high = gr.Image(label="Input Image", type="numpy")
-                        submit_btn_high = gr.Button("Remove Reflection", variant="primary")
-                    with gr.Column():
-                        output_image_high = gr.Image(label="Processed Image")
-
-                # 添加示例
-                gr.Examples(
-                    examples=example_images,
-                    inputs=input_image_high,
-                    outputs=output_image_high,
-                    fn=lambda x: process_image(x, "2560"),
-                    cache_examples=False,  # 缓存结果以加快加载速度
-                    label="Example Images",
-                )
-
-                # 绑定按钮点击事件
-                submit_btn_high.click(
-                    fn=lambda x: process_image(x, "2560"),
-                    inputs=input_image_high,
-                    outputs=output_image_high,
-                )
+        with gr.Row():
+            with gr.Column():
+                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")
+
+        # 添加示例
+        gr.Examples(
+            examples=example_images,
+            inputs=input_image,
+            outputs=output_image,
+            fn=process_image,
+            cache_examples=False,  # 缓存结果以加快加载速度
+            label="Example Images",
+        )
+
+        # 绑定按钮点击事件
+        submit_btn.click(
+            fn=process_image,
+            inputs=input_image,
+            outputs=output_image,
+        )
 
     return demo