Spaces:

intisarhasnain
/

floorplan-detect

Running

File size: 11,927 Bytes

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'<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>"""