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>
This commit is contained in:
co-authored by
Claude Opus 5
parent
154677032f
commit
a4d099e41d
@@ -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)
|
||||
Reference in New Issue
Block a user