mvp

.env.example

@@ -1,3 +0,0 @@
-MODEL_DIR=models
-MODEL_FILE=model.pt
-GRADIO_WATCH_DIRS="planparser,src"

README.md

@@ -0,0 +1 @@
+[Lind for data downloading](https://universe.roboflow.com/research-g8szb/floorplan-details-fork)

mvp/api.py

@@ -0,0 +1,102 @@
+# planparser/api.py
+import os
+import io
+from glob import glob
+
+from ultralytics import YOLO
+from PIL import Image
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from pydantic import BaseModel
+from dotenv import load_dotenv
+
+load_dotenv()
+
+MODEL_DIR = os.getenv("MODEL_DIR")
+if not MODEL_DIR:
+    raise RuntimeError("MODEL_DIR is not set")
+MODEL_DIR = os.path.abspath(os.path.expanduser(MODEL_DIR))
+
+app = FastAPI(title="Model API")
+_models: dict[str, YOLO] = {}
+
+
+def _pick_weights(path: str) -> str:
+    if os.path.isfile(path):
+        return path
+    if not os.path.isdir(path):
+        raise HTTPException(status_code=400, detail=f"model not found: {path}")
+
+    for name in ["model.pt", "best.pt"]:
+        cand = os.path.join(path, name)
+        if os.path.isfile(cand):
+            return cand
+
+    pts = sorted(glob(os.path.join(path, "*.pt")))
+    if pts:
+        return pts[0]
+
+    raise HTTPException(status_code=400, detail=f"no .pt weights in: {path}")
+
+
+def _resolve_weights_path(model_name: str) -> str:
+    if not model_name:
+        raise HTTPException(status_code=400, detail="model_name is required")
+
+    p = os.path.abspath(os.path.join(MODEL_DIR, model_name))
+    if os.path.commonpath([MODEL_DIR, p]) != MODEL_DIR:
+        raise HTTPException(status_code=400, detail="model_name must be inside MODEL_DIR")
+
+    return _pick_weights(p)
+
+
+def load_model(weights_path: str) -> YOLO:
+    weights_path = os.path.abspath(weights_path)
+    m = _models.get(weights_path)
+    if m is None:
+        m = YOLO(weights_path)
+        _models[weights_path] = m
+    return m
+
+
+class Detection(BaseModel):
+    class_id: int
+    class_name: str
+    confidence: float
+    xyxy: list[float]
+
+
+class PredictResponse(BaseModel):
+    detections: list[Detection]
+
+
+@app.get("/health")
+def health():
+    return {"ok": True}
+
+
+@app.post("/predict", response_model=PredictResponse)
+async def predict(
+    file: UploadFile = File(...),
+    model_name: str = Form(...),
+):
+    weights_path = _resolve_weights_path(model_name)
+
+    raw = await file.read()
+    img = Image.open(io.BytesIO(raw)).convert("RGB")
+
+    r0 = load_model(weights_path)(img)[0]
+    boxes = getattr(r0, "boxes", None)
+    if boxes is None:
+        raise HTTPException(status_code=500, detail="model output has no boxes (not a detection model?)")
+
+    names = getattr(r0, "names", {}) or {}
+    dets: list[Detection] = []
+    for cls_t, conf_t, xyxy_t in zip(boxes.cls, boxes.conf, boxes.xyxy):
+        cls_id = int(cls_t.item())
+        conf = float(conf_t.item())
+        xyxy = [float(x) for x in xyxy_t.tolist()]
+        cls_name = names.get(cls_id, str(cls_id)) if isinstance(names, dict) else str(cls_id)
+        dets.append(Detection(class_id=cls_id, class_name=cls_name, confidence=conf, xyxy=xyxy))
+
+    return PredictResponse(detections=dets)
+# uv run uvicorn planparser.api:app --host 0.0.0.0 --port 8000

mvp/app.py

@@ -0,0 +1,103 @@
+# planparser/app.py
+import os
+import io
+import time
+
+import gradio as gr
+import requests
+from PIL import Image, ImageDraw
+from dotenv import load_dotenv
+
+load_dotenv()
+
+API_URL = os.getenv("API_URL")
+if not API_URL:
+    raise RuntimeError("API_URL is not set")
+
+EXAMPLE_1 = os.getenv("EXAMPLE_1")
+EXAMPLE_2 = os.getenv("EXAMPLE_2")
+
+MODEL_MAP = {
+    "yolo11n": os.getenv("MODEL_1"),
+    "yolo11s": os.getenv("MODEL_2"),
+    "yolo11l": os.getenv("MODEL_3"),
+}
+
+MODEL_CHOICES = list(MODEL_MAP.keys())
+DEFAULT_MODEL = MODEL_CHOICES[0]
+
+
+def _draw_detections(img: Image.Image, dets: list[dict]) -> Image.Image:
+    out = img.copy().convert("RGB")
+    d = ImageDraw.Draw(out)
+
+    for det in dets:
+        x1, y1, x2, y2 = det["xyxy"]
+        d.rectangle([x1, y1, x2, y2], width=2)
+        txt = f'{det["class_name"]} {det["confidence"]:.2f}'
+        d.text((x1, max(0, y1 - 12)), txt)
+
+    return out
+
+
+def predict(model_label: str, img: Image.Image):
+    if img is None or not model_label:
+        return None, [], ""
+
+    buf = io.BytesIO()
+    img.convert("RGB").save(buf, format="JPEG")
+    buf.seek(0)
+
+    t0 = time.perf_counter()
+    r = requests.post(
+        f"{API_URL}/predict",
+        files={"file": ("image.jpg", buf, "image/jpeg")},
+        data={"model_name": MODEL_MAP[model_label]},
+        timeout=60,
+    )
+    dt = time.perf_counter() - t0
+
+    if r.status_code != 200:
+        return None, {"error": r.text}, f"{dt:.3f} s"
+
+    data = r.json()
+    dets = data.get("detections", []) or []
+    vis = _draw_detections(img, dets)
+
+    return vis, dets, f"{dt:.3f} s"
+
+
+with gr.Blocks(title="YOLO detection demo") as demo:
+    gr.Markdown("# YOLO detection demo")
+
+    model_dd = gr.Dropdown(
+        choices=MODEL_CHOICES,
+        value=DEFAULT_MODEL,
+        label="Model",
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            img_in = gr.Image(
+                type="pil",
+                label="Image",
+                sources=["upload"],
+                height=200,
+            )
+
+            ex = [p for p in [EXAMPLE_1, EXAMPLE_2] if p]
+            if ex:
+                gr.Examples(examples=ex, inputs=img_in)
+
+            btn = gr.Button("Submit")
+
+        with gr.Column(scale=2):
+            out_img = gr.Image(type="pil", label="Result", height=600)
+            out_json = gr.JSON(label="Detections")
+            out_time = gr.Textbox(label="Processing time")
+
+    btn.click(predict, inputs=[model_dd, img_in], outputs=[out_img, out_json, out_time])
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)
+# uv run gradio planparser/app.py

mvp/utils/images.py

@@ -0,0 +1,17 @@
+from pathlib import Path
+from PIL import Image
+
+def to_white_bg(in_path: str | Path, out_path: str | Path) -> None:
+    in_path = Path(in_path)
+    out_path = Path(out_path)
+
+    im = Image.open(in_path)
+    if im.mode in ("RGBA", "LA") or ("transparency" in im.info):
+        im = im.convert("RGBA")
+        white = Image.new("RGBA", im.size, (255, 255, 255, 255))
+        out = Image.alpha_composite(white, im).convert("RGB")
+    else:
+        out = im.convert("RGB")
+
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out.save(out_path)

planparser/api.py

@@ -0,0 +1,102 @@
+# planparser/api.py
+import os
+import io
+from glob import glob
+
+from ultralytics import YOLO
+from PIL import Image
+from fastapi import FastAPI, UploadFile, File, Form, HTTPException
+from pydantic import BaseModel
+from dotenv import load_dotenv
+
+load_dotenv()
+
+MODEL_DIR = os.getenv("MODEL_DIR")
+if not MODEL_DIR:
+    raise RuntimeError("MODEL_DIR is not set")
+MODEL_DIR = os.path.abspath(os.path.expanduser(MODEL_DIR))
+
+app = FastAPI(title="Model API")
+_models: dict[str, YOLO] = {}
+
+
+def _pick_weights(path: str) -> str:
+    if os.path.isfile(path):
+        return path
+    if not os.path.isdir(path):
+        raise HTTPException(status_code=400, detail=f"model not found: {path}")
+
+    for name in ["model.pt", "best.pt"]:
+        cand = os.path.join(path, name)
+        if os.path.isfile(cand):
+            return cand
+
+    pts = sorted(glob(os.path.join(path, "*.pt")))
+    if pts:
+        return pts[0]
+
+    raise HTTPException(status_code=400, detail=f"no .pt weights in: {path}")
+
+
+def _resolve_weights_path(model_name: str) -> str:
+    if not model_name:
+        raise HTTPException(status_code=400, detail="model_name is required")
+
+    p = os.path.abspath(os.path.join(MODEL_DIR, model_name))
+    if os.path.commonpath([MODEL_DIR, p]) != MODEL_DIR:
+        raise HTTPException(status_code=400, detail="model_name must be inside MODEL_DIR")
+
+    return _pick_weights(p)
+
+
+def load_model(weights_path: str) -> YOLO:
+    weights_path = os.path.abspath(weights_path)
+    m = _models.get(weights_path)
+    if m is None:
+        m = YOLO(weights_path)
+        _models[weights_path] = m
+    return m
+
+
+class Detection(BaseModel):
+    class_id: int
+    class_name: str
+    confidence: float
+    xyxy: list[float]
+
+
+class PredictResponse(BaseModel):
+    detections: list[Detection]
+
+
+@app.get("/health")
+def health():
+    return {"ok": True}
+
+
+@app.post("/predict", response_model=PredictResponse)
+async def predict(
+    file: UploadFile = File(...),
+    model_name: str = Form(...),
+):
+    weights_path = _resolve_weights_path(model_name)
+
+    raw = await file.read()
+    img = Image.open(io.BytesIO(raw)).convert("RGB")
+
+    r0 = load_model(weights_path)(img)[0]
+    boxes = getattr(r0, "boxes", None)
+    if boxes is None:
+        raise HTTPException(status_code=500, detail="model output has no boxes (not a detection model?)")
+
+    names = getattr(r0, "names", {}) or {}
+    dets: list[Detection] = []
+    for cls_t, conf_t, xyxy_t in zip(boxes.cls, boxes.conf, boxes.xyxy):
+        cls_id = int(cls_t.item())
+        conf = float(conf_t.item())
+        xyxy = [float(x) for x in xyxy_t.tolist()]
+        cls_name = names.get(cls_id, str(cls_id)) if isinstance(names, dict) else str(cls_id)
+        dets.append(Detection(class_id=cls_id, class_name=cls_name, confidence=conf, xyxy=xyxy))
+
+    return PredictResponse(detections=dets)
+# uv run uvicorn planparser.api:app --host 0.0.0.0 --port 8000

planparser/app.py

@@ -1,42 +1,103 @@
+# planparser/app.py
 import os
+import io
+import time
+
 import gradio as gr
+import requests
+from PIL import Image, ImageDraw
+from dotenv import load_dotenv
+
+load_dotenv()
+
+API_URL = os.getenv("API_URL")
+if not API_URL:
+    raise RuntimeError("API_URL is not set")
+
+EXAMPLE_1 = os.getenv("EXAMPLE_1")
+EXAMPLE_2 = os.getenv("EXAMPLE_2")
+
+MODEL_MAP = {
+    "yolo11n": os.getenv("MODEL_1"),
+    "yolo11s": os.getenv("MODEL_2"),
+    "yolo11l": os.getenv("MODEL_3"),
+}
+
+MODEL_CHOICES = list(MODEL_MAP.keys())
+DEFAULT_MODEL = MODEL_CHOICES[0]
+
+
+def _draw_detections(img: Image.Image, dets: list[dict]) -> Image.Image:
+    out = img.copy().convert("RGB")
+    d = ImageDraw.Draw(out)
+
+    for det in dets:
+        x1, y1, x2, y2 = det["xyxy"]
+        d.rectangle([x1, y1, x2, y2], width=2)
+        txt = f'{det["class_name"]} {det["confidence"]:.2f}'
+        d.text((x1, max(0, y1 - 12)), txt)
+
+    return out
+
+
+def predict(model_label: str, img: Image.Image):
+    if img is None or not model_label:
+        return None, [], ""
+
+    buf = io.BytesIO()
+    img.convert("RGB").save(buf, format="JPEG")
+    buf.seek(0)
+
+    t0 = time.perf_counter()
+    r = requests.post(
+        f"{API_URL}/predict",
+        files={"file": ("image.jpg", buf, "image/jpeg")},
+        data={"model_name": MODEL_MAP[model_label]},
+        timeout=60,
+    )
+    dt = time.perf_counter() - t0
 
-MODEL_DIR = os.getenv("MODEL_DIR", "models")
-MODEL_FILE = os.getenv("MODEL_FILE", "model.pt")
+    if r.status_code != 200:
+        return None, {"error": r.text}, f"{dt:.3f} s"
 
-_model = None
+    data = r.json()
+    dets = data.get("detections", []) or []
+    vis = _draw_detections(img, dets)
 
+    return vis, dets, f"{dt:.3f} s"
 
-def get_model_path() -> str:
-    path = os.path.join(MODEL_DIR, MODEL_FILE)
-    if not os.path.exists(path):
-        raise FileNotFoundError(
-            f"Model not found: {path}. Put the file there or set MODEL_DIR and MODEL_FILE."
-        )
-    return path
 
+with gr.Blocks(title="YOLO detection demo") as demo:
+    gr.Markdown("# YOLO detection demo")
 
-def load_model():
-    global _model
-    if _model is not None:
-        return _model
-    import torch
+    model_dd = gr.Dropdown(
+        choices=MODEL_CHOICES,
+        value=DEFAULT_MODEL,
+        label="Model",
+    )
 
-    _model = torch.jit.load(get_model_path(), map_location="cpu").eval()
-    return _model
+    with gr.Row():
+        with gr.Column(scale=1):
+            img_in = gr.Image(
+                type="pil",
+                label="Image",
+                sources=["upload"],
+                height=200,
+            )
 
+            ex = [p for p in [EXAMPLE_1, EXAMPLE_2] if p]
+            if ex:
+                gr.Examples(examples=ex, inputs=img_in)
 
-def predict(text: str) -> str:
-    _ = load_model()
-    return f"Loaded local model: {os.path.join(MODEL_DIR, MODEL_FILE)}. Input: {text}"
+            btn = gr.Button("Submit")
 
+        with gr.Column(scale=2):
+            out_img = gr.Image(type="pil", label="Result", height=600)
+            out_json = gr.JSON(label="Detections")
+            out_time = gr.Textbox(label="Processing time")
 
-demo = gr.Interface(
-    fn=predict,
-    inputs=gr.Textbox(label="Input"),
-    outputs=gr.Textbox(label="Output"),
-    title="Minimal demo",
-)
+    btn.click(predict, inputs=[model_dd, img_in], outputs=[out_img, out_json, out_time])
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)
+# uv run gradio planparser/app.py

pyproject.toml

@@ -6,6 +6,14 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "gradio>=6.2.0",
+    "requests>=2.32.5",
     "torch>=2.9.1",
     "torchvision>=0.24.1",
+    "ultralytics>=8.3.243",
+]
+
+[dependency-groups]
+dev = [
+    "jupyterlab>=4.5.1",
+    "python-dotenv>=1.2.1",
 ]

src/data/data.yaml

@@ -0,0 +1,13 @@
+train: ../train/images
+val: ../valid/images
+test: ../test/images
+
+nc: 15
+names: ['bathtub', 'bed', 'bed2', 'chair', 'door', 'door2', 'shower', 'sink', 'sofa1', 'sofa2', 'sofa3', 'stove', 'table', 'toilet', 'vanity']
+
+roboflow:
+  workspace: research-g8szb
+  project: floorplan-details-fork
+  version: 1
+  license: CC BY 4.0
+  url: https://universe.roboflow.com/research-g8szb/floorplan-details-fork/dataset/1