Update app.py

app.py

@@ -43,21 +43,20 @@ def detect_aruco_scale(img, marker_size_cm=10.0):
     corners, ids, _ = ARUCO_DETECTOR.detectMarkers(gray)
     if ids is None:
         return None, None
-    # Use first detected marker
     marker_corners = corners[0][0]
     w_px = np.linalg.norm(marker_corners[0] - marker_corners[1])
     h_px = np.linalg.norm(marker_corners[1] - marker_corners[2])
     pixels_per_cm = (w_px + h_px) / 2 / marker_size_cm
-    return pixels_per_cm, corners
+    return float(pixels_per_cm), corners
 
 @app.get("/")
 def root():
     return {
         "model": MODEL_FILE,
         "classes": list(CLASS_COLORS.keys()),
-        "calibration": "Auto via ArUco marker on hard hat (10cm × 10cm)",
+        "calibration": "Auto via ArUco marker on hard hat (10cm x 10cm)",
         "endpoints": {
-            "POST /detect": "Send image → get detections in cm (if hard hat in frame)",
+            "POST /detect": "Send image → get detections in cm",
             "GET  /health": "Health check"
         }
     }
@@ -70,7 +69,8 @@ def health():
 async def detect(
     file: UploadFile = File(...),
     marker_size_cm: float = Form(10.0),
-    confidence: float = Form(0.4)
+    confidence: float = Form(0.2),
+    iou: float = Form(0.3)
 ):
     start = time.time()
 
@@ -78,7 +78,7 @@ async def detect(
     nparr    = np.frombuffer(contents, np.uint8)
     img      = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
 
-    # Try ArUco auto-calibration
+    # ArUco auto-calibration
     pixels_per_cm, aruco_corners = detect_aruco_scale(img, marker_size_cm)
     calibrated = pixels_per_cm is not None
 
@@ -86,54 +86,58 @@ async def detect(
     if calibrated:
         cv2.aruco.drawDetectedMarkers(img, aruco_corners)
 
-    # Run YOLO
-    results    = model(img, conf=confidence)[0]
+    # Run YOLO with lower confidence + iou for more detections
+    results    = model(img, conf=confidence, iou=iou)[0]
     detections = []
 
     for box in results.boxes:
         x1, y1, x2, y2 = map(int, box.xyxy[0])
         cls   = results.names[int(box.cls[0])]
-        conf  = round(float(box.conf[0]), 3)
+        conf  = round(float(box.conf[0]), 2)
         w_px  = x2 - x1
         h_px  = y2 - y1
         color = CLASS_COLORS.get(cls, (0, 255, 0))
 
-        w_cm = round(w_px / pixels_per_cm, 1) if calibrated else None
-        h_cm = round(h_px / pixels_per_cm, 1) if calibrated else None
+        w_cm = round(float(w_px) / pixels_per_cm, 1) if calibrated else None
+        h_cm = round(float(h_px) / pixels_per_cm, 1) if calibrated else None
 
         # Draw bounding box
         cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
 
-        # Label with cm if calibrated
+        # Label
         label = f"{cls} {conf:.2f}"
         if calibrated:
             label += f" | {w_cm}x{h_cm}cm"
-        cv2.putText(img, label, (x1, y1 - 8),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, color, 2)
+
+        # Background for label text
+        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 2)
+        cv2.rectangle(img, (x1, y1 - th - 10), (x1 + tw, y1), color, -1)
+        cv2.putText(img, label, (x1, y1 - 5),
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 2)
 
         detections.append({
             "class":      cls,
-            "confidence": float(conf),
+            "confidence": conf,
             "bbox":       [int(x1), int(y1), int(x2), int(y2)],
             "width_px":   int(w_px),
             "height_px":  int(h_px),
-            "width_cm":   float(w_cm) if w_cm is not None else None,
-            "height_cm":  float(h_cm) if h_cm is not None else None,
+            "width_cm":   round(float(w_cm), 1) if w_cm is not None else None,
+            "height_cm":  round(float(h_cm), 1) if h_cm is not None else None,
         })
 
     # Encode result image
-    _, buf   = cv2.imencode(".jpg", img)
-    img_b64  = base64.b64encode(buf).decode()
+    _, buf  = cv2.imencode(".jpg", img)
+    img_b64 = base64.b64encode(buf).decode()
 
     return {
-        "success":            True,
-        "calibrated":         bool(calibrated),
-        "pixels_per_cm":      float(round(pixels_per_cm, 2)) if calibrated else None,
-        "marker_size_cm":     float(marker_size_cm),
-        "inference_time_s":   float(round(time.time() - start, 3)),
-        "total":              int(len(detections)),
-        "detections":         detections,
-        "image_base64":       img_b64,
+        "success":          True,
+        "calibrated":       bool(calibrated),
+        "pixels_per_cm":    round(pixels_per_cm, 2) if calibrated else None,
+        "marker_size_cm":   float(marker_size_cm),
+        "inference_time_s": round(float(time.time() - start), 3),
+        "total":            int(len(detections)),
+        "detections":       detections,
+        "image_base64":     img_b64,
     }
 
 if __name__ == "__main__":