from distutils.command.upload import upload
import streamlit as st
from io import StringIO
from PIL import Image
import pandas as pd

import torch
import torch.nn as nn
import torchvision.models as models

import albumentations as A
from albumentations.pytorch.transforms import ToTensorV2

import numpy as np

import cv2

st.title('Dummy')
uploaded_file = st.file_uploader('Select File')

id2class = {0: 'agricultural', 1: 'airplane', 2: 'baseballdiamond', 3: 'beach', 4: 'buildings', 5: 'chaparral', 6: 'denseresidential', 7: 'forest', 8: 'freeway', 9: 'golfcourse', 10: 'intersection', 11: 'mediumresidential', 12: 'mobilehomepark', 13: 'overpass', 14: 'parkinglot', 15: 'river', 16: 'runway', 17: 'sparseresidential', 18: 'storagetanks', 19: 'tenniscourt', 20: 'harbor'}

model = models.resnet50(weights=None)
model.fc = nn.Linear(2048, 21)
model.load_state_dict(torch.load('resnet_best.pth', map_location=torch.device('cpu')), strict=True)
model.eval()

if uploaded_file is not None:
    if '.jpg' in uploaded_file.name.lower() or '.png' in uploaded_file.name.lower():
        st.write(uploaded_file.name)
        img = Image.open(uploaded_file)
        st.image(img)
        img = np.array(img)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # OpenCV images have a different color profile. So remember to switch to RGB, that our resnet model understands.

        cust_transform = A.Compose([A.Resize(height=256, width=256, p=1.0),ToTensorV2(p=1.0)], p=1.0)
        tensor = cust_transform(image=img)
        tensor = tensor['image'].float().resize(1,3,256,256)

        custom_pred = model.forward(tensor).detach().numpy() # Forward is the python method defined inside the resnet.
        custom_pred

        st.write(f'Predicted: {id2class[np.argmax(custom_pred)]}')
    elif '.csv' in uploaded_file.name:
        dataframe = pd.read_csv(uploaded_file)
        st.write(dataframe)