Small Edits

main.py

@@ -1,73 +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 Defect Detection API")
 model = YOLO("yolo_modeln11_1502.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 Defects API works", "time": current_time}
+    return {
+        "message": "PCB Defects API works",
+        "time": current_time
+    }
 
 @app.post("/predict")
 async def predict(request: ImageRequest):
-    # Check if the image string is empty
+    """
+    Process an image to detect PCB defects.
+    
+    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"}
     
     try:
-        # Attempt to decode Base64 string to bytes
         image_bytes = base64.b64decode(request.image, validate=True)
-    except Exception as e:
-        return {"error": "Invalid Image"}
-    
-    # Ensure the decoded bytes are not empty
-    if not image_bytes:
-        return {"error": "Invalid 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"}
-    
-    
-    results = model.predict(image)
-    result = results[0]
-
-    # Response
-    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())
+        
+        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]
 
-        # Counting occurrences of each defect
-        if class_name in class_counters:
-            class_counters[class_name] += 1
+        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
+            statistics_summary = {name: count for name, count in class_counters.items()}
+            
+        return {
+            "statistics": statistics_summary,
+            "detections": json_result
+        }
         
-        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:
-        statistics_summary = {name: count for name, count in class_counters.items()}
-    
-    return {"statistics": statistics_summary, "detections": 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)
+    uvicorn.run(app, host="0.0.0.0", port=7860)