ModelTrainEval/ml_backend/model.py

"""把训练好的 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 自己退回）
    RETRAIN_AFTER    累计多少条新标注后触发重训，默认 0=不触发
    RETRAIN_COMMAND  触发时执行的命令；不设则只记数不执行
"""

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()

    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=[])

        results = model.predict(
            source=paths, 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 等着这个请求返回，长任务会把它卡住。
        所以只累加计数，达到阈值后起一个子进程去跑，立刻返回。
        """
        if event not in self.TRAIN_EVENTS:
            return None
        after = int(os.environ.get("RETRAIN_AFTER", 0))
        command = os.environ.get("RETRAIN_COMMAND", "")

        count = int(self.get("annotations_since_training") or 0) + 1
        self.set("annotations_since_training", str(count))
        logger.info("收到 %s，距上次训练已有 %d 条标注", event, count)

        if not after or count < after or not command:
            return {"annotations_since_training": count}

        self.set("annotations_since_training", "0")
        logger.info("累计 %d 条标注，触发重训: %s", count, command)
        # 不等它结束：训练要几十分钟，这个请求必须立刻返回
        subprocess.Popen(shlex.split(command))  # noqa: S603
        return {"annotations_since_training": 0, "retrain_started": True}