import streamlit as st
import os
from PIL import Image
from train import main

st.title("ComicInator")
st.write("Welcome to the ComicInator! Choose an image and see it displayed below.")

os.makedirs("input", exist_ok=True)
os.makedirs("evaluation", exist_ok=True)

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

if uploaded_file is not None:
    for filename in os.listdir("input/face"):
        file_path = os.path.join("input/face", filename)
        os.remove(file_path)
    with open(os.path.join("input/face", uploaded_file.name), "wb") as f:
        f.write(uploaded_file.getvalue())
    st.sidebar.success("Image successfully uploaded!")
    main()


col1, col2 = st.columns(2)

if uploaded_file is not None:
    col1.image(Image.open(os.path.join("evaluation", "input_0.png")), caption="Uploaded Image", use_column_width=True)
    col2.image(Image.open(os.path.join("evaluation", "y_gen_0.png")), caption="Colorized Image", use_column_width=True)