Add Bounding Box vs Segment Overlay output style option in Gradio UI

app.py

@@ -18,7 +18,7 @@ from describe import describe_change
 
 import gradio as gr
 
-def gradio_compare(img1, img2):
+def gradio_compare(img1, img2, style_option="Segment Overlay"):
     if img1 is None or img2 is None:
         return None, None, "Please upload both images.", {}
 
@@ -39,7 +39,8 @@ def gradio_compare(img1, img2):
         # Run comparison pipeline
         pipeline_used = "AI Pipeline"
         try:
-            result_data = compare_images_ai(path1, path2, output_path, heatmap_path)
+            draw_style = "box" if style_option == "Bounding Boxes Only" else "overlay"
+            result_data = compare_images_ai(path1, path2, output_path, heatmap_path, draw_style=draw_style)
         except Exception as e:
             print(f"[Gradio] AI pipeline failed ({e}); falling back to classical")
             result_data = compare_images(path1, path2, output_path, heatmap_path)
@@ -109,6 +110,11 @@ with gr.Blocks() as demo:
             gr.Markdown("### Input Images")
             img1_input = gr.Image(label="Before Image", type="numpy")
             img2_input = gr.Image(label="After Image", type="numpy")
+            style_input = gr.Radio(
+                choices=["Segment Overlay", "Bounding Boxes Only"],
+                value="Segment Overlay",
+                label="Visualization Output Style"
+            )
             submit_btn = gr.Button("Compare & Analyze", variant="primary", size="lg")
             
         with gr.Column(scale=1):
@@ -123,7 +129,7 @@ with gr.Blocks() as demo:
                     
     submit_btn.click(
         fn=gradio_compare,
-        inputs=[img1_input, img2_input],
+        inputs=[img1_input, img2_input, style_input],
         outputs=[result_img_output, heatmap_img_output, desc_output, metrics_output]
     )
 

diff_ai.py

@@ -458,7 +458,7 @@ def _composite_object(base, object_img, mask, bbox, color):
 # MAIN ENTRY
 # =====================================================================
 
-def _compare_images_ai_impl(img1_path, img2_path, output_path, heatmap_path):
+def _compare_images_ai_impl(img1_path, img2_path, output_path, heatmap_path, draw_style="overlay"):
     """End-to-end change-region-driven AI pipeline. Raises RuntimeError
     if any required model is unavailable."""
     if not _get_loftr() or not _get_sam() or not _get_dinov2():
@@ -611,18 +611,29 @@ def _compare_images_ai_impl(img1_path, img2_path, output_path, heatmap_path):
     n_total   = n_added + n_removed + n_moved
     severity  = "HIGH" if n_total > 0 else "NONE"
 
-    # 5) Render: dark-tint background + highlight changed objects
-    DARK_FACTOR = 0.3
-    result_img = cv2.multiply(aligned, np.array([DARK_FACTOR] * 3, dtype=np.float64))
-    result_img = np.clip(result_img, 0, 255).astype(np.uint8)
-
-    for a in added_objects:
-        _composite_object(result_img, aligned, a["mask"], a["bbox"], COLOR_ADDED)
-    for r in removed_objects:
-        _composite_object(result_img, r["source"], r["mask"], r["bbox"], COLOR_REMOVED)
-    for m in moved_objects:
-        _composite_object(result_img, m["from"]["source"], m["from"]["mask"], m["from"]["bbox"], COLOR_MOVED)
-        _composite_object(result_img, aligned,              m["to"]["mask"],   m["to"]["bbox"],   COLOR_MOVED)
+    # 5) Render output image
+    if draw_style == "box":
+        result_img = aligned.copy()
+        for a in added_objects:
+            _draw_box(result_img, a["bbox"], COLOR_ADDED, "ADDED")
+        for r in removed_objects:
+            _draw_box(result_img, r["bbox"], COLOR_REMOVED, "REMOVED")
+        for m in moved_objects:
+            _draw_box(result_img, m["from"]["bbox"], COLOR_MOVED, "MOVED")
+            _draw_box(result_img, m["to"]["bbox"],   COLOR_MOVED, "MOVED")
+    else:
+        # Dark-tint background + highlight changed objects (overlay)
+        DARK_FACTOR = 0.3
+        result_img = cv2.multiply(aligned, np.array([DARK_FACTOR] * 3, dtype=np.float64))
+        result_img = np.clip(result_img, 0, 255).astype(np.uint8)
+
+        for a in added_objects:
+            _composite_object(result_img, aligned, a["mask"], a["bbox"], COLOR_ADDED)
+        for r in removed_objects:
+            _composite_object(result_img, r["source"], r["mask"], r["bbox"], COLOR_REMOVED)
+        for m in moved_objects:
+            _composite_object(result_img, m["from"]["source"], m["from"]["mask"], m["from"]["bbox"], COLOR_MOVED)
+            _composite_object(result_img, aligned,              m["to"]["mask"],   m["to"]["bbox"],   COLOR_MOVED)
 
     # Draw severity label in top-left
     label = f"{severity}  +{n_added}  -{n_removed}  ~{n_moved}"
@@ -647,8 +658,8 @@ def _compare_images_ai_impl(img1_path, img2_path, output_path, heatmap_path):
 
 if HAS_SPACES:
     @spaces.GPU
-    def compare_images_ai(img1_path, img2_path, output_path, heatmap_path):
-        return _compare_images_ai_impl(img1_path, img2_path, output_path, heatmap_path)
+    def compare_images_ai(img1_path, img2_path, output_path, heatmap_path, draw_style="overlay"):
+        return _compare_images_ai_impl(img1_path, img2_path, output_path, heatmap_path, draw_style=draw_style)
 else:
-    def compare_images_ai(img1_path, img2_path, output_path, heatmap_path):
-        return _compare_images_ai_impl(img1_path, img2_path, output_path, heatmap_path)
+    def compare_images_ai(img1_path, img2_path, output_path, heatmap_path, draw_style="overlay"):
+        return _compare_images_ai_impl(img1_path, img2_path, output_path, heatmap_path, draw_style=draw_style)