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:
co-authored by
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()
|
||||
|
||||
Reference in New Issue
Block a user