Yolo 26 with nms required, argument for cli version passing in tests

Author: dtronmans

tests/constants.py

@@ -76,6 +76,7 @@
     {"name": "yolo26x", "version": "v26"},
     {"name": "yolo26n-seg", "version": "v26"},
     {"name": "yolo26n-pose", "version": "v26"},
+    {"name": "yolo26n", "version": "v26_nms", "cli_version": "yolov26_nms"},
     {"name": "yolov8n-cls", "version": "v8"},
     {"name": "yolov8n-seg", "version": "v8"},
     {"name": "yolov8n-pose", "version": "v8"},

tests/test_end2end.py

@@ -15,7 +15,7 @@
 @pytest.mark.parametrize(
     "model",
     TEST_MODELS,
-    ids=[model["name"] for model in TEST_MODELS],
+    ids=[model.get("cli_version", model["name"]) if model.get("cli_version") else model["name"] for model in TEST_MODELS],
 )
 def test_cli_conversion(model: dict, test_config: dict, subtests):
     """Tests the whole CLI conversion flow with no extra params specified."""
@@ -50,6 +50,8 @@ def test_cli_conversion(model: dict, test_config: dict, subtests):
             pytest.skip("Weights not present and `download_weights` not set")
 
     command = ["tools", model_path]
+    if model.get("cli_version"):
+        command += ["--version", model.get("cli_version")]
     if model.get("size"):  # edge case when stride=64 is needed
         command += ["--imgsz", model.get("size")]
 

tools/main.py

@@ -28,6 +28,7 @@
     YOLOV11_CONVERSION,
     YOLOV12_CONVERSION,
     YOLOV26_CONVERSION,
+    YOLOV26_NMS_CONVERSION,
     detect_version,
 )
 
@@ -50,6 +51,7 @@
     YOLOV11_CONVERSION,
     YOLOV12_CONVERSION,
     YOLOV26_CONVERSION,
+    YOLOV26_NMS_CONVERSION,
 ]
 
 
@@ -176,6 +178,7 @@ def convert(
             YOLOV9_CONVERSION,
             YOLOV11_CONVERSION,
             YOLOV12_CONVERSION,
+            YOLOV26_NMS_CONVERSION,
         ]:
             from tools.yolo.yolov8_exporter import YoloV8Exporter
 

tools/modules/heads.py

@@ -366,11 +366,17 @@ def __init__(self, old_detect, use_rvc2: bool):
 
         self.use_rvc2 = use_rvc2
 
-        self.proj_conv = nn.Conv2d(old_detect.dfl.c1, 1, 1, bias=False).requires_grad_(
-            False
-        )
-        x = torch.arange(old_detect.dfl.c1, dtype=torch.float)
-        self.proj_conv.weight.data[:] = nn.Parameter(x.view(1, old_detect.dfl.c1, 1, 1))
+        # yolo26: dfl will be nn.Identity(), we set proj_conv = None and skip the DFL block in forward
+        if hasattr(old_detect.dfl, "c1"):
+            self.proj_conv = nn.Conv2d(
+                old_detect.dfl.c1, 1, 1, bias=False
+            ).requires_grad_(False)
+            x = torch.arange(old_detect.dfl.c1, dtype=torch.float)
+            self.proj_conv.weight.data[:] = nn.Parameter(
+                x.view(1, old_detect.dfl.c1, 1, 1)
+            )
+        else:
+            self.proj_conv = None
 
     def forward(self, x):
         bs = x[0].shape[0]  # batch size
@@ -382,9 +388,10 @@ def forward(self, x):
 
             # ------------------------------
             # DFL PART
-            box = box.view(bs, 4, self.reg_max, h * w).permute(0, 2, 1, 3)
-            box = self.proj_conv(F.softmax(box, dim=1))[:, 0]
-            box = box.reshape([bs, 4, h, w])
+            if self.proj_conv is not None:
+                box = box.view(bs, 4, self.reg_max, h * w).permute(0, 2, 1, 3)
+                box = self.proj_conv(F.softmax(box, dim=1))[:, 0]
+                box = box.reshape([bs, 4, h, w])
             # ------------------------------
 
             cls = self.cv3[i](x[i])

tools/version_detection/__init__.py

@@ -13,6 +13,7 @@
     YOLOV11_CONVERSION,
     YOLOV12_CONVERSION,
     YOLOV26_CONVERSION,
+    YOLOV26_NMS_CONVERSION,
     detect_version,
 )
 
@@ -30,6 +31,7 @@
     "YOLOV11_CONVERSION",
     "YOLOV12_CONVERSION",
     "YOLOV26_CONVERSION",
+    "YOLOV26_NMS_CONVERSION",
     "GOLD_YOLO_CONVERSION",
     "UNRECOGNIZED",
 ]

tools/version_detection/version_detection.py

@@ -18,6 +18,7 @@
 YOLOV11_CONVERSION = "yolov11"
 YOLOV12_CONVERSION = "yolov12"
 YOLOV26_CONVERSION = "yolov26"
+YOLOV26_NMS_CONVERSION = "yolov26_nms"
 GOLD_YOLO_CONVERSION = "goldyolo"
 UNRECOGNIZED = "none"
 

tools/yolo/yolov8_exporter.py

@@ -119,9 +119,16 @@ def __init__(
     def load_model(self):
         # load the model
         model, _ = load_checkpoint(
-            self.model_path, device="cpu", inplace=True, fuse=True
+            self.model_path, device="cpu", inplace=True, fuse=False
         )
 
+        # for yolo26 end2end has to be disabled before fusing
+        # otherwise cv2/cv3 are removed in the fuse process
+        head = model.model[-1]
+        if getattr(head, "end2end", False):
+            head.end2end = False
+        model.fuse()
+
         self.mode = -1
         if isinstance(model.model[-1], (Segment)) or isinstance(
             model.model[-1], (YOLOESegment)