Update main.py

main.py

@@ -31,8 +31,10 @@ async def get_prediction(file: UploadFile = File(...)):
     results = model(file_name) 
     # View results
     class_name = ""
+    confidence = ""
     for r in results:
         print(names[int(r.probs.top1)]) 
         class_name = names[int(r.probs.top1)]
+        confidence = r.probs.top1conf
 
-    return class_name
+    return {"class_name":class_name,"confidence":confidence}