Spaces:
Running
Running
| 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") | |
| 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}" | |
| }) | |
| def ui(): | |
| classes = list(model.names.values()) | |
| classes_str = ", ".join(f'<span class="tag">{c}</span>' for c in classes) | |
| return f"""<!DOCTYPE html> | |
| <html lang="en"> | |
| <head> | |
| <meta charset="UTF-8"> | |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> | |
| <title>Floor Plan Detector</title> | |
| <style> | |
| *, *::before, *::after {{ box-sizing: border-box; margin: 0; padding: 0; }} | |
| body {{ | |
| font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; | |
| background: #f0f2f5; | |
| min-height: 100vh; | |
| display: flex; | |
| flex-direction: column; | |
| align-items: center; | |
| padding: 32px 16px; | |
| color: #1a1a2e; | |
| }} | |
| h1 {{ font-size: 1.6rem; font-weight: 700; margin-bottom: 4px; }} | |
| .subtitle {{ color: #666; font-size: 0.9rem; margin-bottom: 24px; }} | |
| .card {{ | |
| background: white; | |
| border-radius: 16px; | |
| padding: 28px; | |
| width: 100%; | |
| max-width: 720px; | |
| box-shadow: 0 4px 24px rgba(0,0,0,0.08); | |
| margin-bottom: 20px; | |
| }} | |
| .card h2 {{ font-size: 1rem; font-weight: 600; margin-bottom: 16px; color: #444; text-transform: uppercase; letter-spacing: 0.05em; }} | |
| .drop-zone {{ | |
| border: 2px dashed #c5c9d6; | |
| border-radius: 12px; | |
| padding: 40px 20px; | |
| text-align: center; | |
| cursor: pointer; | |
| transition: all 0.2s; | |
| background: #fafbfc; | |
| position: relative; | |
| }} | |
| .drop-zone:hover, .drop-zone.drag-over {{ border-color: #4f46e5; background: #f5f3ff; }} | |
| .drop-zone input {{ position: absolute; inset: 0; opacity: 0; cursor: pointer; width: 100%; height: 100%; }} | |
| .drop-icon {{ font-size: 2rem; margin-bottom: 8px; }} | |
| .drop-text {{ color: #666; font-size: 0.95rem; }} | |
| .drop-text strong {{ color: #4f46e5; }} | |
| #preview {{ max-width: 100%; border-radius: 8px; margin-top: 16px; display: none; }} | |
| .slider-row {{ | |
| display: flex; | |
| align-items: center; | |
| gap: 12px; | |
| margin-bottom: 14px; | |
| }} | |
| .slider-label {{ width: 110px; font-size: 0.9rem; color: #555; font-weight: 500; }} | |
| input[type=range] {{ flex: 1; accent-color: #4f46e5; cursor: pointer; }} | |
| .slider-val {{ | |
| width: 42px; | |
| text-align: right; | |
| font-size: 0.9rem; | |
| font-weight: 600; | |
| color: #4f46e5; | |
| }} | |
| .btn {{ | |
| width: 100%; | |
| padding: 14px; | |
| background: #4f46e5; | |
| color: white; | |
| border: none; | |
| border-radius: 10px; | |
| font-size: 1rem; | |
| font-weight: 600; | |
| cursor: pointer; | |
| transition: background 0.2s; | |
| margin-top: 8px; | |
| }} | |
| .btn:hover {{ background: #4338ca; }} | |
| .btn:disabled {{ background: #a5b4fc; cursor: not-allowed; }} | |
| #result-img {{ max-width: 100%; border-radius: 10px; display: none; }} | |
| .stats {{ | |
| display: flex; | |
| gap: 16px; | |
| margin-bottom: 16px; | |
| flex-wrap: wrap; | |
| }} | |
| .stat {{ | |
| background: #f5f3ff; | |
| border-radius: 8px; | |
| padding: 10px 18px; | |
| text-align: center; | |
| }} | |
| .stat-val {{ font-size: 1.5rem; font-weight: 700; color: #4f46e5; }} | |
| .stat-lbl {{ font-size: 0.75rem; color: #888; margin-top: 2px; }} | |
| .detections {{ display: flex; flex-direction: column; gap: 8px; }} | |
| .det-row {{ | |
| display: flex; | |
| align-items: center; | |
| justify-content: space-between; | |
| background: #f9fafb; | |
| border-radius: 8px; | |
| padding: 10px 14px; | |
| font-size: 0.9rem; | |
| }} | |
| .det-class {{ font-weight: 600; }} | |
| .det-conf {{ | |
| background: #4f46e5; | |
| color: white; | |
| border-radius: 20px; | |
| padding: 2px 10px; | |
| font-size: 0.8rem; | |
| font-weight: 600; | |
| }} | |
| .tags {{ display: flex; flex-wrap: wrap; gap: 6px; margin-top: 4px; }} | |
| .tag {{ | |
| background: #ede9fe; | |
| color: #5b21b6; | |
| border-radius: 20px; | |
| padding: 3px 10px; | |
| font-size: 0.8rem; | |
| font-weight: 500; | |
| }} | |
| .error {{ color: #dc2626; background: #fef2f2; border-radius: 8px; padding: 12px 16px; font-size: 0.9rem; }} | |
| .hidden {{ display: none; }} | |
| #result-section {{ display: none; }} | |
| .spinner {{ | |
| width: 20px; height: 20px; | |
| border: 3px solid rgba(255,255,255,0.4); | |
| border-top-color: white; | |
| border-radius: 50%; | |
| animation: spin 0.8s linear infinite; | |
| display: inline-block; | |
| vertical-align: middle; | |
| margin-right: 8px; | |
| }} | |
| @keyframes spin {{ to {{ transform: rotate(360deg); }} }} | |
| </style> | |
| </head> | |
| <body> | |
| <h1>🏠 Floor Plan Detector</h1> | |
| <p class="subtitle">YOLOv8s · Trained on 14,274 floor plan images</p> | |
| <div class="card"> | |
| <h2>Upload Image</h2> | |
| <div class="drop-zone" id="dropZone"> | |
| <input type="file" id="fileInput" accept="image/*"> | |
| <div class="drop-icon">📐</div> | |
| <div class="drop-text">Drop a floor plan image or <strong>click to browse</strong></div> | |
| </div> | |
| <img id="preview"> | |
| </div> | |
| <div class="card"> | |
| <h2>Settings</h2> | |
| <div class="slider-row"> | |
| <span class="slider-label">Confidence</span> | |
| <input type="range" id="conf" min="1" max="99" value="50"> | |
| <span class="slider-val" id="confVal">0.50</span> | |
| </div> | |
| <div class="slider-row"> | |
| <span class="slider-label">Overlap (IoU)</span> | |
| <input type="range" id="overlap" min="1" max="99" value="50"> | |
| <span class="slider-val" id="overlapVal">0.50</span> | |
| </div> | |
| <div style="margin-top:8px"> | |
| <span style="font-size:0.85rem;color:#888;">Detectable classes: </span> | |
| <div class="tags" style="margin-top:6px">{classes_str}</div> | |
| </div> | |
| <button class="btn" id="detectBtn" onclick="detect()" disabled>Select an image first</button> | |
| </div> | |
| <div id="result-section"> | |
| <div class="card"> | |
| <h2>Annotated Result</h2> | |
| <img id="result-img"> | |
| </div> | |
| <div class="card"> | |
| <h2>Detections</h2> | |
| <div class="stats"> | |
| <div class="stat"><div class="stat-val" id="countVal">0</div><div class="stat-lbl">Objects Found</div></div> | |
| <div class="stat"><div class="stat-val" id="confSetting">0.50</div><div class="stat-lbl">Confidence Used</div></div> | |
| </div> | |
| <div class="detections" id="detList"></div> | |
| <div class="error hidden" id="errBox"></div> | |
| </div> | |
| </div> | |
| <script> | |
| const confSlider = document.getElementById('conf'); | |
| const overlapSlider = document.getElementById('overlap'); | |
| const confVal = document.getElementById('confVal'); | |
| const overlapVal = document.getElementById('overlapVal'); | |
| const fileInput = document.getElementById('fileInput'); | |
| const preview = document.getElementById('preview'); | |
| const btn = document.getElementById('detectBtn'); | |
| const dropZone = document.getElementById('dropZone'); | |
| confSlider.oninput = () => confVal.textContent = (confSlider.value / 100).toFixed(2); | |
| overlapSlider.oninput = () => overlapVal.textContent = (overlapSlider.value / 100).toFixed(2); | |
| dropZone.addEventListener('dragover', e => {{ e.preventDefault(); dropZone.classList.add('drag-over'); }}); | |
| dropZone.addEventListener('dragleave', () => dropZone.classList.remove('drag-over')); | |
| dropZone.addEventListener('drop', e => {{ | |
| e.preventDefault(); | |
| dropZone.classList.remove('drag-over'); | |
| if (e.dataTransfer.files[0]) loadFile(e.dataTransfer.files[0]); | |
| }}); | |
| fileInput.onchange = () => {{ if (fileInput.files[0]) loadFile(fileInput.files[0]); }}; | |
| function loadFile(file) {{ | |
| const reader = new FileReader(); | |
| reader.onload = e => {{ | |
| preview.src = e.target.result; | |
| preview.style.display = 'block'; | |
| btn.disabled = false; | |
| btn.textContent = 'Run Detection'; | |
| }}; | |
| reader.readAsDataURL(file); | |
| }} | |
| async function detect() {{ | |
| const file = fileInput.files[0]; | |
| if (!file) return; | |
| btn.disabled = true; | |
| btn.innerHTML = '<span class="spinner"></span>Detecting...'; | |
| document.getElementById('result-section').style.display = 'none'; | |
| document.getElementById('errBox').classList.add('hidden'); | |
| const fd = new FormData(); | |
| fd.append('file', file); | |
| fd.append('confidence', confSlider.value / 100); | |
| fd.append('overlap', overlapSlider.value / 100); | |
| try {{ | |
| const res = await fetch('/detect', {{ method: 'POST', body: fd }}); | |
| const data = await res.json(); | |
| if (data.error) throw new Error(data.error); | |
| // Show annotated image | |
| document.getElementById('result-img').src = data.annotated_image; | |
| document.getElementById('result-img').style.display = 'block'; | |
| // Stats | |
| document.getElementById('countVal').textContent = data.count; | |
| document.getElementById('confSetting').textContent = data.settings.confidence.toFixed(2); | |
| // Detection list | |
| const list = document.getElementById('detList'); | |
| if (data.detections.length === 0) {{ | |
| list.innerHTML = '<div class="det-row" style="color:#888">No objects detected — try lowering the confidence threshold</div>'; | |
| }} else {{ | |
| list.innerHTML = data.detections | |
| .sort((a, b) => b.confidence - a.confidence) | |
| .map(d => ` | |
| <div class="det-row"> | |
| <span class="det-class">${{d.class}}</span> | |
| <span class="det-conf">${{(d.confidence * 100).toFixed(1)}}%</span> | |
| </div>`) | |
| .join(''); | |
| }} | |
| document.getElementById('result-section').style.display = 'block'; | |
| }} catch (err) {{ | |
| const errBox = document.getElementById('errBox'); | |
| errBox.textContent = 'Error: ' + err.message; | |
| errBox.classList.remove('hidden'); | |
| document.getElementById('result-section').style.display = 'block'; | |
| }} finally {{ | |
| btn.disabled = false; | |
| btn.textContent = 'Run Detection'; | |
| }} | |
| }} | |
| </script> | |
| </body> | |
| </html>""" |