import streamlit as st
import numpy as np
from colortransfer import transfer_color, convert_bytes_to_pil
import io
from PIL import Image

with open("references.md") as file:
    refrences = file.read()

with st.container():
    st.write(refrences)

with st.container():
    col1, col2 = st.columns(2)

    with col1:
        target_img = st.file_uploader("Choose Target Image", type=["png", "jpg"])
        if target_img is not None:
            st.image(target_img, width=256)

    with col2:
        style_img = st.file_uploader("Choose Style Image", type=["png", "jpg"])
        if style_img is not None:
            st.image(style_img, width=256)


with st.container():
    if st.button("Transfer!"):
        if target_img is not None and style_img is not None:
            result = transfer_color(
                np.asarray(convert_bytes_to_pil(style_img.read())),
                np.asarray(convert_bytes_to_pil(target_img.read())),
            )
            st.image(result, width=700)
        else:
            st.error("You need to specify the target and style images first")
    else:
        st.write("Upload images and press Transfer button")