diff --git a/ml_backend/README.md b/ml_backend/README.md index 800c1b3..c8cd7ce 100644 --- a/ml_backend/README.md +++ b/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` | 空 | 触发时执行的命令 | diff --git a/ml_backend/model.py b/ml_backend/model.py index 506ac5a..c3912dd 100644 --- a/ml_backend/model.py +++ b/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( - source=paths, imgsz=imgsz, conf=conf, device=device, verbose=False - ) + # 必须自己分批:ultralytics 不理会 batch 参数,会把整个列表当成一个 batch。 + # 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 = [ diff --git a/scripts/push_predictions.py b/scripts/push_predictions.py new file mode 100644 index 0000000..baa91eb --- /dev/null +++ b/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//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()) diff --git a/tests/test_push_predictions.py b/tests/test_push_predictions.py new file mode 100644 index 0000000..17310d4 --- /dev/null +++ b/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