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
+28
-24
@@ -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()
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user