Update app.py

app.py

@@ -20,6 +20,7 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 class_mapping = {'tb': 0, 'healthy': 1, 'sick_but_no_tb': 2}
 reverse_mapping = {v: k for k, v in class_mapping.items()}
+labels = class_mapping.keys()
 
 def load_model():
     # config = read_params(config_path)
@@ -99,9 +100,9 @@ def predict(inp):
   inp = transforms.ToTensor()(inp).unsqueeze(0)
   with torch.no_grad():
     prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
-    # confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
-    prediction = reverse_mapping[prediction]
-  return prediction 
+    confidences = {labels[i]: float(prediction[i]) for i in range(3)}
+    # prediction = reverse_mapping[prediction]
+  return confidences 
 
 
 import gradio as gr