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
2026-09-23 15:12:35 +08:00
co-authored by Claude Opus 5
parent 2a92275317
commit 474b7b50a0
9 changed files with 412 additions and 75 deletions
+28 -24
View File
@@ -40,6 +40,7 @@ sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "ml_backend"))
from convert import to_prediction # noqa: E402
from engine.tiled_inference import load_image, predict_image # noqa: E402
from utils import setup_logging # noqa: E402
from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402
@@ -186,7 +187,12 @@ def main() -> int:
)
parser.add_argument("--model-version", default=None, help="缺省按权重文件名与修改时间生成")
parser.add_argument("--conf", type=float, default=0.25)
parser.add_argument("--imgsz", type=int, default=640)
parser.add_argument(
"--tile-size", type=int, default=160,
help="模型输入边长;按它切片推理再拼回,全程不缩放。0=整图推理(会 letterbox)",
)
parser.add_argument("--tile-overlap", type=int, default=32, help="切片重叠像素")
parser.add_argument("--imgsz", type=int, default=640, help="仅 --tile-size 0 时用")
parser.add_argument("--device", default="auto")
parser.add_argument("--batch", type=int, default=16, help="每次推理几张")
parser.add_argument("--chunk", type=int, default=500, help="每次上传几条预测")
@@ -256,12 +262,15 @@ def main() -> int:
removed = client.delete_predictions(args.project, model_version)
logger.info("删除同版本(%s)旧预测 %d 条", model_version, removed)
from PIL import Image
from ultralytics import YOLO
model = YOLO(str(weights))
classes = model.names
device_arg = device.index or 0 if device.type == "cuda" else device.type
if args.tile_size > 0:
logger.info("切片推理 %dx%d,重叠 %d(不缩放)", args.tile_size, args.tile_size, args.tile_overlap)
else:
logger.warning("整图推理 imgsz=%d —— ultralytics 会 letterbox 缩放", args.imgsz)
pending: list[dict] = []
uploaded = with_boxes = total_boxes = 0
@@ -275,35 +284,30 @@ def main() -> int:
uploaded += len(pending)
pending = []
for start in range(0, len(pairs), args.batch):
chunk = pairs[start:start + args.batch]
results = model.predict(
source=[str(path) for _, path in chunk], imgsz=args.imgsz, conf=args.conf,
device=device_arg, verbose=False,
for processed, (task_id, path) in enumerate(pairs, start=1):
image = load_image(path)
height, width = image.shape[:2]
detections = predict_image(
model, image, size=args.tile_size, overlap=args.tile_overlap, conf=args.conf,
device=device_arg, batch=args.batch, whole_image_imgsz=args.imgsz,
)
for (task_id, path), result in zip(chunk, results, strict=True):
with Image.open(path) as image:
width, height = image.size
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,
)
]
prediction = to_prediction(boxes, width, height, "label", "image", model_version)
prediction["task"] = task_id
pending.append(prediction)
total_boxes += len(boxes)
with_boxes += bool(boxes)
boxes = [
(classes[box.class_id], (box.x1, box.y1, box.x2, box.y2), box.score)
for box in detections
]
prediction = to_prediction(boxes, width, height, "label", "image", model_version)
prediction["task"] = task_id
pending.append(prediction)
total_boxes += len(boxes)
with_boxes += bool(boxes)
if len(pending) >= args.chunk:
try:
flush()
except RuntimeError as exc:
logger.error("%s", exc)
return 1
processed = start + len(chunk)
if processed % 1000 < args.batch:
if processed % 1000 == 0:
logger.info("处理 [%d/%d],已上传 %d", processed, len(pairs), uploaded)
try:
flush()