fix(ml-backend): 推理分批，避免 Retrieve Predictions 撑爆显存；新增批量推送脚本

Retrieve Predictions 会把选中的任务一次发给后端，而后端把整个列表交给
model.predict(source=paths)——ultralytics 不理会 batch 参数，会当成一个 batch，
于是显存爆掉。这与之前 visualize_detections 踩的是同一个坑。后端现在按
YOLO_BATCH（默认 16）自己分批。

另加 scripts/push_predictions.py：重训完在本地分批推理，再用批量导入接口
POST /api/projects/<id>/import/predictions 提交，不经过后端服务，显存与请求体都可控。
支持 --replace 先删同版本旧预测（预测是追加的，重复跑会堆两份）、--limit 与 --dry-run。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

ml_backend/README.md

@@ -64,6 +64,7 @@ Label Studio 里：Settings → **Model** → Connect Model，URL 填 `http://lo
 | `YOLO_CONF` | `0.25` | 置信度阈值。预标注可以调低些，宁可多给几个框让人删 |
 | `YOLO_IMGSZ` | `640` | 推理尺寸，与训练保持一致 |
 | `YOLO_DEVICE` | `cuda` | `cuda` / `cpu` / `mps` |
+| `YOLO_BATCH` | `16` | 一次推理几张。Retrieve Predictions 会一次发来上千个任务，不分批会撑爆显存 |
 | `RETRAIN_AFTER` | `0` | 累计多少条新标注触发重训，0 = 不触发 |
 | `RETRAIN_COMMAND` | 空 | 触发时执行的命令 |
 

ml_backend/model.py

@@ -6,6 +6,8 @@
     YOLO_CONF        置信度阈值，默认 0.25
     YOLO_IMGSZ       推理尺寸，默认 640
     YOLO_DEVICE      cuda / cpu / mps，默认 cuda（不可用时 ultralytics 自己退回）
+    YOLO_BATCH       一次推理几张，默认 16。Retrieve Predictions 会一次发来成百上千个
+                     任务，不分批会直接把显存撑爆
     RETRAIN_AFTER    累计多少条新标注后自动触发重训，默认 0=只靠手动触发
     RETRAIN_COMMAND  重训命令；不设则只记数不执行
     RETRAIN_LOG      重训命令的输出写到哪，默认 /tmp/retrain.log
@@ -82,9 +84,16 @@ class FilmDefectDetector(LabelStudioMLBase):
         if not paths:
             return ModelResponse(model_version=model_version, predictions=[])
 
-        results = model.predict(
+        # 必须自己分批：ultralytics 不理会 batch 参数，会把整个列表当成一个 batch。
-            source=paths, imgsz=imgsz, conf=conf, device=device, verbose=False
+        # Label Studio 的 Retrieve Predictions 一次能发来上千个任务，不分批必然 OOM
-        )
+        batch = int(os.environ.get("YOLO_BATCH", 16))
+        results = []
+        for start in range(0, len(paths), batch):
+            results.extend(model.predict(
+                source=paths[start:start + batch], imgsz=imgsz, conf=conf,
+                device=device, verbose=False,
+            ))
+
         predictions = []
         for result, (width, height) in zip(results, sizes, strict=True):
             boxes = [

scripts/push_predictions.py

@@ -0,0 +1,273 @@
+#!/usr/bin/env python
+"""用指定权重给 Label Studio 项目里的所有任务批量生成预测。
+
+    uv run python scripts/push_predictions.py --project 39 \\
+        --weights runs/detect/current.pt
+
+适合重训完成后，一次性把全部样本的预标注刷新成新模型的结果。
+
+与网页端的 **Retrieve Predictions** 相比：那条路要经过 ML backend，Label Studio 会把选中的
+任务一股脑发过去；本脚本在本地按 `--batch` 分批推理，再用批量导入接口
+（`POST /api/projects/<id>/import/predictions`）一次提交一批，显存和请求体都可控，
+也不依赖后端服务在跑。
+
+需要一个 API token（网页端 Account & Settings → Access Token）：
+
+    export LABEL_STUDIO_API_KEY=xxxx
+
+图片路径从任务的 `/data/local-files/?d=...` 还原，所以 `--document-root` 必须和
+Label Studio 的 `LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT` 一致。
+
+**预测是追加的，不会覆盖。** 同一个任务可以有多版预测，Label Studio 在项目设置里按
+model_version 选用哪一版。重复跑同一个版本会堆重复预测，用 `--replace` 先删掉同名版本。
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import urllib.parse
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "ml_backend"))
+
+from convert import to_prediction  # noqa: E402
+
+from utils import setup_logging  # noqa: E402
+from utils.device import DeviceNotAvailableError, describe_device, resolve_device  # noqa: E402
+
+logger = None
+
+
+class LabelStudio:
+    """Label Studio REST API 的最小封装。"""
+
+    def __init__(self, url: str, token: str) -> None:
+        import requests
+
+        self.url = url.rstrip("/")
+        self.session = requests.Session()
+        self.session.headers["Authorization"] = f"Token {token}"
+
+    def tasks(self, project: int, page_size: int = 500) -> list[dict]:
+        """列出项目里的全部任务（只取 id 与 data，分页拉完）。"""
+        collected, page = [], 1
+        while True:
+            response = self.session.get(
+                f"{self.url}/api/tasks/",
+                params={"project": project, "page": page, "page_size": page_size,
+                        "fields": "task_only"},
+                timeout=120,
+            )
+            response.raise_for_status()
+            payload = response.json()
+            # 不同版本的分页字段不一样：data manager 用 tasks/total，DRF 默认用 results/count
+            batch = payload.get("tasks", payload.get("results", payload)) if isinstance(
+                payload, dict
+            ) else payload
+            if not batch:
+                break
+            collected.extend(batch)
+            total = payload.get("total", payload.get("count")) if isinstance(payload, dict) else None
+            if total is not None and len(collected) >= int(total):
+                break
+            page += 1
+        return collected
+
+    def delete_predictions(self, project: int, model_version: str) -> int:
+        """删掉该项目下指定版本的预测，返回删除条数。"""
+        removed = 0
+        page = 1
+        while True:
+            response = self.session.get(
+                f"{self.url}/api/predictions/",
+                params={"project": project, "page": page, "page_size": 500},
+                timeout=120,
+            )
+            if response.status_code == 404:      # 翻到尾页，某些版本直接返回 404
+                break
+            response.raise_for_status()
+            payload = response.json()
+            items = payload.get("results", payload) if isinstance(payload, dict) else payload
+            if not items:
+                break
+            for item in items:
+                if item.get("model_version") == model_version:
+                    self.session.delete(
+                        f"{self.url}/api/predictions/{item['id']}/", timeout=60
+                    ).raise_for_status()
+                    removed += 1
+            if isinstance(payload, dict) and not payload.get("next"):
+                break
+            page += 1
+        return removed
+
+    def import_predictions(self, project: int, predictions: list[dict]) -> None:
+        """批量导入，服务端自己再分批入库。"""
+        response = self.session.post(
+            f"{self.url}/api/projects/{project}/import/predictions",
+            json=predictions,
+            timeout=600,
+        )
+        response.raise_for_status()
+
+
+def local_path(image_url: str, document_root: Path) -> Path:
+    """把任务里的图片 URL 还原成磁盘路径。
+
+    只处理本地存储（`/data/local-files/?d=相对路径`）——上传或云存储的图得走下载，
+    那种情况用网页端的 Retrieve Predictions 更省事。
+    """
+    if "?d=" not in image_url:
+        raise ValueError(
+            f"不是本地存储的图片 URL：{image_url}；本脚本只支持 /data/local-files/?d=..."
+        )
+    relative = urllib.parse.unquote(image_url.split("?d=", 1)[1])
+    return document_root / relative
+
+
+def main() -> int:
+    global logger
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
+    )
+    parser.add_argument("--project", type=int, required=True, help="Label Studio 项目 id")
+    parser.add_argument(
+        "--weights", default="runs/detect/current.pt",
+        help="权重；缺省用重训脚本维护的 current.pt 软链接",
+    )
+    parser.add_argument("--url", default=os.environ.get("LABEL_STUDIO_URL", "http://localhost:8080"))
+    parser.add_argument(
+        "--token", default=os.environ.get("LABEL_STUDIO_API_KEY", ""),
+        help="API token，缺省读环境变量 LABEL_STUDIO_API_KEY",
+    )
+    parser.add_argument(
+        "--document-root",
+        default=os.environ.get("LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT", "/home/deep/labeling"),
+        help="必须与 Label Studio 的同名环境变量一致",
+    )
+    parser.add_argument("--model-version", default=None, help="缺省按权重文件名与修改时间生成")
+    parser.add_argument("--conf", type=float, default=0.25)
+    parser.add_argument("--imgsz", type=int, default=640)
+    parser.add_argument("--device", default="auto")
+    parser.add_argument("--batch", type=int, default=16, help="每次推理几张")
+    parser.add_argument("--chunk", type=int, default=500, help="每次上传几条预测")
+    parser.add_argument("--limit", type=int, default=0, help="只处理前 N 个任务（试运行用）")
+    parser.add_argument(
+        "--replace", action="store_true",
+        help="先删掉同一 model_version 的旧预测；否则会追加，同版本堆两份",
+    )
+    parser.add_argument("--dry-run", action="store_true", help="只推理不上传")
+    args = parser.parse_args()
+
+    logger = setup_logging()
+    if not args.token:
+        logger.error("缺少 API token：export LABEL_STUDIO_API_KEY=... 或用 --token 传入")
+        return 1
+
+    weights = Path(args.weights)
+    if not weights.exists():
+        logger.error("找不到权重 %s", weights)
+        return 1
+    model_version = args.model_version or f"{weights.resolve().stem}-{int(weights.stat().st_mtime)}"
+
+    try:
+        device = resolve_device(args.device)
+    except DeviceNotAvailableError as exc:
+        logger.error("%s", exc)
+        return 1
+
+    client = LabelStudio(args.url, args.token)
+    logger.info("拉取项目 %d 的任务列表…", args.project)
+    tasks = client.tasks(args.project)
+    if args.limit:
+        tasks = tasks[: args.limit]
+    logger.info(
+        "任务 %d 个 | 权重 %s | 版本 %s | 设备 %s",
+        len(tasks), weights, model_version, describe_device(device),
+    )
+    if not tasks:
+        return 0
+
+    document_root = Path(args.document_root)
+    try:
+        pairs = [(task["id"], local_path(task["data"]["image"], document_root)) for task in tasks]
+    except (KeyError, ValueError) as exc:
+        logger.error("%s", exc)
+        return 1
+    missing = [path for _, path in pairs if not path.exists()]
+    if missing:
+        logger.error(
+            "有 %d 张图在磁盘上找不到，例如 %s；--document-root 是否与 Label Studio 一致（当前 %s）",
+            len(missing), missing[0], document_root,
+        )
+        return 1
+
+    if args.replace and not args.dry_run:
+        removed = client.delete_predictions(args.project, model_version)
+        logger.info("删除同版本旧预测 %d 条", removed)
+
+    from PIL import Image
+    from ultralytics import YOLO
+
+    model = YOLO(str(weights))
+    classes = model.names
+    device_arg = device.index or 0 if device.type == "cuda" else device.type
+
+    pending: list[dict] = []
+    uploaded = with_boxes = total_boxes = 0
+
+    def flush() -> None:
+        nonlocal pending, uploaded
+        if not pending or args.dry_run:
+            pending = []
+            return
+        client.import_predictions(args.project, pending)
+        uploaded += len(pending)
+        pending = []
+
+    for start in range(0, len(pairs), args.batch):
+        chunk = pairs[start:start + args.batch]
+        results = model.predict(
+            source=[str(path) for _, path in chunk], imgsz=args.imgsz, conf=args.conf,
+            device=device_arg, verbose=False,
+        )
+        for (task_id, path), result in zip(chunk, results, strict=True):
+            with Image.open(path) as image:
+                width, height = image.size
+            boxes = [
+                (classes[int(class_id)], xyxyn, float(score))
+                for xyxyn, class_id, score in zip(
+                    result.boxes.xyxyn.tolist(), result.boxes.cls.tolist(),
+                    result.boxes.conf.tolist(), strict=True,
+                )
+            ]
+            prediction = to_prediction(boxes, width, height, "label", "image", model_version)
+            prediction["task"] = task_id
+            pending.append(prediction)
+            total_boxes += len(boxes)
+            with_boxes += bool(boxes)
+        if len(pending) >= args.chunk:
+            flush()
+        processed = start + len(chunk)
+        if processed % 1000 < args.batch:
+            logger.info("处理 [%d/%d]，已上传 %d", processed, len(pairs), uploaded)
+    flush()
+
+    logger.info(
+        "完成：%d 个任务，%d 个有框（%d 个框），上传 %d 条预测，版本 %s",
+        len(pairs), with_boxes, total_boxes, uploaded, model_version,
+    )
+    if args.dry_run:
+        logger.info("--dry-run：未上传")
+    else:
+        logger.info(
+            "在项目设置 Annotation → Prelabeling 里把预测版本选成 %s 即可用于预标注", model_version
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/test_push_predictions.py

@@ -0,0 +1,88 @@
+"""批量给 Label Studio 推预测：URL 还原与分页拉取。"""
+
+import sys
+
+import pytest
+
+pytest.importorskip("requests")
+
+sys.path.insert(0, "scripts")
+import push_predictions as pp  # noqa: E402
+
+
+def test_local_path_from_local_files_url(tmp_path):
+    url = "/data/local-files/?d=boxes_shap_label_studio/images/blackdot_0489.png"
+    assert pp.local_path(url, tmp_path) == (
+        tmp_path / "boxes_shap_label_studio/images/blackdot_0489.png"
+    )
+
+
+def test_local_path_unquotes_escaped_names(tmp_path):
+    """文件名里有中文或空格时 URL 是转义过的，得还原回来。"""
+    url = "/data/local-files/?d=proj/images/%E5%8F%8D%E9%9D%A2%20a.png"
+    assert pp.local_path(url, tmp_path) == tmp_path / "proj/images/反面 a.png"
+
+
+def test_local_path_rejects_non_local_urls(tmp_path):
+    with pytest.raises(ValueError, match="本地存储"):
+        pp.local_path("https://example.com/a.png", tmp_path)
+
+
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self._payload = payload
+        self.status_code = status
+
+    def json(self):
+        return self._payload
+
+    def raise_for_status(self):
+        if self.status_code >= 400:
+            raise AssertionError(f"HTTP {self.status_code}")
+
+
+class _FakeSession:
+    """按页返回任务，记录请求参数。"""
+
+    def __init__(self, pages):
+        self.pages = pages
+        self.headers = {}
+        self.calls = []
+
+    def get(self, url, params=None, timeout=None):
+        self.calls.append(params)
+        page = params["page"]
+        return _FakeResponse(self.pages[page - 1] if page <= len(self.pages) else {"tasks": []})
+
+
+@pytest.fixture
+def client(monkeypatch):
+    monkeypatch.setattr(pp, "LabelStudio", pp.LabelStudio)
+    instance = pp.LabelStudio.__new__(pp.LabelStudio)
+    instance.url = "http://ls"
+    return instance
+
+
+def test_tasks_pagination_stops_at_total(client):
+    """分页要拉完再停，不能只拿第一页。"""
+    client.session = _FakeSession([
+        {"tasks": [{"id": 1}, {"id": 2}], "total": 3},
+        {"tasks": [{"id": 3}], "total": 3},
+    ])
+    assert [t["id"] for t in client.tasks(39, page_size=2)] == [1, 2, 3]
+
+
+def test_tasks_handles_drf_style_pagination(client):
+    """有的版本返回 results/count 而不是 tasks/total。"""
+    client.session = _FakeSession([
+        {"results": [{"id": 1}], "count": 2},
+        {"results": [{"id": 2}], "count": 2},
+    ])
+    assert [t["id"] for t in client.tasks(39, page_size=1)] == [1, 2]
+
+
+def test_tasks_stops_on_empty_page(client):
+    """没有 total 字段时，靠空页收尾，不能无限翻。"""
+    client.session = _FakeSession([{"tasks": [{"id": 1}]}, {"tasks": []}])
+    assert [t["id"] for t in client.tasks(39)] == [1]
+    assert len(client.session.calls) == 2