Shaun ZhouandClaude Opus 5
dce2eef0a7
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>
164 lines
7.3 KiB
Python
164 lines
7.3 KiB
Python
"""把训练好的 YOLO 接进 Label Studio，边标边用模型预标注。
|
||
|
||
启动见 ml_backend/README.md。行为由环境变量控制：
|
||
|
||
YOLO_WEIGHTS 必填，权重路径，如 runs/detect/film/weights/best.pt
|
||
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
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import os
|
||
import shlex
|
||
import subprocess
|
||
import threading
|
||
from pathlib import Path
|
||
|
||
from convert import to_prediction
|
||
from label_studio_ml.model import LabelStudioMLBase
|
||
from label_studio_ml.response import ModelResponse
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class FilmDefectDetector(LabelStudioMLBase):
|
||
"""薄膜缺陷检测的预标注后端。"""
|
||
|
||
_model = None
|
||
_lock = threading.Lock()
|
||
_retraining = None # 正在跑的重训子进程
|
||
|
||
def setup(self) -> None:
|
||
"""记录模型版本。
|
||
|
||
版本带上权重文件的修改时间：重训换了权重之后版本号会变，Label Studio 就能把
|
||
新旧预测区分开（项目设置里按 model_version 选用哪一版预标注）。
|
||
"""
|
||
weights = Path(os.environ.get("YOLO_WEIGHTS", "yolo26n.pt"))
|
||
stamp = int(weights.stat().st_mtime) if weights.exists() else 0
|
||
self.set("model_version", f"{weights.stem}-{stamp}")
|
||
|
||
def _load(self):
|
||
"""首次预测时才加载权重，且只加载一次——Flask 是多线程的，加锁。"""
|
||
with self._lock:
|
||
if FilmDefectDetector._model is None:
|
||
from ultralytics import YOLO
|
||
|
||
weights = os.environ.get("YOLO_WEIGHTS", "yolo26n.pt")
|
||
logger.info("加载权重 %s", weights)
|
||
FilmDefectDetector._model = YOLO(weights)
|
||
return FilmDefectDetector._model
|
||
|
||
def predict(self, tasks: list[dict], context: dict | None = None, **kwargs) -> ModelResponse:
|
||
"""对一批任务推理，返回 Label Studio 格式的预测。"""
|
||
from PIL import Image
|
||
|
||
# 控件名从项目的 labeling config 里读，不写死：改了配置这里要跟着变
|
||
from_name, to_name, value_key = self.get_first_tag_occurence(
|
||
"RectangleLabels", "Image"
|
||
)
|
||
model = self._load()
|
||
classes = model.names
|
||
conf = float(os.environ.get("YOLO_CONF", 0.25))
|
||
imgsz = int(os.environ.get("YOLO_IMGSZ", 640))
|
||
device = os.environ.get("YOLO_DEVICE", "cuda")
|
||
|
||
paths, sizes = [], []
|
||
for task in tasks:
|
||
# 本地存储的图要靠 LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT 解析；
|
||
# 上传或云存储的图则会被下载到本地缓存
|
||
path = self.get_local_path(task["data"][value_key], task_id=task.get("id"))
|
||
with Image.open(path) as image:
|
||
sizes.append(image.size)
|
||
paths.append(path)
|
||
|
||
model_version = self.get("model_version")
|
||
if not paths:
|
||
return ModelResponse(model_version=model_version, predictions=[])
|
||
|
||
# 必须自己分批：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 = [
|
||
(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,
|
||
)
|
||
]
|
||
predictions.append(
|
||
to_prediction(boxes, width, height, from_name, to_name, model_version)
|
||
)
|
||
return ModelResponse(model_version=model_version, predictions=predictions)
|
||
|
||
def fit(self, event, data, **kwargs):
|
||
"""Label Studio 的标注事件与 Start Training 按钮都走这里。
|
||
|
||
这里**不能**直接训练：Label Studio 等着这个请求返回，长任务会把它卡住。
|
||
所以只决定"要不要训"，真正的训练交给子进程，立刻返回。
|
||
"""
|
||
if event not in self.TRAIN_EVENTS:
|
||
return None
|
||
after = int(os.environ.get("RETRAIN_AFTER", 0))
|
||
|
||
count = int(self.get("annotations_since_training") or 0)
|
||
if event == "ANNOTATION_CREATED":
|
||
# 只数新增：UPDATED / DELETED 是对已有标注的修改，不该让阈值提前到点
|
||
count += 1
|
||
self.set("annotations_since_training", str(count))
|
||
logger.info("收到 %s，距上次训练已新增 %d 条标注", event, count)
|
||
|
||
# Start Training 是人明确点的，无条件训；攒够阈值则自动训
|
||
manual = event == "START_TRAINING"
|
||
if not manual and not (after and count >= after):
|
||
return {"annotations_since_training": count}
|
||
return self._start_retraining(count, manual)
|
||
|
||
def _start_retraining(self, count: int, manual: bool) -> dict:
|
||
"""起一个子进程跑重训命令，立刻返回。"""
|
||
command = os.environ.get("RETRAIN_COMMAND", "")
|
||
if not command:
|
||
logger.warning(
|
||
"%s，但没有设置 RETRAIN_COMMAND，不会真的训练"
|
||
"（在 scripts/start_ml_backend.sh 里配置）",
|
||
"收到 Start Training" if manual else f"已累计 {count} 条标注",
|
||
)
|
||
return {"annotations_since_training": count, "retrain_started": False,
|
||
"detail": "RETRAIN_COMMAND 未设置"}
|
||
|
||
running = FilmDefectDetector._retraining
|
||
if running is not None and running.poll() is None:
|
||
# 重复点 Start Training 不该叠加启动多个训练
|
||
logger.warning("已有重训在跑（pid=%s），这次忽略", running.pid)
|
||
return {"annotations_since_training": count, "retrain_started": False,
|
||
"detail": f"已有重训在跑 pid={running.pid}"}
|
||
|
||
log_path = os.environ.get("RETRAIN_LOG", "/tmp/retrain.log")
|
||
self.set("annotations_since_training", "0")
|
||
logger.info("触发重训: %s（输出见 %s）", command, log_path)
|
||
with open(log_path, "a", encoding="utf-8") as log: # noqa: PTH123
|
||
FilmDefectDetector._retraining = subprocess.Popen( # noqa: S603
|
||
shlex.split(command), stdout=log, stderr=subprocess.STDOUT, stdin=subprocess.DEVNULL,
|
||
)
|
||
return {"annotations_since_training": 0, "retrain_started": True,
|
||
"pid": FilmDefectDetector._retraining.pid, "log": log_path}
|