feat(ml-backend): 把训练好的 YOLO 接进 Label Studio 做预标注

ml_backend/ 是一个 label-studio-ml 后端：Label Studio 打开任务时调 /predict，
后端跑 YOLO 并把框转成 Label Studio 的百分比坐标返回，人工只需修正。

- model.py：FilmDefectDetector。权重懒加载且只加载一次（Flask 多线程，加锁）；
  控件名从项目的 labeling config 读取，不写死；model_version 带权重的修改时间，
  换权重后 Label Studio 能区分新旧预测。
- convert.py：坐标转换，刻意不 import label_studio_ml，好在主环境里单测。
- fit()：每条标注累加计数，达到 RETRAIN_AFTER 就起子进程执行 RETRAIN_COMMAND
  并立刻返回——训练几十分钟，同步做会卡住 Label Studio。
- _wsgi.py 由官方模板生成，ruff 跳过，便于跟随上游。

必须用独立虚拟环境：label-studio-ml 会把 jinja2 降到 3.0.3，而 label-studio 1.23
需要 3.1，装进主环境会直接弄坏 Label Studio。PyPI 上的 1.0.9 是 2023 年的旧版，
要装 GitHub 上的现行版本。

本地用 COCO 权重端到端验证过：/health、/predict（bus.jpg 出 10 个框、坐标换算正确）、
/webhook（第 3 条标注触发重训命令并清零计数）。

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

ml_backend/README.md

@@ -0,0 +1,94 @@
+# Label Studio ML Backend：用训练好的 YOLO 做预标注
+
+把 `runs/detect/<name>/weights/best.pt` 接进 Label Studio，打开任务时自动给出框，
+人工只需修正而不是从零画。标注攒够数量还能自动触发重训，形成"边标边训"的闭环。
+
+## 为什么要单独的虚拟环境
+
+`label-studio-ml` 会把 jinja2 降到 3.0.3、带进 Flask 1.1.2，而 label-studio 1.23 需要
+jinja2 3.1。装进主环境会直接弄坏 Label Studio。它本来就是个独立服务，单独建环境最省事。
+
+另外 **PyPI 上的 `label-studio-ml` 是 2023 年的 1.0.9，不要用**，要装 GitHub 上的现行版本
+（本文档基于 2.0.1.dev0 验证）。
+
+```bash
+uv venv .venv-ml --python 3.12
+VIRTUAL_ENV=.venv-ml uv pip install \
+    "label-studio-ml @ git+https://github.com/HumanSignal/label-studio-ml-backend.git" \
+    ultralytics
+```
+
+## 启动
+
+```bash
+cd ml_backend
+YOLO_WEIGHTS=$PWD/../runs/detect/film/weights/best.pt \
+YOLO_DEVICE=cuda \
+LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT=/home/deep/labeling \
+../.venv-ml/bin/python _wsgi.py -p 9090
+```
+
+**用 `python _wsgi.py` 而不是 `label-studio-ml start .`**：后者用 PATH 上的 python 起子进程，
+在独立环境里会 `ModuleNotFoundError`。
+
+`LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT` 必须和启动 Label Studio 时的值一致——任务里的图片是
+`/data/local-files/?d=...`，后端要靠它把 URL 还原成磁盘路径。图片若来自上传或云存储，
+则改为设置 `LABEL_STUDIO_URL` 与 `LABEL_STUDIO_API_KEY`，后端会自行下载。
+
+健康检查：
+
+```bash
+curl http://localhost:9090/health
+# {"model_class":"FilmDefectDetector","status":"UP"}
+```
+
+## 接入项目
+
+Label Studio 里：Settings → **Model** → Connect Model，URL 填 `http://localhost:9090`。
+
+- 勾选 **Interactive preannotations** 后，打开任务就会实时调用模型；
+- 也可以在 Data Manager 里选中若干任务 → Actions → **Retrieve Predictions** 批量拉取。
+
+**类别名必须和 labeling config 里的 `<Label value="...">` 完全一致。** 后端用的是权重里的
+`model.names`（也就是训练时 `data.yaml` 的 names），对不上的话 Label Studio 会拒绝这条预测。
+我们的模型是 9 类：black、blackdot、grease、mosquito、normal、obstacle、ripple、uneven、wrinkle。
+
+## 环境变量
+
+| 变量 | 默认 | 说明 |
+| --- | --- | --- |
+| `YOLO_WEIGHTS` | `yolo26n.pt` | 权重路径。版本号里带它的修改时间，换权重后 Label Studio 能区分新旧预测 |
+| `YOLO_CONF` | `0.25` | 置信度阈值。预标注可以调低些，宁可多给几个框让人删 |
+| `YOLO_IMGSZ` | `640` | 推理尺寸，与训练保持一致 |
+| `YOLO_DEVICE` | `cuda` | `cuda` / `cpu` / `mps` |
+| `RETRAIN_AFTER` | `0` | 累计多少条新标注触发重训，0 = 不触发 |
+| `RETRAIN_COMMAND` | 空 | 触发时执行的命令 |
+
+## 边标边训
+
+Label Studio 每产生一条标注就会 POST `/webhook`，后端在 `fit()` 里累加计数，
+达到 `RETRAIN_AFTER` 就起一个**子进程**执行 `RETRAIN_COMMAND` 并立刻返回——
+训练要几十分钟，绝不能在这个请求里同步做，否则会把 Label Studio 卡住。
+
+```bash
+RETRAIN_AFTER=200 \
+RETRAIN_COMMAND="/home/deep/ModelTrainEval/scripts/retrain.sh" \
+../.venv-ml/bin/python _wsgi.py -p 9090
+```
+
+`retrain.sh` 需要自己写，大致是：导出项目 → `build_yolo_dataset.py` 装配 → `detect.py` 训练 →
+把新权重放到 `YOLO_WEIGHTS` 指向的位置 → 重启本服务（或让它下次启动时加载新权重）。
+**目前权重只在进程启动时加载一次**，重训完要重启服务才会生效。
+
+## 已验证的行为
+
+本地用 COCO 预训练权重实测过：
+
+- `/health` 返回 UP；
+- `/predict` 用 Label Studio 的真实请求体（tasks + label_config + params）调用，
+  `bus.jpg` 返回 10 个框，坐标换算正确（百分比、`original_width/height` 与原图一致）；
+- 图片无检出时返回一条空 `result` 的预测，表示"这张图没问题"；
+- `/webhook` 连发 3 条 `ANNOTATION_CREATED`，在第 3 条触发 `RETRAIN_COMMAND` 并把计数清零。
+
+转换逻辑的单元测试在 `tests/test_ml_backend_convert.py`，跑在主环境里
+（`convert.py` 刻意不 import `label_studio_ml`，就是为了能这样测）。

ml_backend/_wsgi.py

@@ -0,0 +1,125 @@
+import os
+import argparse
+import json
+import logging
+import logging.config
+
+# Set a default log level if LOG_LEVEL is not defined
+log_level = os.getenv("LOG_LEVEL", "INFO")
+
+logging.config.dictConfig({
+  "version": 1,
+  "disable_existing_loggers": False,  # Prevent overriding existing loggers
+  "formatters": {
+    "standard": {
+      "format": "[%(asctime)s] [%(levelname)s] [%(name)s::%(funcName)s::%(lineno)d] %(message)s"
+    }
+  },
+  "handlers": {
+    "console": {
+      "class": "logging.StreamHandler",
+      "level": log_level,
+      "stream": "ext://sys.stdout",
+      "formatter": "standard"
+    }
+  },
+  "root": {
+    "level": log_level,
+    "handlers": [
+      "console"
+    ],
+    "propagate": True
+  }
+})
+
+from label_studio_ml.api import init_app
+from model import FilmDefectDetector
+
+
+_DEFAULT_CONFIG_PATH = os.path.join(os.path.dirname(__file__), 'config.json')
+
+
+def get_kwargs_from_config(config_path=_DEFAULT_CONFIG_PATH):
+    if not os.path.exists(config_path):
+        return dict()
+    with open(config_path) as f:
+        config = json.load(f)
+    assert isinstance(config, dict)
+    return config
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Label studio')
+    parser.add_argument(
+        '-p', '--port', dest='port', type=int, default=9090,
+        help='Server port')
+    parser.add_argument(
+        '--host', dest='host', type=str, default='0.0.0.0',
+        help='Server host')
+    parser.add_argument(
+        '--kwargs', '--with', dest='kwargs', metavar='KEY=VAL', nargs='+', type=lambda kv: kv.split('='),
+        help='Additional LabelStudioMLBase model initialization kwargs')
+    parser.add_argument(
+        '-d', '--debug', dest='debug', action='store_true',
+        help='Switch debug mode')
+    parser.add_argument(
+        '--log-level', dest='log_level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], default=log_level,
+        help='Logging level')
+    parser.add_argument(
+        '--model-dir', dest='model_dir', default=os.path.dirname(__file__),
+        help='Directory where models are stored (relative to the project directory)')
+    parser.add_argument(
+        '--check', dest='check', action='store_true',
+        help='Validate model instance before launching server')
+    parser.add_argument('--basic-auth-user',
+                        default=os.environ.get('ML_SERVER_BASIC_AUTH_USER', None),
+                        help='Basic auth user')
+    
+    parser.add_argument('--basic-auth-pass',
+                        default=os.environ.get('ML_SERVER_BASIC_AUTH_PASS', None),
+                        help='Basic auth pass')    
+    
+    args = parser.parse_args()
+
+    # setup logging level
+    if args.log_level:
+        logging.root.setLevel(args.log_level)
+
+    def isfloat(value):
+        try:
+            float(value)
+            return True
+        except ValueError:
+            return False
+
+    def parse_kwargs():
+        param = dict()
+        for k, v in args.kwargs:
+            if v.isdigit():
+                param[k] = int(v)
+            elif v == 'True' or v == 'true':
+                param[k] = True
+            elif v == 'False' or v == 'false':
+                param[k] = False
+            elif isfloat(v):
+                param[k] = float(v)
+            else:
+                param[k] = v
+        return param
+
+    kwargs = get_kwargs_from_config()
+
+    if args.kwargs:
+        kwargs.update(parse_kwargs())
+
+    if args.check:
+        print('Check "' + FilmDefectDetector.__name__ + '" instance creation..')
+        model = FilmDefectDetector(**kwargs)
+
+    app = init_app(model_class=FilmDefectDetector, basic_auth_user=args.basic_auth_user, basic_auth_pass=args.basic_auth_pass)
+
+    app.run(host=args.host, port=args.port, debug=args.debug)
+
+else:
+    # for uWSGI use
+    app = init_app(model_class=FilmDefectDetector)

ml_backend/convert.py

@@ -0,0 +1,73 @@
+"""YOLO 结果 <-> Label Studio 结果格式的互转。
+
+单独成一个模块、不 import label_studio_ml，这样主环境里就能直接测：
+ML backend 跑在自己的虚拟环境里（见 ml_backend/README.md），主环境装不了它。
+
+Label Studio 的框用**百分比**表示左上角与宽高（0~100），YOLO 用归一化的中心点与宽高。
+`original_width` / `original_height` 必须是图片的真实像素尺寸，Label Studio 据此换算。
+"""
+
+from __future__ import annotations
+
+import uuid
+from collections.abc import Sequence
+
+
+def box_to_result(
+    class_name: str,
+    xyxyn: Sequence[float],
+    width: int,
+    height: int,
+    from_name: str,
+    to_name: str,
+    score: float | None = None,
+) -> dict:
+    """一个归一化 xyxy 框 -> 一条 Label Studio 的 rectanglelabels 结果。"""
+    x1, y1, x2, y2 = (min(max(float(v), 0.0), 1.0) for v in xyxyn)
+    item = {
+        "id": uuid.uuid4().hex[:10],
+        "type": "rectanglelabels",
+        "from_name": from_name,
+        "to_name": to_name,
+        "original_width": int(width),
+        "original_height": int(height),
+        "image_rotation": 0,
+        "value": {
+            "x": x1 * 100,
+            "y": y1 * 100,
+            "width": (x2 - x1) * 100,
+            "height": (y2 - y1) * 100,
+            "rotation": 0,
+            "rectanglelabels": [class_name],
+        },
+    }
+    if score is not None:
+        item["score"] = float(score)
+    return item
+
+
+def to_prediction(
+    boxes: Sequence[tuple[str, Sequence[float], float]],
+    width: int,
+    height: int,
+    from_name: str,
+    to_name: str,
+    model_version: str,
+) -> dict:
+    """一张图的若干框 -> 一条 Label Studio 预测。
+
+    Args:
+        boxes: (类别名, 归一化 xyxy, 置信度) 的序列。
+
+    预测整体的 `score` 取所有框里最高的置信度：Label Studio 用它给任务排序，
+    取最大值才能让"模型最有把握的图"排在前面；没有框时为 0，表示这张图没找到缺陷。
+    """
+    results = [
+        box_to_result(name, xyxyn, width, height, from_name, to_name, score)
+        for name, xyxyn, score in boxes
+    ]
+    return {
+        "model_version": model_version,
+        "score": max((score for _, _, score in boxes), default=0.0),
+        "result": results,
+    }

ml_backend/model.py

@@ -0,0 +1,124 @@
+"""把训练好的 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}

pyproject.toml

@@ -54,6 +54,8 @@ addopts = "-q"
 [tool.ruff]
 line-length = 110
 target-version = "py312"
+# ml_backend/_wsgi.py 由 label-studio-ml 的官方模板生成，保持原样便于跟随上游更新
+extend-exclude = ["ml_backend/_wsgi.py"]
 
 [tool.ruff.lint]
 select = ["E", "F", "I", "UP", "B"]

tests/test_ml_backend_convert.py

@@ -0,0 +1,71 @@
+"""YOLO 结果 -> Label Studio 预测格式。
+
+只测转换模块：ML backend 本体跑在独立虚拟环境里（label_studio_ml 与主环境的 jinja2
+版本冲突），主环境 import 不了 model.py。
+"""
+
+import sys
+
+import pytest
+
+sys.path.insert(0, "ml_backend")
+from convert import box_to_result, to_prediction  # noqa: E402
+
+FROM_NAME, TO_NAME = "label", "image"
+
+
+def test_box_converts_to_percent_coordinates():
+    """Label Studio 用左上角+宽高的百分比，YOLO 给的是归一化 xyxy。"""
+    item = box_to_result("blackdot", [0.2, 0.3, 0.6, 0.5], 200, 100, FROM_NAME, TO_NAME, 0.9)
+
+    assert item["value"]["x"] == pytest.approx(20.0)
+    assert item["value"]["y"] == pytest.approx(30.0)
+    assert item["value"]["width"] == pytest.approx(40.0)
+    assert item["value"]["height"] == pytest.approx(20.0)
+    assert item["value"]["rectanglelabels"] == ["blackdot"]
+    assert item["type"] == "rectanglelabels"
+    assert item["from_name"] == FROM_NAME and item["to_name"] == TO_NAME
+
+
+def test_original_size_is_the_real_pixel_size():
+    """Label Studio 按 original_width/height 把百分比换算回像素，填错框就会偏。"""
+    item = box_to_result("grease", [0, 0, 1, 1], 810, 1080, FROM_NAME, TO_NAME)
+    assert (item["original_width"], item["original_height"]) == (810, 1080)
+
+
+def test_ids_are_unique():
+    boxes = [box_to_result("a", [0, 0, 0.1, 0.1], 10, 10, FROM_NAME, TO_NAME) for _ in range(20)]
+    assert len({item["id"] for item in boxes}) == 20
+
+
+def test_out_of_range_coordinates_are_clamped():
+    """推理偶尔会给出略微越界的坐标；越界的框 Label Studio 画不出来。"""
+    item = box_to_result("x", [-0.01, 0.0, 1.02, 1.0], 100, 100, FROM_NAME, TO_NAME)
+    value = item["value"]
+    assert value["x"] == 0.0
+    assert value["x"] + value["width"] == pytest.approx(100.0)
+
+
+def test_prediction_score_is_the_highest_confidence():
+    """Label Studio 用 score 给任务排序，取最大值才能把最有把握的排在前面。"""
+    prediction = to_prediction(
+        [("blackdot", [0.1, 0.1, 0.2, 0.2], 0.4), ("grease", [0.5, 0.5, 0.7, 0.7], 0.85)],
+        128, 128, FROM_NAME, TO_NAME, "v1",
+    )
+    assert prediction["score"] == pytest.approx(0.85)
+    assert prediction["model_version"] == "v1"
+    assert len(prediction["result"]) == 2
+
+
+def test_empty_prediction_is_wellformed():
+    """没检出缺陷也要返回一条预测：空 result 表示"这张图没问题"。"""
+    prediction = to_prediction([], 128, 128, FROM_NAME, TO_NAME, "v1")
+    assert prediction["result"] == []
+    assert prediction["score"] == 0.0
+
+
+def test_score_is_kept_per_box():
+    prediction = to_prediction(
+        [("blackdot", [0.1, 0.1, 0.2, 0.2], 0.42)], 64, 64, FROM_NAME, TO_NAME, "v1"
+    )
+    assert prediction["result"][0]["score"] == pytest.approx(0.42)