Update main.py

main.py

@@ -21,14 +21,26 @@ async def root():
 
 @app.post("/predict")
 async def predict(request: ImageRequest):
-    # Decode Base64
-    image_bytes = base64.b64decode(request.image)
+    # Check if the image string is empty
+    if not request.image:
+        return {"error": "Empty image string"}
+    
+    try:
+        # Attempt to decode Base64 string to bytes
+        image_bytes = base64.b64decode(request.image, validate=True)
+    except Exception as e:
+        return {"error": "Invalid Base64 string", "details": str(e)}
+    
+    # Ensure the decoded bytes are not empty
+    if not image_bytes:
+        return {"error": "Decoded data is empty"}
+    
     np_arr = np.frombuffer(image_bytes, np.uint8)
     image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
     if image is None:
         return {"error": "Invalid image"}
     
-    # Inference
+    
     results = model.predict(image)
     result = results[0]
 
@@ -44,18 +56,21 @@ async def predict(request: ImageRequest):
         '''
         Create unique key for each detection
         EX: if multiple "missing_hole", keys will be:
-        "missing_hole", "missing_hole_2", ...
+        "missing_hole", "missing_hole2", ...
         '''
         if class_name in class_counters:
             class_counters[class_name] += 1
-            key = f"{class_name}_{class_counters[class_name]}"
+            key = f"{class_name}{class_counters[class_name]}"
         else:
             class_counters[class_name] = 1
             key = class_name
         
         json_result[key] = [x1, y1, x2, y2]
     
-    return json_result
+    # Creating statistics summary
+    statistics_summary = {f"{count} {name}{'s' if count > 1 else ''}": count for name, count in class_counters.items()}
+    
+    return {"statistics": statistics_summary, "detections": json_result}
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)