Update app.py

app.py

@@ -1,190 +1,141 @@
-"""
-Construction Detection API — Hugging Face Space
-Loads model from HF Hub, serves REST API for mobile app
-"""
-
-from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
 from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
+import cv2
 import numpy as np
-import cv2, base64, time, os
-
-# ── CONFIG ──────────────────────────────────────────────────────────────────
-HF_REPO_ID  = "newtechdevng/construction_detection_fine_tune"
-MODEL_FILE  = "best_v2_finetune.pt"
-CONF        = 0.25
-IOU         = 0.45
-
-COLORS = {
-    "beam":   [255, 0,   0  ],
-    "column": [0,   255, 255],
-    "door":   [255, 0,   255],
-    "floor":  [0,   165, 255],
-    "stairs": [0,   255, 0  ],
-    "wall":   [255, 255, 0  ],
-    "window": [0,   0,   255],
-}
+import base64
+import time
+import os
 
-# ── APP ─────────────────────────────────────────────────────────────────────
-app = FastAPI(
-    title       = "Construction Detection API",
-    description = "Detects construction elements and measures dimensions",
-    version     = "1.0.0",
-)
+app = FastAPI(title="Construction Detection API")
 
 app.add_middleware(
     CORSMiddleware,
-    allow_origins  = ["*"],
-    allow_methods  = ["*"],
-    allow_headers  = ["*"],
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
-# ── GLOBAL STATE ─────────────────────────────────────────────────────────────
-model         = None
-pixels_per_cm = None
-
-# ── STARTUP ──────────────────────────────────────────────────────────────────
-@app.on_event("startup")
-async def load_model():
-    global model
-    print(f"Downloading {MODEL_FILE} from {HF_REPO_ID}...")
-    path = hf_hub_download(repo_id=HF_REPO_ID, filename=MODEL_FILE)
-    model = YOLO(path)
-    print("✅ Model loaded!")
-
-# ── HELPERS ──────────────────────────────────────────────────────────────────
-def bytes_to_image(data: bytes) -> np.ndarray:
-    arr = np.frombuffer(data, np.uint8)
-    return cv2.imdecode(arr, cv2.IMREAD_COLOR)
-
-def image_to_base64(img: np.ndarray) -> str:
-    _, buf = cv2.imencode(".jpg", img, [cv2.IMWRITE_JPEG_QUALITY, 85])
-    return base64.b64encode(buf).decode("utf-8")
-
-def px_to_cm(pixels: float) -> float | None:
-    if pixels_per_cm is None:
-        return None
-    return round(pixels / pixels_per_cm, 1)
-
-def draw_boxes(img: np.ndarray, detections: list) -> np.ndarray:
-    for det in detections:
-        x1, y1, x2, y2 = det["bbox"]
-        cls   = det["class"]
-        conf  = det["confidence"]
-        color = COLORS.get(cls, [255, 255, 255])
-
-        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
-
-        if det.get("width_cm"):
-            label = f"{cls} {conf:.2f} W:{det['width_cm']}cm H:{det['height_cm']}cm"
-        else:
-            label = f"{cls} {conf:.2f} W:{det['width_px']}px H:{det['height_px']}px"
+# Load YOLO model
+HF_REPO_ID = "newtechdevng/construction_detection_fine_tune"
+MODEL_FILE  = "best_v2_finetune.pt"
+model_path  = hf_hub_download(repo_id=HF_REPO_ID, filename=MODEL_FILE)
+model       = YOLO(model_path)
+
+# ArUco setup
+ARUCO_DICT     = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
+ARUCO_PARAMS   = cv2.aruco.DetectorParameters()
+ARUCO_DETECTOR = cv2.aruco.ArucoDetector(ARUCO_DICT, ARUCO_PARAMS)
+
+CLASS_COLORS = {
+    "beam":    (255, 100,   0),
+    "column":  (  0, 255, 255),
+    "door":    (255,   0, 255),
+    "floor":   (  0, 255,   0),
+    "stairs":  (255, 255,   0),
+    "wall":    (  0, 100, 255),
+    "window":  (100,   0, 255),
+}
 
-        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
-        cv2.rectangle(img, (x1, y1 - th - 8), (x1 + tw + 4, y1), color, -1)
-        cv2.putText(img, label, (x1 + 2, y1 - 5),
-                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 0), 1)
-    return img
+def detect_aruco_scale(img, marker_size_cm=10.0):
+    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+    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
 
-def run_detection(img: np.ndarray) -> list:
-    results    = model.predict(img, conf=CONF, iou=IOU, task="detect", verbose=False)
-    detections = []
-    for result in results:
-        for box in result.boxes:
-            x1, y1, x2, y2 = map(int, box.xyxy[0])
-            cls  = model.names[int(box.cls)]
-            conf = round(float(box.conf), 3)
-            w_px = x2 - x1
-            h_px = y2 - y1
-            detections.append({
-                "class":      cls,
-                "confidence": conf,
-                "bbox":       [x1, y1, x2, y2],
-                "width_px":   w_px,
-                "height_px":  h_px,
-                "width_cm":   px_to_cm(w_px),
-                "height_cm":  px_to_cm(h_px),
-                "color":      COLORS.get(cls, [255, 255, 255]),
-            })
-    return detections
-
-# ── ROUTES ───────────────────────────────────────────────────────────────────
 @app.get("/")
-async def root():
+def root():
     return {
-        "status":    "running",
-        "model":     MODEL_FILE,
-        "classes":   list(COLORS.keys()),
+        "model": MODEL_FILE,
+        "classes": list(CLASS_COLORS.keys()),
+        "calibration": "Auto via ArUco marker on hard hat (10cm × 10cm)",
         "endpoints": {
-            "POST /detect":    "Upload image → detections + dimensions",
-            "POST /calibrate": "Set reference object for real-world units",
-            "GET  /health":    "Health check",
+            "POST /detect": "Send image → get detections in cm (if hard hat in frame)",
+            "GET  /health": "Health check"
         }
     }
 
 @app.get("/health")
-async def health():
-    return {"status": "ok", "model_loaded": model is not None}
-
-@app.post("/calibrate")
-async def calibrate(
-    file:        UploadFile = File(...),
-    bbox_x1:     int   = 0,
-    bbox_y1:     int   = 0,
-    bbox_x2:     int   = 210,
-    bbox_y2:     int   = 297,
-    real_width:  float = 21.0,
-    real_height: float = 29.7,
+def health():
+    return {"status": "ok", "model": MODEL_FILE}
+
+@app.post("/detect")
+async def detect(
+    file: UploadFile = File(...),
+    marker_size_cm: float = Form(10.0),
+    confidence: float = Form(0.4)
 ):
-    """
-    Calibrate using a reference object (e.g. A4 paper = 21cm x 29.7cm).
-    Provide bounding box of the reference object in pixels.
-    """
-    global pixels_per_cm
-    data = await file.read()
-    img  = bytes_to_image(data)
-    if img is None:
-        raise HTTPException(400, "Invalid image")
-
-    ref_px_w  = bbox_x2 - bbox_x1
-    ref_px_h  = bbox_y2 - bbox_y1
-    px_per_w  = ref_px_w / real_width
-    px_per_h  = ref_px_h / real_height
-    pixels_per_cm = round((px_per_w + px_per_h) / 2, 4)
+    start = time.time()
+
+    contents = await file.read()
+    nparr    = np.frombuffer(contents, np.uint8)
+    img      = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+
+    # Try ArUco auto-calibration
+    pixels_per_cm, aruco_corners = detect_aruco_scale(img, marker_size_cm)
+    calibrated = pixels_per_cm is not None
+
+    # Draw ArUco marker highlight
+    if calibrated:
+        cv2.aruco.drawDetectedMarkers(img, aruco_corners)
+
+    # Run YOLO
+    results    = model(img, conf=confidence)[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)
+        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
+
+        # Draw bounding box
+        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
+
+        # Label with cm if calibrated
+        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)
+
+        detections.append({
+            "class":      cls,
+            "confidence": conf,
+            "bbox":       [x1, y1, x2, y2],
+            "width_px":   w_px,
+            "height_px":  h_px,
+            "width_cm":   w_cm,
+            "height_cm":  h_cm,
+        })
+
+    # Encode result image
+    _, buf   = cv2.imencode(".jpg", img)
+    img_b64  = base64.b64encode(buf).decode()
 
     return {
-        "message":       "✅ Calibration successful",
-        "pixels_per_cm": pixels_per_cm,
+        "success":         True,
+        "calibrated":      calibrated,
+        "pixels_per_cm":   round(pixels_per_cm, 2) if calibrated else None,
+        "marker_size_cm":  marker_size_cm,
+        "inference_time_s": round(time.time() - start, 3),
+        "total":           len(detections),
+        "detections":      detections,
+        "image_base64":    img_b64,
     }
 
-@app.post("/detect")
-async def detect(file: UploadFile = File(...)):
-    """
-    Upload a construction site image.
-    Returns all detected objects with bounding boxes and dimensions.
-    """
-    if model is None:
-        raise HTTPException(503, "Model not loaded")
-
-    data = await file.read()
-    img  = bytes_to_image(data)
-    if img is None:
-        raise HTTPException(400, "Invalid image")
-
-    start      = time.time()
-    detections = run_detection(img)
-    elapsed    = round(time.time() - start, 3)
-
-    annotated  = draw_boxes(img.copy(), detections)
-    img_b64    = image_to_base64(annotated)
-
-    return JSONResponse({
-        "success":          True,
-        "total":            len(detections),
-        "inference_time_s": elapsed,
-        "calibrated":       pixels_per_cm is not None,
-        "image_base64":     img_b64,
-        "detections":       detections,
-    })
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)