Shaun ZhouandClaude Opus 5
474b7b50a0
训练数据已经是固定边长的片，推理还把整图交给 ultralytics 就会被 letterbox 缩放， 与模型见过的形态不符。三个调用方（可视化、批量推送预测、ML backend）统一改成 切片推理再拼回原图。 - engine/tiling.merge_tiles / tiled_predict：映射回原图、裁掉镜像填充区的框、跨片去重。 推理框架留在调用方，几何部分可单测。 - engine/detection.nms：标准 NMS，另加按包含比例抑制。缺陷被切割线切开时，相邻片 只看到一角，给出的小残框与完整框 IoU 很小（面积悬殊），光靠 IoU 压不掉。 - engine/tiled_inference：ultralytics 的接线，内部分批，避免片数多时 OOM。 - 三处调用点加 --tile-size / --tile-overlap（后端为 YOLO_TILE_SIZE / YOLO_TILE_OVERLAP）， 默认 160 / 32；设 0 退回整图推理，仅用于和旧模型对比。 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
157 lines
7.4 KiB
Python
157 lines
7.4 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 会一次发来成百上千个
|
||
任务，不分批会直接把显存撑爆
|
||
YOLO_TILE_SIZE 模型输入边长，默认 160。按它把图切片推理再拼回，全程不缩放；
|
||
0 = 整图推理，交给 ultralytics letterbox（会缩放）
|
||
YOLO_TILE_OVERLAP 切片重叠像素，默认 32
|
||
RETRAIN_AFTER 累计多少条新标注后自动触发重训，默认 0=只靠手动触发
|
||
RETRAIN_COMMAND 重训命令；不设则只记数不执行
|
||
RETRAIN_LOG 重训命令的输出写到哪，默认 /tmp/retrain.log
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import os
|
||
import shlex
|
||
import subprocess
|
||
import sys
|
||
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
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||
|
||
from engine.tiled_inference import load_image, predict_image # noqa: E402
|
||
|
||
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 格式的预测。"""
|
||
# 控件名从项目的 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))
|
||
device = os.environ.get("YOLO_DEVICE", "cuda")
|
||
batch = int(os.environ.get("YOLO_BATCH", 16))
|
||
tile_size = int(os.environ.get("YOLO_TILE_SIZE", 160))
|
||
overlap = int(os.environ.get("YOLO_TILE_OVERLAP", 32))
|
||
imgsz = int(os.environ.get("YOLO_IMGSZ", 640))
|
||
|
||
model_version = self.get("model_version")
|
||
predictions = []
|
||
for task in tasks:
|
||
# 本地存储的图要靠 LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT 解析；
|
||
# 上传或云存储的图则会被下载到本地缓存
|
||
path = self.get_local_path(task["data"][value_key], task_id=task.get("id"))
|
||
image = load_image(path)
|
||
height, width = image.shape[:2]
|
||
# 切片推理：模型的输入尺寸是固定的，整图直接送进去会被 letterbox 缩放。
|
||
# predict_image 内部也会分批，Retrieve Predictions 一次发来上千个任务也不会 OOM
|
||
detections = predict_image(
|
||
model, image, size=tile_size, overlap=overlap, conf=conf,
|
||
device=device, batch=batch, whole_image_imgsz=imgsz,
|
||
)
|
||
boxes = [
|
||
(classes[box.class_id], (box.x1, box.y1, box.x2, box.y2), box.score)
|
||
for box in detections
|
||
]
|
||
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}
|