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>
86 lines
2.7 KiB
Python
86 lines
2.7 KiB
Python
"""用 ultralytics 做切片推理的公共封装。
|
||
|
||
三个调用方（可视化、批量推送预测、ML backend）都要"切片 -> 逐片推理 -> 拼回原图"，
|
||
逻辑一致，放在一处免得各写一份走样。
|
||
|
||
几何在 `engine.tiling` 里，与推理框架无关且单独可测；这里只负责把 ultralytics 接上去。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from collections.abc import Sequence
|
||
|
||
import numpy as np
|
||
|
||
from engine.detection import Box
|
||
from engine.tiling import tiled_predict
|
||
|
||
|
||
def predict_image(
|
||
model,
|
||
image: np.ndarray,
|
||
size: int,
|
||
overlap: int = 0,
|
||
conf: float = 0.25,
|
||
device: str | int = "cpu",
|
||
batch: int = 16,
|
||
iou_threshold: float = 0.5,
|
||
whole_image_imgsz: int | None = None,
|
||
) -> list[Box]:
|
||
"""对一张任意尺寸的图做切片推理，返回整图上的归一化框。
|
||
|
||
`size <= 0` 时退化为整图推理，交给 ultralytics 按 `whole_image_imgsz` 做 letterbox
|
||
——那会缩放，只用于和旧模型对比。
|
||
"""
|
||
if size <= 0:
|
||
return _predict_batch(
|
||
model, [image], whole_image_imgsz, conf=conf, device=device, batch=batch
|
||
)[0]
|
||
|
||
def run(patches: Sequence[np.ndarray]) -> list[list[Box]]:
|
||
return _predict_batch(model, patches, size, conf, device, batch)
|
||
|
||
return tiled_predict(run, image, size, overlap, iou_threshold)
|
||
|
||
|
||
def _predict_batch(
|
||
model,
|
||
images: Sequence[np.ndarray],
|
||
size: int | None,
|
||
conf: float,
|
||
device: str | int,
|
||
batch: int,
|
||
) -> list[list[Box]]:
|
||
"""分批送进 ultralytics，返回每张图的归一化框。
|
||
|
||
必须自己分批：ultralytics 不理会 batch 参数，会把整个列表当成一个 batch，
|
||
片数一多就会撑爆显存。
|
||
"""
|
||
results = []
|
||
for start in range(0, len(images), batch):
|
||
chunk = list(images[start:start + batch])
|
||
kwargs = {"conf": conf, "device": device, "verbose": False}
|
||
if size:
|
||
kwargs["imgsz"] = size
|
||
results.extend(model.predict(source=chunk, **kwargs))
|
||
|
||
out = []
|
||
for result in results:
|
||
out.append([
|
||
Box(int(class_id), *(float(v) for v in xyxyn), float(score))
|
||
for xyxyn, class_id, score in zip(
|
||
result.boxes.xyxyn.tolist(), result.boxes.cls.tolist(),
|
||
result.boxes.conf.tolist(), strict=True,
|
||
)
|
||
])
|
||
return out
|
||
|
||
|
||
def load_image(path) -> np.ndarray:
|
||
"""读图为 numpy 数组，保持原始通道数（灰度图不强行转 RGB，免得多占三倍内存）。"""
|
||
from PIL import Image
|
||
|
||
with Image.open(path) as handle:
|
||
image = handle.convert("L") if handle.mode in ("L", "1", "I;16") else handle.convert("RGB")
|
||
return np.asarray(image)
|