Adding YOLOX native support

.gitmodules

@@ -25,3 +25,7 @@
 [submodule "YOLOv6R3"]
 	path = tools/yolov6r3/YOLOv6R3
 	url = https://github.com/meituan/YOLOv6
+[submodule "tools/yolox/YOLOX"]
+	path = tools/yolox/YOLOX
+	url = https://github.com/Megvii-BaseDetection/YOLOX.git
+	branch = main

README.md

@@ -120,6 +120,7 @@ Currently, the following models are supported:
 | `yolov11`     | **Detection, Instance Segmentation, Pose, Oriented Detection, Classification:** YOLO11n, YOLO11s, YOLO11m, YOLO11l, YOLO11x                                                                                                                                                         |
 | `yolov12`     | **Detection:** YOLO12n, YOLO12s, YOLO12m, YOLO12l, YOLO12x                                                                                                                                                                                                                          |
 | `yolo26`      | **Detection, Instance Segmentation, Pose, Semantic Segmentation:** YOLO26n, YOLO26s, YOLO26m, YOLO26l, YOLO26x                                                                                                                                                                      |
+| `yolox`       | **Detection:** YOLOX-Nano, YOLOX-Tiny, YOLOX-S, YOLOX-M, YOLOX-L, YOLOX-X                                                                                                                                                                                                           |
 | `yoloe`       | **Detection, Instance Segmentation:** YOLOE-11s, YOLOE-11m, YOLOE-11l; YOLOE-v8s, YOLOE-v8m, YOLOE-v8l                                                                                                                                                                              |
 | `goldyolo`    | Gold-YOLO-N, Gold-YOLO-S, Gold-YOLO-M, Gold-YOLO-L                                                                                                                                                                                                                                  |
 
@@ -131,7 +132,7 @@ If you don't find your model in the list, it is possible that it can be converte
 
 ## 📝 Credits
 
-This application uses source code of the following repositories: [YOLOv5](https://github.com/ultralytics/yolov5), [YOLOv6](https://github.com/meituan/YOLOv6), [GoldYOLO](https://github.com/huawei-noah/Efficient-Computing) [YOLOv7](https://github.com/WongKinYiu/yolov7), and [Ultralytics](https://github.com/ultralytics/ultralytics) (see each of them for more information).
+This application uses source code of the following repositories: [YOLOv5](https://github.com/ultralytics/yolov5), [YOLOv6](https://github.com/meituan/YOLOv6), [GoldYOLO](https://github.com/huawei-noah/Efficient-Computing), [YOLOv7](https://github.com/WongKinYiu/yolov7), [YOLOX](https://github.com/Megvii-BaseDetection/YOLOX), and [Ultralytics](https://github.com/ultralytics/ultralytics) (see each of them for more information).
 
 <a name="license"></a>
 

pyproject.toml

@@ -57,7 +57,7 @@ version = { attr = "tools.__version__" }
 
 [tool.ruff]
 target-version = "py310"
-exclude = ["tools/yolov7/yolov7/"]
+exclude = ["tools/yolov7/yolov7/", "tools/yolox/YOLOX/"]
 
 [tool.ruff.lint]
 ignore = ["F403", "B028", "B905", "D1"]

tests/conftest.py

@@ -37,6 +37,7 @@ def pytest_addoption(parser):
             "v11",
             "v12",
             "v26",
+            "yolox",
         ],
         default=None,
         help="If set then test only that specific yolo version",

tests/constants.py

@@ -14,6 +14,12 @@
     {"name": "yolov6sr3", "version": "v6r2"},
     {"name": "yolov6nr1", "version": "v6"},
     {"name": "yolov6tr1", "version": "v6"},
+    {"name": "yoloxnano", "version": "yolox", "size": "416"},
+    {"name": "yoloxtiny", "version": "yolox", "size": "416"},
+    {"name": "yoloxs", "version": "yolox", "size": "640"},
+    {"name": "yoloxm", "version": "yolox", "size": "640"},
+    {"name": "yoloxl", "version": "yolox", "size": "640"},
+    {"name": "yoloxx", "version": "yolox", "size": "640"},
     {"name": "yolov5n", "version": "v5", "size": "320"},
     {"name": "yolov5s", "version": "v5", "size": "320"},
     {"name": "yolov5n6", "version": "v5", "size": "320"},
@@ -107,6 +113,12 @@
     "yolov6sr3": "https://github.com/meituan/YOLOv6/releases/download/0.3.0/yolov6s.pt",
     "yolov6nr1": "https://github.com/meituan/YOLOv6/releases/download/0.1.0/yolov6n.pt",
     "yolov6tr1": "https://github.com/meituan/YOLOv6/releases/download/0.1.0/yolov6t.pt",
+    "yoloxnano": "https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_nano.pth",
+    "yoloxtiny": "https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_tiny.pth",
+    "yoloxs": "https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_s.pth",
+    "yoloxm": "https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_m.pth",
+    "yoloxl": "https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_l.pth",
+    "yoloxx": "https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/yolox_x.pth",
     "yolov5n": "https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5n.pt",
     "yolov5s": "https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5s.pt",
     "yolov5m": "https://github.com/ultralytics/yolov5/releases/download/v7.0/yolov5m.pt",

tests/e2e_shards.py

@@ -38,6 +38,9 @@
                 "tests/test_end2end.py::test_cli_conversion[yolov8n-obb]",
                 "tests/test_end2end.py::test_cli_conversion[yolov8n-pose]",
                 "tests/test_end2end.py::test_cli_conversion[yolov8s]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxm]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxnano]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxx]",
                 "tests/test_end2end.py::test_cli_conversion[yolov9c]",
                 "tests/test_end2end.py::test_cli_conversion[yolov9e]",
                 "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yolo26n-pose]",
@@ -78,6 +81,7 @@
                 "tests/test_end2end.py::test_cli_conversion[yolov8m]",
                 "tests/test_end2end.py::test_cli_conversion[yolov8n]",
                 "tests/test_end2end.py::test_cli_conversion[yolov8x]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxtiny]",
                 "tests/test_end2end.py::test_cli_conversion[yolov9s]",
                 "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yolo26n]",
                 "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yolov11n-seg]",
@@ -114,11 +118,14 @@
                 "tests/test_end2end.py::test_cli_conversion[yolov7x]",
                 "tests/test_end2end.py::test_cli_conversion[yolov8l]",
                 "tests/test_end2end.py::test_cli_conversion[yolov8n-seg]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxl]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxs]",
                 "tests/test_end2end.py::test_cli_conversion[yolov9m]",
                 "tests/test_end2end.py::test_cli_conversion[yolov9t]",
                 "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yolov11n-pose]",
                 "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yolov12n]",
                 "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yolov8n-seg]",
+                "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yoloxnano]",
                 "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yolov9t]",
             }
         ),
@@ -149,6 +156,8 @@
                 "tests/test_end2end.py::test_cli_conversion[yolov6nr21]",
                 "tests/test_end2end.py::test_cli_conversion[yolov8n-pose]",
                 "tests/test_end2end.py::test_cli_conversion[yolov8s]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxm]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxnano]",
                 "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yolov8n-pose]",
             }
         ),
@@ -162,6 +171,8 @@
                 "tests/test_end2end.py::test_cli_conversion[yolov5m6u]",
                 "tests/test_end2end.py::test_cli_conversion[yolov6mr2]",
                 "tests/test_end2end.py::test_cli_conversion[yolov6sr1]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxl]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxtiny]",
                 "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yolo26n-pose]",
             }
         ),
@@ -205,6 +216,7 @@
                 "tests/test_end2end.py::test_cli_conversion[yolov6nr1]",
                 "tests/test_end2end.py::test_cli_conversion[yolov6sr21]",
                 "tests/test_end2end.py::test_cli_conversion[yolov8n-seg]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxx]",
                 "tests/test_end2end.py::test_cli_conversion[yolov9e]",
                 "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yolov8n-seg]",
             }
@@ -220,6 +232,7 @@
                 "tests/test_end2end.py::test_cli_conversion[yolov6nr4]",
                 "tests/test_end2end.py::test_cli_conversion[yolov7t]",
                 "tests/test_end2end.py::test_cli_conversion[yolov8x]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxs]",
             }
         ),
         frozenset(
@@ -263,6 +276,7 @@
                 "tests/test_end2end.py::test_cli_conversion[yolov6sr4]",
                 "tests/test_end2end.py::test_cli_conversion[yolov7]",
                 "tests/test_end2end.py::test_cli_conversion[yolov9c]",
+                "tests/test_end2end.py::test_n_variant_nnarchive_outputs[yoloxnano]",
             }
         ),
     ),
@@ -286,6 +300,7 @@
                 "tests/test_end2end.py::test_cli_conversion[yolov8n-seg]",
                 "tests/test_end2end.py::test_cli_conversion[yolov8n]",
                 "tests/test_end2end.py::test_cli_conversion[yolov9t]",
+                "tests/test_end2end.py::test_cli_conversion[yoloxnano]",
             }
         ),
         frozenset(

tests/helper_functions.py

@@ -51,10 +51,21 @@ def download_private_model(model_name: str, filename: str, folder: str) -> str:
 def nn_archive_checker(
     extra_keys_to_check: Optional[list] = None,
     output_dir: str = "shared_with_container/outputs",
+    expected_scale: Optional[list[float]] = None,
+    expected_dai_type: Optional[str] = "RGB888p",
 ):
-    """Tests the content of the exported NNArchive."""
+    """Tests the content of the exported NNArchive.
+
+    Args:
+        extra_keys_to_check: Additional archive configuration values to check.
+        output_dir: Directory containing generated NNArchive output folders.
+        expected_scale: Expected per-channel input scale. Defaults to 255.
+        expected_dai_type: Expected DepthAI input type. Defaults to RGB888p.
+    """
     if extra_keys_to_check is None:
         extra_keys_to_check = []
+    if expected_scale is None:
+        expected_scale = [255.0, 255.0, 255.0]
     subdirs = [
         d for d in os.listdir(output_dir) if os.path.isdir(os.path.join(output_dir, d))
     ]
@@ -105,23 +116,21 @@ def nn_archive_checker(
         ), (
             f"Inputs mean `{config_data['model']['inputs'][0]['preprocessing']['mean']}` doesn't match the expected mean `{expected_mean}`"
         )
-        expected_scale = [255.0, 255.0, 255.0]
         assert (
             config_data["model"]["inputs"][0]["preprocessing"]["scale"]
             == expected_scale
         ), (
             f"Inputs scale `{config_data['model']['inputs'][0]['preprocessing']['scale']}` doesn't match the expected scale `{expected_scale}`"
         )
 
-        if len(extra_keys_to_check) and not any(
-            ["dai_type" in i for i in extra_keys_to_check[0]]
-        ):  # only check if we are not already checking though "extra_keys_to_check"
-            dai_type = "RGB888p"
+        if expected_dai_type is not None and not any(
+            "dai_type" in keys for keys, _ in extra_keys_to_check
+        ):
             assert (
                 config_data["model"]["inputs"][0]["preprocessing"]["dai_type"]
-                == dai_type
+                == expected_dai_type
             ), (
-                f"Inputs dai_type `{config_data['model']['inputs'][0]['preprocessing']['dai_type']}` doesn't match the expected dai_type `{dai_type}`"
+                f"Inputs dai_type `{config_data['model']['inputs'][0]['preprocessing']['dai_type']}` doesn't match the expected dai_type `{expected_dai_type}`"
             )
 
         if extra_keys_to_check:

tests/nnarchive_output_checks.py

@@ -98,4 +98,11 @@ def _clone_check(base_case: dict, *, name: str, version: str) -> dict:
         "yolo_outputs": ["output_yolo26"],
         "keypoints_outputs": ["kpt_output"],
     },
+    {
+        "name": "yoloxnano",
+        "version": "yolox",
+        "model_outputs": ["output1_yolov6", "output2_yolov6", "output3_yolov6"],
+        "head_outputs": ["output1_yolov6", "output2_yolov6", "output3_yolov6"],
+        "yolo_outputs": ["output1_yolov6", "output2_yolov6", "output3_yolov6"],
+    },
 ]

tests/test_end2end.py

@@ -89,7 +89,12 @@ def test_cli_conversion(model: dict, test_config: dict, subtests):
                 if model.get("size")
                 else []
             )
-            nn_archive_checker(extra_keys_to_check=extra_keys_to_check)
+            is_yolox = model["version"] == "yolox"
+            nn_archive_checker(
+                extra_keys_to_check=extra_keys_to_check,
+                expected_scale=[1.0, 1.0, 1.0] if is_yolox else None,
+                expected_dai_type="BGR888p" if is_yolox else "RGB888p",
+            )
 
 
 @pytest.mark.parametrize(

tests/test_unittests.py

@@ -269,8 +269,8 @@ def test_e2e_shard_count_three_manifest_matches_profile():
 
     assert e2e_shards.DEFAULT_E2E_SHARD_COUNT == 3
     assert e2e_shards.supported_e2e_shard_counts() == (3, 10)
-    assert counts == (33, 35, 33)
-    assert len(public_union) == 101
+    assert counts == (36, 36, 36)
+    assert len(public_union) == 108
     assert public_overlap == set()
 
 
@@ -294,9 +294,9 @@ def test_e2e_shard_count_ten_manifest_matches_selected_profile():
         seen.update(shard)
 
     assert len(assignment) == 10
-    assert counts == (11, 9, 9, 11, 11, 10, 9, 11, 11, 9)
+    assert counts == (11, 11, 11, 11, 11, 11, 10, 11, 11, 10)
     assert all(assignment)
-    assert len(public_union) == 101
+    assert len(public_union) == 108
     assert public_union == count_three_union
     assert overlap == set()
 
@@ -319,10 +319,11 @@ def test_e2e_representative_shard_manifest():
         "representative",
     )
     assert e2e_shards.supported_e2e_shard_counts("representative") == (2,)
-    assert counts == (14, 16)
-    assert len(representative_nodeids) == 30
+    assert counts == (15, 16)
+    assert len(representative_nodeids) == 31
     assert representative_nodeids < full_nodeids
     assert assignment[0].isdisjoint(assignment[1])
+    assert "tests/test_end2end.py::test_cli_conversion[yoloxnano]" in assignment[0]
 
 
 def test_e2e_shard_assignment_detects_collection_drift():

tests/test_yolox_exporter.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+import zipfile
+
+import pytest
+import torch
+
+from tools.utils.constants import Encoding
+from tools.version_detection import YOLOX_CONVERSION, detect_version
+from tools.yolox.yolox_exporter import YoloXExporter, _infer_standard_architecture
+
+
+def test_detects_yolox_state_dict_checkpoint(tmp_path) -> None:
+    """Recognize YOLOX's state-dict markers without relying on its extension."""
+    checkpoint = tmp_path / "weights.pth"
+    markers = (
+        "backbone.backbone.stem.conv.conv backbone.lateral_conv0 "
+        "head.stems head.cls_convs head.reg_convs head.cls_preds "
+        "head.reg_preds head.obj_preds"
+    )
+    with zipfile.ZipFile(checkpoint, "w") as archive:
+        archive.writestr("archive/data.pkl", markers)
+
+    assert detect_version(str(checkpoint)) == YOLOX_CONVERSION
+
+
+@pytest.mark.parametrize(
+    ("head_channels", "expected"),
+    [
+        (64, (0.33, 0.25, True)),
+        (96, (0.33, 0.375, False)),
+        (128, (0.33, 0.50, False)),
+        (192, (0.67, 0.75, False)),
+        (256, (1.00, 1.00, False)),
+        (320, (1.33, 1.25, False)),
+    ],
+)
+def test_infers_standard_yolox_architecture(head_channels, expected):
+    state_dict = {"head.cls_preds.0.weight": torch.empty(80, head_channels, 1, 1)}
+
+    assert _infer_standard_architecture(state_dict) == expected
+
+
+def test_yolox_exporter_uses_unscaled_bgr_preprocessing(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    archive_kwargs: dict = {}
+
+    def record_make_nn_archive(_exporter, **kwargs) -> None:
+        archive_kwargs.update(kwargs)
+
+    monkeypatch.setattr(
+        YoloXExporter,
+        "make_nn_archive",
+        record_make_nn_archive,
+    )
+
+    exporter = object.__new__(YoloXExporter)
+    exporter.names = ["person"]
+    exporter.nc = 1
+
+    exporter.export_nn_archive()
+
+    assert archive_kwargs["encoding"] == Encoding.BGR
+    assert archive_kwargs["mean"] == [0, 0, 0]
+    assert archive_kwargs["scale"] == [1, 1, 1]

tools/conversion_registry.py

@@ -19,6 +19,7 @@
     YOLOV26_CONVERSION,
     YOLOV26_NMS_CONVERSION,
     YOLOV26_SEM_CONVERSION,
+    YOLOX_CONVERSION,
 )
 
 ExporterFactory = Callable[[str, tuple[int, int], bool], Any]
@@ -102,6 +103,14 @@ def _build_yolov10_exporter(
     return YoloV10Exporter(model_path, imgsz, use_rvc2)
 
 
+def _build_yolox_exporter(
+    model_path: str, imgsz: tuple[int, int], use_rvc2: bool
+) -> Any:
+    from tools.yolox.yolox_exporter import YoloXExporter
+
+    return YoloXExporter(model_path, imgsz, use_rvc2)
+
+
 CONVERSION_SPECS: dict[str, ConversionSpec] = {
     GOLD_YOLO_CONVERSION: ConversionSpec("goldyolo", _build_goldyolo_exporter),
     YOLOV5_CONVERSION: ConversionSpec("yolov5", _build_yolov5_exporter),
@@ -118,6 +127,7 @@ def _build_yolov10_exporter(
     YOLOV26_CONVERSION: ConversionSpec("yolo26", _build_yolo26_exporter),
     YOLOV26_NMS_CONVERSION: ConversionSpec("yolov8", _build_yolov8_exporter),
     YOLOV26_SEM_CONVERSION: ConversionSpec("yolo26", _build_yolo26_exporter),
+    YOLOX_CONVERSION: ConversionSpec("yolox", _build_yolox_exporter),
 }