| from fastapi import FastAPI, File, UploadFile, Form |
| from fastapi.middleware.cors import CORSMiddleware |
| from huggingface_hub import hf_hub_download |
| from ultralytics import YOLO |
| import cv2 |
| import numpy as np |
| import base64 |
| import time |
| import os |
|
|
| app = FastAPI(title="Construction Detection API") |
|
|
| app.add_middleware( |
| CORSMiddleware, |
| allow_origins=["*"], |
| allow_methods=["*"], |
| allow_headers=["*"], |
| ) |
|
|
| |
| 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_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), |
| } |
|
|
| 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 |
| 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 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 x 10cm)", |
| "endpoints": { |
| "POST /detect": "Send image → get detections in cm", |
| "GET /health": "Health check" |
| } |
| } |
|
|
| @app.get("/health") |
| 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.2), |
| iou: float = Form(0.3) |
| ): |
| start = time.time() |
|
|
| contents = await file.read() |
| nparr = np.frombuffer(contents, np.uint8) |
| img = cv2.imdecode(nparr, cv2.IMREAD_COLOR) |
|
|
| |
| pixels_per_cm, aruco_corners = detect_aruco_scale(img, marker_size_cm) |
| calibrated = pixels_per_cm is not None |
|
|
| |
| if calibrated: |
| cv2.aruco.drawDetectedMarkers(img, aruco_corners) |
|
|
| |
| 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]), 2) |
| w_px = x2 - x1 |
| h_px = y2 - y1 |
| color = CLASS_COLORS.get(cls, (0, 255, 0)) |
|
|
| 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 |
|
|
| |
| cv2.rectangle(img, (x1, y1), (x2, y2), color, 2) |
|
|
| |
| label = f"{cls} {conf:.2f}" |
| if calibrated: |
| label += f" | {w_cm}x{h_cm}cm" |
|
|
| |
| (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": conf, |
| "bbox": [int(x1), int(y1), int(x2), int(y2)], |
| "width_px": int(w_px), |
| "height_px": int(h_px), |
| "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, |
| }) |
|
|
| |
| _, buf = cv2.imencode(".jpg", img) |
| img_b64 = base64.b64encode(buf).decode() |
|
|
| return { |
| "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__": |
| import uvicorn |
| uvicorn.run(app, host="0.0.0.0", port=7860) |
