feat: export classifier with YOLO OpenVINO interface

README.md

@@ -801,6 +801,34 @@ ultralytics 自己还会在 `runs/detect/film/` 下写：
 配置支持 `_base_` 继承与命令行 `--set a.b.c=value` 覆盖，见 [utils/config.py](utils/config.py)。
 新增骨架 / 数据集只需用注册表装饰器注册，再在 yaml 里按 `name` 引用。
 
+## 分类模型复用 YOLO OpenVINO 接口
+
+生产端只使用类别时，可以将分类器导出成与既有 YOLO IR 相同的输入输出接口：
+
+```bash
+uv pip install openvino
+uv run python scripts/export_classifier_openvino.py \
+    --checkpoint runs/classify/best.pth \
+    --template runs/detect/models/film_128_20261013_0902_openvino_model/film_128_20261013_0902.xml \
+    --out runs/classify/models/resnet18_128_yolo_adapter
+```
+
+`--template` 请指定实际部署的 YOLO `.xml`，其同目录需要有 `metadata.yaml`。
+脚本读取分类检查点的配置和同目录 `eval_report.json` 中的类别顺序；也可用
+`--classes classes.txt` 或 `--classes eval_report.json` 明确指定分类器的类别顺序。
+两者类别集合不一致、训练预处理不受支持或输入尺寸不匹配时会拒绝导出。
+
+产出的 XML/BIN 使用模板文件名，输入为 RGB float32 BCHW、值域 `[0,1]`。
+逐图标准化按训练配置嵌入模型，分类概率按 YOLO 类别顺序排列。
+输出保持 `[1,4+nc,anchors]`：仅第一个候选携带全图伪框及 Softmax 概率，其他候选为零。
+伪框是像素 `cx,cy,w,h`，不表示缺陷位置。脚本不进行 resize，图片必须已切为目标尺寸。
+
+部署时成对替换 XML/BIN；使用元数据的程序也需同步 `metadata.yaml`。
+仓库的 `deploy/cpp` 解码器可直接读取适配模型，生产代码若有额外硬编码则需核对。
+Softmax 分数不同于 YOLO 置信度，需在实际数据上确认过滤阈值。
+`normal` 也作为普通分类结果返回；原图被切成多片时，每片各返回一个分类结果。
+输出目录必须是新目录，原有模型不会被覆盖。
+
 ## 开发
 
 ```bash

scripts/export_classifier_openvino.py

@@ -0,0 +1,182 @@
+#!/usr/bin/env python
+"""Export a classifier with the input/output contract of a YOLO OpenVINO IR.
+
+The sole active candidate is a full-image pseudo box, not a localization result.
+Input must already be RGB float32 BCHW in [0, 1], exactly as in deploy/cpp.
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import yaml
+from torch import nn
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from models import build_model  # noqa: E402
+
+
+def ordered_names(names) -> list[str]:
+    if isinstance(names, dict):
+        indexed = {int(k): v for k, v in names.items()}
+        if set(indexed) != set(range(len(indexed))):
+            raise ValueError("Class IDs must be contiguous from zero")
+        names = [indexed[i] for i in range(len(indexed))]
+    if not isinstance(names, list) or not names or not all(isinstance(n, str) for n in names):
+        raise ValueError("Expected a nonempty class-name list or indexed mapping")
+    if len(set(names)) != len(names):
+        raise ValueError("Duplicate class names")
+    return names
+
+
+def class_order(source: list[str], target: list[str]) -> list[int]:
+    source, target = ordered_names(source), ordered_names(target)
+    if set(source) != set(target):
+        raise ValueError(f"Class sets differ: classifier={source}, YOLO={target}")
+    return [source.index(name) for name in target]
+
+
+def preprocessing(config: dict, size: int) -> float | None:
+    """Reject preprocessing that cannot be reproduced without resampling the input."""
+    transforms = config["data"]["val_transforms"]
+    sequence = [t["name"] for t in transforms]
+    if sequence not in (["Resize", "ToTensor", "PerImageStandardize"],
+                        ["ToTensor", "PerImageStandardize"], ["Resize", "ToTensor"], ["ToTensor"]):
+        raise ValueError(f"Unsupported validation preprocessing: {sequence}")
+    for transform in transforms:
+        if transform["name"] == "Resize":
+            if transform["size"] != [size, size]:
+                raise ValueError("Training size differs from template; retrain at the target size")
+        if transform["name"] == "PerImageStandardize":
+            return float(transform.get("eps", 1e-5))
+    return None
+
+
+class YoloClassifierAdapter(nn.Module):
+    def __init__(self, classifier: nn.Module, size: int, anchors: int,
+                 order: list[int], eps: float | None = 1e-5):
+        super().__init__()
+        if size <= 0 or anchors <= 0:
+            raise ValueError("size and anchors must be positive")
+        self.classifier = classifier
+        self.eps = eps
+        self.register_buffer("order", torch.tensor(order, dtype=torch.long))
+        self.register_buffer("box", torch.tensor([size / 2, size / 2, size, size])
+                             .reshape(1, 4, 1))
+        # Only candidate zero is active; do not duplicate predictions across anchors.
+        self.register_buffer("inactive", torch.zeros(1, 4 + len(order), anchors - 1))
+
+    def normalize(self, x: torch.Tensor) -> torch.Tensor:
+        if self.eps is None:
+            return x
+        dims = (1, 2, 3)
+        mean = x.mean(dim=dims, keepdim=True)
+        # Match datasets.transforms.PerImageStandardize: torch.std correction=1.
+        variance = ((x - mean) ** 2).sum(dim=dims, keepdim=True) / (x[0].numel() - 1)
+        return (x - mean) / (torch.sqrt(variance) + self.eps)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        scores = self.classifier(self.normalize(x)).softmax(dim=1)[:, self.order]
+        active = torch.cat((self.box.expand(x.shape[0], -1, -1), scores.unsqueeze(2)), dim=1)
+        return torch.cat((active, self.inactive.expand(x.shape[0], -1, -1)), dim=2)
+
+
+def main() -> int:
+    import openvino as ov
+
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", default="runs/classify/best.pth")
+    parser.add_argument("--template", required=True, help="Production YOLO .xml (metadata.yaml beside it)")
+    parser.add_argument("--classes", help="Classifier class order: classes.txt or eval_report.json")
+    parser.add_argument("--out", required=True, help="New directory; existing directories are rejected")
+    args = parser.parse_args()
+    checkpoint, template, out = Path(args.checkpoint), Path(args.template), Path(args.out)
+    if out.exists():
+        raise FileExistsError(f"Refusing to overwrite {out}")
+    state = torch.load(checkpoint, map_location="cpu", weights_only=False)
+    config = state["config"]
+    if args.classes:
+        path = Path(args.classes)
+        source = (json.loads(path.read_text())["classes"] if path.suffix == ".json"
+                  else path.read_text().splitlines())
+    else:
+        source = config["data"].get("classes")
+        if not source:
+            report = checkpoint.parent / "eval_report.json"
+            if not report.exists():
+                raise ValueError("No recorded class order; provide --classes (do not guess folder order)")
+            source = json.loads(report.read_text())["classes"]
+    source = ordered_names(source)
+    metadata = yaml.safe_load((template.parent / "metadata.yaml").read_text())
+    if metadata.get("end2end") or metadata.get("args", {}).get("nms"):
+        raise ValueError("Expected raw [1, 4+nc, anchors] output, not end-to-end/NMS output")
+    target = ordered_names(metadata["names"])
+    order = class_order(source, target)
+    core = ov.Core()
+    reference = core.read_model(str(template))
+    if len(reference.inputs) != 1 or len(reference.outputs) != 1:
+        raise ValueError("Expected one input and one output")
+    input_shape, output_shape = list(reference.input().shape), list(reference.output().shape)
+    if (len(input_shape) != 4 or input_shape[:2] != [1, 3]
+            or input_shape[2] != input_shape[3] or len(output_shape) != 3
+            or output_shape[:2] != [1, 4 + len(target)] or output_shape[2] < 1):
+        raise ValueError(f"Unsupported template shapes: {input_shape}, {output_shape}")
+    if (reference.input().element_type != ov.Type.f32
+            or reference.output().element_type != ov.Type.f32):
+        raise ValueError("Only float32 input/output are supported")
+    size = input_shape[2]
+    eps = preprocessing(config, size)
+    model_config = copy.deepcopy(config["model"])
+    model_config["backbone"]["pretrained"] = False
+    model_config["init_from"] = None
+    model_config["num_classes"] = len(source)
+    model = build_model(model_config).eval()
+    model.load_state_dict(state["model"], strict=True)
+    adapter = YoloClassifierAdapter(model, size, output_shape[2], order, eps).eval()
+    torch.manual_seed(42)
+    example = torch.rand(input_shape)
+    converted = ov.convert_model(adapter, example_input=example, input=input_shape)
+    converted.input().get_tensor().set_names(reference.input().get_names())
+    converted.output().get_tensor().set_names(reference.output().get_names())
+    # Preserve any interface metadata used by OpenVINO consumers, then mark pseudo boxes.
+    for key, value in reference.get_rt_info().items():
+        if key != "conversion_parameters":
+            converted.set_rt_info(value, [key])
+    converted.set_rt_info("classifier_pseudo_box", ["adapter_kind"])
+    out.mkdir(parents=True)
+    xml = out / template.name
+    ov.save_model(converted, str(xml), compress_to_fp16=True)
+    compiled = core.compile_model(str(xml), "CPU", {"PERFORMANCE_HINT": "LATENCY"})
+    errors = []
+    with torch.inference_mode():
+        for sample in (example, torch.zeros_like(example), torch.ones_like(example)):
+            expected = adapter(sample).numpy()
+            actual = compiled(sample.numpy())[compiled.output()]
+            np.testing.assert_allclose(actual, expected, atol=0.02, rtol=0.02)
+            assert np.argmax(actual[0, 4:, 0]) == np.argmax(expected[0, 4:, 0])
+            assert np.count_nonzero(actual[:, :, 1:]) == 0
+            errors.append(float(np.max(np.abs(expected - actual))))
+    metadata.update({"names": dict(enumerate(target)), "adapter_kind": "classifier_pseudo_box",
+                     "description": "Classifier with one full-image pseudo box; no localization",
+                     "end2end": False, "args": {"nms": False}, "imgsz": [size, size]})
+    (out / "metadata.yaml").write_text(yaml.safe_dump(metadata, allow_unicode=True))
+    (out / "classes.txt").write_text("\n".join(target) + "\n")
+    summary = {"checkpoint": str(checkpoint.resolve()), "template": str(template.resolve()),
+               "classifier_classes": source, "output_classes": target, "class_order": order,
+               "input_shape": input_shape, "output_shape": output_shape,
+               "standardize_eps": eps, "max_abs_errors": errors,
+               "pseudo_box": [size / 2, size / 2, size, size]}
+    (out / "adapter.json").write_text(json.dumps(summary, ensure_ascii=False, indent=2))
+    print(json.dumps({"model": str(xml), **summary}, ensure_ascii=False, indent=2))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

tests/test_classifier_adapter.py

@@ -0,0 +1,48 @@
+import pytest
+import torch
+from torch import nn
+
+from scripts.export_classifier_openvino import YoloClassifierAdapter, class_order, preprocessing
+
+
+class FixedClassifier(nn.Module):
+    def forward(self, x):
+        return torch.tensor([1.0, 3.0, -1.0]).expand(x.shape[0], -1)
+
+
+def test_yolo_contract_and_class_reordering():
+    adapter = YoloClassifierAdapter(FixedClassifier(), 128, 336, [2, 0, 1])
+    result = adapter(torch.rand(1, 3, 128, 128))
+    assert result.shape == (1, 7, 336)
+    torch.testing.assert_close(result[0, :4, 0], torch.tensor([64., 64., 128., 128.]))
+    expected = torch.tensor([1., 3., -1.]).softmax(0)[[2, 0, 1]]
+    torch.testing.assert_close(result[0, 4:, 0], expected)
+    assert result[0, 4:, 0].argmax().item() == 2
+    assert torch.count_nonzero(result[:, :, 1:]) == 0
+
+
+def test_normalization_matches_training_for_each_image():
+    adapter = YoloClassifierAdapter(FixedClassifier(), 128, 1, [0, 1, 2])
+    images = torch.stack((torch.rand(3, 128, 128), torch.ones(3, 128, 128)))
+    expected = torch.stack([(x - x.mean()) / (x.std() + 1e-5) for x in images])
+    torch.testing.assert_close(adapter.normalize(images), expected)
+    assert torch.isfinite(adapter(images)).all()
+
+
+def test_class_sets_must_match():
+    assert class_order(["normal", "mosquito"], ["mosquito", "normal"]) == [1, 0]
+    with pytest.raises(ValueError, match="Class sets differ"):
+        class_order(["normal", "mosquito"], ["normal", "blackdot"])
+    with pytest.raises(ValueError, match="Duplicate"):
+        class_order(["normal", "normal"], ["normal"])
+
+
+def test_refuses_resampling_and_unknown_preprocessing():
+    config = {"data": {"val_transforms": [
+        {"name": "Resize", "size": [160, 160]}, {"name": "ToTensor"},
+        {"name": "PerImageStandardize"}]}}
+    with pytest.raises(ValueError, match="Training size differs"):
+        preprocessing(config, 128)
+    config["data"]["val_transforms"] = [{"name": "Normalize"}]
+    with pytest.raises(ValueError, match="Unsupported"):
+        preprocessing(config, 128)