Update app.py

app.py

@@ -23,9 +23,7 @@ model.load_state_dict(
     )
 )
 
-
 def predict(img) -> Tuple[Dict, float]:
-
     start_time = timer()
 
     preprocess = transforms.Compose([
@@ -34,16 +32,25 @@ def predict(img) -> Tuple[Dict, float]:
         transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
     ])
 
-    img = preprocess(img).unsqueeze(0)  # Add batch dimension
+    image = preprocess(img).unsqueeze(0)  # Add batch dimension
 
+    # Make predictions
     model.eval()
-    with torch.inference_mode():
-        pred_probs = torch.softmax(model(img), dim=1)
+    with torch.no_grad():
+        outputs = model(image).logits
+        predicted_probs = torch.softmax(outputs, dim=1)
 
-    pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
+    # Get the class name and its associated probability
+    predicted_class_idx = torch.argmax(predicted_probs).item()
+    predicted_class = class_names[predicted_class_idx]
+    predicted_probability = predicted_probs[0][predicted_class_idx].item()
 
+    # Calculate the prediction time
     pred_time = round(timer() - start_time, 5)
 
+    # Create a prediction label and prediction probability dictionary for the predicted class
+    pred_labels_and_probs = {predicted_class: predicted_probability}
+
     return pred_labels_and_probs, pred_time