source: convert_lib — a Detect head must hold the box branch (cv2.* / one2one_cv2.*)

source/convert_lib.py

@@ -115,15 +115,23 @@ _HEAD_MODULE = re.compile(r"^(/model\.\d+)/")
 def detect_head_prefix(ops: "Iterable[tuple[str, str]]") -> "Optional[str]":
     """The Ultralytics Detect module (``/model.22`` for YOLOv9, ``/model.23``
     for YOLO26): the module holding the graph's LAST Sigmoid, which is the
-    class-score sigmoid (the Ultralytics exporter finds its head the same way).
+    class-score sigmoid (the Ultralytics exporter finds its head the same way),
+    provided that module holds the box branch (``cv2.*`` / ``one2one_cv2.*``).
     ``ops`` is ``(friendly_name, type_name)`` in graph order. None when the
-    model is not an Ultralytics export — it is then quantised whole."""
+    model is not an Ultralytics detector — it then keeps NNCF's defaults."""
+    ops = list(ops)
     last = None
     for name, kind in ops:
         if kind == "Sigmoid":
             last = name
     match = _HEAD_MODULE.match(last) if last else None
-    return match.group(1) if match else None
+    if not match:
+        return None
+    head = match.group(1)
+    # `/model.N/` is also PyTorch's naming for any `self.model = nn.Sequential`:
+    # only a module that holds the box branch is a Detect head.
+    box_branch = (f"{head}/cv2.", f"{head}/one2one_cv2.")
+    return head if any(name.startswith(box_branch) for name, _ in ops) else None
 
 
 def is_end_to_end_head(names: "Iterable[str]", head: str) -> bool: