Update app.py

app.py

@@ -80,9 +80,15 @@ def predict(image):
         
         # Get the predicted class with the highest score
         _, predicted_idx = torch.max(output, 1)
-        predicted_class = labels[str(predicted_idx.item())]
-        print(f"Predicted class index: {predicted_idx.item()}")
-        print(f"Predicted class: {predicted_class}")
+        predicted_idx = predicted_idx.item()
+        print(f"Predicted class index: {predicted_idx}")
+        
+        # Check if the predicted index exists in labels
+        if str(predicted_idx) in labels:
+            predicted_class = labels[str(predicted_idx)]
+        else:
+            predicted_class = f"Unknown class index: {predicted_idx}. Please check the label mapping."
+            print(predicted_class)
 
         return f"Predicted class: {predicted_class}"