From 63bab36147ae0d411bda66ca6d75284b5acdb829 Mon Sep 17 00:00:00 2001 From: Shaun Zhou Date: Wed, 14 Oct 2026 15:31:52 +0800 Subject: [PATCH] fix: suppress normal pseudo boxes in classifier exports --- README.md | 5 +++- scripts/export_classifier_openvino.py | 21 ++++++++++++++-- tests/test_classifier_adapter.py | 35 +++++++++++++++++++++++++++ 3 files changed, 58 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index cc28d98..6d758bb 100644 --- a/README.md +++ b/README.md @@ -826,7 +826,10 @@ uv run python scripts/export_classifier_openvino.py \ 部署时成对替换 XML/BIN；使用元数据的程序也需同步 `metadata.yaml`。 仓库的 `deploy/cpp` 解码器可直接读取适配模型，生产代码若有额外硬编码则需核对。 Softmax 分数不同于 YOLO 置信度，需在实际数据上确认过滤阈值。 -`normal` 也作为普通分类结果返回；原图被切成多片时，每片各返回一个分类结果。 +默认 top1 为 `normal` 时，所有候选类别分数清零，C++ 在正数置信度阈值下返回空列表， +与 YOLO 的背景行为一致。类别编号保持不变，`normal` 仍占其原通道。 +用 `--background-class` 指定其他背景类别，或 `--keep-background-box` 恢复普通分类输出。 +原图被切成多片时，每片独立分类，仅前景片返回伪框。 输出目录必须是新目录，原有模型不会被覆盖。 ## 开发 diff --git a/scripts/export_classifier_openvino.py b/scripts/export_classifier_openvino.py index a22164e..0490763 100644 --- a/scripts/export_classifier_openvino.py +++ b/scripts/export_classifier_openvino.py @@ -61,12 +61,16 @@ def preprocessing(config: dict, size: int) -> float | None: class YoloClassifierAdapter(nn.Module): def __init__(self, classifier: nn.Module, size: int, anchors: int, - order: list[int], eps: float | None = 1e-5): + order: list[int], eps: float | None = 1e-5, + background_index: int | None = None): super().__init__() if size <= 0 or anchors <= 0: raise ValueError("size and anchors must be positive") self.classifier = classifier self.eps = eps + if background_index is not None and not 0 <= background_index < len(order): + raise ValueError("background_index must be an output class index") + self.background_index = background_index 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)) @@ -84,6 +88,9 @@ class YoloClassifierAdapter(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: scores = self.classifier(self.normalize(x)).softmax(dim=1)[:, self.order] + if self.background_index is not None: + foreground = (scores.argmax(dim=1) != self.background_index).unsqueeze(1) + scores = scores * foreground.to(scores.dtype) 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) @@ -95,6 +102,9 @@ def main() -> int: 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("--background-class", default="normal", help="Top1 class to emit no boxes for") + parser.add_argument("--keep-background-box", action="store_true", + help="Opt out of background suppression (legacy classification behavior)") 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) @@ -119,6 +129,9 @@ def main() -> int: 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) + if not args.keep_background_box and args.background_class not in target: + raise ValueError(f"Background class {args.background_class!r} is not in {target}") + background_index = None if args.keep_background_box else target.index(args.background_class) core = ov.Core() reference = core.read_model(str(template)) if len(reference.inputs) != 1 or len(reference.outputs) != 1: @@ -139,7 +152,8 @@ def main() -> int: 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() + adapter = YoloClassifierAdapter(model, size, output_shape[2], order, eps, + background_index=background_index).eval() torch.manual_seed(42) example = torch.rand(input_shape) converted = ov.convert_model(adapter, example_input=example, input=input_shape) @@ -165,6 +179,7 @@ def main() -> int: 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", + "background_class": None if background_index is None else args.background_class, "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") @@ -172,6 +187,8 @@ def main() -> int: "classifier_classes": source, "output_classes": target, "class_order": order, "input_shape": input_shape, "output_shape": output_shape, "standardize_eps": eps, "max_abs_errors": errors, + "background_class": None if background_index is None else args.background_class, + "background_index": background_index, "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)) diff --git a/tests/test_classifier_adapter.py b/tests/test_classifier_adapter.py index 7b40f2a..a1868dd 100644 --- a/tests/test_classifier_adapter.py +++ b/tests/test_classifier_adapter.py @@ -37,6 +37,41 @@ def test_class_sets_must_match(): class_order(["normal", "normal"], ["normal"]) +def test_background_suppression_uses_reordered_output_index(): + # Source class 1 wins, but it is output class 2 after reordering. + adapter = YoloClassifierAdapter(FixedClassifier(), 128, 336, [2, 0, 1], background_index=2) + result = adapter(torch.rand(1, 3, 128, 128)) + assert result.shape == (1, 7, 336) + assert torch.count_nonzero(result[:, 4:, :]) == 0 + # The existing C++ decoder returns zero candidates at any positive threshold. + assert not (result[0, 4:, :].max(dim=0).values >= 0.25).any() + + +def test_foreground_class_indices_and_scores_are_preserved(): + foreground = YoloClassifierAdapter(FixedClassifier(), 128, 336, [0, 1, 2], + background_index=0) + legacy = YoloClassifierAdapter(FixedClassifier(), 128, 336, [0, 1, 2]) + x = torch.rand(1, 3, 128, 128) + torch.testing.assert_close(foreground(x), legacy(x)) + assert foreground(x)[0, 4:, 0].argmax() == 1 + + +def test_background_gate_is_per_image_in_batch(): + class BatchClassifier(nn.Module): + def forward(self, x): + return x[:, :, 0, 0] + + x = torch.zeros(2, 3, 4, 4) + x[0, 0, 0, 0] = 3 # normal + x[1, 1, 0, 0] = 3 # blackdot + adapter = YoloClassifierAdapter(BatchClassifier(), 4, 10, [0, 1, 2], + eps=None, background_index=0) + result = adapter(x) + assert not torch.count_nonzero(result[0, 4:, :]) + assert result[1, 4:, 0].argmax() == 1 + assert result[1, 5, 0] > 0.9 + + def test_refuses_resampling_and_unknown_preprocessing(): config = {"data": {"val_transforms": [ {"name": "Resize", "size": [160, 160]}, {"name": "ToTensor"},