import streamlit as st
from super_image import EdsrModel, ImageLoader
from PIL import Image
import cv2

st.title('Take your picture to the next level')


col1, col2 = st.columns([3, 1])
preds = None

with col2:
    choice = st.selectbox('Scale to', [2, 3, 4])
    model_choice = st.selectbox('Model', ['EDSR', 'PAN', 'CARN'])

with col1:
    file = st.file_uploader('Post your picture', type=['png', 'jpeg'])
    if file is not None:
        image = Image.open(file)

        col1, col2 = st.columns(2)
        with col1:
            st.subheader('Input image')
            st.image(image)

        if model_choice == "EDSR":

            model = EdsrModel.from_pretrained("eugenesiow/edsr-base", scale=int(choice))
            inputs = ImageLoader.load_image(image)
            preds = model(inputs)

            ImageLoader.save_image(preds, './scaled_2x.png')
            img = Image.open('./scaled_2x.png')
            with col2:
                st.subheader('New image')
                st.image(img, caption=f'Upscaled {choice} times image')
    
                with open('./scaled_2x.png', 'rb') as f:
                    st.download_button(label="Download image", data=f, file_name="upscaled.png", mime="image/png")
            
        else:
            st.text('Not supported yet. Comeback soon for fun releases')