saliency_gradio.py

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

# 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()  # Return grayscale map

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_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_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 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 Maps for Two Images",
    description="Upload two images to compare their saliency maps and total saliency values.",
)

iface.launch(share=True)