Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from streamlit_back_camera_input import back_camera_input | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import tensorflow as tf | |
| from PIL import Image | |
| loaded_model = tf.saved_model.load("model/") | |
| loaded_model = loaded_model.signatures["serving_default"] | |
| 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: | |
| target_shape = (320, 320) | |
| elif best_mode == landscape_mode: | |
| target_shape = (240, 320) | |
| else: | |
| target_shape = (320, 240) | |
| return target_shape | |
| 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) | |
| output_array = output_tensor.numpy().squeeze() | |
| output_array = plt.cm.inferno(output_array)[..., :3] | |
| return output_array | |
| def compute_saliency(input_image, alpha=0.65): | |
| if input_image is not None: | |
| 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 | |
| ) | |
| saliency_map = loaded_model(input_tensor)["output"] | |
| saliency_map = postprocess_output( | |
| saliency_map, vertical_padding, horizontal_padding, original_shape | |
| ) | |
| blended_image = alpha * saliency_map + (1 - alpha) * input_image / 255 | |
| return blended_image | |
| def main(): | |
| st.title("Visual Saliency Estimator") | |
| # Let the user upload an image | |
| uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) | |
| if uploaded_file is not None: | |
| # Open the uploaded image and display it | |
| image = Image.open(uploaded_file) | |
| st.image(image, caption="Uploaded Image", use_column_width=True) | |
| if st.button("Compute"): | |
| # Convert the image to a numpy array for processing | |
| img_array = np.array(image) | |
| # Compute the saliency map (replace this with actual computation) | |
| saliency_map = compute_saliency(img_array) | |
| # Display the saliency map | |
| st.image(saliency_map, caption="Saliency Map", use_column_width=True) | |
| if __name__ == "__main__": | |
| main() | |