Update saliency_gradio.py

saliency_gradio.py

@@ -4,6 +4,7 @@ import tensorflow as tf
 from huggingface_hub import snapshot_download, from_pretrained_keras
 import gradio as gr
 
+# Load the model
 model = from_pretrained_keras("alexanderkroner/MSI-Net")
 hf_dir = snapshot_download(repo_id="alexanderkroner/MSI-Net")
 
@@ -17,13 +18,11 @@ def get_target_shape(original_shape):
     best_mode = min(square_mode, landscape_mode, portrait_mode)
 
     if best_mode == square_mode:
-        target_shape = (320, 320)
+        return (320, 320)
     elif best_mode == landscape_mode:
-        target_shape = (240, 320)
+        return (240, 320)
     else:
-        target_shape = (320, 240)
-
-    return target_shape
+        return (320, 240)
 
 def preprocess_input(input_image, target_shape):
     input_tensor = tf.expand_dims(input_image, axis=0)
@@ -37,7 +36,6 @@ def preprocess_input(input_image, target_shape):
 
     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
 
@@ -66,39 +64,43 @@ def postprocess_output(output_tensor, vertical_padding, horizontal_padding, orig
     ]
 
     output_tensor = tf.image.resize(output_tensor, original_shape)
+    return output_tensor.numpy().squeeze()  # Return grayscale map
 
-    output_array = output_tensor.numpy().squeeze()
-    return output_array  # Keep as grayscale
-
-def predict_saliency(image):
-    input_image = np.array(image, dtype=np.float32)
+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)
-
-    # Compute the sum of grayscale values
     total_saliency = np.sum(saliency_gray)
 
-    # Convert to colormap for visualization
-    saliency_map_rgb = plt.cm.inferno(saliency_gray)[..., :3]
-
-    # Blend with original image
+    saliency_rgb = plt.cm.inferno(saliency_gray)[..., :3]
     alpha = 0.9
-    blended_image = alpha * saliency_map_rgb + (1 - alpha) * input_image / 255
+    blended_image = alpha * saliency_rgb + (1 - alpha) * input_image / 255
 
     return blended_image, f"Total grayscale saliency: {total_saliency:.2f}"
 
+def predict_two_images(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
+
 iface = gr.Interface(
-    fn=predict_saliency,
-    inputs=gr.Image(type="pil"),
+    fn=predict_two_images,
+    inputs=[
+        gr.Image(type="pil", label="Input Image 1"),
+        gr.Image(type="pil", label="Input Image 2"),
+    ],
     outputs=[
-        gr.Image(type="numpy", label="Saliency Map"),
-        gr.Textbox(label="Grayscale Pixel Sum")
+        gr.Image(type="numpy", label="Saliency Map 1"),
+        gr.Textbox(label="Grayscale Sum 1"),
+        gr.Image(type="numpy", label="Saliency Map 2"),
+        gr.Textbox(label="Grayscale Sum 2"),
     ],
-    title="MSI-Net Saliency Map",
-    description="Upload an image to generate its saliency map and view the total intensity.",
+    title="MSI-Net Saliency Maps for Two Images",
+    description="Upload two images to compare their saliency maps and total saliency values.",
 )
 
 iface.launch(share=True)