Update saliency_gradio.py

saliency_gradio.py

@@ -1,9 +1,101 @@
+import matplotlib.pyplot as plt
+import numpy as np
+import tensorflow as tf
+from huggingface_hub import snapshot_download, from_pretrained_keras
 import gradio as gr
-from gradio_image_prompter import ImagePrompter
-
-demo = gr.Interface(
-    lambda prompts: (prompts["image"], prompts["points"]),
-    ImagePrompter(show_label=False),
-    [gr.Image(show_label=False), gr.Dataframe(label="Points")],
-)
-demo.launch()
+
+# Load the 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:
+        return (320, 320)
+    elif best_mode == landscape_mode:
+        return (240, 320)
+    else:
+        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]
+
+def postprocess_output(output_tensor, vertical_padding, horizontal_padding, 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],
+        :,
+    ]
+    output_tensor = tf.image.resize(output_tensor, original_shape)
+    return output_tensor.numpy().squeeze()
+
+def process_image(input_image):
+    input_image = np.array(input_image, dtype=np.float32)
+    original_shape = input_image.shape[:2]
+    target_shape = get_target_shape(original_shape)
+    input_tensor, vertical_padding, horizontal_padding = preprocess_input(input_image, target_shape)
+    output_tensor = model(input_tensor)["output"]
+    saliency_gray = postprocess_output(output_tensor, vertical_padding, horizontal_padding, original_shape)
+    total_saliency = np.sum(saliency_gray)
+    saliency_rgb = plt.cm.inferno(saliency_gray)[..., :3]
+    alpha = 0.9
+    blended_image = alpha * saliency_rgb + (1 - alpha) * input_image / 255
+    return blended_image, f"Total grayscale saliency: {total_saliency:.2f}"
+
+def predict_single(image):
+    return process_image(image)
+
+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
+
+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.")
+            with gr.Row():
+                input_image_single = gr.Image(type="pil", label="Input Image")
+            with gr.Row():
+                output_image_single = gr.Image(type="numpy", label="Saliency Map")
+                output_text_single = gr.Textbox(label="Grayscale Sum")
+            submit_single = gr.Button("Generate Saliency")
+            submit_single.click(fn=predict_single, inputs=input_image_single, 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():
+                input_image1 = gr.Image(type="pil", label="Image 1")
+                input_image2 = gr.Image(type="pil", label="Image 2")
+            with gr.Row():
+                output_image1 = gr.Image(type="numpy", label="Saliency Map 1")
+                output_text1 = gr.Textbox(label="Grayscale Sum 1")
+                output_image2 = gr.Image(type="numpy", label="Saliency Map 2")
+                output_text2 = gr.Textbox(label="Grayscale Sum 2")
+            submit_dual = gr.Button("Compare Saliency")
+            submit_dual.click(fn=predict_dual, inputs=[input_image1, input_image2], outputs=[output_image1, output_text1, output_image2, output_text2])
+
+demo.launch(share=True)