import streamlit as st
import cv2
from saliency_map import avg_saliency 
import numpy as np
from PIL import Image
import io

st.set_page_config(
    page_title="Thumbnail Analyser",
    page_icon="thumbnail_logo.png",
    layout="wide",  # wide layout reduces left and right margins
    initial_sidebar_state="expanded"
)

def saliency_analysis(image,saliency_image,col2):
    with saliency_image:
        with st.spinner('Loading...'):
            img, saliency_output = avg_saliency(image)
            if saliency_output is not None:
                pil_image = Image.fromarray(saliency_output)
                if pil_image.mode != 'RGB':
                    pil_image = pil_image.convert('L')
    saliency_image.image(pil_image, caption='Attention Spots', use_column_width=True)
    return pil_image
       

def main():
    st.markdown('<h1 style="text-align: center;">Thumbnail Analyser</h1>', unsafe_allow_html=True)


    # Create two image containers
    col1, col2 = st.columns(2)

    # Placeholder images
    image_placeholder = 'Upload_image.png'
    original_image = col1.image(image_placeholder, use_column_width=True)
    saliency_image = col2.image(image_placeholder, use_column_width=True)

    # Upload images and replace placeholders
    original_image_upload = col1.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
    

    if original_image_upload is not None:
        original_image.image(original_image_upload, use_column_width=True)
        file_bytes = np.asarray(bytearray(original_image_upload.read()), dtype=np.uint8)
        cv2_image = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
        pil_image = saliency_analysis(cv2_image,saliency_image,col2)
        buffered = io.BytesIO()
        pil_image.save(buffered, format="PNG")
        buffered.seek(0)
    
        col2.download_button(
                label="Download Image",
                data=buffered,
                file_name="AttentionSpots.png",
                mime="image/png",
                key="download_button"
            )

if __name__ == "__main__":
    main()