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

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)
    
    custom_pred = model.forward(tensor).detach().numpy()
    custom_pred
    
  elif ".csv" in uploaded_file.name:
    dataframe = pd.read_csv(uploaded_file)
    st.write(dataframe)