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

st.title("Convert Horse To Zebra")
st.write("Welcome to the Horse to Zebra Convertor! Choose a horse image and see it displayed below.")

os.makedirs("input", exist_ok=True)
os.makedirs("saved_images", 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"):
        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.sidebar.success("Image successfully uploaded!")
    main()


col1, col2 = st.columns(2)

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