from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from ultralytics import YOLO
import numpy as np
import cv2
import base64
app = FastAPI()
# Allow your Vercel app to call this API
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # replace * with your vercel domain in production
allow_methods=["*"],
allow_headers=["*"],
)
model = YOLO("best.pt")
@app.post("/detect")
async def detect(
file: UploadFile = File(...),
confidence: float = Form(default=0.5), # 0.0 - 1.0, matches Roboflow slider default
overlap: float = Form(default=0.5), # NMS IoU threshold, same as Roboflow overlap slider
):
# Clamp values to valid range
confidence = max(0.01, min(1.0, confidence))
overlap = max(0.01, min(1.0, overlap))
contents = await file.read()
nparr = np.frombuffer(contents, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
if img is None:
return JSONResponse({"error": "Could not decode image"}, status_code=400)
results = model(img, conf=confidence, iou=overlap)[0]
detections = []
for box in results.boxes:
detections.append({
"class": model.names[int(box.cls)],
"confidence": round(float(box.conf), 3),
"bbox": {
# Normalized coords (0-1), easy to draw on any canvas size
"x": round(float(box.xywhn[0][0]), 4),
"y": round(float(box.xywhn[0][1]), 4),
"w": round(float(box.xywhn[0][2]), 4),
"h": round(float(box.xywhn[0][3]), 4),
},
"bbox_pixels": {
# Absolute pixel coords in the original image
"x1": int(box.xyxy[0][0]),
"y1": int(box.xyxy[0][1]),
"x2": int(box.xyxy[0][2]),
"y2": int(box.xyxy[0][3]),
}
})
# Draw boxes on image and return as base64
annotated = results.plot()
_, buffer = cv2.imencode(".jpg", annotated)
annotated_b64 = base64.b64encode(buffer).decode("utf-8")
return JSONResponse({
"detections": detections,
"count": len(detections),
"image_shape": {
"width": img.shape[1],
"height": img.shape[0],
},
"settings": {
"confidence": confidence,
"overlap": overlap,
},
"annotated_image": f"data:image/jpeg;base64,{annotated_b64}"
})
@app.get("/", response_class=HTMLResponse)
def ui():
classes = list(model.names.values())
classes_str = ", ".join(f'{c}' for c in classes)
return f"""
Floor Plan Detector
🏠 Floor Plan Detector
YOLOv8s · Trained on 14,274 floor plan images
Upload Image
![]()
Annotated Result
![]()
"""