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()