Update app.py

app.py

@@ -147,7 +147,7 @@ for key in list(state_dict.keys()):
 
 model.load_state_dict(state_dict)
 model.eval()
-
+"""
 class_names = {
     'akk': 'Actinic Keratosis',
     'bcc': 'Basal Cell Carcinoma',
@@ -157,11 +157,14 @@ class_names = {
     'nv': 'Melanocytic Nevi',
     'vasc': 'Vascular Lesion'
 }
-
+"""
+class_names = ['akk', 'bcc', 'bkl', 'df', 'mel', 'nv', 'vasc']
+class_names.sort()
 examples_dir = "sample"
 
 
 
+
 transformation_pipeline = transforms.Compose([
     transforms.ToPILImage(),
     transforms.Grayscale(num_output_channels=3),
@@ -217,7 +220,7 @@ def image_classifier(inp):
     # postprocess
     result = torch.nn.functional.softmax(result, dim=1) # apply softmax
     result = result[0].detach().numpy().tolist() # take the first batch
-    labeled_result = {class_names[key]: score for key, score in result.items()}
+    labeled_result = {name:score for name, score in zip(class_names, result)}
 
     return labeled_result