import streamlit as st
from neural_style_transfer import generate_img

def display_img(org_img, style_img):
    img = generate_img(org_img, style_img)
    st.image(img, caption='Result after Neural Style Transfer', use_column_width=True)

def main():
    markdown = ""
    with open('README.md', 'r') as f:
        markdown = f.read()
    st.markdown(markdown, unsafe_allow_html=True)    
    with st.container(border=True):
        col1, col2 = st.columns(2)
        with col1:
            st.header('Original Image')
            org_img = st.file_uploader(label='Upload original image', type=['png', 'jpg', 'jpeg'])
        with col2:
            st.header('Style Image')
            style_img = st.file_uploader(label='Upload style image', type=['png', 'jpg', 'jpeg'])

        if org_img is not None and style_img is not None:
            if st.button("Submit"):
                col1, col2, col3 = st.columns([1,2.5,1])
                with col2:
                    display_img(org_img.read(), style_img.read())


if __name__ == '__main__':
    main()