Update saliency_gradio.py

saliency_gradio.py

@@ -51,26 +51,36 @@ def postprocess_output(output_tensor, vertical_padding, horizontal_padding, orig
     output_tensor = tf.image.resize(output_tensor, original_shape)
     return output_tensor.numpy().squeeze()
 
-def process_image(input_image, threshold=0.5):
-    input_image = np.array(input_image, dtype=np.float32)
-    original_shape = input_image.shape[:2]
+def process_image_with_bbox(input_image, bbox, threshold=0.0):
+    input_image_np = np.array(input_image, dtype=np.float32)
+    original_shape = input_image_np.shape[:2]
     target_shape = get_target_shape(original_shape)
 
-    input_tensor, vertical_padding, horizontal_padding = preprocess_input(input_image, target_shape)
+    input_tensor, vertical_padding, horizontal_padding = preprocess_input(input_image_np, target_shape)
     output_tensor = model(input_tensor)["output"]
     saliency_gray = postprocess_output(output_tensor, vertical_padding, horizontal_padding, original_shape)
 
+    # Total saliency
     total_saliency = np.sum(saliency_gray)
-    above_threshold_saliency = np.sum(saliency_gray[saliency_gray > threshold])
-    ratio = above_threshold_saliency / total_saliency if total_saliency > 0 else 0.0
 
+    # Bounding box: bbox = [x_min, y_min, x_max, y_max] in image pixel coordinates
+    if bbox is not None:
+        x_min, y_min, x_max, y_max = map(int, bbox)
+        saliency_crop = saliency_gray[y_min:y_max, x_min:x_max]
+        bbox_sum = np.sum(saliency_crop)
+        bbox_ratio = bbox_sum / total_saliency if total_saliency > 0 else 0.0
+    else:
+        bbox_sum = 0
+        bbox_ratio = 0.0
+
+    # Heatmap overlay
     saliency_rgb = plt.cm.inferno(saliency_gray)[..., :3]
     alpha = 0.9
-    blended_image = alpha * saliency_rgb + (1 - alpha) * input_image / 255
+    blended_image = alpha * saliency_rgb + (1 - alpha) * input_image_np / 255
 
-    summary_text = f"{ratio:.4f}"
+    summary = f"{bbox_ratio:.4f}"
 
-    return blended_image, summary_text
+    return blended_image, summary
 
 def predict_single(image):
     return process_image(image)
@@ -84,15 +94,18 @@ with gr.Blocks(title="MSI-Net Saliency App") as demo:
     gr.Markdown("## MSI-Net Saliency Map Viewer")
     with gr.Tabs():
         with gr.Tab("Single Image"):
-            gr.Markdown("### Upload an image to see its saliency map and total grayscale saliency value.")
+            gr.Markdown("### Upload an image and draw a bounding box to measure saliency inside it.")
             with gr.Row():
-                input_image_single = gr.Image(type="pil", label="Input Image")
+                input_image_single = gr.Image(type="pil", tool="select", label="Input Image with Bounding Box")
             with gr.Row():
                 output_image_single = gr.Image(type="numpy", label="Saliency Map")
-                output_text_single = gr.Textbox(label="Grayscale Sum")
+                output_text_single = gr.Textbox(label="Saliency Stats")
             submit_single = gr.Button("Generate Saliency")
-            submit_single.click(fn=predict_single, inputs=input_image_single, outputs=[output_image_single, output_text_single])
-
+            submit_single.click(
+                fn=process_image_with_bbox,
+                inputs=[input_image_single, input_image_single.select_region],
+                outputs=[output_image_single, output_text_single],
+            )
         with gr.Tab("Compare Two Images"):
             gr.Markdown("### Upload two images to compare their saliency maps and grayscale saliency values.")
             with gr.Row():