Update saliency_gradio.py

saliency_gradio.py

@@ -1,19 +1,12 @@
-# prompt: save the gradio app into  a file
-
 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
 
-# Note: 'keras<3.x' or 'tf_keras' must be installed (legacy)
-# See https://github.com/keras-team/tf-keras for more details.
-
 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]
 
@@ -32,7 +25,6 @@ def get_target_shape(original_shape):
 
     return target_shape
 
-
 def preprocess_input(input_image, target_shape):
     input_tensor = tf.expand_dims(input_image, axis=0)
 
@@ -65,10 +57,7 @@ def preprocess_input(input_image, target_shape):
         [horizontal_padding_1, horizontal_padding_2],
     )
 
-
-def postprocess_output(
-    output_tensor, vertical_padding, horizontal_padding, original_shape
-):
+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],
@@ -79,10 +68,7 @@ def postprocess_output(
     output_tensor = tf.image.resize(output_tensor, original_shape)
 
     output_array = output_tensor.numpy().squeeze()
-    output_array = plt.cm.inferno(output_array)[..., :3]
-
-    return output_array
-
+    return output_array  # Keep as grayscale
 
 def predict_saliency(image):
     input_image = np.array(image, dtype=np.float32)
@@ -90,18 +76,29 @@ def predict_saliency(image):
     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_map = postprocess_output(output_tensor, vertical_padding, horizontal_padding, original_shape)
+    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
     alpha = 0.9
-    blended_image = alpha * saliency_map + (1 - alpha) * input_image / 255
-    return blended_image
+    blended_image = alpha * saliency_map_rgb + (1 - alpha) * input_image / 255
 
+    return blended_image, f"Total grayscale saliency: {total_saliency:.2f}"
 
 iface = gr.Interface(
     fn=predict_saliency,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="numpy"),
+    outputs=[
+        gr.Image(type="numpy", label="Saliency Map"),
+        gr.Textbox(label="Grayscale Pixel Sum")
+    ],
     title="MSI-Net Saliency Map",
-    description="Upload an image to generate its saliency map.",
+    description="Upload an image to generate its saliency map and view the total intensity.",
 )
 
-iface.launch(share=True)
+iface.launch(share=True)