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。
|
||||
|
||||
|
||||
@@ -65,6 +65,8 @@ Label Studio 里：Settings → **Model** → Connect Model，URL 填 `http://lo
|
||||
| `YOLO_IMGSZ` | `640` | 推理尺寸，与训练保持一致 |
|
||||
| `YOLO_DEVICE` | `cuda` | `cuda` / `cpu` / `mps` |
|
||||
| `YOLO_BATCH` | `16` | 一次推理几张。Retrieve Predictions 会一次发来上千个任务，不分批会撑爆显存 |
|
||||
| `YOLO_TILE_SIZE` | `160` | 模型输入边长。按它切片推理再拼回，全程不缩放；0 = 整图推理（会 letterbox） |
|
||||
| `YOLO_TILE_OVERLAP` | `32` | 切片重叠像素。大于缺陷尺寸时，总有一片能完整看到它 |
|
||||
| `RETRAIN_AFTER` | `0` | 累计多少条新标注触发重训，0 = 不触发 |
|
||||
| `RETRAIN_COMMAND` | 空 | 触发时执行的命令 |
|
||||
|
||||
|
||||
+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):
|
||||
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,
|
||||
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)
|
||||
|
||||
+21
-17
@@ -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,
|
||||
)
|
||||
(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()
|
||||
|
||||
@@ -25,6 +25,9 @@ train / val 不分目录，但 `index.csv` 里有 `split` 列——训练集的
|
||||
同样地，"命中 / 漏检 / 误检"记在 `index.csv` 的 `outcome` 列里，按它排序筛选即可。
|
||||
|
||||
`--only errors` 只输出有漏检或误检的图，扫起来快得多。
|
||||
|
||||
模型的输入固定为 `--tile-size`（默认 160，与训练一致）时，图会按同样规则切片推理再拼回，
|
||||
全程不缩放。`--tile-size 0` 退回整图推理（ultralytics 自己 letterbox），用于和旧模型对比。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -43,6 +46,7 @@ import yaml # noqa: E402
|
||||
|
||||
from datasets.yolo import IMAGE_SUFFIXES # noqa: E402
|
||||
from engine.detection import Box, draw_boxes, match_boxes, parse_label # 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
|
||||
|
||||
@@ -124,6 +128,11 @@ def main() -> int:
|
||||
parser.add_argument("--conf", type=float, default=0.25, help="置信度阈值")
|
||||
parser.add_argument("--iou-threshold", type=float, default=0.5, help="判定命中的 IoU")
|
||||
parser.add_argument("--batch", type=int, default=16)
|
||||
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(
|
||||
"--only", choices=("all", "errors"), default="all",
|
||||
help="errors=只输出有漏检或误检的图，统计仍按全量计",
|
||||
@@ -150,7 +159,16 @@ def main() -> int:
|
||||
"数据集: %s（%d 类，%s）", args.data, len(classes),
|
||||
"，".join(f"{name} {len(images)} 张" for name, images in splits.items()),
|
||||
)
|
||||
logger.info("conf=%.2f IoU=%.2f imgsz=%d", args.conf, args.iou_threshold, args.imgsz)
|
||||
if args.tile_size > 0:
|
||||
logger.info(
|
||||
"conf=%.2f 命中 IoU=%.2f 切片 %dx%d 重叠 %d（不缩放）",
|
||||
args.conf, args.iou_threshold, args.tile_size, args.tile_size, args.tile_overlap,
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"conf=%.2f 命中 IoU=%.2f 整图推理 imgsz=%d —— ultralytics 会 letterbox 缩放",
|
||||
args.conf, args.iou_threshold, args.imgsz,
|
||||
)
|
||||
|
||||
if out_dir.exists() and any(out_dir.iterdir()):
|
||||
if not args.overwrite:
|
||||
@@ -172,26 +190,13 @@ def main() -> int:
|
||||
rows = []
|
||||
processed = 0
|
||||
|
||||
def predict_in_batches(images):
|
||||
"""自己按批切分：ultralytics 的 predict 不理会 batch 参数，
|
||||
一次性把整个列表当成一个 batch，上万张图必然 OOM。"""
|
||||
for start in range(0, len(images), args.batch):
|
||||
chunk = images[start:start + args.batch]
|
||||
results = model.predict(
|
||||
source=[str(p) for p in chunk], imgsz=args.imgsz, conf=args.conf,
|
||||
device=device_arg, verbose=False,
|
||||
)
|
||||
yield from zip(chunk, results, strict=True)
|
||||
|
||||
for split, images in splits.items():
|
||||
for image_path, result in predict_in_batches(images):
|
||||
predictions = [
|
||||
Box(int(cls), *(float(v) for v in xyxy), float(score))
|
||||
for xyxy, cls, score in zip(
|
||||
result.boxes.xyxyn.tolist(), result.boxes.cls.tolist(),
|
||||
result.boxes.conf.tolist(), strict=True,
|
||||
for image_path in images:
|
||||
predictions = predict_image(
|
||||
model, load_image(image_path), size=args.tile_size, overlap=args.tile_overlap,
|
||||
conf=args.conf, device=device_arg, batch=args.batch,
|
||||
whole_image_imgsz=args.imgsz,
|
||||
)
|
||||
]
|
||||
label_file = label_path_for(image_path)
|
||||
truth = parse_label(label_file.read_text(encoding="utf-8")) if label_file.exists() else []
|
||||
match = match_boxes(predictions, truth, args.iou_threshold)
|
||||
|
||||
@@ -13,6 +13,7 @@ from engine.detection import ( # noqa: E402
|
||||
from_yolo,
|
||||
iou,
|
||||
match_boxes,
|
||||
nms,
|
||||
parse_label,
|
||||
)
|
||||
|
||||
@@ -129,3 +130,41 @@ def test_draw_handles_boxes_at_the_border():
|
||||
prediction = Box(0, 0.0, 0.0, 0.5, 0.05, 0.9)
|
||||
draw_boxes(image, [], [prediction], ["blackdot"], match_boxes([prediction], []))
|
||||
assert COLOR_UNMATCHED in _colors(image)
|
||||
|
||||
|
||||
# ---------------- NMS ----------------
|
||||
|
||||
|
||||
def test_nms_removes_duplicates_of_the_same_defect():
|
||||
high = Box(0, 0.30, 0.30, 0.50, 0.50, 0.9)
|
||||
duplicate = Box(0, 0.31, 0.31, 0.51, 0.51, 0.7)
|
||||
assert nms([duplicate, high]) == [1]
|
||||
|
||||
|
||||
def test_nms_keeps_neighbours_below_the_threshold():
|
||||
"""真正分开的两个缺陷，轻微相邻也要都留住。"""
|
||||
a = Box(0, 0.0, 0.0, 0.5, 0.5, 0.9)
|
||||
b = Box(0, 0.4, 0.4, 0.9, 0.9, 0.8) # IoU 约 0.02
|
||||
assert nms([a, b]) == [0, 1]
|
||||
|
||||
|
||||
def test_nms_is_per_class_by_default():
|
||||
"""不同类别在同一位置的分歧是有意义的，不该互相压制。"""
|
||||
a = Box(0, 0.3, 0.3, 0.6, 0.6, 0.9)
|
||||
b = Box(4, 0.31, 0.31, 0.61, 0.61, 0.8)
|
||||
assert nms([a, b]) == [0, 1]
|
||||
assert nms([a, b], per_class=False) == [0]
|
||||
|
||||
|
||||
def test_nms_handles_empty_input():
|
||||
assert nms([]) == []
|
||||
|
||||
|
||||
def test_nms_can_suppress_contained_fragments():
|
||||
"""切片接缝上的残框：与完整框 IoU 很小，却几乎整个落在里面。"""
|
||||
full = Box(0, 0.30, 0.50, 0.40, 0.60, 0.9)
|
||||
fragment = Box(0, 0.39, 0.50, 0.40, 0.60, 0.8) # 只剩一条边
|
||||
assert iou(full, fragment) < 0.15
|
||||
|
||||
assert nms([full, fragment]) == [0, 1] # 只看 IoU 压不掉
|
||||
assert nms([full, fragment], contained_threshold=0.8) == [0] # 按包含比例能压掉
|
||||
|
||||
@@ -184,3 +184,106 @@ def test_score_survives_both_directions():
|
||||
mapped, = remap_boxes([box], tile, 100, 100, SIZE)
|
||||
assert mapped.score == 0.83
|
||||
assert to_source(mapped, tile, SIZE).score == 0.83
|
||||
|
||||
|
||||
# ---------------- 切片推理的拼回 ----------------
|
||||
|
||||
|
||||
def test_merge_deduplicates_across_the_seam():
|
||||
"""重叠区里同一处缺陷会被两片各检出一次，合并后应只剩一个。"""
|
||||
from engine.tiling import merge_tiles
|
||||
|
||||
width = height = 300
|
||||
tiles = plan_tiles(width, height, SIZE, overlap=64)
|
||||
# 让每片都"检出"同一个真实位置的框（像素 150..170）
|
||||
per_tile = []
|
||||
for tile in tiles:
|
||||
px1, py1, px2, py2 = 150.0, 150.0, 170.0, 170.0
|
||||
bx1 = (px1 - tile.x + tile.pad_left) / SIZE
|
||||
by1 = (py1 - tile.y + tile.pad_top) / SIZE
|
||||
bx2 = (px2 - tile.x + tile.pad_left) / SIZE
|
||||
by2 = (py2 - tile.y + tile.pad_top) / SIZE
|
||||
inside = 0 <= bx1 and bx2 <= 1 and 0 <= by1 and by2 <= 1
|
||||
per_tile.append([Box(0, bx1, by1, bx2, by2, 0.9)] if inside else [])
|
||||
|
||||
assert sum(len(b) for b in per_tile) >= 2 # 确实有多片检出
|
||||
merged = merge_tiles(per_tile, tiles, SIZE, width, height)
|
||||
assert len(merged) == 1
|
||||
assert merged[0].x1 * width == pytest.approx(150, abs=1e-6)
|
||||
assert merged[0].x2 * width == pytest.approx(170, abs=1e-6)
|
||||
|
||||
|
||||
def test_merge_keeps_genuinely_separate_defects():
|
||||
from engine.tiling import merge_tiles
|
||||
|
||||
width = height = 320
|
||||
tiles = plan_tiles(width, height, SIZE, overlap=0)
|
||||
per_tile = [[] for _ in tiles]
|
||||
per_tile[0] = [Box(0, 0.1, 0.1, 0.2, 0.2, 0.9)] # 左上片
|
||||
per_tile[3] = [Box(0, 0.8, 0.8, 0.9, 0.9, 0.8)] # 右下片
|
||||
assert len(merge_tiles(per_tile, tiles, SIZE, width, height)) == 2
|
||||
|
||||
|
||||
def test_merge_clips_boxes_that_reach_into_mirrored_padding():
|
||||
"""小图填充后，模型可能在镜像区出框；那里不是真实区域，要裁掉。"""
|
||||
from engine.tiling import merge_tiles
|
||||
|
||||
width = height = 80
|
||||
tiles = plan_tiles(width, height, SIZE)
|
||||
# 一个横跨整片的框，两端都落在镜像填充里
|
||||
merged = merge_tiles([[Box(0, 0.0, 0.4, 1.0, 0.6, 0.9)]], tiles, SIZE, width, height)
|
||||
|
||||
assert len(merged) == 1
|
||||
assert merged[0].x1 == pytest.approx(0.0)
|
||||
assert merged[0].x2 == pytest.approx(1.0) # 裁到原图边界，不超出
|
||||
assert 0.0 <= merged[0].y1 <= merged[0].y2 <= 1.0
|
||||
|
||||
|
||||
def test_merge_drops_boxes_entirely_inside_the_mirror():
|
||||
from engine.tiling import merge_tiles
|
||||
|
||||
width = height = 80
|
||||
tiles = plan_tiles(width, height, SIZE) # pad 40，原图占 40..120
|
||||
ghost = Box(0, 0.05, 0.05, 0.2, 0.2, 0.9) # 像素 8..32，全在镜像区
|
||||
assert merge_tiles([[ghost]], tiles, SIZE, width, height) == []
|
||||
|
||||
|
||||
def test_tiled_predict_wires_geometry_to_the_model():
|
||||
"""端到端：切片 -> 逐片推理 -> 拼回，位置要落在真实缺陷上。"""
|
||||
from engine.tiling import tiled_predict
|
||||
|
||||
image = np.zeros((300, 400), dtype=np.uint8)
|
||||
image[200:220, 100:130] = 255 # 一块亮斑
|
||||
|
||||
def predict_batch(patches):
|
||||
"""假模型：把片里最亮的连通区域当成检出。"""
|
||||
out = []
|
||||
for patch in patches:
|
||||
ys, xs = np.nonzero(patch == 255)
|
||||
if len(xs) == 0:
|
||||
out.append([])
|
||||
continue
|
||||
out.append([Box(0, xs.min() / SIZE, ys.min() / SIZE,
|
||||
(xs.max() + 1) / SIZE, (ys.max() + 1) / SIZE, 0.9)])
|
||||
return out
|
||||
|
||||
merged = tiled_predict(predict_batch, image, SIZE, overlap=32)
|
||||
assert len(merged) == 1
|
||||
assert merged[0].x1 * 400 == pytest.approx(100, abs=1)
|
||||
assert merged[0].y1 * 300 == pytest.approx(200, abs=1)
|
||||
assert merged[0].x2 * 400 == pytest.approx(130, abs=1)
|
||||
assert merged[0].y2 * 300 == pytest.approx(220, abs=1)
|
||||
|
||||
|
||||
def test_tiled_predict_on_a_small_image_uses_one_tile():
|
||||
from engine.tiling import tiled_predict
|
||||
|
||||
seen = []
|
||||
|
||||
def predict_batch(patches):
|
||||
seen.append(len(patches))
|
||||
assert all(p.shape[:2] == (SIZE, SIZE) for p in patches)
|
||||
return [[] for _ in patches]
|
||||
|
||||
tiled_predict(predict_batch, np.zeros((83, 116), dtype=np.uint8), SIZE)
|
||||
assert seen == [1]
|
||||
|
||||
Reference in New Issue
Block a user