fix: suppress normal pseudo boxes in classifier exports

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