Update main.py

main.py

@@ -30,7 +30,9 @@ async def get_prediction(file: UploadFile = File(...)):
     # Predict with the model
     results = model(file_name) 
     # View results
+    class_name = ""
     for r in results:
         print(names[int(r.probs.top1)]) 
+        class_name = names[int(r.probs.top1)]
 
-    return r.probs
+    return class_name