app.py

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

st.title("Anime Drawing Colorizer")
st.write("Welcome to the Anime Colorizer App! Choose an image and see it displayed below.")

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

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

if uploaded_file is not None:
    for filename in os.listdir("input"):
        file_path = os.path.join("input", filename)
        os.remove(file_path)
    with open(os.path.join("input", uploaded_file.name), "wb") as f:
        f.write(uploaded_file.getvalue())
    st.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)