fix: suppress normal pseudo boxes in classifier exports

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` 恢复普通分类输出。
+原图被切成多片时，每片独立分类，仅前景片返回伪框。
 输出目录必须是新目录，原有模型不会被覆盖。
 
 ## 开发

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))

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"},