from distutils.command.upload import upload
import streamlit as st
from io import StringIO
from PIL import Image
import pandas as pd
import numpy as np
#import glob
#import os

import torch
import torch.nn as nn
#from torch.utils.data import Dataset, DataLoader
import torchvision.models as models
#from torchinfo import summary

#from sklearn.model_selection import StratifiedKFold
#from sklearn.metrics import accuracy_score

#from tqdm import tqdm
#import opencv-python
import cv2
import albumentations as A # Albumentations is a computer vision tool that boosts the performance of deep convolutional neural networks. (https://albumentations.ai/)
#import matplotlib.pyplot as plt
#import seaborn as sns
from albumentations.pytorch.transforms import ToTensorV2

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(pretrained=False)
model.fc = nn.Linear(2048, 21)
model.load_state_dict(torch.load('resnet_best.pth', map_location=torch.device('cpu')), strict=True)

st.title("some big ML function")

uploaded_file = st.file_uploader("Choose a file")

if uploaded_file is not None:
  if ".jpg" in uploaded_file.name or ".png" in uploaded_file.name:
    img = Image.open(uploaded_file)
    st.image(img)
    img = np.array(img)
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    
    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)
    
    model.eval()
    custom_pred = model.forward(tensor).detach().numpy()
    custom_pred
    id2class[np.argmax(custom_pred)]
    
  elif ".csv" in uploaded_file.name:
    dataframe = pd.read_csv(uploaded_file)
    st.write(dataframe)