import streamlit as st
from PIL import Image
import torch

from model import ModelColorization

from utils import process_gs_image, inverse_transform_cs

# create the model
model = ModelColorization().from_pretrained("sebastiansarasti/AutoEncoderImageColorization")

# create the streamlit app
st.title("Image Colorization App")
st.write("This is an app to colorize black and white images.")

# create a botton to upload the image
uploaded_file = st.file_uploader("Choose an image...", type="jpg")

# check if the image is uploaded
if uploaded_file is not None:
    # display the image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image.", use_container_width=True)

    # create a button to colorize the image
    if st.button("Colorize"):
        # process the grayscale image
        image, original_size = process_gs_image(image)
        # run the model
        model.eval()
        with torch.no_grad():
            result = model(image)
        # colorize the image
        colorized_image = inverse_transform_cs(result.squeeze(0), original_size)
        # display the colorized image
        st.image(colorized_image, caption="Colorized Image.", use_container_width=True)