diff --git a/engine/tiled_inference.py b/engine/tiled_inference.py index 187d7c8..7343147 100644 --- a/engine/tiled_inference.py +++ b/engine/tiled_inference.py @@ -16,6 +16,41 @@ from engine.detection import Box from engine.tiling import tiled_predict +def training_imgsz(model) -> int | None: + """权重里记录的训练输入尺寸；读不到返回 None。 + + ultralytics 把训练参数存在 checkpoint 里，推理时可以取回来核对。 + """ + for source in ( + getattr(model, "ckpt", None) or {}, + getattr(model, "overrides", None) or {}, + ): + value = source.get("train_args", source).get("imgsz") if isinstance(source, dict) else None + if isinstance(value, int): + return value + return None + + +def check_size_matches(model, size: int, logger=None) -> str | None: + """推理尺寸与训练尺寸不一致时给出警告文本（一致或读不到则返回 None）。 + + 这一条是有代价换来的：用 640 整图训练的模型去跑 160 的切片推理，结果**看起来正常** + ——照样出框、照样有置信度，但类别系统性地错（实测 mosquito 全部变成 blackdot， + TP 为 0）。不主动核对的话，这种结果会被当成模型退化去查很久。 + """ + trained = training_imgsz(model) + if trained is None or size <= 0 or trained == size: + return None + message = ( + f"权重是按 imgsz={trained} 训练的，现在按 {size} 切片推理——两者不一致，" + f"结果不可信（会照常出框，但类别系统性地错）。" + f"要么换成按 {size} 训练的权重，要么用 --tile-size 0 退回整图推理" + ) + if logger is not None: + logger.warning("%s", message) + return message + + def predict_image( model, image: np.ndarray, diff --git a/ml_backend/model.py b/ml_backend/model.py index 4a3bac3..b4dabb6 100644 --- a/ml_backend/model.py +++ b/ml_backend/model.py @@ -32,7 +32,7 @@ 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 +from engine.tiled_inference import check_size_matches, load_image, predict_image # noqa: E402 logger = logging.getLogger(__name__) @@ -63,6 +63,9 @@ class FilmDefectDetector(LabelStudioMLBase): weights = os.environ.get("YOLO_WEIGHTS", "yolo26n.pt") logger.info("加载权重 %s", weights) FilmDefectDetector._model = YOLO(weights) + check_size_matches( + FilmDefectDetector._model, int(os.environ.get("YOLO_TILE_SIZE", 160)), logger + ) return FilmDefectDetector._model def predict(self, tasks: list[dict], context: dict | None = None, **kwargs) -> ModelResponse: diff --git a/scripts/push_predictions.py b/scripts/push_predictions.py index f1b4693..5d4dcc8 100644 --- a/scripts/push_predictions.py +++ b/scripts/push_predictions.py @@ -40,7 +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 engine.tiled_inference import check_size_matches, load_image, predict_image # noqa: E402 from utils import setup_logging # noqa: E402 from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402 @@ -267,6 +267,7 @@ def main() -> int: model = YOLO(str(weights)) classes = model.names device_arg = device.index or 0 if device.type == "cuda" else device.type + check_size_matches(model, args.tile_size, logger) if args.tile_size > 0: logger.info("切片推理 %dx%d，重叠 %d（不缩放）", args.tile_size, args.tile_size, args.tile_overlap) else: diff --git a/scripts/visualize_detections.py b/scripts/visualize_detections.py index d4f1b22..ffc2a53 100644 --- a/scripts/visualize_detections.py +++ b/scripts/visualize_detections.py @@ -46,7 +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 engine.tiled_inference import check_size_matches, load_image, predict_image # noqa: E402 from utils import setup_logging # noqa: E402 from utils.device import DeviceNotAvailableError, describe_device, resolve_device # noqa: E402 @@ -180,6 +180,7 @@ def main() -> int: from ultralytics import YOLO model = YOLO(args.weights) + check_size_matches(model, args.tile_size, logger) device_arg = device.index or 0 if device.type == "cuda" else device.type outcomes: dict[str, Counter[str]] = {name: Counter() for name in splits} @@ -246,6 +247,8 @@ def main() -> int: "conf": args.conf, "iou_threshold": args.iou_threshold, "imgsz": args.imgsz, + "tile_size": args.tile_size, + "tile_overlap": args.tile_overlap, "outcomes": {split: dict(counter) for split, counter in outcomes.items()}, "per_class": { split: { diff --git a/tests/test_size_guard.py b/tests/test_size_guard.py new file mode 100644 index 0000000..cca6293 --- /dev/null +++ b/tests/test_size_guard.py @@ -0,0 +1,61 @@ +"""推理尺寸与训练尺寸的一致性校验。 + +回归：用 640 整图训练的权重跑 160 切片推理，结果看起来完全正常——照常出框、 +照常有置信度，但类别系统性地错（实测 mosquito 的 TP 归零、599/760 被判成 blackdot）。 +这种错误不会自己暴露，必须主动核对。 +""" + + +from engine.tiled_inference import check_size_matches, training_imgsz + + +class _Model: + def __init__(self, ckpt=None, overrides=None): + if ckpt is not None: + self.ckpt = ckpt + if overrides is not None: + self.overrides = overrides + + +def test_reads_imgsz_from_checkpoint(): + assert training_imgsz(_Model(ckpt={"train_args": {"imgsz": 640}})) == 640 + + +def test_reads_imgsz_from_overrides(): + assert training_imgsz(_Model(overrides={"imgsz": 160})) == 160 + + +def test_missing_imgsz_returns_none(): + assert training_imgsz(_Model()) is None + assert training_imgsz(_Model(ckpt={})) is None + + +def test_mismatch_is_reported(): + message = check_size_matches(_Model(ckpt={"train_args": {"imgsz": 640}}), 160) + assert message is not None + assert "640" in message and "160" in message + + +def test_matching_sizes_are_silent(): + assert check_size_matches(_Model(ckpt={"train_args": {"imgsz": 160}}), 160) is None + + +def test_whole_image_mode_is_not_checked(): + """--tile-size 0 是显式退回整图推理，交给 ultralytics letterbox，不该报警。""" + assert check_size_matches(_Model(ckpt={"train_args": {"imgsz": 640}}), 0) is None + + +def test_unknown_training_size_is_not_reported(): + """读不到就不报——宁可不说，也不误导。""" + assert check_size_matches(_Model(), 160) is None + + +def test_logger_receives_the_warning(): + seen = [] + + class _Logger: + def warning(self, fmt, *args): + seen.append(fmt % args) + + check_size_matches(_Model(overrides={"imgsz": 640}), 160, _Logger()) + assert seen and "不一致" in seen[0]