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>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
8fa7a9d142
commit
dce2eef0a7
@@ -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` | 空 | 触发时执行的命令 |
|
||||
|
||||
|
||||
+12
-3
@@ -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 = [
|
||||
|
||||
@@ -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())
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user