Update app.py

app.py

@@ -24,6 +24,8 @@ import albumentations as A # Albumentations is a computer vision tool that boost
 #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)
@@ -45,6 +47,7 @@ if uploaded_file is not None:
     
     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)