import streamlit as st
from PIL import Image
from transformers import pipeline
import numpy as np
import cv2
import matplotlib.cm as cm
import base64
from io import BytesIO

st.set_page_config(layout="wide")

with open("styles.css") as f:
    st.markdown('<style>{}</style>'.format(f.read()), unsafe_allow_html=True)

st.markdown("<h1 class='title'>Segformer Semantic Segmentation</h1>", unsafe_allow_html=True)
st.markdown("""
<div class='text-center'>
This app uses the Segformer deep learning model to perform semantic segmentation on <b style='color: red; font-weight: 40px;'>road images</b>. The Transformer-based model is 
trained on the CityScapes dataset which contains images of urban road scenes. Upload a 
road scene and the app will return the image with semantic segmentation applied.
</div>
""", unsafe_allow_html=True)

group_members = ["Ang Ngo Ching, Josh Darren W.", "Bautista, Ryan Matthew M.", "Lacuesta, Angelo Giuseppe M.", "Reyes, Kenwin Hans", "Ting, Sidney Mitchell O."]

st.markdown("""
            <h3 class='text-center' style='margin-top: 0.5rem;'>
ℹ️ You can get sample images of road scenes in this <a href='https://drive.google.com/drive/folders/1202EMeXAHnN18NuhJKWWme34vg0V-svY?fbclid=IwAR3kyjGS895nOBKi9aGT_P4gLX9jvSNrV5b5y3GH49t2Pvg2sZSRA58LLxs' target='_blank'>link</a>.
</h3>""", unsafe_allow_html=True)

st.markdown("""
            <h3 class='text-center' style='margin-top: 0.5rem;'>
📜 Read more about the paper <a href='https://arxiv.org/pdf/2105.15203.pdf' target='_blank'>here</a>.
</h3>""", unsafe_allow_html=True)

label_colors = {}

def draw_masks_fromDict(image, results):
    masked_image = image.copy()

    colormap = cm.get_cmap('nipy_spectral')
    
    for i, result in enumerate(results):
        mask = np.array(result['mask'])  
        mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)  
        
        color = colormap(i / len(results))[:3] 
        color = tuple(int(c * 255) for c in color)  
        
        masked_image = np.where(mask, color, masked_image)

        label_colors[color] = result['label']

    masked_image = masked_image.astype(np.uint8)
    return cv2.addWeighted(image, 0.3, masked_image, 0.7, 0)

uploaded_file = st.file_uploader("", type=["jpg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    
    col1, col2 = st.columns(2)

    with col1:
        st.image(image, caption='Uploaded Image.', use_column_width=True)

    with st.spinner('Processing...'):  
        semantic_segmentation = pipeline("image-segmentation", f"nvidia/segformer-b1-finetuned-cityscapes-1024-1024")
        segmentation_results = semantic_segmentation(image)
        image_with_masks = draw_masks_fromDict(np.array(image)[:, :, :3], segmentation_results)
        image_with_masks_pil = Image.fromarray(image_with_masks, 'RGB')

    with col2:
        st.image(image_with_masks_pil, caption='Segmented Image.', use_column_width=True)

        html_segment = "<div class='container'><h3>Labels:</h3>"
        
        for color, label in label_colors.items():
            html_segment += f"<div style='display: flex; align-items: center; margin-bottom: 0.5rem;'><span style='display: inline-block; width: 20px; height: 20px; background-color: rgb{color}; margin-right: 1rem; border-radius: 10px;'></span><p style='margin: 0;'>{label}</p></div>"
        
        buffered = BytesIO()
        image_with_masks_pil.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        html_segment += f'<a href="data:file/png;base64,{img_str}" download="segmented_{uploaded_file.name}">Download Segmented Image</a>'
        st.markdown(html_segment + "</div>", unsafe_allow_html=True)



html_members = "<hr><div style='display: flex; justify-content: center;'><h3>Group 6 - Members:</h3><ul>"
for member in group_members:
    html_members += "<li>" + member + "</li>"

st.markdown(html_members + "</ul></div>", unsafe_allow_html=True)