Update app.py

app.py

@@ -1,14 +1,15 @@
-#import streamlit as st
+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 glob
+#import os
 
 import torch
 import torch.nn as nn
-from torch.utils.data import Dataset, DataLoader
+#from torch.utils.data import Dataset, DataLoader
 import torchvision.models as models
 #from torchinfo import summary
 
@@ -18,14 +19,14 @@ import torchvision.models as models
 #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 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
+from albumentations.pytorch.transforms import ToTensorV2
 
 model = models.resnet50(pretrained=False)
 model.fc = nn.Linear(2048, num_classes)
-model.load_state_dict(torch.load('resnet_best.pth'), strict=True)
+model.load_state_dict(torch.load('resnet_best.pth', map_location=torch.device('cpu')), strict=True)
 
 st.title("some big ML function")
 
@@ -33,8 +34,9 @@ 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:
-    image = Image.open(uploaded_file)
-    st.image(image)  #show it
+    img = Image.open(uploaded_file)
+    st.image(img)
+    img = np.array(img)
     
   elif ".csv" in uploaded_file.name:
     dataframe = pd.read_csv(uploaded_file)