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  重训命令；不设则只记数不执行
    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=[])

        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 的标注事件与 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}