From 03f667276bff00c19be55fce53e808464026a45f Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Wed, 14 Oct 2026 12:00:36 +0800 Subject: [PATCH] feat: export classifier with YOLO OpenVINO interface --- README.md | 28 ++++ scripts/export_classifier_openvino.py | 182 ++++++++++++++++++++++++++ tests/test_classifier_adapter.py | 48 +++++++ 3 files changed, 258 insertions(+) create mode 100644 scripts/export_classifier_openvino.py create mode 100644 tests/test_classifier_adapter.py diff --git a/README.md b/README.md index 6b4c4e6..cc28d98 100644 --- a/README.md +++ b/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 diff --git a/scripts/export_classifier_openvino.py b/scripts/export_classifier_openvino.py new file mode 100644 index 0000000..a22164e --- /dev/null +++ b/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()) diff --git a/tests/test_classifier_adapter.py b/tests/test_classifier_adapter.py new file mode 100644 index 0000000..7b40f2a --- /dev/null +++ b/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)