Update saliency_gradio.py

saliency_gradio.py

@@ -1,111 +1,103 @@
-import matplotlib.pyplot as plt
+import gradio as gr
+from gradio_image_annotation import ImageAnnotation
 import numpy as np
+import matplotlib.pyplot as plt
 import tensorflow as tf
-from huggingface_hub import snapshot_download, from_pretrained_keras
-import gradio as gr
+from huggingface_hub import from_pretrained_keras, snapshot_download
 
-# Load the model
+# Load model
 model = from_pretrained_keras("alexanderkroner/MSI-Net")
 hf_dir = snapshot_download(repo_id="alexanderkroner/MSI-Net")
 
 def get_target_shape(original_shape):
-    original_aspect_ratio = original_shape[0] / original_shape[1]
-    square_mode = abs(original_aspect_ratio - 1.0)
-    landscape_mode = abs(original_aspect_ratio - 240 / 320)
-    portrait_mode = abs(original_aspect_ratio - 320 / 240)
-    best_mode = min(square_mode, landscape_mode, portrait_mode)
-    if best_mode == square_mode:
+    aspect = original_shape[0] / original_shape[1]
+    if abs(aspect - 1.0) < min(abs(aspect - 240/320), abs(aspect - 320/240)):
         return (320, 320)
-    elif best_mode == landscape_mode:
+    elif abs(aspect - 240/320) < abs(aspect - 320/240):
         return (240, 320)
-    else:
-        return (320, 240)
+    return (320, 240)
 
 def preprocess_input(input_image, target_shape):
     input_tensor = tf.expand_dims(input_image, axis=0)
     input_tensor = tf.image.resize(input_tensor, target_shape, preserve_aspect_ratio=True)
     vertical_padding = target_shape[0] - input_tensor.shape[1]
     horizontal_padding = target_shape[1] - input_tensor.shape[2]
-    vertical_padding_1 = vertical_padding // 2
-    vertical_padding_2 = vertical_padding - vertical_padding_1
-    horizontal_padding_1 = horizontal_padding // 2
-    horizontal_padding_2 = horizontal_padding - horizontal_padding_1
-    input_tensor = tf.pad(
-        input_tensor,
-        [
-            [0, 0],
-            [vertical_padding_1, vertical_padding_2],
-            [horizontal_padding_1, horizontal_padding_2],
-            [0, 0],
-        ],
-    )
-    return input_tensor, [vertical_padding_1, vertical_padding_2], [horizontal_padding_1, horizontal_padding_2]
+    vpad1, vpad2 = vertical_padding // 2, vertical_padding - vertical_padding // 2
+    hpad1, hpad2 = horizontal_padding // 2, horizontal_padding - horizontal_padding // 2
+    input_tensor = tf.pad(input_tensor, [[0,0], [vpad1, vpad2], [hpad1, hpad2], [0,0]])
+    return input_tensor, [vpad1, vpad2], [hpad1, hpad2]
 
-def postprocess_output(output_tensor, vertical_padding, horizontal_padding, original_shape):
+def postprocess_output(output_tensor, vpad, hpad, original_shape):
     output_tensor = output_tensor[
-        :,
-        vertical_padding[0] : output_tensor.shape[1] - vertical_padding[1],
-        horizontal_padding[0] : output_tensor.shape[2] - horizontal_padding[1],
-        :,
+        :, vpad[0]:output_tensor.shape[1]-vpad[1],
+        hpad[0]:output_tensor.shape[2]-hpad[1], :
     ]
     output_tensor = tf.image.resize(output_tensor, original_shape)
     return output_tensor.numpy().squeeze()
 
-def process_image_with_bbox(input_image, bbox, threshold=0.0):
-    input_image_np = np.array(input_image, dtype=np.float32)
+def process_image(image):
+    input_image_np = np.array(image, dtype=np.float32)
     original_shape = input_image_np.shape[:2]
     target_shape = get_target_shape(original_shape)
+    input_tensor, vpad, hpad = preprocess_input(input_image_np, target_shape)
+    output_tensor = model(input_tensor)["output"]
+    saliency_gray = postprocess_output(output_tensor, vpad, hpad, original_shape)
+    total_saliency = np.sum(saliency_gray)
+    saliency_rgb = plt.cm.inferno(saliency_gray)[..., :3]
+    alpha = 0.9
+    blended = alpha * saliency_rgb + (1 - alpha) * input_image_np / 255
+    return blended, f"Total grayscale saliency: {total_saliency:.2f}"
 
-    input_tensor, vertical_padding, horizontal_padding = preprocess_input(input_image_np, target_shape)
+def process_image_with_bbox(image, annotations):
+    input_image_np = np.array(image, dtype=np.float32)
+    original_shape = input_image_np.shape[:2]
+    target_shape = get_target_shape(original_shape)
+    input_tensor, vpad, hpad = 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)
+    saliency_gray = postprocess_output(output_tensor, vpad, hpad, original_shape)
 
-    # Total saliency
     total_saliency = np.sum(saliency_gray)
-
-    # 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 = 0.0
+    if annotations and isinstance(annotations, list) and "coordinates" in annotations[0]:
+        x, y, w, h = map(int, annotations[0]["coordinates"])
+        saliency_crop = saliency_gray[y:y+h, x:x+w]
         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
+    ratio = bbox_sum / total_saliency if total_saliency > 0 else 0
     saliency_rgb = plt.cm.inferno(saliency_gray)[..., :3]
     alpha = 0.9
-    blended_image = alpha * saliency_rgb + (1 - alpha) * input_image_np / 255
+    blended = alpha * saliency_rgb + (1 - alpha) * input_image_np / 255
 
-    summary = f"{bbox_ratio:.4f}"
-
-    return blended_image, summary
+    summary = (
+        f"Total saliency: {total_saliency:.2f}\n"
+        f"Bounding box saliency: {bbox_sum:.2f}\n"
+        f"Ratio: {ratio:.4f}"
+    )
 
-def predict_single(image):
-    return process_image(image)
+    return blended, summary
 
 def predict_dual(image1, image2):
     result1_img, result1_val = process_image(image1)
     result2_img, result2_val = process_image(image2)
     return result1_img, result1_val, result2_img, result2_val
 
+# Interface
 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 and draw a bounding box to measure saliency inside it.")
             with gr.Row():
-                input_image_single = gr.Image(type="pil", tool="select", label="Input Image with Bounding Box")
+                image_input = gr.Image(type="pil", label="Upload Image")
+                annotator = ImageAnnotation()
             with gr.Row():
-                output_image_single = gr.Image(type="numpy", label="Saliency Map")
-                output_text_single = gr.Textbox(label="Saliency Stats")
-            submit_single = gr.Button("Generate Saliency")
-            submit_single.click(
-                fn=process_image_with_bbox,
-                inputs=[input_image_single, input_image_single.select_region],
-                outputs=[output_image_single, output_text_single],
-            )
+                image_output = gr.Image(type="numpy", label="Saliency Map")
+                stats_output = gr.Textbox(label="Saliency Stats")
+            button = gr.Button("Generate Saliency")
+            button.click(fn=process_image_with_bbox, inputs=[image_input, annotator], outputs=[image_output, stats_output])
+            image_input.change(fn=lambda img: img, inputs=image_input, outputs=annotator)
+
         with gr.Tab("Compare Two Images"):
             gr.Markdown("### Upload two images to compare their saliency maps and grayscale saliency values.")
             with gr.Row():