diff --git a/.gitmodules b/.gitmodules index acb42ef..bc01f80 100644 --- a/.gitmodules +++ b/.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 diff --git a/README.md b/README.md index cb2f457..dd6c0da 100644 --- a/README.md +++ b/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). diff --git a/media/coverage_badge.svg b/media/coverage_badge.svg index 5d787db..b4a82e6 100644 --- a/media/coverage_badge.svg +++ b/media/coverage_badge.svg @@ -15,7 +15,7 @@ coverage coverage - 58% - 58% + 56% + 56% diff --git a/pyproject.toml b/pyproject.toml index ba45c82..3d3c8a1 100644 --- a/pyproject.toml +++ b/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"] diff --git a/tests/conftest.py b/tests/conftest.py index 5047ea1..c8ca918 100644 --- a/tests/conftest.py +++ b/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", diff --git a/tests/constants.py b/tests/constants.py index 84b99be..e27ecab 100644 --- a/tests/constants.py +++ b/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", diff --git a/tests/e2e_shards.py b/tests/e2e_shards.py index 4c1eba5..87637c1 100644 --- a/tests/e2e_shards.py +++ b/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( diff --git a/tests/helper_functions.py b/tests/helper_functions.py index 922cc12..a61bf8a 100644 --- a/tests/helper_functions.py +++ b/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,7 +116,6 @@ 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 @@ -113,15 +123,14 @@ def nn_archive_checker( 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: diff --git a/tests/nnarchive_output_checks.py b/tests/nnarchive_output_checks.py index 129ac5f..1522fbc 100644 --- a/tests/nnarchive_output_checks.py +++ b/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"], + }, ] diff --git a/tests/test_end2end.py b/tests/test_end2end.py index af88136..57745e9 100644 --- a/tests/test_end2end.py +++ b/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( diff --git a/tests/test_unittests.py b/tests/test_unittests.py index 8ed1710..cdedb1a 100644 --- a/tests/test_unittests.py +++ b/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(): diff --git a/tests/test_yolox_exporter.py b/tests/test_yolox_exporter.py new file mode 100644 index 0000000..9a14e78 --- /dev/null +++ b/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] diff --git a/tools/conversion_registry.py b/tools/conversion_registry.py index 3380a64..96afa11 100644 --- a/tools/conversion_registry.py +++ b/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), } diff --git a/tools/main.py b/tools/main.py index 4301439..f81fd18 100644 --- a/tools/main.py +++ b/tools/main.py @@ -39,6 +39,7 @@ ) from tools.version_detection import ( YOLOV5U_CONVERSION, + YOLOX_CONVERSION, detect_version, ) @@ -61,9 +62,9 @@ def convert( Parameter(show_default=True), ] = None, encoding: Annotated[ - Encoding, + Encoding | None, Parameter(show_default=True), - ] = Encoding.RGB, + ] = None, use_rvc2: Annotated[ bool, Parameter(show_default=True), @@ -94,7 +95,7 @@ def convert( version: YOLO variant to force, such as ``"yolov8"``. When omitted, the command runs automatic version detection. encoding: Color encoding used by the input model. Must be ``RGB`` or - ``BGR``. + ``BGR``. When omitted it is selected based on version. use_rvc2: Whether to target RVC2 instead of RVC3. class_names: Comma-separated class names recognized by the model. output_remote_url: Remote destination URL for uploading the generated NN @@ -147,26 +148,9 @@ def convert( else: class_names_list = class_names - try: - config = Config.get_config( - { - "model": model, - "imgsz": imgsz_list, - "encoding": encoding, - "use_rvc2": use_rvc2, - "class_names": class_names_list, - "output_remote_url": output_remote_url, - "put_file_plugin": put_file_plugin, - } - ) - except Exception as e: - logger.error(f"Invalid configuration: {e}") - raise SystemExit(ExitCode.VALIDATION_FAILED.value) from e - exporter_imgsz = cast(tuple[int, int], tuple(config.imgsz)) - phase = Phase.PATH_RESOLUTION try: - model_path = resolve_path(config.model, MISC_DIR) + model_path = resolve_path(model, MISC_DIR) except Exception as e: logger.error(f"Error resolving model path: {e}") raise SystemExit(ExitCode.PATH_RESOLUTION_FAILED.value) from e @@ -188,8 +172,25 @@ def convert( ) logger.info(f"Detected version: {version} {version_note}") - if version is None: - raise RuntimeError("Version must be resolved before telemetry capture.") + if encoding is None: + encoding = Encoding.BGR if version == YOLOX_CONVERSION else Encoding.RGB + + try: + config = Config.get_config( + { + "model": model, + "imgsz": imgsz_list, + "encoding": encoding, + "use_rvc2": use_rvc2, + "class_names": class_names_list, + "output_remote_url": output_remote_url, + "put_file_plugin": put_file_plugin, + } + ) + except Exception as e: + logger.error(f"Invalid configuration: {e}") + raise SystemExit(ExitCode.VALIDATION_FAILED.value) from e + exporter_imgsz = cast(tuple[int, int], tuple(config.imgsz)) conversion_summary = build_conversion_summary( config=config, diff --git a/tools/modules/exporter.py b/tools/modules/exporter.py index c7ed4a4..484abdc 100644 --- a/tools/modules/exporter.py +++ b/tools/modules/exporter.py @@ -164,6 +164,8 @@ def make_nn_archive( anchors: list[list[list[float]]] | None = None, output_kwargs: dict | None = None, encoding: Encoding = Encoding.RGB, + mean: list[float] | None = None, + scale: list[float] | None = None, ): """Create an NN archive for a detection-style model. @@ -182,6 +184,8 @@ def make_nn_archive( anchors: Optional anchor definitions for the detection head. output_kwargs: Additional output metadata forwarded to the head. encoding: Color encoding used by the input model. + mean: Per-channel input mean. Defaults to ``[0, 0, 0]``. + scale: Per-channel input scale. Defaults to ``[255, 255, 255]``. """ self.f_nn_archive = (self.output_folder / f"{self.model_name}.tar.xz").resolve() if stage2_executable_path is not None: @@ -191,6 +195,10 @@ def make_nn_archive( if output_kwargs is None: output_kwargs = {} + if mean is None: + mean = [0, 0, 0] + if scale is None: + scale = [255, 255, 255] output_specs = self.get_output_specs() archive = ArchiveGenerator( @@ -210,8 +218,8 @@ def make_nn_archive( "input_type": InputType.IMAGE, "shape": [1, self.number_of_channels, *self.imgsz[::-1]], "preprocessing": { - "mean": [0, 0, 0], - "scale": [255, 255, 255], + "mean": mean, + "scale": scale, "dai_type": encoding.get_dai_type(), }, } diff --git a/tools/version_detection/__init__.py b/tools/version_detection/__init__.py index 79ee6ba..95cbaf1 100644 --- a/tools/version_detection/__init__.py +++ b/tools/version_detection/__init__.py @@ -15,6 +15,7 @@ YOLOV26_CONVERSION, YOLOV26_NMS_CONVERSION, YOLOV26_SEM_CONVERSION, + YOLOX_CONVERSION, detect_version, ) @@ -34,6 +35,7 @@ "YOLOV26_CONVERSION", "YOLOV26_NMS_CONVERSION", "YOLOV26_SEM_CONVERSION", + "YOLOX_CONVERSION", "GOLD_YOLO_CONVERSION", "UNRECOGNIZED", ] diff --git a/tools/version_detection/version_detection.py b/tools/version_detection/version_detection.py index a0fd03f..d7fa1f2 100644 --- a/tools/version_detection/version_detection.py +++ b/tools/version_detection/version_detection.py @@ -20,6 +20,7 @@ YOLOV26_CONVERSION = "yolov26" YOLOV26_NMS_CONVERSION = "yolov26_nms" YOLOV26_SEM_CONVERSION = "yolov26_sem" +YOLOX_CONVERSION = "yolox" GOLD_YOLO_CONVERSION = "goldyolo" UNRECOGNIZED = "none" @@ -126,6 +127,20 @@ def detect_version(path: str, debug: bool = False) -> str: elif "gold_yolo" in content: return GOLD_YOLO_CONVERSION return YOLOV6R3_CONVERSION + elif all( + marker in content + for marker in ( + "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", + ) + ): + return YOLOX_CONVERSION elif "yolov7" in content: return YOLOV7_CONVERSION elif ( diff --git a/tools/yolox/YOLOX b/tools/yolox/YOLOX new file mode 160000 index 0000000..6ddff48 --- /dev/null +++ b/tools/yolox/YOLOX @@ -0,0 +1 @@ +Subproject commit 6ddff4824372906469a7fae2dc3206c7aa4bbaee diff --git a/tools/yolox/__init__.py b/tools/yolox/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tools/yolox/yolox_exporter.py b/tools/yolox/yolox_exporter.py new file mode 100644 index 0000000..b28bdff --- /dev/null +++ b/tools/yolox/yolox_exporter.py @@ -0,0 +1,244 @@ +# ruff: noqa: I001 +from __future__ import annotations + +import os +import sys +from collections.abc import Mapping + +import torch +import torch.nn as nn + +from tools.modules import Exporter +from tools.utils.constants import Encoding + +current_dir = os.path.dirname(os.path.abspath(__file__)) +yolox_path = os.path.join(current_dir, "YOLOX") +if yolox_path not in sys.path: + sys.path.insert(0, yolox_path) + +from yolox.models import YOLOX, YOLOPAFPN, YOLOXHead # noqa: E402 +from yolox.models.network_blocks import SiLU # noqa: E402 +from yolox.utils import replace_module # noqa: E402 + + +COCO_CLASSES = [ + "person", + "bicycle", + "car", + "motorcycle", + "airplane", + "bus", + "train", + "truck", + "boat", + "traffic light", + "fire hydrant", + "stop sign", + "parking meter", + "bench", + "bird", + "cat", + "dog", + "horse", + "sheep", + "cow", + "elephant", + "bear", + "zebra", + "giraffe", + "backpack", + "umbrella", + "handbag", + "tie", + "suitcase", + "frisbee", + "skis", + "snowboard", + "sports ball", + "kite", + "baseball bat", + "baseball glove", + "skateboard", + "surfboard", + "tennis racket", + "bottle", + "wine glass", + "cup", + "fork", + "knife", + "spoon", + "bowl", + "banana", + "apple", + "sandwich", + "orange", + "broccoli", + "carrot", + "hot dog", + "pizza", + "donut", + "cake", + "chair", + "couch", + "potted plant", + "bed", + "dining table", + "toilet", + "tv", + "laptop", + "mouse", + "remote", + "keyboard", + "cell phone", + "microwave", + "oven", + "toaster", + "sink", + "refrigerator", + "book", + "clock", + "vase", + "scissors", + "teddy bear", + "hair drier", + "toothbrush", +] + + +class DetectYOLOX(nn.Module): + """Expose YOLOX's three raw detection maps for archive postprocessing.""" + + def __init__(self, head: YOLOXHead): + super().__init__() + self.nc = head.num_classes + self.nl = len(head.strides) + self.stride = head.strides + self.stems = head.stems + self.cls_convs = head.cls_convs + self.reg_convs = head.reg_convs + self.cls_preds = head.cls_preds + self.reg_preds = head.reg_preds + self.obj_preds = head.obj_preds + + def forward(self, features: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, ...]: + outputs = [] + for index in range(self.nl): + x = self.stems[index](features[index]) + cls_output = self.cls_preds[index](self.cls_convs[index](x)) + reg_features = self.reg_convs[index](x) + reg_output = self.reg_preds[index](reg_features) + obj_output = self.obj_preds[index](reg_features) + outputs.append( + torch.cat( + [reg_output, obj_output.sigmoid(), cls_output.sigmoid()], dim=1 + ) + ) + return tuple(outputs) + + +def _infer_standard_architecture( + state_dict: Mapping[str, torch.Tensor], +) -> tuple[float, float, bool]: + """Infer a standard YOLOX architecture from its head width.""" + try: + head_channels = state_dict["head.cls_preds.0.weight"].shape[1] + except KeyError as exc: + raise ValueError("Checkpoint does not contain a YOLOX detection head.") from exc + + # Maps detection-head channel count to (depth multiplier, width multiplier, depthwise). + architectures = { + 64: (0.33, 0.25, True), # YOLOX-Nano + 96: (0.33, 0.375, False), # YOLOX-Tiny + 128: (0.33, 0.50, False), # YOLOX-S + 192: (0.67, 0.75, False), # YOLOX-M + 256: (1.00, 1.00, False), # YOLOX-L + 320: (1.33, 1.25, False), # YOLOX-X + } + try: + return architectures[head_channels] + except KeyError as exc: + raise ValueError( + "Unsupported YOLOX head width. Only standard YOLOX-Nano/Tiny/S/M/L/X " + "checkpoints are currently supported." + ) from exc + + +class YoloXExporter(Exporter): + """Export standard YOLOX detection checkpoints as Luxonis NNArchives.""" + + strides = [8, 16, 32] + output_names = ["output1_yolov6", "output2_yolov6", "output3_yolov6"] + + def __init__( + self, + model_path: str, + imgsz: tuple[int, int], + use_rvc2: bool, + ): + super().__init__( + model_path, + imgsz, + use_rvc2, + # YOLOX uses the same grid decode as this already-supported subtype. + subtype="yolov6r1", + output_names=self.output_names, + ) + self.load_model() + + def load_model(self) -> None: + checkpoint = torch.load(self.model_path, map_location="cpu", weights_only=True) + state_dict = checkpoint.get("model", checkpoint) + if not isinstance(state_dict, Mapping): + raise ValueError("YOLOX checkpoint must contain a model state dictionary.") + + depth, width, depthwise = _infer_standard_architecture(state_dict) + try: + num_classes = state_dict["head.cls_preds.0.weight"].shape[0] + except KeyError as exc: + raise ValueError( + "Checkpoint does not contain YOLOX class predictions." + ) from exc + + backbone = YOLOPAFPN(depth=depth, width=width, act="silu", depthwise=depthwise) + head = YOLOXHead( + num_classes=num_classes, width=width, act="silu", depthwise=depthwise + ) + model = YOLOX(backbone, head) + model.load_state_dict(state_dict, strict=True) + model = replace_module(model, nn.SiLU, SiLU) + model.head = DetectYOLOX(model.head) + model.eval() + + if any(size % max(self.strides) != 0 for size in self.imgsz): + raise ValueError("Image size must be divisible by the maximum stride (32).") + + self.number_of_channels = 3 + self.nc = num_classes + self.names = ( + COCO_CLASSES + if num_classes == len(COCO_CLASSES) + else [f"Class_{index}" for index in range(num_classes)] + ) + self.model = model + + def export_nn_archive( + self, class_names: list[str] | None = None, encoding: Encoding = Encoding.BGR + ) -> None: + """Create an NNArchive with the existing YOLOv6 R1 grid decoder.""" + names = self.names + if class_names is not None: + assert len(class_names) == self.nc, ( + f"Number of given class names {len(class_names)} does not match " + f"the model class count {self.nc}." + ) + names = class_names + + self.make_nn_archive( + class_list=names, + n_classes=self.nc, + parser="YOLOExtendedParser", + output_kwargs={"strides": self.strides}, + encoding=encoding, + mean=[0, 0, 0], + scale=[1, 1, 1], + )