Update app.py

app.py

@@ -201,29 +201,6 @@ def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str
     except Exception as e:
         raise gr.Error(f"Processing failed: {str(e)}")
 
-def process_bbox(img: Image.Image, box_input: str) -> tuple[list[Image.Image], str]:
-    try:
-        if img is None or box_input.strip() == "":
-            raise gr.Error("Please provide both image and bounding box coordinates")
-        
-        # Parse box coordinates
-        try:
-            coords = eval(box_input)
-            if not isinstance(coords, list) or len(coords) != 4:
-                raise ValueError("Invalid box format")
-            bbox = tuple(int(x) for x in coords)
-        except:
-            raise gr.Error("Invalid box format. Please provide [xmin, ymin, xmax, ymax]")
-        
-        # Process the image
-        results, download_path = _process(img, bbox)
-        
-        # Convert results to list for gallery
-        gallery_images = list(results)
-        
-        return gallery_images, download_path
-    except Exception as e:
-        raise gr.Error(str(e))
 
 
 def on_change_bbox(prompts: dict[str, Any] | None):
@@ -233,21 +210,39 @@ def on_change_bbox(prompts: dict[str, Any] | None):
 def on_change_prompt(img: Image.Image | None, prompt: str | None, bg_prompt: str | None = None):
     return gr.update(interactive=bool(img and prompt))
 
-def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None) -> tuple[list[Image.Image], str]:
+def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None) -> tuple[Image.Image, Image.Image]:
     try:
         if img is None or prompt.strip() == "":
             raise gr.Error("Please provide both image and prompt")
         
         # Process the image
-        results, download_path = _process(img, prompt, bg_prompt)
+        results, _ = _process(img, prompt, bg_prompt)
         
-        # Convert results to list for gallery
-        gallery_images = list(results)
-        
-        return gallery_images, download_path
+        # 합성된 이미지와 추출된 이미지만 반환
+        return results[1], results[2]
     except Exception as e:
         raise gr.Error(str(e))
 
+def process_bbox(img: Image.Image, box_input: str) -> tuple[Image.Image, Image.Image]:
+    try:
+        if img is None or box_input.strip() == "":
+            raise gr.Error("Please provide both image and bounding box coordinates")
+        
+        try:
+            coords = eval(box_input)
+            if not isinstance(coords, list) or len(coords) != 4:
+                raise ValueError("Invalid box format")
+            bbox = tuple(int(x) for x in coords)
+        except:
+            raise gr.Error("Invalid box format. Please provide [xmin, ymin, xmax, ymax]")
+        
+        # Process the image
+        results, _ = _process(img, bbox)
+        
+        # 합성된 이미지와 추출된 이미지만 반환
+        return results[1], results[2]
+    except Exception as e:
+        raise gr.Error(str(e))
 
 # Event handler functions 수정
 def update_process_button(img, prompt):
@@ -322,6 +317,7 @@ button.primary:hover {
 
 # UI 부분만 수정
 # Main Gradio app
+# UI 부분 수정
 with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     gr.HTML("""
         <div class="main-title">
@@ -357,15 +353,20 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     )
 
                 with gr.Column(scale=1):
-                    output_display = gr.Gallery(
-                        label="Results",
-                        show_download_button=False,
-                        visible=True
-                    )
-                    download_btn = gr.DownloadButton(
-                        "Download Result",
-                        visible=True
-                    )
+                    with gr.Row():
+                        combined_image = gr.Image(
+                            label="Combined Result",
+                            show_download_button=True,
+                            type="pil",
+                            height=512
+                        )
+                    with gr.Row():
+                        extracted_image = gr.Image(
+                            label="Extracted Object",
+                            show_download_button=True,
+                            type="pil",
+                            height=256
+                        )
 
         # Box-based extraction tab
         with gr.TabItem("Extract by Box"):
@@ -388,15 +389,20 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                     )
 
                 with gr.Column(scale=1):
-                    box_output = gr.Gallery(
-                        label="Results",
-                        show_download_button=False,
-                        visible=True
-                    )
-                    box_download = gr.DownloadButton(
-                        "Download Result",
-                        visible=True
-                    )
+                    with gr.Row():
+                        box_combined_image = gr.Image(
+                            label="Combined Result",
+                            show_download_button=True,
+                            type="pil",
+                            height=512
+                        )
+                    with gr.Row():
+                        box_extracted_image = gr.Image(
+                            label="Extracted Object",
+                            show_download_button=True,
+                            type="pil",
+                            height=256
+                        )
 
     # Event bindings
     input_image.change(
@@ -416,7 +422,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     process_btn.click(
         fn=process_prompt,
         inputs=[input_image, text_prompt, bg_prompt],
-        outputs=[output_display, download_btn],
+        outputs=[combined_image, extracted_image],  # 출력 변경
         queue=True
     )
 
@@ -437,7 +443,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     box_btn.click(
         fn=process_bbox,
         inputs=[box_image, box_input],
-        outputs=[box_output, box_download],
+        outputs=[box_combined_image, box_extracted_image],  # 출력 변경
         queue=True
     )