ai-training/tests/test_export.py

import shutil
import sys
from pathlib import Path

import cv2
import numpy as np
import onnxruntime as ort
import pytest
import torch
from ultralytics import YOLO

import constants as c
import exports as exports_mod

_HAS_TENSORRT = torch.cuda.is_available()
try:
    import tensorrt
except ImportError:
    _HAS_TENSORRT = False

_TESTS_DIR = Path(__file__).resolve().parent
_CONFIG_TEST = _TESTS_DIR.parent / "config.test.yaml"
_DATASET_IMAGES = _TESTS_DIR / "root" / "data" / "images"


class TestOnnxExport:
    def test_onnx_file_created(self, exported_models):
        # Assert
        p = Path(exported_models["onnx"])
        assert p.exists()
        assert p.stat().st_size > 0

    def test_onnx_batch_dimension_is_dynamic(self, exported_models):
        # Arrange
        session = ort.InferenceSession(exported_models["onnx"], providers=["CPUExecutionProvider"])
        batch_dim = session.get_inputs()[0].shape[0]

        # Assert
        assert isinstance(batch_dim, str) or batch_dim == -1

    def test_onnx_inference_batch_1(self, exported_models):
        # Arrange
        session = ort.InferenceSession(exported_models["onnx"], providers=["CPUExecutionProvider"])
        meta = session.get_inputs()[0]
        imgsz = exported_models["imgsz"]
        imgs = sorted(_DATASET_IMAGES.glob("*.jpg"))
        if not imgs:
            pytest.skip("no test images")
        blob = cv2.dnn.blobFromImage(
            cv2.imread(str(imgs[0])), 1.0 / 255.0, (imgsz, imgsz), (0, 0, 0), swapRB=True, crop=False,
        )

        # Act
        out = session.run(None, {meta.name: blob})

        # Assert
        assert out[0].shape[0] == 1

    def test_onnx_inference_batch_multiple(self, exported_models):
        # Arrange
        session = ort.InferenceSession(exported_models["onnx"], providers=["CPUExecutionProvider"])
        meta = session.get_inputs()[0]
        imgsz = exported_models["imgsz"]
        imgs = sorted(_DATASET_IMAGES.glob("*.jpg"))
        if not imgs:
            pytest.skip("no test images")
        single = cv2.dnn.blobFromImage(
            cv2.imread(str(imgs[0])), 1.0 / 255.0, (imgsz, imgsz), (0, 0, 0), swapRB=True, crop=False,
        )
        batch = np.concatenate([single] * 4, axis=0)

        # Act
        out = session.run(None, {meta.name: batch})

        # Assert
        assert out[0].shape[0] == 4


@pytest.mark.skipif(not _HAS_TENSORRT, reason="TensorRT requires NVIDIA GPU and tensorrt package")
class TestTensorrtExport:
    @pytest.fixture(scope="class")
    def tensorrt_model(self, exported_models):
        # Arrange
        model_dir = exported_models["model_dir"]
        pt_path = exported_models["pt_path"]
        old_config = c.config
        c.config = c.Config.from_yaml(str(_CONFIG_TEST), root=str(model_dir.parent))

        # Act
        exports_mod.export_tensorrt(pt_path)

        c.config = old_config
        engines = list(model_dir.glob("*.engine"))
        yield {
            "engine": str(engines[0]) if engines else None,
            "model_dir": model_dir,
            "imgsz": exported_models["imgsz"],
        }

        for e in model_dir.glob("*.engine"):
            e.unlink(missing_ok=True)

    def test_tensorrt_engine_created(self, tensorrt_model):
        # Assert
        assert tensorrt_model["engine"] is not None
        p = Path(tensorrt_model["engine"])
        assert p.exists()
        assert p.stat().st_size > 0

    def test_tensorrt_inference_batch_1(self, tensorrt_model):
        # Arrange
        assert tensorrt_model["engine"] is not None
        imgs = sorted(_DATASET_IMAGES.glob("*.jpg"))
        if not imgs:
            pytest.skip("no test images")
        model = YOLO(tensorrt_model["engine"])

        # Act
        results = model.predict(source=str(imgs[0]), imgsz=tensorrt_model["imgsz"], verbose=False)

        # Assert
        assert len(results) == 1
        assert results[0].boxes is not None

    def test_tensorrt_inference_batch_multiple(self, tensorrt_model):
        # Arrange
        assert tensorrt_model["engine"] is not None
        imgs = sorted(_DATASET_IMAGES.glob("*.jpg"))
        if len(imgs) < 4:
            pytest.skip("need at least 4 test images")
        model = YOLO(tensorrt_model["engine"])

        # Act
        results = model.predict(source=[str(p) for p in imgs[:4]], imgsz=tensorrt_model["imgsz"], verbose=False)

        # Assert
        assert len(results) == 4

    def test_tensorrt_inference_batch_max(self, tensorrt_model):
        # Arrange
        assert tensorrt_model["engine"] is not None
        imgs = sorted(_DATASET_IMAGES.glob("*.jpg"))
        if not imgs:
            pytest.skip("no test images")
        model = YOLO(tensorrt_model["engine"])
        sources = [str(imgs[0])] * 8

        # Act
        results = model.predict(source=sources, imgsz=tensorrt_model["imgsz"], verbose=False)

        # Assert
        assert len(results) == 8


@pytest.mark.skipif(sys.platform != "darwin", reason="CoreML requires macOS")
class TestCoremlExport:
    def test_coreml_package_created(self, exported_models):
        # Assert
        pkgs = list(exported_models["model_dir"].glob("*.mlpackage"))
        assert len(pkgs) >= 1

    def test_coreml_package_has_model(self, exported_models):
        # Assert
        pkgs = list(exported_models["model_dir"].glob("*.mlpackage"))
        assert len(pkgs) >= 1
        model_file = pkgs[0] / "Data" / "com.apple.CoreML" / "model.mlmodel"
        assert model_file.exists()

    def test_coreml_inference_produces_detections(self, exported_models):
        # Arrange
        pkgs = list(exported_models["model_dir"].glob("*.mlpackage"))
        assert len(pkgs) >= 1
        imgs = sorted(_DATASET_IMAGES.glob("*.jpg"))
        if not imgs:
            pytest.skip("no test images")
        model = YOLO(str(pkgs[0]))

        # Act
        results = model.predict(source=str(imgs[0]), imgsz=exported_models["imgsz"], verbose=False)

        # Assert
        assert len(results) == 1
        assert results[0].boxes is not None


_INPUT_DATA = _TESTS_DIR.parent / "_docs" / "00_problem" / "input_data"


@pytest.mark.skipif(sys.platform != "darwin", reason="CoreML requires macOS")
class TestCoremlExportRealModel:
    def test_export_azaion_pt_to_coreml(self, fixture_pt_model):
        # Arrange
        output_dir = _INPUT_DATA / "azaion.mlpackage"
        if output_dir.exists():
            shutil.rmtree(output_dir)

        # Act
        model = YOLO(fixture_pt_model)
        model.export(format="coreml", imgsz=1280)
        exported = Path(fixture_pt_model).with_suffix(".mlpackage")
        if exported != output_dir:
            shutil.move(str(exported), str(output_dir))

        # Assert
        assert output_dir.exists()
        model_file = output_dir / "Data" / "com.apple.CoreML" / "model.mlmodel"
        assert model_file.exists()