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
@@ -102,6 +102,53 @@ def match_boxes(
|
||||
)
|
||||
|
||||
|
||||
def containment(inner: Box, outer: Box) -> float:
|
||||
"""inner 有多大比例落在 outer 里（0~1）。面积悬殊时 IoU 会很小，但这个值接近 1。"""
|
||||
left, top = max(inner.x1, outer.x1), max(inner.y1, outer.y1)
|
||||
right, bottom = min(inner.x2, outer.x2), min(inner.y2, outer.y2)
|
||||
if right <= left or bottom <= top:
|
||||
return 0.0
|
||||
area = max(inner.x2 - inner.x1, 0) * max(inner.y2 - inner.y1, 0)
|
||||
return (right - left) * (bottom - top) / area if area > 0 else 0.0
|
||||
|
||||
|
||||
def nms(
|
||||
boxes: Sequence[Box],
|
||||
iou_threshold: float = 0.5,
|
||||
per_class: bool = True,
|
||||
contained_threshold: float = 0.0,
|
||||
) -> list[int]:
|
||||
"""非极大值抑制，返回保留下来的下标（按输入顺序）。
|
||||
|
||||
用于把切片推理的结果拼回整图：相邻片在重叠区会把同一处缺陷各检出一次，
|
||||
合并后必须去重。注意这与"禁止框重叠"是两回事——这里只压掉重复，
|
||||
真正分开的两个缺陷即使轻微相邻也都保留。
|
||||
|
||||
Args:
|
||||
per_class: 只在同类之间抑制。不同类别的框重叠往往是有意义的分歧，不该互相压制。
|
||||
contained_threshold: 额外按**包含比例**抑制，0 = 关闭。切片时缺陷被切割线切开，
|
||||
相邻片只看到一角，会给出一个小残框；它与完整框的 IoU 很小（面积悬殊），
|
||||
光靠 IoU 压不掉，但它几乎完全落在完整框里。
|
||||
"""
|
||||
order = sorted(range(len(boxes)), key=lambda i: -(boxes[i].score or 0.0))
|
||||
kept: list[int] = []
|
||||
for index in order:
|
||||
box = boxes[index]
|
||||
duplicate = False
|
||||
for other in kept:
|
||||
if per_class and boxes[other].class_id != box.class_id:
|
||||
continue
|
||||
if iou(box, boxes[other]) > iou_threshold:
|
||||
duplicate = True
|
||||
break
|
||||
if contained_threshold and containment(box, boxes[other]) >= contained_threshold:
|
||||
duplicate = True
|
||||
break
|
||||
if not duplicate:
|
||||
kept.append(index)
|
||||
return sorted(kept)
|
||||
|
||||
|
||||
def draw_boxes(image, ground_truth: Sequence[Box], predictions: Sequence[Box],
|
||||
classes: Sequence[str], match: Match, width: int = 2):
|
||||
"""把真值与预测画在图上（就地修改传入的 PIL 图像并返回）。
|
||||
|
||||
@@ -0,0 +1,85 @@
|
||||
"""用 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)
|
||||
@@ -132,6 +132,65 @@ def remap_boxes(
|
||||
return kept
|
||||
|
||||
|
||||
def merge_tiles(
|
||||
per_tile: Sequence[Sequence[Box]],
|
||||
tiles: Sequence[Tile],
|
||||
size: int,
|
||||
width: int,
|
||||
height: int,
|
||||
iou_threshold: float = 0.5,
|
||||
contained_threshold: float = 0.8,
|
||||
) -> list[Box]:
|
||||
"""把每片各自的框拼回整图，返回相对原图归一化的框。
|
||||
|
||||
三步：映射回原图像素 → 裁掉落在镜像填充区里的部分（那里是复制品，不是真实区域）
|
||||
→ 跨片去重。
|
||||
|
||||
去重同时看 IoU 与包含比例：缺陷被切割线切开时，相邻片只看到一角，给出的小残框
|
||||
与完整框的 IoU 很小（面积悬殊），但几乎整个落在完整框里。只看 IoU 会把残框留下来。
|
||||
|
||||
残框与完整框互不包含的情况（缺陷正好被切成两半、两片各见一半）只能靠 `overlap`
|
||||
避免——重叠大于缺陷尺寸时，总有一片能完整看到它。
|
||||
"""
|
||||
from engine.detection import nms
|
||||
|
||||
collected = []
|
||||
for tile, boxes in zip(tiles, per_tile, strict=True):
|
||||
for box in boxes:
|
||||
mapped = to_source(box, tile, size)
|
||||
x1, y1 = max(mapped.x1, 0.0), max(mapped.y1, 0.0)
|
||||
x2, y2 = min(mapped.x2, float(width)), min(mapped.y2, float(height))
|
||||
if x2 <= x1 or y2 <= y1:
|
||||
continue
|
||||
collected.append(
|
||||
Box(box.class_id, x1 / width, y1 / height, x2 / width, y2 / height, box.score)
|
||||
)
|
||||
return [
|
||||
collected[i]
|
||||
for i in nms(collected, iou_threshold, contained_threshold=contained_threshold)
|
||||
]
|
||||
|
||||
|
||||
def tiled_predict(
|
||||
predict_batch,
|
||||
image: np.ndarray,
|
||||
size: int,
|
||||
overlap: int = 0,
|
||||
iou_threshold: float = 0.5,
|
||||
) -> list[Box]:
|
||||
"""按固定边长切片推理再拼回，返回整图上的归一化框。
|
||||
|
||||
Args:
|
||||
predict_batch: 接收一批 `size × size` 的图、返回每张图的归一化框列表。
|
||||
推理框架留在调用方，这里只管几何，便于单测。
|
||||
image: HxW 或 HxWxC 的原图。
|
||||
"""
|
||||
height, width = image.shape[:2]
|
||||
tiles = plan_tiles(width, height, size, overlap)
|
||||
patches = [cut(image, tile, size) for tile in tiles]
|
||||
return merge_tiles(predict_batch(patches), tiles, size, width, height, iou_threshold)
|
||||
|
||||
|
||||
def to_source(box: Box, tile: Tile, size: int) -> Box:
|
||||
"""反向：把片上的归一化框映射回原图的**像素** xyxy。
|
||||
|
||||
|
||||
Reference in New Issue
Block a user