feat(detect): 推理侧改为切片，与固定 160 输入一致
训练数据已经是固定边长的片，推理还把整图交给 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>
This commit is contained in:
Shaun Zhou
co-authored by
Claude Opus 5
Claude Opus 5
parent
2a92275317
commit
474b7b50a0
+24
-31
@@ -8,6 +8,9 @@
|
||||
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
|
||||
@@ -19,6 +22,7 @@ import logging
|
||||
import os
|
||||
import shlex
|
||||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
@@ -26,6 +30,10 @@ 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__)
|
||||
|
||||
|
||||
@@ -59,8 +67,6 @@ class FilmDefectDetector(LabelStudioMLBase):
|
||||
|
||||
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"
|
||||
@@ -68,42 +74,29 @@ class FilmDefectDetector(LabelStudioMLBase):
|
||||
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")
|
||||
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))
|
||||
|
||||
paths, sizes = [], []
|
||||
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"))
|
||||
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=[])
|
||||
|
||||
# 必须自己分批：ultralytics 不理会 batch 参数，会把整个列表当成一个 batch。
|
||||
# Label Studio 的 Retrieve Predictions 一次能发来上千个任务，不分批必然 OOM
|
||||
batch = int(os.environ.get("YOLO_BATCH", 16))
|
||||
results = []
|
||||
for start in range(0, len(paths), batch):
|
||||
results.extend(model.predict(
|
||||
source=paths[start:start + batch], imgsz=imgsz, conf=conf,
|
||||
device=device, verbose=False,
|
||||
))
|
||||
|
||||
predictions = []
|
||||
for result, (width, height) in zip(results, sizes, strict=True):
|
||||
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[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,
|
||||
)
|
||||
(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)
|
||||
|
||||
Reference in New Issue
Block a user