Spaces:
Sleeping
Sleeping
File size: 3,693 Bytes
4a76860 f32fc1c 4a76860 f32fc1c 4a76860 f32fc1c 4a76860 f32fc1c 4a76860 7f649e2 4a76860 f32fc1c 4a76860 f32fc1c 4a76860 f32fc1c 4a76860 7f649e2 f32fc1c d102922 f32fc1c 4a76860 7f649e2 4a76860 f32fc1c 4a76860 f32fc1c 7f649e2 f32fc1c 7f649e2 f32fc1c 4a76860 7f649e2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
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)
|