Shaun ZhouandClaude Opus 5
744e5fd18a
两个在真实环境暴露的 bug： 1. /api/predictions/ 根本不分页——Label Studio 的 REST_FRAMEWORK 没设 DEFAULT_PAGINATION_CLASS，PAGE_SIZE 形同虚设，它一次返回全部预测并忽略 page。 原先按页翻永远不终止，实测翻到第 1798 页才因别的原因报错。改成只请求一次。 2. access token 默认只有 5 分钟（LS 未覆盖 SimpleJWT 默认值），而推上万条预测远超 5 分钟。所有请求统一走 request()，碰到 401 自动用 refresh token 续期后重试一次。 另加 --replace-all：用 Data Manager 的批量操作一次删光项目里所有预测， 比逐条删同版本快得多（适合版本号变了、旧预测留着没用的情况）。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
316 lines
13 KiB
Python
316 lines
13 KiB
Python
#!/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
|
||
|
||
新版 Label Studio 默认只认 **JWT 个人访问令牌**（组织设置里 `legacy_api_tokens_enabled=false`）。
|
||
脚本会自动识别：令牌形如 JWT（两个点分三段）时先用 `/api/token/refresh/` 换取 access token
|
||
并以 `Bearer` 发送，否则按旧式 `Token` 发送。
|
||
|
||
图片路径从任务的 `/data/local-files/?d=...` 还原，所以 `--document-root` 必须和
|
||
Label Studio 的 `LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT` 一致。
|
||
|
||
**预测是追加的，不会覆盖。** 同一个任务可以有多版预测，Label Studio 在项目设置里按
|
||
model_version 选用哪一版。重复跑同一个版本会堆重复预测：`--replace` 先删掉同名版本
|
||
（换了权重时一般删不到东西，很快），`--replace-all` 则一次删光项目里的所有预测。
|
||
"""
|
||
|
||
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.token = token
|
||
self._authorize()
|
||
|
||
def _authorize(self) -> None:
|
||
"""设置 Authorization 头；JWT 令牌会换成 access token。
|
||
|
||
新版默认关闭了旧式 `Token xxx`（组织设置 legacy_api_tokens_enabled=false），
|
||
个人访问令牌是一个 JWT refresh token，要先换成 access token 再以 Bearer 发送。
|
||
"""
|
||
if self.token.count(".") != 2: # 不是 JWT，按旧式 Token 处理
|
||
self.session.headers["Authorization"] = f"Token {self.token}"
|
||
return
|
||
response = self.session.post(
|
||
f"{self.url}/api/token/refresh/", json={"refresh": self.token}, timeout=60
|
||
)
|
||
if response.status_code >= 400:
|
||
raise RuntimeError(
|
||
f"用个人访问令牌换取 access token 失败（HTTP {response.status_code}）："
|
||
f"{response.text[:200]}"
|
||
)
|
||
self.session.headers["Authorization"] = f"Bearer {response.json()['access']}"
|
||
|
||
def request(self, method: str, path: str, **kwargs):
|
||
"""带自动续期的请求。
|
||
|
||
access token 默认只有 5 分钟有效期（Label Studio 没有覆盖 SimpleJWT 的默认值），
|
||
而推完上万条预测远不止 5 分钟，所以碰到 401 就用 refresh token 再换一个再试一次。
|
||
"""
|
||
kwargs.setdefault("timeout", 120)
|
||
response = self.session.request(method, f"{self.url}{path}", **kwargs)
|
||
if response.status_code == 401:
|
||
self._authorize()
|
||
response = self.session.request(method, f"{self.url}{path}", **kwargs)
|
||
response.raise_for_status()
|
||
return response
|
||
|
||
def tasks(self, project: int, page_size: int = 500) -> list[dict]:
|
||
"""列出项目里的全部任务（只取 id 与 data，分页拉完）。"""
|
||
collected, page = [], 1
|
||
while True:
|
||
response = self.request(
|
||
"GET", "/api/tasks/",
|
||
params={"project": project, "page": page, "page_size": page_size,
|
||
"fields": "task_only"},
|
||
)
|
||
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:
|
||
"""删掉该项目下指定版本的预测，返回删除条数。
|
||
|
||
`/api/predictions/` **不分页**：Label Studio 的 REST_FRAMEWORK 没设
|
||
DEFAULT_PAGINATION_CLASS，PAGE_SIZE 形同虚设，它一次返回全部预测并忽略 page 参数。
|
||
早先按页翻会翻不到头（实测翻到第 1798 页才因令牌过期报错），所以这里只请求一次。
|
||
"""
|
||
payload = self.request("GET", "/api/predictions/", params={"project": project}).json()
|
||
items = payload.get("results", []) if isinstance(payload, dict) else payload
|
||
removed = 0
|
||
for item in items:
|
||
if item.get("model_version") == model_version:
|
||
self.request("DELETE", f"/api/predictions/{item['id']}/")
|
||
removed += 1
|
||
return removed
|
||
|
||
def delete_all_predictions(self, project: int) -> dict:
|
||
"""用 Data Manager 的批量操作一次删光项目里的所有预测（不分版本）。"""
|
||
return self.request(
|
||
"POST", "/api/dm/actions/",
|
||
params={"id": "delete_tasks_predictions", "project": project},
|
||
json={"selectedItems": {"all": True, "excluded": []}},
|
||
).json()
|
||
|
||
def import_predictions(self, project: int, predictions: list[dict]) -> None:
|
||
"""批量导入，服务端自己再分批入库。"""
|
||
self.request(
|
||
"POST", f"/api/projects/{project}/import/predictions",
|
||
json=predictions, timeout=600,
|
||
)
|
||
|
||
|
||
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(
|
||
"--replace-all", action="store_true",
|
||
help="先删光项目里的**所有**预测（不分版本），用 Data Manager 的批量操作，一次请求",
|
||
)
|
||
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
|
||
|
||
try:
|
||
client = LabelStudio(args.url, args.token)
|
||
except RuntimeError as exc:
|
||
logger.error("%s", exc)
|
||
return 1
|
||
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_all and not args.dry_run:
|
||
detail = client.delete_all_predictions(args.project)
|
||
logger.info("删除项目里的全部旧预测：%s", detail.get("detail", detail))
|
||
elif args.replace and not args.dry_run:
|
||
removed = client.delete_predictions(args.project, model_version)
|
||
logger.info("删除同版本（%s）旧预测 %d 条", model_version, 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())
|