Update app.py

app.py

@@ -1,13 +1,60 @@
+Hugging Face's logo
+Hugging Face
+Search models, datasets, users...
+Models
+Datasets
+Spaces
+Docs
+Solutions
+Pricing
+
+
+
+Spaces:
+
+jvahala
+/
+dummy Copied
+like
+0
+App
+Files and versions
+Community
+dummy
+/
+app.py
+jvahala's picture
+jvahala
+fix imread issue
+011b03a
+about 1 hour ago
+raw
+history
+blame
+contribute
+delete
+Safe
+1.89 kB
+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 torch
 import torch.nn as nn
-import cv2
+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)
@@ -15,36 +62,22 @@ 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)
 
-#Load model torch.load.state_dict(PATH)
-#model = torch.load.state_dict('resnet_best.pth')
-#model = models.resnet50(pretrained=False)
-#model.fc = nn.Linear(2048, 21)
-#model.load_state_dict(torch.load('resnet_best.pth'), strict=True)#load weights from training run
-#model.eval()
-
-st.title("My awesome ML Function")
-
-uploaded_file = st.file_uploader("Choose a file")
+        custom_pred = model.forward(tensor).detach().numpy() # Forward is the python method defined inside the resnet.
+        custom_pred
 
-if uploaded_file is not None:
-  if ".jpg" in uploaded_file.name or ".png" in uploaded_file.name: #image check
-    image = Image.open(uploaded_file)
-    st.image(image) #showing the image
-    np_img = np.array(image)
-    cv_image = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    
-    # Let's resize the image and make a pytorch tensor from the image np array. 
-    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)
-    tensor,tensor.shape
-    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: #csv check
-    dataframe = pd.read_csv(uploaded_file)
-    st.write(dataframe)
- 
+        st.write(f'Predicted: {id2class[np.argmax(custom_pred)]}')
+    elif '.csv' in uploaded_file.name:
+        dataframe = pd.read_csv(uploaded_file)
+        st.write(dataframe)