Update main.py

main.py

@@ -1,56 +1,83 @@
 from fastapi import FastAPI
+from pydantic import BaseModel
 import uvicorn
 import base64
 import cv2
 import numpy as np
 from ultralytics import YOLO
 from datetime import datetime
-from pydantic import BaseModel
-
-app = FastAPI()
 
+app = FastAPI(title="PCB Component Detection API")
 model = YOLO("pcb_component_detection_best.pt")
 
 class ImageRequest(BaseModel):
+    """Request model for image processing endpoint."""
     image: str
 
 @app.get("/")
 async def root():
+    """Root endpoint to verify API status."""
     current_time = datetime.now().isoformat()
-    return {"message": "PCB components API works", "time": current_time}
+    return {
+        "message": "PCB Components API works",
+        "time": current_time
+    }
 
 @app.post("/predict")
 async def predict(request: ImageRequest):
-    # Decode Base64
-    image_bytes = base64.b64decode(request.image)
-    np_arr = np.frombuffer(image_bytes, np.uint8)
-    image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
-    if image is None:
-        return {"error": "Invalid image"}
+    """
+    Process an image to detect PCB components.
     
-    # Inference
-    results = model.predict(image)
-    result = results[0]
-
-    # Response
-    json_result = {}
-    class_counters = {}
+    Args:
+        request: Contains base64 encoded image
+        
+    Returns:
+        JSON with detection statistics and bounding boxes
+    """
+    # Validate image input
+    if not request.image:
+        return {"error": "Invalid Image"}
     
-    for box in result.boxes:
-        class_id = int(box.cls[0])
-        class_name = result.names[class_id]
-        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
-
-        if class_name in class_counters:
-            class_counters[class_name] += 1
-            key = f"{class_name}_{class_counters[class_name]}"
+    try:
+        image_bytes = base64.b64decode(request.image, validate=True)
+        
+        np_arr = np.frombuffer(image_bytes, np.uint8)
+        image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
+        
+        if image is None:
+            return {"error": "Invalid image"}
+            
+        results = model.predict(image)
+        result = results[0]
+        
+        json_result = {}
+        class_counters = {}
+        
+        for box in result.boxes:
+            class_id = int(box.cls[0])
+            class_name = result.names[class_id]
+            x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
+            
+            if class_name in class_counters:
+                class_counters[class_name] += 1
+            else:
+                class_counters[class_name] = 1
+                
+            key = f"{class_name}{class_counters[class_name]}" if class_counters[class_name] > 1 else class_name
+            json_result[key] = [x1, y1, x2, y2]
+        
+        if hasattr(result, "summary") and isinstance(result.summary, dict):
+            statistics_summary = result.summary
         else:
-            class_counters[class_name] = 1
-            key = class_name
+            statistics_summary = {name: count for name, count in class_counters.items()}
+            
+        return {
+            "statistics": statistics_summary,
+            "components": json_result
+        }
         
-        json_result[key] = [x1, y1, x2, y2]
-    
-    return json_result
+    except Exception as e:
+        return {"error": f"Invalid Image: {str(e)}"}
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)