Model path check in App

planparser/api.py

@@ -1,56 +1,18 @@
-# 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")
+app = FastAPI(title="Planparser 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)
+    weights_path = os.path.abspath(os.path.expanduser(weights_path))
     m = _models.get(weights_path)
     if m is None:
         m = YOLO(weights_path)
@@ -77,9 +39,11 @@ def health():
 @app.post("/predict", response_model=PredictResponse)
 async def predict(
     file: UploadFile = File(...),
-    model_name: str = Form(...),
+    weights_path: str = Form(...),
 ):
-    weights_path = _resolve_weights_path(model_name)
+    weights_path = os.path.abspath(os.path.expanduser(weights_path))
+    if not (os.path.isfile(weights_path) and weights_path.lower().endswith(".pt")):
+        raise HTTPException(status_code=400, detail=f"weights_path must be an existing .pt file: {weights_path}")
 
     raw = await file.read()
     img = Image.open(io.BytesIO(raw)).convert("RGB")
@@ -99,4 +63,3 @@ async def predict(
         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

@@ -21,15 +21,27 @@ if not API_URL:
 
 EXAMPLES_DIR = os.getenv("EXAMPLES_DIR")
 
+MODEL_DIR = os.getenv("MODEL_DIR")
+MODEL_1 = os.getenv("MODEL_1")
+MODEL_2 = os.getenv("MODEL_2")
+
+def join_pt(folder: str | None, name: str | None) -> str | None:
+    if not folder or not name:
+        return None
+    p = (Path(folder).expanduser() / name).resolve()
+    if p.is_file() and p.suffix.lower() == ".pt":
+        return str(p)
+    return None
+
 MODEL_MAP = {
-    "yolo11l_custom": os.getenv("MODEL_1"),
-    "custom": os.getenv("MODEL_2"),
+    "yolo11l_custom": join_pt(MODEL_DIR, MODEL_1),
+    "custom": join_pt(MODEL_DIR, MODEL_2),
 }
 MODEL_MAP = {k: v for k, v in MODEL_MAP.items() if v}
 
 MODEL_CHOICES = list(MODEL_MAP.keys())
-if len(MODEL_CHOICES) != 2:
-    raise RuntimeError("Exactly two models must be set: MODEL_1 and MODEL_2")
+if not MODEL_CHOICES:
+    raise RuntimeError("No valid .pt models found via MODEL_DIR + MODEL_1/MODEL_2")
 
 DEFAULT_MODEL = MODEL_CHOICES[0]
 
@@ -80,8 +92,8 @@ def _collect_example_images(max_n: int = 30) -> list[str]:
     if not files:
         return []
 
-    k = min(max_n, len(files))
-    return [str(x) for x in random.sample(files, k=k)]
+    random.shuffle(files)
+    return [str(x) for x in files[:min(max_n, len(files))]]
 
 
 def _hex2rgb(h: str) -> tuple[int, int, int]:
@@ -154,7 +166,7 @@ def _request_predict(model_label: str, img: Image.Image) -> tuple[list[dict], fl
     r = requests.post(
         f"{API_URL}/predict",
         files={"file": ("image.jpg", buf, "image/jpeg")},
-        data={"model_name": MODEL_MAP[model_label]},
+        data={"weights_path": MODEL_MAP[model_label]},
         timeout=60,
     )
     dt = time.perf_counter() - t0
@@ -220,7 +232,7 @@ def maybe_autorun(model_label: str, img: Image.Image, auto_run: bool):
     return run_predict(model_label, img)
 
 
-with gr.Blocks(title="Architectural plan elements detection") as demo:
+with gr.Blocks(title="Planparser") as demo:
     gr.Markdown("# Architectural plan elements detection")
 
     with gr.Row():
@@ -281,5 +293,3 @@ with gr.Blocks(title="Architectural plan elements detection") as demo:
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)
-
-# uv run gradio planparser/app.py