import streamlit as st
from PIL import Image

import style
st.title(":rainbow[Stylize app is so kewl]:sun_with_face::sunglasses:")
#st.title(":rainbow[using _Streamlit_ is so cool] :sunglasses:")
col1, col2 = st.columns(2,gap="large")
img = st.sidebar.selectbox(
    'Select image',
    ('amber.jpg','cat.jpg','owl.jpg','dipika.jpg','mayanti.jpg','model1.jpg','model2.jpg','taapse.jpg','tamanna.jpg')
)

style_name = st.sidebar.selectbox(
    'Select style',
    ('candy','mosaic','rain_princess','udnie')
)

model = "./saved_models/" +style_name+ ".pth"
#print(model)
input_image = "./images/content-images/" + img
output_image = "./images/output-images/" + style_name + "-" + img
with col1:
    st.write("### Source Image:")
    image = Image.open(input_image)
    st.image(image,width=300)

clicked = st.button("Stylize",type="primary")

if clicked:
    with col2:
        model = style.load_model(model)
        style.stylize(model,input_image,output_image)

        st.write('### Output Image:')
        image=Image.open(output_image)
        st.image(image,width=300)