import streamlit as st
from io import BytesIO
from PIL import Image
from utils import convert_to_bw, load_colorization_model, colorize_bw_image
import os
os.environ["STREAMLIT_SERVER_HEADLESS"] = "true"

st.set_page_config(page_title="Image Colorizer", layout="centered")

st.title("Convert B&W images to Colored and vice versa")

uploaded_file = st.sidebar.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])



if uploaded_file:
    # Open and convert the uploaded image to RGB format
    image = Image.open(uploaded_file).convert("RGB")

    option = st.sidebar.selectbox("Choose an action", ("Convert to Black & White", "Colorize Your B&W images"))

    if st.sidebar.button("Process"):
        #Convert the uploaded image to black and white
        if option == "Convert to Black & White":
            result_img = convert_to_bw(image)
        #Colorize a black and white image using a pre-trained model
        elif option == "Colorize Your B&W images":
            with st.spinner("Colorizing..."):
                net = load_colorization_model()
                result_img = colorize_bw_image(image, net)
            

        # Display both images in columns
        col1, col2 = st.columns(2)
        with col1:
            st.image(image, caption="Original Image")
        with col2:
            st.image(result_img, caption="Processed Image")
            
            
        # Create a buffer to store image bytes
        buffer = BytesIO()
        result_img.save(buffer, format="JPEG")  
        buffer.seek(0)  # Reset cursor to the beginning
        
        #Download Image in Jpeg
        st.download_button(
            label="Download Output Image",
            data=buffer ,#result_img.tobytes()
            file_name="output.jpeg",
            mime="image/jpeg"
        )