wehub-resource-sync
16031aae96
CPU tests Workflow / Testing (ubuntu-latest, 3.12) (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.13) (push) Failing after 0s
Mypy Type Check / Type Check (push) Failing after 0s
Docs/Test WorkFlow / Test docs build (push) Failing after 1s
PR Conflict Labeler / labeling (push) Failing after 1s
Dependency resolution / Resolve [tflite] extra — Python 3.12 (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.10) (push) Failing after 0s
Smoke Tests / try-all-models (ubuntu-latest, 3.13) (push) Failing after 1s
CPU tests Workflow / build-pkg (push) Failing after 1s
CPU tests Workflow / Testing (ubuntu-latest, 3.10) (push) Failing after 0s
CPU tests Workflow / Testing (ubuntu-latest, 3.11) (push) Failing after 0s
Smoke Tests / try-all-models (macos-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (macos-latest, 3.13) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.10) (push) Has been cancelled
Smoke Tests / try-all-models (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (macos-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.10) (push) Has been cancelled
CPU tests Workflow / Testing (windows-latest, 3.13) (push) Has been cancelled
CPU tests Workflow / testing-guardian (push) Has been cancelled
GPU tests Workflow / Testing (push) Has been cancelled
1111 lines
42 KiB
Python
1111 lines
42 KiB
Python
# ------------------------------------------------------------------------
|
|
# RF-DETR
|
|
# Copyright (c) 2025 Roboflow. All Rights Reserved.
|
|
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
|
|
# ------------------------------------------------------------------------
|
|
"""Tests for the ONNX → TFLite export pipeline.
|
|
|
|
Tests cover:
|
|
* ``export_tflite()`` — the main conversion function (mocked ``onnx2tf``)
|
|
* ``_check_onnx2tf_available()`` — import-based availability check
|
|
* ``_numpy_allow_pickle()`` — NumPy monkey-patch context manager
|
|
* ``_patch_validation_download()`` — validation download redirect
|
|
* ``_get_onnx_input_info()`` — ONNX model input metadata reader
|
|
* ``_prepare_calibration_data()`` — calibration data preparation
|
|
* ``format="tflite"`` parameter wiring through ``RFDETR.export()``
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
from typing import Any, Generator
|
|
from unittest import mock
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from rfdetr.export._tflite import _IS_ONNX2TF_AVAILABLE
|
|
from rfdetr.export._tflite.converter import (
|
|
_DEFAULT_CALIB_SAMPLES,
|
|
_DEFAULT_DIR_CALIB_SAMPLES,
|
|
_IMAGE_EXTENSIONS,
|
|
_NUMPY_LOAD_PATCH_LOCK,
|
|
_VALID_QUANTIZATIONS,
|
|
_check_onnx2tf_available,
|
|
_get_onnx_input_info,
|
|
_load_calibration_images,
|
|
_numpy_allow_pickle,
|
|
_patch_validation_download,
|
|
_prepare_calibration_data,
|
|
export_tflite,
|
|
)
|
|
|
|
onnx2tf_available = pytest.mark.skipif(not _IS_ONNX2TF_AVAILABLE, reason="onnx2tf not installed")
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers — fake onnx2tf module injected into sys.modules
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _FakeOnnx2tfModule:
|
|
"""Namespace that mimics ``onnx2tf`` for testing."""
|
|
|
|
def __init__(self) -> None:
|
|
self.convert = mock.MagicMock()
|
|
|
|
|
|
_ONNX2TF_KEYS = ("onnx2tf", "onnx2tf.onnx2tf", "onnx2tf.utils", "onnx2tf.utils.common_functions")
|
|
|
|
|
|
def _install_fake_onnx2tf() -> tuple[_FakeOnnx2tfModule, mock.MagicMock, dict[str, object]]:
|
|
"""Insert a fake ``onnx2tf`` package into ``sys.modules``.
|
|
|
|
Saves any pre-existing real modules under the same keys so they can be restored by ``_remove_fake_onnx2tf``
|
|
(Copilot: do not silently clobber real modules that a prior test may have imported).
|
|
|
|
Returns:
|
|
Tuple of (fake_module, convert_mock, saved_originals).
|
|
"""
|
|
# Snapshot originals before overwriting (None means the key was absent).
|
|
saved: dict[str, object] = {k: sys.modules.get(k) for k in _ONNX2TF_KEYS}
|
|
|
|
fake = _FakeOnnx2tfModule()
|
|
pkg = types.ModuleType("onnx2tf")
|
|
pkg.convert = fake.convert # type: ignore[attr-defined]
|
|
pkg.__version__ = "2.4.0" # type: ignore[attr-defined]
|
|
|
|
# onnx2tf.onnx2tf — force-imported by export_tflite() before patching
|
|
inner_mod = types.ModuleType("onnx2tf.onnx2tf")
|
|
inner_mod.download_test_image_data = mock.MagicMock( # type: ignore[attr-defined]
|
|
return_value=np.zeros((20, 128, 128, 3), dtype=np.float32),
|
|
)
|
|
|
|
# onnx2tf.utils and onnx2tf.utils.common_functions
|
|
utils_mod = types.ModuleType("onnx2tf.utils")
|
|
cf_mod = types.ModuleType("onnx2tf.utils.common_functions")
|
|
cf_mod.download_test_image_data = mock.MagicMock( # type: ignore[attr-defined]
|
|
return_value=np.zeros((20, 128, 128, 3), dtype=np.float32),
|
|
)
|
|
|
|
# Wire up module hierarchy
|
|
pkg.onnx2tf = inner_mod # type: ignore[attr-defined]
|
|
pkg.utils = utils_mod # type: ignore[attr-defined]
|
|
utils_mod.common_functions = cf_mod # type: ignore[attr-defined]
|
|
|
|
sys.modules["onnx2tf"] = pkg
|
|
sys.modules["onnx2tf.onnx2tf"] = inner_mod
|
|
sys.modules["onnx2tf.utils"] = utils_mod
|
|
sys.modules["onnx2tf.utils.common_functions"] = cf_mod
|
|
return fake, fake.convert, saved
|
|
|
|
|
|
def _remove_fake_onnx2tf(saved: dict[str, object] | None = None) -> None:
|
|
"""Remove fake ``onnx2tf`` entries from ``sys.modules`` and restore originals.
|
|
|
|
Args:
|
|
saved: Snapshot returned by ``_install_fake_onnx2tf``. If a key was
|
|
present before installation its original value is restored; if it was absent it is deleted. When *saved* is
|
|
``None`` all ``onnx2tf*`` keys are simply deleted (legacy behaviour).
|
|
"""
|
|
if saved is not None:
|
|
for key in _ONNX2TF_KEYS:
|
|
original = saved.get(key)
|
|
if original is None:
|
|
sys.modules.pop(key, None)
|
|
else:
|
|
sys.modules[key] = original # type: ignore[assignment]
|
|
else:
|
|
for key in list(sys.modules):
|
|
if key == "onnx2tf" or key.startswith("onnx2tf."):
|
|
del sys.modules[key]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture()
|
|
def fake_onnx2tf():
|
|
"""Provide a fake ``onnx2tf`` that records *convert()* calls.
|
|
|
|
Also patches ``_replace_gridsample_for_tflite`` to return the input path unchanged so tests that supply stub ONNX
|
|
bytes do not depend on ``onnx.load`` tolerating those bytes.
|
|
"""
|
|
fake, convert_mock, saved = _install_fake_onnx2tf()
|
|
with mock.patch(
|
|
"rfdetr.export._tflite.converter._replace_gridsample_for_tflite",
|
|
side_effect=lambda path, _dir: path,
|
|
):
|
|
yield fake, convert_mock
|
|
_remove_fake_onnx2tf(saved)
|
|
|
|
|
|
@pytest.fixture()
|
|
def onnx_model(tmp_path: Path) -> Path:
|
|
"""Create a dummy ``.onnx`` file."""
|
|
p = tmp_path / "model.onnx"
|
|
p.write_bytes(b"\x08\x06") # minimal bytes
|
|
return p
|
|
|
|
|
|
@pytest.fixture()
|
|
def mock_prepare_calib(tmp_path: Path) -> Generator:
|
|
"""Mock ``_prepare_calibration_data`` so dummy ONNX files work.
|
|
|
|
``export_tflite`` calls ``_prepare_calibration_data`` which calls ``_get_onnx_input_info`` (requiring a real ONNX
|
|
file). Since the ``onnx_model`` fixture writes only stub bytes, this mock prevents the ONNX parse from being
|
|
attempted.
|
|
"""
|
|
dummy_npy = tmp_path / "_rfdetr_calib_data.npy"
|
|
np.save(str(dummy_npy), np.zeros((1, 64, 64, 3), dtype=np.float32))
|
|
with mock.patch(
|
|
"rfdetr.export._tflite.converter._prepare_calibration_data",
|
|
return_value=dummy_npy,
|
|
) as m:
|
|
yield m
|
|
|
|
|
|
@pytest.fixture()
|
|
def tflite_output(tmp_path: Path, onnx_model: Path) -> Path:
|
|
"""Create expected TFLite output file so export_tflite finds it."""
|
|
out = tmp_path / "output"
|
|
out.mkdir()
|
|
(out / f"{onnx_model.stem}_float32.tflite").write_bytes(b"tflite")
|
|
return out
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestExportTfliteConverter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@onnx2tf_available
|
|
class TestExportTfliteConverter:
|
|
"""Tests for ``export_tflite()``."""
|
|
|
|
def test_missing_onnx_raises_file_not_found(self, tmp_path: Path, fake_onnx2tf: Any) -> None:
|
|
with pytest.raises(FileNotFoundError, match="ONNX model not found"):
|
|
export_tflite(tmp_path / "nope.onnx", tmp_path / "out")
|
|
|
|
def test_invalid_quantization_raises_value_error(self, onnx_model: Path, tmp_path: Path, fake_onnx2tf: Any) -> None:
|
|
with pytest.raises(ValueError, match="Unsupported quantization"):
|
|
export_tflite(onnx_model, tmp_path / "out", quantization="q4")
|
|
|
|
@pytest.mark.parametrize(
|
|
"static_mode",
|
|
[
|
|
pytest.param("int8_static", id="int8_static"),
|
|
pytest.param("full_int8", id="full_int8"),
|
|
pytest.param("integer_quant", id="integer_quant"),
|
|
],
|
|
)
|
|
def test_static_int8_raises(self, onnx_model: Path, tmp_path: Path, fake_onnx2tf: Any, static_mode: str) -> None:
|
|
"""A static / full-integer INT8 request must raise a ValueError.
|
|
|
|
Static INT8 is intentionally unsupported; only dynamic-range 'int8' is offered.
|
|
"""
|
|
with pytest.raises(ValueError, match="[Ss]tatic / full-integer INT8 is not supported"):
|
|
export_tflite(onnx_model, tmp_path / "out", quantization=static_mode)
|
|
|
|
def test_default_quantization_calls_convert(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
_, convert_mock = fake_onnx2tf
|
|
result = export_tflite(onnx_model, tflite_output)
|
|
|
|
convert_mock.assert_called_once()
|
|
kwargs = convert_mock.call_args.kwargs
|
|
assert kwargs["input_onnx_file_path"] == str(onnx_model)
|
|
assert kwargs["output_folder_path"] == str(tflite_output)
|
|
assert kwargs["output_signaturedefs"] is True
|
|
assert kwargs["non_verbose"] is True
|
|
assert "output_integer_quantized_tflite" not in kwargs
|
|
assert result == tflite_output / "model_float32.tflite"
|
|
|
|
def test_custom_input_not_passed_to_convert(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
"""custom_input_op_name_np_data_path must NOT be passed to convert().
|
|
|
|
The onnx2tf custom_input code path triggers a tf.tile rank mismatch with DINOv2-style backbones when N > 1. We
|
|
rely on patching download_test_image_data() instead.
|
|
"""
|
|
_, convert_mock = fake_onnx2tf
|
|
export_tflite(onnx_model, tflite_output)
|
|
|
|
kwargs = convert_mock.call_args.kwargs
|
|
assert "custom_input_op_name_np_data_path" not in kwargs
|
|
|
|
def test_output_signaturedefs_always_enabled(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
"""output_signaturedefs must always be True.
|
|
|
|
Segmentation models produce ONNX node names with leading "/" characters that violate the TF saved_model naming
|
|
pattern. Enabling signature defs bypasses this restriction.
|
|
"""
|
|
_, convert_mock = fake_onnx2tf
|
|
export_tflite(onnx_model, tflite_output)
|
|
|
|
kwargs = convert_mock.call_args.kwargs
|
|
assert kwargs["output_signaturedefs"] is True
|
|
|
|
def test_tflite_backend_forced_to_tf_converter(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
"""tflite_backend must always be 'tf_converter' to avoid the TFLite TopK_V2 kernel check.
|
|
|
|
onnx2tf 2.x defaults to flatbuffer_direct, which trips a "k > internal dimension" error at AllocateTensors()
|
|
time on RF-DETR's encoder TopK node. tf_converter is forced unconditionally.
|
|
"""
|
|
_, convert_mock = fake_onnx2tf
|
|
export_tflite(onnx_model, tflite_output)
|
|
|
|
assert convert_mock.call_args.kwargs["tflite_backend"] == "tf_converter"
|
|
|
|
def test_replace_to_pseudo_operators_contains_erf_and_gelu(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
"""replace_to_pseudo_operators must include Erf and GeLU.
|
|
|
|
Without this, AllocateTensors() fails with "FlexErf failed to prepare" because the TFLite runtime lacks native
|
|
Erf / GeLU kernels.
|
|
"""
|
|
_, convert_mock = fake_onnx2tf
|
|
export_tflite(onnx_model, tflite_output)
|
|
|
|
pseudo_ops = convert_mock.call_args.kwargs.get("replace_to_pseudo_operators", [])
|
|
assert "Erf" in pseudo_ops
|
|
assert "GeLU" in pseudo_ops
|
|
|
|
def test_fp32_quantization_no_int8_flag(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
_, convert_mock = fake_onnx2tf
|
|
export_tflite(onnx_model, tflite_output, quantization="fp32")
|
|
assert "output_integer_quantized_tflite" not in convert_mock.call_args.kwargs
|
|
|
|
def test_fp16_quantization_no_int8_flag(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
_, convert_mock = fake_onnx2tf
|
|
export_tflite(onnx_model, tflite_output, quantization="fp16")
|
|
assert "output_integer_quantized_tflite" not in convert_mock.call_args.kwargs
|
|
|
|
def test_int8_quantization_produces_dynamic_range(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
"""Int8 export derives a dynamic-range model and avoids onnx2tf's -oiqt path.
|
|
|
|
onnx2tf's ``output_integer_quantized_tflite`` (-oiqt) only yields static quantization, which RF-DETR's
|
|
transformer activations do not survive. The converter instead builds dynamic-range INT8 from the SavedModel via
|
|
``_quantize_dynamic_range``, so the onnx2tf call must NOT carry the ``output_integer_quantized_tflite`` flag.
|
|
"""
|
|
_, convert_mock = fake_onnx2tf
|
|
dyn_path = tflite_output / "model_dynamic_range_quant.tflite"
|
|
with mock.patch(
|
|
"rfdetr.export._tflite.converter._quantize_dynamic_range",
|
|
return_value=dyn_path,
|
|
) as quant_mock:
|
|
result = export_tflite(onnx_model, tflite_output, quantization="int8")
|
|
assert "output_integer_quantized_tflite" not in convert_mock.call_args.kwargs
|
|
quant_mock.assert_called_once()
|
|
assert result == dyn_path
|
|
|
|
def test_verbosity_forwarded(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
_, convert_mock = fake_onnx2tf
|
|
export_tflite(onnx_model, tflite_output, verbosity="debug")
|
|
assert convert_mock.call_args.kwargs["verbosity"] == "debug"
|
|
|
|
def test_convert_failure_raises_runtime_error(
|
|
self,
|
|
onnx_model: Path,
|
|
tmp_path: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
_, convert_mock = fake_onnx2tf
|
|
convert_mock.side_effect = RuntimeError("boom")
|
|
with pytest.raises(RuntimeError, match="onnx2tf conversion failed"):
|
|
export_tflite(onnx_model, tmp_path / "out")
|
|
|
|
def test_fallback_when_primary_tflite_missing(
|
|
self,
|
|
onnx_model: Path,
|
|
tmp_path: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
"""Fallback returns a stem-scoped file when the primary *_float32.tflite is absent."""
|
|
out = tmp_path / "out"
|
|
out.mkdir()
|
|
# Scoped fallback: must match {stem}_*.tflite (stem == "model" here).
|
|
(out / "model_float16.tflite").write_bytes(b"fb")
|
|
result = export_tflite(onnx_model, out)
|
|
assert result.name == "model_float16.tflite"
|
|
|
|
def test_fallback_does_not_return_unrelated_tflite(
|
|
self,
|
|
onnx_model: Path,
|
|
tmp_path: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
"""Stale artifacts from a different export are never returned as fallback."""
|
|
out = tmp_path / "out"
|
|
out.mkdir()
|
|
# Unrelated file — does NOT match model_*.tflite.
|
|
(out / "other_model.tflite").write_bytes(b"stale")
|
|
with pytest.raises(RuntimeError, match="no .tflite file matching"):
|
|
export_tflite(onnx_model, out)
|
|
|
|
def test_no_tflite_output_raises_runtime_error(
|
|
self,
|
|
onnx_model: Path,
|
|
tmp_path: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
"""Empty output directory raises RuntimeError after conversion."""
|
|
out = tmp_path / "empty_out"
|
|
out.mkdir()
|
|
with pytest.raises(RuntimeError, match="no .tflite file matching"):
|
|
export_tflite(onnx_model, out)
|
|
|
|
def test_returns_path_object(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
result = export_tflite(onnx_model, tflite_output)
|
|
assert isinstance(result, Path)
|
|
|
|
def test_calibration_data_forwarded_to_prepare(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
"""Verify that calibration_data is passed to _prepare_calibration_data."""
|
|
calib_path = "/some/calib.npy"
|
|
export_tflite(onnx_model, tflite_output, calibration_data=calib_path)
|
|
call_args = mock_prepare_calib.call_args
|
|
assert call_args[0][1] == calib_path # second positional arg
|
|
|
|
def test_max_images_forwarded_to_prepare(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
"""Verify that max_images is passed to _prepare_calibration_data."""
|
|
export_tflite(onnx_model, tflite_output, max_images=42)
|
|
call_kwargs = mock_prepare_calib.call_args
|
|
assert call_kwargs.kwargs.get("max_images") == 42
|
|
|
|
def test_max_images_default_is_100(
|
|
self,
|
|
onnx_model: Path,
|
|
tflite_output: Path,
|
|
fake_onnx2tf: Any,
|
|
mock_prepare_calib: Any,
|
|
) -> None:
|
|
"""Verify that max_images defaults to 100 when not specified."""
|
|
export_tflite(onnx_model, tflite_output)
|
|
call_kwargs = mock_prepare_calib.call_args
|
|
assert call_kwargs.kwargs.get("max_images") == 100
|
|
|
|
def test_valid_quantizations_set(self) -> None:
|
|
assert _VALID_QUANTIZATIONS == {None, "fp32", "fp16", "int8"}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestExportFormatParameter
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestExportFormatParameter:
|
|
"""Tests for ``format="tflite"`` wiring through ``RFDETR.export()``."""
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _patch_export_deps(self, tmp_path: Path) -> None:
|
|
"""Mock the heavy export dependencies so ``RFDETR.export()`` is fast."""
|
|
self._tmp_path = tmp_path
|
|
onnx_out = tmp_path / "inference_model.onnx"
|
|
onnx_out.write_bytes(b"onnx")
|
|
|
|
import contextlib
|
|
|
|
self._mock_stack = contextlib.ExitStack()
|
|
|
|
# Mock export_onnx to return a fake ONNX file path
|
|
self._mock_export_onnx = self._mock_stack.enter_context(
|
|
mock.patch("rfdetr.export.main.export_onnx", return_value=str(onnx_out))
|
|
)
|
|
# Mock make_infer_image to return a small tensor
|
|
import torch
|
|
|
|
self._mock_stack.enter_context(
|
|
mock.patch(
|
|
"rfdetr.export.main.make_infer_image",
|
|
return_value=torch.zeros(1, 3, 560, 560),
|
|
)
|
|
)
|
|
# Mock export_tflite
|
|
self._mock_export_tflite = self._mock_stack.enter_context(
|
|
mock.patch(
|
|
"rfdetr.export._tflite.converter.export_tflite",
|
|
return_value=tmp_path / "inference_model_float32.tflite",
|
|
)
|
|
)
|
|
yield
|
|
self._mock_stack.close()
|
|
|
|
@staticmethod
|
|
def _make_rfdetr() -> Any:
|
|
"""Create a minimal RFDETR instance with mocked internals."""
|
|
from rfdetr.detr import RFDETR
|
|
|
|
obj = RFDETR.__new__(RFDETR)
|
|
obj.model = mock.MagicMock()
|
|
obj.model.resolution = 560
|
|
obj.model.device = "cpu"
|
|
obj.model.model.to.return_value = obj.model.model
|
|
obj.model_config = mock.MagicMock()
|
|
obj.model_config.segmentation_head = False
|
|
obj.model_config.patch_size = 14
|
|
obj.model_config.num_windows = 1
|
|
return obj
|
|
|
|
def test_tflite_format_calls_export_tflite(self) -> None:
|
|
obj = self._make_rfdetr()
|
|
obj.export(format="tflite", output_dir=str(self._tmp_path / "out"))
|
|
self._mock_export_tflite.assert_called_once()
|
|
|
|
def test_onnx_format_does_not_call_export_tflite(self) -> None:
|
|
obj = self._make_rfdetr()
|
|
obj.export(format="onnx", output_dir=str(self._tmp_path / "out"))
|
|
self._mock_export_tflite.assert_not_called()
|
|
|
|
def test_quantization_forwarded(self) -> None:
|
|
obj = self._make_rfdetr()
|
|
obj.export(
|
|
format="tflite",
|
|
output_dir=str(self._tmp_path / "out"),
|
|
quantization="int8",
|
|
)
|
|
call_kwargs = self._mock_export_tflite.call_args
|
|
assert call_kwargs[1].get("quantization") == "int8" or call_kwargs.kwargs.get("quantization") == "int8"
|
|
|
|
@pytest.mark.parametrize(
|
|
"quant",
|
|
[
|
|
pytest.param(None, id="none"),
|
|
pytest.param("fp32", id="fp32"),
|
|
pytest.param("fp16", id="fp16"),
|
|
pytest.param("int8", id="int8"),
|
|
],
|
|
)
|
|
def test_all_quantization_modes_accepted(self, quant: str | None) -> None:
|
|
obj = self._make_rfdetr()
|
|
obj.export(
|
|
format="tflite",
|
|
output_dir=str(self._tmp_path / "out"),
|
|
quantization=quant,
|
|
)
|
|
self._mock_export_tflite.assert_called_once()
|
|
|
|
def test_unsupported_format_raises(self) -> None:
|
|
obj = self._make_rfdetr()
|
|
with pytest.raises(ValueError, match="[Uu]nsupported.*format"):
|
|
obj.export(format="banana", output_dir=str(self._tmp_path / "out"))
|
|
|
|
def test_calibration_data_forwarded(self) -> None:
|
|
"""Verify calibration_data kwarg reaches export_tflite."""
|
|
obj = self._make_rfdetr()
|
|
calib = "/my/calib.npy"
|
|
obj.export(
|
|
format="tflite",
|
|
output_dir=str(self._tmp_path / "out"),
|
|
calibration_data=calib,
|
|
)
|
|
call_kwargs = self._mock_export_tflite.call_args
|
|
assert call_kwargs[1].get("calibration_data") == calib or call_kwargs.kwargs.get("calibration_data") == calib
|
|
|
|
def test_max_images_forwarded(self) -> None:
|
|
"""Verify max_images kwarg reaches export_tflite."""
|
|
obj = self._make_rfdetr()
|
|
obj.export(
|
|
format="tflite",
|
|
output_dir=str(self._tmp_path / "out"),
|
|
max_images=50,
|
|
)
|
|
call_kwargs = self._mock_export_tflite.call_args
|
|
assert call_kwargs[1].get("max_images") == 50 or call_kwargs.kwargs.get("max_images") == 50
|
|
|
|
def test_max_images_default_is_100(self) -> None:
|
|
"""Verify max_images defaults to 100 when not specified."""
|
|
obj = self._make_rfdetr()
|
|
obj.export(
|
|
format="tflite",
|
|
output_dir=str(self._tmp_path / "out"),
|
|
)
|
|
call_kwargs = self._mock_export_tflite.call_args
|
|
assert call_kwargs[1].get("max_images") == 100 or call_kwargs.kwargs.get("max_images") == 100
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestNumpyAllowPickle
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestNumpyAllowPickle:
|
|
"""Tests for ``_numpy_allow_pickle()`` context manager."""
|
|
|
|
def test_patches_np_load(self) -> None:
|
|
original = np.load
|
|
with _numpy_allow_pickle():
|
|
assert np.load is not original
|
|
assert np.load is original
|
|
|
|
def test_sets_allow_pickle_default(self) -> None:
|
|
calls: list[dict[str, Any]] = []
|
|
original = np.load
|
|
|
|
def _spy(*args: Any, **kwargs: Any) -> Any:
|
|
calls.append(kwargs.copy())
|
|
raise ValueError("stop")
|
|
|
|
np.load = _spy # type: ignore[assignment]
|
|
try:
|
|
with _numpy_allow_pickle():
|
|
with pytest.raises(ValueError, match="stop"):
|
|
np.load("dummy.npy")
|
|
assert calls[0].get("allow_pickle") is True
|
|
finally:
|
|
np.load = original # type: ignore[assignment]
|
|
|
|
def test_restores_on_exception(self) -> None:
|
|
"""Patch is restored even when an exception propagates out of the context."""
|
|
original = np.load
|
|
try:
|
|
with _numpy_allow_pickle():
|
|
raise RuntimeError("boom")
|
|
except RuntimeError:
|
|
pass
|
|
assert np.load is original
|
|
|
|
def test_concurrent_access_lock_not_reentrant(self) -> None:
|
|
"""Lock prevents a second context from patching np.load while the first is still active."""
|
|
import threading
|
|
|
|
entered = threading.Event()
|
|
released = threading.Event()
|
|
|
|
def _hold_context() -> None:
|
|
with _numpy_allow_pickle():
|
|
entered.set()
|
|
released.wait(timeout=2.0)
|
|
|
|
holder = threading.Thread(target=_hold_context)
|
|
holder.start()
|
|
entered.wait(timeout=2.0)
|
|
|
|
acquired = _NUMPY_LOAD_PATCH_LOCK.acquire(blocking=False)
|
|
released.set()
|
|
holder.join(timeout=2.0)
|
|
|
|
assert not acquired, "Lock must remain held while first context is active"
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestPatchValidationDownload
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPatchValidationDownload:
|
|
"""Tests for ``_patch_validation_download()`` context manager."""
|
|
|
|
def test_patches_download_in_common_functions(self, tmp_path: Path, fake_onnx2tf: Any) -> None:
|
|
"""The patch replaces download_test_image_data in common_functions."""
|
|
data = np.random.rand(5, 128, 128, 3).astype(np.float32)
|
|
npy_path = tmp_path / "calib.npy"
|
|
np.save(str(npy_path), data)
|
|
|
|
cf_mod = sys.modules["onnx2tf.utils.common_functions"]
|
|
original_fn = cf_mod.download_test_image_data
|
|
|
|
with _patch_validation_download(str(npy_path)):
|
|
assert cf_mod.download_test_image_data is not original_fn
|
|
result = cf_mod.download_test_image_data()
|
|
np.testing.assert_array_equal(result, data)
|
|
|
|
# Restored after exit
|
|
assert cf_mod.download_test_image_data is original_fn
|
|
|
|
def test_patches_download_in_onnx2tf_module(self, tmp_path: Path, fake_onnx2tf: Any) -> None:
|
|
"""The patch also covers onnx2tf.onnx2tf if it has the function."""
|
|
data = np.random.rand(3, 64, 64, 3).astype(np.float32)
|
|
npy_path = tmp_path / "calib.npy"
|
|
np.save(str(npy_path), data)
|
|
|
|
# Create onnx2tf.onnx2tf submodule with download_test_image_data
|
|
onnx2tf_inner = types.ModuleType("onnx2tf.onnx2tf")
|
|
original_fn = mock.MagicMock()
|
|
onnx2tf_inner.download_test_image_data = original_fn # type: ignore[attr-defined]
|
|
sys.modules["onnx2tf.onnx2tf"] = onnx2tf_inner
|
|
|
|
try:
|
|
with _patch_validation_download(str(npy_path)):
|
|
assert onnx2tf_inner.download_test_image_data is not original_fn
|
|
result = onnx2tf_inner.download_test_image_data()
|
|
np.testing.assert_array_equal(result, data)
|
|
|
|
assert onnx2tf_inner.download_test_image_data is original_fn
|
|
finally:
|
|
sys.modules.pop("onnx2tf.onnx2tf", None)
|
|
|
|
def test_restores_on_exception(self, tmp_path: Path, fake_onnx2tf: Any) -> None:
|
|
"""Functions are restored even when an exception occurs."""
|
|
npy_path = tmp_path / "calib.npy"
|
|
np.save(str(npy_path), np.zeros((1, 8, 8, 3), dtype=np.float32))
|
|
|
|
cf_mod = sys.modules["onnx2tf.utils.common_functions"]
|
|
original_fn = cf_mod.download_test_image_data
|
|
|
|
try:
|
|
with _patch_validation_download(str(npy_path)):
|
|
raise RuntimeError("boom")
|
|
except RuntimeError:
|
|
pass
|
|
|
|
assert cf_mod.download_test_image_data is original_fn
|
|
|
|
def test_skips_missing_modules(self, tmp_path: Path) -> None:
|
|
"""Does not raise when onnx2tf modules are not in sys.modules."""
|
|
npy_path = tmp_path / "calib.npy"
|
|
np.save(str(npy_path), np.zeros((1, 8, 8, 3), dtype=np.float32))
|
|
|
|
# Ensure onnx2tf is NOT in sys.modules
|
|
keys = [k for k in sys.modules if k == "onnx2tf" or k.startswith("onnx2tf.")]
|
|
saved = {k: sys.modules.pop(k) for k in keys}
|
|
|
|
try:
|
|
with _patch_validation_download(str(npy_path)):
|
|
pass # should not raise
|
|
finally:
|
|
sys.modules.update(saved)
|
|
|
|
|
|
class TestGetOnnxInputInfo:
|
|
"""Tests for ``_get_onnx_input_info()``."""
|
|
|
|
def test_reads_input_name_and_shape(self, tmp_path: Path) -> None:
|
|
"""Build a minimal ONNX model and verify we read back its metadata."""
|
|
onnx = pytest.importorskip("onnx", reason="onnx not installed")
|
|
TensorProto, helper = onnx.TensorProto, onnx.helper # noqa: N806
|
|
|
|
inp = helper.make_tensor_value_info("images", TensorProto.FLOAT, [1, 3, 560, 560])
|
|
out = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 100, 4])
|
|
node = helper.make_node("Identity", inputs=["images"], outputs=["output"])
|
|
graph = helper.make_graph([node], "test", [inp], [out])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
onnx_path = tmp_path / "test_model.onnx"
|
|
onnx.save(model, str(onnx_path))
|
|
|
|
name, dims = _get_onnx_input_info(onnx_path)
|
|
assert name == "images"
|
|
assert dims == [1, 3, 560, 560]
|
|
|
|
def test_different_input_shape(self, tmp_path: Path) -> None:
|
|
"""Verify non-square resolution reads correctly."""
|
|
onnx = pytest.importorskip("onnx", reason="onnx not installed")
|
|
TensorProto, helper = onnx.TensorProto, onnx.helper # noqa: N806
|
|
|
|
inp = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 448, 640])
|
|
out = helper.make_tensor_value_info("output", TensorProto.FLOAT, [1, 10, 4])
|
|
node = helper.make_node("Identity", inputs=["input"], outputs=["output"])
|
|
graph = helper.make_graph([node], "test", [inp], [out])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)])
|
|
onnx_path = tmp_path / "test_model.onnx"
|
|
onnx.save(model, str(onnx_path))
|
|
|
|
name, dims = _get_onnx_input_info(onnx_path)
|
|
assert name == "input"
|
|
assert dims == [1, 3, 448, 640]
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestPrepareCalibrationData
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestPrepareCalibrationData:
|
|
"""Tests for ``_prepare_calibration_data()``."""
|
|
|
|
@pytest.fixture()
|
|
def _mock_onnx_info(self) -> Generator:
|
|
"""Mock ``_get_onnx_input_info`` to return a known shape."""
|
|
with mock.patch(
|
|
"rfdetr.export._tflite.converter._get_onnx_input_info",
|
|
return_value=("input", [1, 3, 256, 256]),
|
|
):
|
|
yield
|
|
|
|
def test_none_generates_random_data(self, tmp_path: Path, _mock_onnx_info: None) -> None:
|
|
onnx_path = tmp_path / "model.onnx"
|
|
onnx_path.write_bytes(b"\x00")
|
|
|
|
npy_path = _prepare_calibration_data(onnx_path, None, tmp_path, "fp32")
|
|
|
|
assert isinstance(npy_path, Path)
|
|
assert npy_path.is_file()
|
|
|
|
data = np.load(str(npy_path))
|
|
assert data.shape == (_DEFAULT_CALIB_SAMPLES, 256, 256, 3)
|
|
assert data.dtype == np.float32
|
|
|
|
def test_none_int8_emits_warning(self, tmp_path: Path, _mock_onnx_info: None) -> None:
|
|
onnx_path = tmp_path / "model.onnx"
|
|
onnx_path.write_bytes(b"\x00")
|
|
|
|
with mock.patch("rfdetr.export._tflite.converter.logger") as mock_logger:
|
|
_prepare_calibration_data(onnx_path, None, tmp_path, "int8")
|
|
mock_logger.warning.assert_called_once()
|
|
assert "INT8" in mock_logger.warning.call_args[0][0]
|
|
|
|
def test_ndarray_saves_to_npy(self, tmp_path: Path, _mock_onnx_info: None) -> None:
|
|
onnx_path = tmp_path / "model.onnx"
|
|
onnx_path.write_bytes(b"\x00")
|
|
calib = np.random.rand(10, 256, 256, 3).astype(np.float32)
|
|
|
|
npy_path = _prepare_calibration_data(onnx_path, calib, tmp_path, "fp32")
|
|
|
|
loaded = np.load(str(npy_path))
|
|
np.testing.assert_array_equal(loaded, calib)
|
|
|
|
def test_path_string_used_directly(self, tmp_path: Path, _mock_onnx_info: None) -> None:
|
|
onnx_path = tmp_path / "model.onnx"
|
|
onnx_path.write_bytes(b"\x00")
|
|
npy_file = tmp_path / "my_calib.npy"
|
|
np.save(str(npy_file), np.zeros((5, 256, 256, 3), dtype=np.float32))
|
|
|
|
npy_path = _prepare_calibration_data(onnx_path, str(npy_file), tmp_path, "fp32")
|
|
|
|
assert npy_path == npy_file
|
|
|
|
def test_directory_loads_images(self, tmp_path: Path, _mock_onnx_info: None) -> None:
|
|
"""A directory path triggers image loading and .npy creation."""
|
|
from PIL import Image
|
|
|
|
onnx_path = tmp_path / "model.onnx"
|
|
onnx_path.write_bytes(b"\x00")
|
|
img_dir = tmp_path / "images"
|
|
img_dir.mkdir()
|
|
for i in range(5):
|
|
img = Image.new("RGB", (100, 80), color=(i * 50, 0, 0))
|
|
img.save(img_dir / f"img_{i:03d}.jpg")
|
|
|
|
npy_path = _prepare_calibration_data(onnx_path, str(img_dir), tmp_path, "int8")
|
|
|
|
assert npy_path.is_file()
|
|
data = np.load(str(npy_path))
|
|
# _mock_onnx_info returns [1, 3, 256, 256] → H=256, W=256
|
|
assert data.shape == (5, 256, 256, 3)
|
|
assert data.dtype == np.float32
|
|
|
|
def test_directory_respects_max_images(self, tmp_path: Path, _mock_onnx_info: None) -> None:
|
|
"""Verify that max_images limits the number of images loaded from a directory."""
|
|
from PIL import Image
|
|
|
|
onnx_path = tmp_path / "model.onnx"
|
|
onnx_path.write_bytes(b"\x00")
|
|
img_dir = tmp_path / "images"
|
|
img_dir.mkdir()
|
|
for i in range(10):
|
|
img = Image.new("RGB", (100, 80), color=(i * 25, 0, 0))
|
|
img.save(img_dir / f"img_{i:03d}.jpg")
|
|
|
|
npy_path = _prepare_calibration_data(onnx_path, str(img_dir), tmp_path, "int8", max_images=3)
|
|
|
|
assert npy_path.is_file()
|
|
data = np.load(str(npy_path))
|
|
assert data.shape[0] == 3 # only 3 of 10 images loaded
|
|
|
|
def test_missing_path_raises_file_not_found(self, tmp_path: Path, _mock_onnx_info: None) -> None:
|
|
onnx_path = tmp_path / "model.onnx"
|
|
onnx_path.write_bytes(b"\x00")
|
|
|
|
with pytest.raises(FileNotFoundError, match="Calibration data path not found"):
|
|
_prepare_calibration_data(onnx_path, "/nonexistent/calib.npy", tmp_path, "fp32")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestLoadCalibrationImages
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class TestLoadCalibrationImages:
|
|
"""Tests for ``_load_calibration_images()``."""
|
|
|
|
@staticmethod
|
|
def _make_images(directory: Path, count: int = 5, size: tuple[int, int] = (100, 80)) -> list[Path]:
|
|
"""Create *count* small JPEG images in *directory*."""
|
|
from PIL import Image
|
|
|
|
paths: list[Path] = []
|
|
for i in range(count):
|
|
p = directory / f"img_{i:04d}.jpg"
|
|
Image.new("RGB", size, color=(i * 40 % 256, 0, 0)).save(p)
|
|
paths.append(p)
|
|
return paths
|
|
|
|
def test_loads_images_with_correct_shape(self, tmp_path: Path) -> None:
|
|
self._make_images(tmp_path, count=3)
|
|
result = _load_calibration_images(tmp_path, height=128, width=256)
|
|
|
|
assert isinstance(result, np.ndarray)
|
|
assert result.shape == (3, 128, 256, 3)
|
|
assert result.dtype == np.float32
|
|
assert result.min() >= 0.0
|
|
assert result.max() <= 1.0
|
|
|
|
def test_respects_max_images(self, tmp_path: Path) -> None:
|
|
self._make_images(tmp_path, count=10)
|
|
result = _load_calibration_images(tmp_path, height=64, width=64, max_images=4)
|
|
assert result.shape[0] == 4
|
|
|
|
def test_empty_directory_raises(self, tmp_path: Path) -> None:
|
|
empty = tmp_path / "empty"
|
|
empty.mkdir()
|
|
with pytest.raises(FileNotFoundError, match="No supported image files"):
|
|
_load_calibration_images(empty, height=64, width=64)
|
|
|
|
def test_nonexistent_directory_raises(self, tmp_path: Path) -> None:
|
|
with pytest.raises(FileNotFoundError, match="Calibration image directory not found"):
|
|
_load_calibration_images(tmp_path / "does_not_exist", height=64, width=64)
|
|
|
|
def test_unsupported_extensions_ignored(self, tmp_path: Path) -> None:
|
|
"""Only image extensions are loaded; .txt files are skipped."""
|
|
self._make_images(tmp_path, count=2)
|
|
(tmp_path / "readme.txt").write_text("hello")
|
|
(tmp_path / "data.csv").write_text("a,b")
|
|
|
|
result = _load_calibration_images(tmp_path, height=32, width=32)
|
|
assert result.shape[0] == 2
|
|
|
|
def test_skips_unreadable_files(self, tmp_path: Path) -> None:
|
|
"""Corrupt images are skipped without raising."""
|
|
self._make_images(tmp_path, count=3)
|
|
# Write a corrupt file with a supported extension
|
|
(tmp_path / "corrupt.jpg").write_bytes(b"not-a-jpeg")
|
|
|
|
result = _load_calibration_images(tmp_path, height=32, width=32)
|
|
assert result.shape[0] == 3 # only the 3 valid images
|
|
|
|
def test_all_unreadable_raises(self, tmp_path: Path) -> None:
|
|
"""If all files are unreadable, raises FileNotFoundError."""
|
|
(tmp_path / "bad1.jpg").write_bytes(b"garbage")
|
|
(tmp_path / "bad2.png").write_bytes(b"garbage")
|
|
|
|
with pytest.raises(FileNotFoundError, match="No readable images"):
|
|
_load_calibration_images(tmp_path, height=32, width=32)
|
|
|
|
def test_png_and_jpeg_both_loaded(self, tmp_path: Path) -> None:
|
|
"""Both .jpg and .png formats are loaded."""
|
|
from PIL import Image
|
|
|
|
Image.new("RGB", (50, 50), "red").save(tmp_path / "a.jpg")
|
|
Image.new("RGB", (50, 50), "blue").save(tmp_path / "b.png")
|
|
|
|
result = _load_calibration_images(tmp_path, height=32, width=32)
|
|
assert result.shape[0] == 2
|
|
|
|
def test_constants_are_reasonable(self) -> None:
|
|
"""Sanity-check the module-level constants."""
|
|
assert _DEFAULT_DIR_CALIB_SAMPLES > 0
|
|
assert ".jpg" in _IMAGE_EXTENSIONS
|
|
assert ".png" in _IMAGE_EXTENSIONS
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestCheckOnnx2tfAvailable
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@onnx2tf_available
|
|
class TestCheckOnnx2tfAvailable:
|
|
"""Tests for ``_check_onnx2tf_available()``."""
|
|
|
|
def test_available_when_importable(self, fake_onnx2tf: Any) -> None:
|
|
_check_onnx2tf_available() # should not raise
|
|
|
|
def test_raises_when_not_importable(self) -> None:
|
|
"""ImportError is raised with install hint when onnx2tf is absent."""
|
|
_remove_fake_onnx2tf()
|
|
with mock.patch.dict(sys.modules, {"onnx2tf": None}):
|
|
with pytest.raises(ImportError, match="onnx2tf is not installed"):
|
|
_check_onnx2tf_available()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# TestGridSampleOnnxRewrite
|
|
# ---------------------------------------------------------------------------
|
|
|
|
onnx_gs_available = pytest.mark.skipif(
|
|
not all(
|
|
__import__("importlib").util.find_spec(p) is not None for p in ("onnx", "onnx_graphsurgeon", "onnxruntime")
|
|
),
|
|
reason="onnx, onnx_graphsurgeon, and onnxruntime required",
|
|
)
|
|
|
|
|
|
def _build_gridsample_onnx(
|
|
path: Path,
|
|
*,
|
|
n: int = 1,
|
|
c: int = 4,
|
|
h: int = 8,
|
|
w: int = 8,
|
|
h_out: int = 4,
|
|
w_out: int = 4,
|
|
) -> None:
|
|
"""Write a minimal ONNX model with one GridSample node to *path*."""
|
|
import onnx
|
|
from onnx import TensorProto, helper
|
|
|
|
im = helper.make_tensor_value_info("im", TensorProto.FLOAT, [n, c, h, w])
|
|
grid = helper.make_tensor_value_info("grid", TensorProto.FLOAT, [n, h_out, w_out, 2])
|
|
out = helper.make_tensor_value_info("out", TensorProto.FLOAT, [n, c, h_out, w_out])
|
|
node = helper.make_node(
|
|
"GridSample",
|
|
inputs=["im", "grid"],
|
|
outputs=["out"],
|
|
mode="bilinear",
|
|
padding_mode="zeros",
|
|
align_corners=0,
|
|
)
|
|
graph = helper.make_graph([node], "gs_test", [im, grid], [out])
|
|
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 16)])
|
|
onnx.save(model, str(path))
|
|
|
|
|
|
@pytest.fixture()
|
|
def gridsample_onnx(tmp_path: Path) -> Path:
|
|
"""Build a tiny ONNX file with a single GridSample node."""
|
|
p = tmp_path / "gs_model.onnx"
|
|
_build_gridsample_onnx(p)
|
|
return p
|
|
|
|
|
|
class TestGridSampleOnnxRewrite:
|
|
"""Tests for the ONNX-level GridSample → Gather(axis=0) rewrite."""
|
|
|
|
def test_module_import_does_not_raise(self) -> None:
|
|
"""Importing the converter module must succeed regardless of onnx2tf version."""
|
|
import rfdetr.export._tflite.converter # noqa: F401
|
|
|
|
@onnx_gs_available
|
|
def test_no_gridsample_nodes_after_rewrite(self, gridsample_onnx: Path, tmp_path: Path) -> None:
|
|
"""_replace_gridsample_for_tflite removes all GridSample nodes from the graph."""
|
|
import onnx
|
|
import onnx_graphsurgeon as gs
|
|
|
|
from rfdetr.export._tflite.converter import _replace_gridsample_for_tflite
|
|
|
|
patched_path = _replace_gridsample_for_tflite(gridsample_onnx, tmp_path)
|
|
|
|
model = onnx.load(str(patched_path))
|
|
graph = gs.import_onnx(model)
|
|
remaining = [n for n in graph.nodes if n.op == "GridSample"]
|
|
assert remaining == [], f"Expected no GridSample nodes; found {len(remaining)}"
|
|
|
|
@onnx_gs_available
|
|
def test_gather_nodes_present_after_rewrite(self, gridsample_onnx: Path, tmp_path: Path) -> None:
|
|
"""Rewritten graph contains Gather nodes (the TFLite-safe replacement ops)."""
|
|
import onnx
|
|
import onnx_graphsurgeon as gs
|
|
|
|
from rfdetr.export._tflite.converter import _replace_gridsample_for_tflite
|
|
|
|
patched_path = _replace_gridsample_for_tflite(gridsample_onnx, tmp_path)
|
|
|
|
model = onnx.load(str(patched_path))
|
|
graph = gs.import_onnx(model)
|
|
gather_nodes = [n for n in graph.nodes if n.op == "Gather"]
|
|
assert len(gather_nodes) >= 4, f"Expected ≥4 Gather nodes (one per bilinear corner); found {len(gather_nodes)}"
|
|
|
|
@onnx_gs_available
|
|
def test_numerical_equivalence_vs_pytorch(self, gridsample_onnx: Path, tmp_path: Path) -> None:
|
|
"""Rewritten ONNX output matches torch.nn.functional.grid_sample within 1e-5."""
|
|
import onnxruntime as ort
|
|
import torch
|
|
import torch.nn.functional as F # noqa: N812
|
|
|
|
from rfdetr.export._tflite.converter import _replace_gridsample_for_tflite
|
|
|
|
rng = np.random.default_rng(0)
|
|
im_np = rng.standard_normal((1, 4, 8, 8)).astype(np.float32)
|
|
grid_np = rng.uniform(-1, 1, (1, 4, 4, 2)).astype(np.float32)
|
|
|
|
ref = F.grid_sample(
|
|
torch.from_numpy(im_np),
|
|
torch.from_numpy(grid_np),
|
|
mode="bilinear",
|
|
padding_mode="zeros",
|
|
align_corners=False,
|
|
).numpy()
|
|
|
|
patched_path = _replace_gridsample_for_tflite(gridsample_onnx, tmp_path)
|
|
sess = ort.InferenceSession(str(patched_path), providers=["CPUExecutionProvider"])
|
|
(result,) = sess.run(None, {"im": im_np, "grid": grid_np})
|
|
|
|
np.testing.assert_allclose(
|
|
result, ref, atol=1e-5, rtol=0, err_msg="Gather(axis=0) rewrite output diverges from F.grid_sample"
|
|
)
|