Update app.py

app.py

@@ -8,13 +8,22 @@ import torch.nn as nn
 import cv2
 from albumentations.pytorch.transforms import ToTensorV2
 
-#Load model torch.load.state_dict(PATH)
-model = torch.load.state_dict('resnet_best.pth')
-#model = models.resnet50(pretrained=False)
+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'), strict=True)#load weights from training run
+model.load_state_dict(torch.load('resnet_best.pth', map_location=torch.device('cpu')), strict=True)
 model.eval()
 
+
+
+#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")