newtechdevng/construction-detection-api

1from fastapi import FastAPI, File, UploadFile, Form2from fastapi.middleware.cors import CORSMiddleware3from huggingface_hub import hf_hub_download4from ultralytics import YOLO5import cv26import numpy as np7import base648import time9import os10 11app = FastAPI(title="Construction Detection API")12 13app.add_middleware(14    CORSMiddleware,15    allow_origins=["*"],16    allow_methods=["*"],17    allow_headers=["*"],18)19 20# Load YOLO model21HF_REPO_ID = "newtechdevng/construction_detection_fine_tune"22MODEL_FILE  = "best_v2_finetune.pt"23model_path  = hf_hub_download(repo_id=HF_REPO_ID, filename=MODEL_FILE)24model       = YOLO(model_path)25 26# ArUco setup27ARUCO_DICT     = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)28ARUCO_PARAMS   = cv2.aruco.DetectorParameters()29ARUCO_DETECTOR = cv2.aruco.ArucoDetector(ARUCO_DICT, ARUCO_PARAMS)30 31CLASS_COLORS = {32    "beam":    (255, 100,   0),33    "column":  (  0, 255, 255),34    "door":    (255,   0, 255),35    "floor":   (  0, 255,   0),36    "stairs":  (255, 255,   0),37    "wall":    (  0, 100, 255),38    "window":  (100,   0, 255),39}40 41def detect_aruco_scale(img, marker_size_cm=10.0):42    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)43    corners, ids, _ = ARUCO_DETECTOR.detectMarkers(gray)44    if ids is None:45        return None, None46    marker_corners = corners[0][0]47    w_px = np.linalg.norm(marker_corners[0] - marker_corners[1])48    h_px = np.linalg.norm(marker_corners[1] - marker_corners[2])49    pixels_per_cm = (w_px + h_px) / 2 / marker_size_cm50    return float(pixels_per_cm), corners51 52@app.get("/")53def root():54    return {55        "model": MODEL_FILE,56        "classes": list(CLASS_COLORS.keys()),57        "calibration": "Auto via ArUco marker on hard hat (10cm x 10cm)",58        "endpoints": {59            "POST /detect": "Send image → get detections in cm",60            "GET  /health": "Health check"61        }62    }63 64@app.get("/health")65def health():66    return {"status": "ok", "model": MODEL_FILE}67 68@app.post("/detect")69async def detect(70    file: UploadFile = File(...),71    marker_size_cm: float = Form(10.0),72    confidence: float = Form(0.2),73    iou: float = Form(0.3)74):75    start = time.time()76 77    contents = await file.read()78    nparr    = np.frombuffer(contents, np.uint8)79    img      = cv2.imdecode(nparr, cv2.IMREAD_COLOR)80 81    # ArUco auto-calibration82    pixels_per_cm, aruco_corners = detect_aruco_scale(img, marker_size_cm)83    calibrated = pixels_per_cm is not None84 85    # Draw ArUco marker highlight86    if calibrated:87        cv2.aruco.drawDetectedMarkers(img, aruco_corners)88 89    # Run YOLO with lower confidence + iou for more detections90    results    = model(img, conf=confidence, iou=iou)[0]91    detections = []92 93    for box in results.boxes:94        x1, y1, x2, y2 = map(int, box.xyxy[0])95        cls   = results.names[int(box.cls[0])]96        conf  = round(float(box.conf[0]), 2)97        w_px  = x2 - x198        h_px  = y2 - y199        color = CLASS_COLORS.get(cls, (0, 255, 0))100 101        w_cm = round(float(w_px) / pixels_per_cm, 1) if calibrated else None102        h_cm = round(float(h_px) / pixels_per_cm, 1) if calibrated else None103 104        # Draw bounding box105        cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)106 107        # Label108        label = f"{cls} {conf:.2f}"109        if calibrated:110            label += f" | {w_cm}x{h_cm}cm"111 112        # Background for label text113        (tw, th), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.55, 2)114        cv2.rectangle(img, (x1, y1 - th - 10), (x1 + tw, y1), color, -1)115        cv2.putText(img, label, (x1, y1 - 5),116                    cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 2)117 118        detections.append({119            "class":      cls,120            "confidence": conf,121            "bbox":       [int(x1), int(y1), int(x2), int(y2)],122            "width_px":   int(w_px),123            "height_px":  int(h_px),124            "width_cm":   round(float(w_cm), 1) if w_cm is not None else None,125            "height_cm":  round(float(h_cm), 1) if h_cm is not None else None,126        })127 128    # Encode result image129    _, buf  = cv2.imencode(".jpg", img)130    img_b64 = base64.b64encode(buf).decode()131 132    return {133        "success":          True,134        "calibrated":       bool(calibrated),135        "pixels_per_cm":    round(pixels_per_cm, 2) if calibrated else None,136        "marker_size_cm":   float(marker_size_cm),137        "inference_time_s": round(float(time.time() - start), 3),138        "total":            int(len(detections)),139        "detections":       detections,140        "image_base64":     img_b64,141    }142 143if __name__ == "__main__":144    import uvicorn145    uvicorn.run(app, host="0.0.0.0", port=7860)