import streamlit as st
# import fastai
from pathlib import Path
import pathlib
temp = pathlib.PosixPath
# pathlib.PosixPath = pathlib.WindowsPath
from fastai.learner import load_learner
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt

path = Path('/content/')
def label_func(x): 
    num = x.stem.split('_')[1]
    return path/'labels'/f'label_{num}.png'

st.title('Flood Segmentation App')

uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg"])
print(uploaded_file)
model = load_learner('fastai_tfl_253img.pkl', cpu = True)
if uploaded_file is not None:
    img = Image.open(uploaded_file)
    x = np.asarray(img)
    tm, tb, tbp = model.predict(x)   
    
    # img = img.resize((300,300))
    # print(img)
    # img = image.load_img(im, target_size=(300,300))
    # x = image.img_to_array(img)
    st.image(img, caption='Uploaded Image.', use_column_width=True)
    st.write("")
    st.write("Generating Segmentation Mask...")
    # x = np.expand_dims(x, axis=0)
    # prob = model.predict(x)
    # print(tm.numpy())
    fig, ax = plt.subplots()
    plt.axis('off')
    ax.imshow(tm)
    # mask = Image.fromarray(np.uint8(tm))
    # st.image(mask, caption='Segmentation Mask', use_column_width=True)
    st.pyplot(fig)