fasterrcnn_resnet50 model

planparser/api.py

@@ -1,25 +1,50 @@
 import os
 import io
+from typing import Any
 
+import torch
 from ultralytics import YOLO
-from PIL import Image
+from PIL import Image, UnidentifiedImageError
 from fastapi import FastAPI, UploadFile, File, Form, HTTPException
 from pydantic import BaseModel
+from torchvision.transforms.functional import pil_to_tensor as tv_pil_to_tensor
 
 
 app = FastAPI(title="Planparser API")
-_models: dict[str, YOLO] = {}
 
+_MODEL_CACHE: dict[tuple[str, str], Any] = {}
 
-def load_model(weights_path: str) -> YOLO:
-    weights_path = os.path.abspath(os.path.expanduser(weights_path))
-    m = _models.get(weights_path)
-    if m is None:
-        m = YOLO(weights_path)
-        _models[weights_path] = m
+
+def _abs_pt(p: str) -> str:
+    return os.path.abspath(os.path.expanduser(p))
+
+
+def get_model(model_type: str, weights_path: str):
+    wp = _abs_pt(weights_path)
+    key = (model_type, wp)
+
+    m = _MODEL_CACHE.get(key)
+    if m is not None:
+        return m
+
+    if model_type == "yolo":
+        m = YOLO(wp)
+
+    elif model_type == "fasterrcnn":
+        m = torch.jit.load(wp, map_location="cpu")
+        m.eval()
+
+    else:
+        raise HTTPException(status_code=400, detail="model_type must be yolo or fasterrcnn")
+
+    _MODEL_CACHE[key] = m
     return m
 
 
+def pil_to_tensor(img: Image.Image) -> torch.Tensor:
+    return tv_pil_to_tensor(img.convert("RGB")).float().div(255.0)
+
+
 class Detection(BaseModel):
     class_id: int
     class_name: str
@@ -40,26 +65,85 @@ def health():
 async def predict(
     file: UploadFile = File(...),
     weights_path: str = Form(...),
+    model_type: str = Form("yolo"),
+    conf: float = Form(0.25),
 ):
-    weights_path = os.path.abspath(os.path.expanduser(weights_path))
+    weights_path = _abs_pt(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")
-
-    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)
+    if not raw:
+        raise HTTPException(status_code=400, detail="empty file")
+
+    try:
+        img = Image.open(io.BytesIO(raw)).convert("RGB")
+    except UnidentifiedImageError:
+        raise HTTPException(status_code=400, detail="invalid image file")
+
+    if model_type == "yolo":
+        model = get_model("yolo", weights_path)
+        with torch.inference_mode():
+            r0 = model.predict(img, device="cpu", verbose=False)[0]
+
+        boxes = getattr(r0, "boxes", None)
+        if boxes is None:
+            raise HTTPException(status_code=500, detail="model output has no boxes")
+
+        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())
+            c = float(conf_t.item())
+            if c < float(conf):
+                continue
+            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=c, xyxy=xyxy))
+
+        return PredictResponse(detections=dets)
+
+    if model_type == "fasterrcnn":
+        try:
+            x = pil_to_tensor(img)
+            model = get_model("fasterrcnn", weights_path)
+
+            get_names = getattr(model, "get_class_names", None)
+            class_names = list(get_names()) if callable(get_names) else []
+
+            with torch.inference_mode():
+                out = model([x])
+
+            boxes = out["boxes"].detach().cpu()
+            scores = out["scores"].detach().cpu()
+            labels = out["labels"].detach().cpu()
+
+            dets: list[Detection] = []
+            for b, s, l in zip(boxes, scores, labels):
+                c = float(s)
+                if c < float(conf):
+                    continue
+
+                cid = int(l)
+                if cid == 0:
+                    continue
+
+                cname = class_names[cid] if 0 <= cid < len(class_names) else str(cid)
+
+                dets.append(
+                    Detection(
+                        class_id=cid,
+                        class_name=cname,
+                        confidence=c,
+                        xyxy=[float(v) for v in b.tolist()],
+                    )
+                )
+
+            return PredictResponse(detections=dets)
+
+        except HTTPException:
+            raise
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=f"fasterrcnn failed: {type(e).__name__}: {e}")
+
+    raise HTTPException(status_code=400, detail="model_type must be yolo or fasterrcnn")

planparser/app.py

@@ -1,4 +1,3 @@
-# planparser/app.py
 import os
 import io
 import time
@@ -61,10 +60,13 @@ RESOLVED_MODEL_DIR = _resolve_models_dir()
 
 MODEL_MAP = {
     "yolo11l_custom": join_pt(RESOLVED_MODEL_DIR, MODEL_1),
-    "custom": join_pt(RESOLVED_MODEL_DIR, MODEL_2),
+    "fasterrcnn_resnet50.pt": join_pt(RESOLVED_MODEL_DIR, MODEL_2),
 }
 MODEL_MAP = {k: v for k, v in MODEL_MAP.items() if v}
 
+def _model_type(model_label: str) -> str:
+    return "fasterrcnn" if "fasterrcnn" in model_label.lower() else "yolo"
+
 MODEL_CHOICES = list(MODEL_MAP.keys())
 if not MODEL_CHOICES:
     raise RuntimeError("No valid .pt models found via MODEL_DIR + MODEL_1/MODEL_2")
@@ -130,21 +132,24 @@ def _pretty_name(raw: str) -> str:
     return CLASS_NAME_MAP.get(raw, raw)
 
 
-def _color_for_det(det: dict) -> tuple[int, int, int]:
+def _color_for_det(det: dict, model_type: str) -> tuple[int, int, int]:
     cls_id = det.get("class_id", None)
     if cls_id is None:
         name = det.get("class_name", "")
         cls_id = abs(hash(name))
+    else:
+        if model_type == "fasterrcnn":
+            cls_id = cls_id - 1
     return _hex2rgb(_PALETTE[int(cls_id) % len(_PALETTE)])
 
 
-def _draw_detections(img: Image.Image, dets: list[dict]) -> Image.Image:
+def _draw_detections(img: Image.Image, dets: list[dict], model_type: str) -> Image.Image:
     out = img.copy().convert("RGB")
     d = ImageDraw.Draw(out)
 
     for det in dets:
         x1, y1, x2, y2 = det["xyxy"]
-        color = _color_for_det(det)
+        color = _color_for_det(det, model_type)
 
         d.rectangle([x1, y1, x2, y2], outline=color, width=2)
 
@@ -193,13 +198,16 @@ 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={"weights_path": MODEL_MAP[model_label]},
+        data={
+            "weights_path": MODEL_MAP[model_label],
+            "model_type": _model_type(model_label),
+        },
         timeout=60,
     )
     dt = time.perf_counter() - t0
 
     if r.status_code != 200:
-        return [{"error": r.text}], dt, "error"
+        return [{"error": r.text}], dt, r.text
 
     data = r.json()
     dets = data.get("detections", []) or []
@@ -225,14 +233,16 @@ def run_predict(model_label: str, img: Image.Image):
     if err is not None:
         return (
             None,
-            time_md,
+            f"_processing time: {dt:.3f} s_\n\n**API error:**\n{err}",
             empty_df,
             dets,
             gr.update(value=None, visible=False),
-            gr.update(visible=False),
+            gr.update(visible=True, open=True),
         )
 
-    vis = _draw_detections(img, dets)
+    mt = _model_type(model_label)
+    vis = _draw_detections(img, dets, mt)
+
     df_counts = _counts_df(dets)
     csv_path = export_df(df_counts)
 

pyproject.toml

@@ -17,5 +17,7 @@ dependencies = [
 
 [dependency-groups]
 dev = [
+    "albumentations>=2.0.8",
+    "ipywidgets>=8.1.8",
     "jupyterlab>=4.5.1",
 ]

src/models/fasterrcnn_resnet50.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:26a8b77e25357e2af9f66a6d7ade101a2f0204d48fe3b567759d6ec007f7f3be
+size 166312561