diff --git a/CHANGELOG.md b/CHANGELOG.md index aa483f2f2..b957b568a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,16 @@ before 1.4.0 are documented in the ### Added +- **TFLite INT8 export.** `export(format="tflite", int8=True, data=...)` now + runs post-training quantization instead of rejecting the request, and + returns the fully integer artifact (int8 input and output tensors, which is + what EdgeTPU and int8-only MCU runtimes require). `TFLiteBackend` already + quantized inputs and dequantized outputs from the tensor scales, so + `predict()` on the artifact takes float images as before. Calibration data + is mandatory: there is no eight-image default for this format. Accuracy is + not parity-validated per family yet, so measure with `val()` before + deploying. + - **LibreGround** sibling factory: screenshot + instruction → `Results.points`. Shipped adapters are Florence-2-base (MIT), ShowUI-2B (MIT weights; Apache-2.0 code/base), and Qwen3-VL-2B (Apache-2.0). diff --git a/libreyolo/export/exporter.py b/libreyolo/export/exporter.py index deac189c6..36d35ce69 100644 --- a/libreyolo/export/exporter.py +++ b/libreyolo/export/exporter.py @@ -2306,9 +2306,13 @@ class TFLiteExporter(BaseExporter): format_name = "tflite" suffix = ".tflite" requires_onnx = True - supports_int8 = False + supports_int8 = True supports_fp16 = False apply_model_half = False + # Deliberately no default calibration set, matching TensorRT and OpenVINO. + # A full-integer TFLite graph calibrated on eight images is quietly wrong + # rather than loudly missing, so data= stays mandatory. + default_int8_calibration_data = False def __call__(self, *args, dynamic: bool = False, **kwargs) -> str: if dynamic: @@ -2326,11 +2330,6 @@ def _validate(self, half: bool, int8: bool, data: Optional[str]): raise ValueError( "TFLite FP16 export is not supported yet. Omit half=True for FP32." ) - if int8: - raise ValueError( - "TFLite INT8 quantization is not supported yet. " - "Omit int8=True for FP32." - ) return super()._validate(half, int8, data) def _preflight(self, **kwargs): @@ -2348,6 +2347,8 @@ def _export( metadata, onnx_path, half, + int8, + calibration_data, verbose, onnx2tf_args=None, **kwargs, @@ -2364,6 +2365,8 @@ def _export( onnx_path=onnx_path, output_path=output_path, half=half, + int8=int8, + calibration_data=calibration_data, verbose=verbose, onnx2tf_args=onnx2tf_args, metadata=metadata, diff --git a/libreyolo/export/tflite.py b/libreyolo/export/tflite.py index b6f3ef3da..29107b460 100644 --- a/libreyolo/export/tflite.py +++ b/libreyolo/export/tflite.py @@ -674,7 +674,31 @@ def _onnx2tf_command() -> list[str]: return [sys.executable, "-m", "onnx2tf"] -def _find_converted_tflite(output_dir: Path, onnx_path: Path) -> Path: +# onnx2tf emits every quantized variant it can build under -oiqt. Prefer the +# fully integer artifact: it is the only one EdgeTPU and int8-only MCU runtimes +# accept, and TFLiteBackend already quantizes the input and dequantizes the +# outputs from the tensor scales, so float32 callers see no difference. +_INT8_ARTIFACT_SUFFIXES = ("_full_integer_quant", "_integer_quant") + + +def _find_converted_tflite( + output_dir: Path, onnx_path: Path, *, int8: bool = False +) -> Path: + if int8: + for suffix in _INT8_ARTIFACT_SUFFIXES: + exact = output_dir / f"{onnx_path.stem}{suffix}.tflite" + if exact.exists(): + return exact + matches = sorted(output_dir.rglob(f"*{suffix}.tflite")) + if matches: + return matches[0] + # Never fall back to a float artifact here: returning one would hand + # back an FP32 model under int8 filename and precision metadata. + produced = sorted(str(f.relative_to(output_dir)) for f in output_dir.rglob("*")) + raise RuntimeError( + f"onnx2tf did not produce an INT8 TFLite file. Files found: {produced[:20]}" + ) + exact = output_dir / f"{onnx_path.stem}_float32.tflite" if exact.exists(): return exact @@ -693,6 +717,73 @@ def _find_converted_tflite(output_dir: Path, onnx_path: Path) -> Path: ) +def _onnx_input_name(onnx_path: Path) -> str: + """Return the name of the first ONNX graph input.""" + import onnx + + model = onnx.load(str(onnx_path), load_external_data=False) + if not model.graph.input: + raise ValueError(f"ONNX graph has no inputs: {onnx_path}") + return model.graph.input[0].name + + +def _write_int8_calibration_npy(calibration_data: Any, output_dir: Path) -> Path: + """Write the calibration batches as one NHWC float32 ``.npy`` for onnx2tf. + + ``CalibrationDataLoader`` yields already-preprocessed NCHW batches and pads + the final batch by repeating its last image so TensorRT sees a full batch. + onnx2tf instead reads the whole file as a representative dataset, so the + padding is trimmed here and every image is weighted once. The array is + built through a memmap because a few hundred 640x640 float32 images are + larger than the process should hold at once. + """ + npy_path = output_dir / "_int8_calibration.npy" + total = int(calibration_data.num_samples) + array = None + written = 0 + + for batch in calibration_data: + if written >= total: + break + if batch.ndim != 4: + raise ValueError( + "TFLite INT8 calibration requires rank-4 NCHW batches; " + f"got shape {batch.shape}." + ) + chunk = np.ascontiguousarray(np.transpose(batch, (0, 2, 3, 1))) + chunk = chunk[: total - written] + if array is None: + array = np.lib.format.open_memmap( + npy_path, + mode="w+", + dtype=np.float32, + shape=(total, *chunk.shape[1:]), + ) + array[written : written + len(chunk)] = chunk + written += len(chunk) + + if array is None or written == 0: + raise RuntimeError( + "The calibration dataset produced no usable images for TFLite " + "INT8 export. Check that the images in data= can be read." + ) + array.flush() + del array + + if written < total: + # Unreadable images are skipped by the loader, so the memmap can be + # longer than what was filled. Copy the used prefix into a second file + # rather than truncating one that may still be mapped on Windows. + trimmed_path = output_dir / "_int8_calibration_trimmed.npy" + source = np.load(npy_path, mmap_mode="r") + np.save(trimmed_path, np.asarray(source[:written])) + del source + npy_path = trimmed_path + + logger.info("INT8 calibration tensor: %d images -> %s", written, npy_path) + return npy_path + + def _write_metadata_sidecar(output_path: Path, metadata: dict) -> None: sidecar_path = Path(str(output_path) + ".json") with open(sidecar_path, "w") as f: @@ -705,12 +796,19 @@ def export_tflite( output_path: str, *, half: bool = False, + int8: bool = False, + calibration_data: Any = None, verbose: bool = False, onnx2tf_args: Iterable[str] | None = None, metadata: dict | None = None, ) -> str: """Convert a static ONNX model to TensorFlow Lite using onnx2tf. + With ``int8=True`` the converter runs post-training quantization against + ``calibration_data`` and the fully integer artifact is returned. The + calibration batches are the model's own preprocessed tensors, so onnx2tf is + told not to normalize them a second time (mean 0, std 1). + Note: ``onnx2tf_args`` is forwarded only on the YOLO9 CLI path. It is not applicable to the RF-DETR Python-API path and will be ignored there. """ @@ -718,11 +816,23 @@ def export_tflite( raise ValueError( "TFLite FP16 export is not supported yet. Omit half=True for FP32." ) + if int8 and calibration_data is None: + raise ValueError( + "TFLite INT8 export requires calibration data. " + "Pass data= to export()." + ) check_tflite_export_available() model_family = ((metadata or {}).get("model_family") or "").lower() if model_family == "rfdetr": + if int8: + raise NotImplementedError( + "TFLite INT8 export is not implemented for RF-DETR. That family " + "converts through the onnx2tf Python API with a bespoke " + "GridSample and position-embedding fixup that the quantization " + "path has not been run against." + ) if onnx2tf_args is not None: logger.warning( "onnx2tf_args is not supported on the RF-DETR TFLite path " @@ -747,10 +857,40 @@ def export_tflite( "-o", str(tmp_output), "-tb", - "flatbuffer_direct", + # flatbuffer_direct is the validated FP32 lowering, but its integer + # path is strict: it aborts on any op it cannot represent in int8 + # end to end, which YOLO9 hits on EQUAL. tf_converter runs the + # standard TFLite representative-dataset quantization instead and + # leaves such ops in float. + "tf_converter" if int8 else "flatbuffer_direct", "-v", "info" if verbose else "warn", ] + if int8: + calib_npy = _write_int8_calibration_npy(calibration_data, tmp_output) + # mean 0 / std 1: onnx2tf applies (value - mean) / std to the + # calibration tensor, and these batches already went through the + # model's own preprocessing. + cmd += [ + "-oiqt", + "-cind", + _onnx_input_name(onnx_file), + str(calib_npy), + "0.0", + "1.0", + ] + logger.warning( + "TFLite INT8 is post-training quantization; accuracy is not " + "parity-validated per family. Measure the exported artifact " + "with val() before deploying it." + ) + if shutil.which("onnxsim") is None: + logger.warning( + "The onnxsim executable is not on PATH. onnx2tf shells out " + "to it by name, and the tf_converter backend needs the " + "shapes it propagates. Install onnx-simplifier into the " + "active environment if the conversion below fails." + ) if onnx2tf_args is not None: cmd.extend(str(arg) for arg in onnx2tf_args) @@ -762,14 +902,22 @@ def export_tflite( if result.returncode != 0: stdout = result.stdout or "" stderr = result.stderr or "" + hint = "" + if int8 and shutil.which("onnxsim") is None: + hint = ( + "\nHint: onnxsim was not found on PATH. onnx2tf invokes it " + "as a bare command, and without the shapes it propagates " + "the INT8 backend fails while building the Keras graph." + ) raise RuntimeError( - f"onnx2tf failed with exit code {result.returncode}.\n" + f"onnx2tf failed with exit code {result.returncode}.{hint}\n" f"Command: {' '.join(cmd)}\n" f"stdout: {stdout}\n" f"stderr: {stderr}" ) - converted = _find_converted_tflite(tmp_output, onnx_file) + converted = _find_converted_tflite(tmp_output, onnx_file, int8=int8) + logger.info("Selected converted artifact: %s", converted.name) shutil.copy2(converted, dst) if metadata is not None: diff --git a/tests/unit/test_export_tflite.py b/tests/unit/test_export_tflite.py index 8ade158b1..c06129e26 100644 --- a/tests/unit/test_export_tflite.py +++ b/tests/unit/test_export_tflite.py @@ -9,6 +9,7 @@ from pathlib import Path from unittest.mock import MagicMock +import cv2 import numpy as np import pytest import torch @@ -67,9 +68,10 @@ def test_tflite_format_registered(): assert "tflite" in BaseExporter._registry assert TFLiteExporter.suffix == ".tflite" assert TFLiteExporter.requires_onnx is True - assert TFLiteExporter.supports_int8 is False + assert TFLiteExporter.supports_int8 is True assert TFLiteExporter.supports_fp16 is False assert TFLiteExporter.apply_model_half is False + assert TFLiteExporter.default_int8_calibration_data is False def test_tflite_family_support_scaffold(): @@ -103,11 +105,11 @@ def test_tflite_rejects_dynamic_export(): exporter(dynamic=True) -def test_tflite_rejects_int8_export(): +def test_tflite_int8_requires_calibration_data(): exporter = TFLiteExporter(_make_wrapper()) - with pytest.raises(ValueError, match="INT8"): - exporter(output_path="unused.tflite", int8=True, data="coco8") + with pytest.raises(ValueError, match="requires calibration data"): + exporter(output_path="unused.tflite", int8=True) def test_tflite_rejects_fp16_export(): @@ -185,7 +187,10 @@ def test_tflite_export_copies_float32_output(monkeypatch, tmp_path): monkeypatch.setattr(tflite_module, "check_tflite_export_available", lambda: None) monkeypatch.setattr(tflite_module, "_onnx2tf_command", lambda: ["onnx2tf"]) + captured = {} + def fake_run(cmd, capture_output, text): + captured["cmd"] = list(cmd) output_dir = Path(cmd[cmd.index("-o") + 1]) output_dir.mkdir(parents=True, exist_ok=True) (output_dir / "model_float32.tflite").write_bytes(b"fp32") @@ -203,6 +208,8 @@ def fake_run(cmd, capture_output, text): assert result == str(fp32_dst) assert fp32_dst.read_bytes() == b"fp32" + assert captured["cmd"][captured["cmd"].index("-tb") + 1] == "flatbuffer_direct" + assert "-oiqt" not in captured["cmd"] sidecar = json.loads(Path(str(fp32_dst) + ".json").read_text()) assert sidecar["model_family"] == "yolo9" @@ -302,6 +309,262 @@ def fake_export_tflite(**kwargs): assert not Path(captured["tflite"]["onnx_path"]).exists() +class _FakeCalibration: + """Stand-in for CalibrationDataLoader: NCHW batches, padded final batch.""" + + def __init__(self, batches, num_samples): + self._batches = batches + self.num_samples = num_samples + + def __iter__(self): + return iter(self._batches) + + +def _constant_images(*values): + return np.concatenate( + [np.full((1, 3, 2, 2), value, dtype=np.float32) for value in values] + ) + + +def _int8_converter(monkeypatch, produced): + """Patch onnx2tf away and record the command it would have been given.""" + from libreyolo.export import tflite as tflite_module + + monkeypatch.setattr(tflite_module, "check_tflite_export_available", lambda: None) + monkeypatch.setattr(tflite_module, "_onnx2tf_command", lambda: ["onnx2tf"]) + monkeypatch.setattr(tflite_module, "_onnx_input_name", lambda _path: "images") + captured = {} + + def fake_run(cmd, capture_output, text): + captured["cmd"] = list(cmd) + output_dir = Path(cmd[cmd.index("-o") + 1]) + output_dir.mkdir(parents=True, exist_ok=True) + for name, payload in produced.items(): + (output_dir / name).write_bytes(payload) + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + return captured + + +def test_int8_calibration_npy_is_nhwc_and_drops_batch_padding(tmp_path): + from libreyolo.export import tflite as tflite_module + + # Second batch repeats its last image, the way CalibrationDataLoader pads. + loader = _FakeCalibration( + [_constant_images(1.0, 2.0), _constant_images(3.0, 3.0)], + num_samples=3, + ) + + npy_path = tflite_module._write_int8_calibration_npy(loader, tmp_path) + data = np.load(npy_path) + + assert data.shape == (3, 2, 2, 3) + assert data.dtype == np.float32 + np.testing.assert_array_equal(data[:, 0, 0, 0], [1.0, 2.0, 3.0]) + + +def test_int8_calibration_npy_trims_when_images_are_skipped(tmp_path): + from libreyolo.export import tflite as tflite_module + + # num_samples counts files on disk; unreadable ones never reach a batch. + loader = _FakeCalibration([_constant_images(1.0, 2.0)], num_samples=4) + + data = np.load(tflite_module._write_int8_calibration_npy(loader, tmp_path)) + + assert data.shape == (2, 2, 2, 3) + np.testing.assert_array_equal(data[:, 0, 0, 0], [1.0, 2.0]) + + +def test_int8_calibration_npy_rejects_empty_dataset(tmp_path): + from libreyolo.export import tflite as tflite_module + + with pytest.raises(RuntimeError, match="no usable images"): + tflite_module._write_int8_calibration_npy( + _FakeCalibration([], num_samples=4), tmp_path + ) + + +def test_tflite_int8_helper_requires_calibration_data(tmp_path): + from libreyolo.export import tflite as tflite_module + + with pytest.raises(ValueError, match="requires calibration data"): + tflite_module.export_tflite( + str(tmp_path / "model.onnx"), + str(tmp_path / "model.tflite"), + int8=True, + ) + + +def test_tflite_int8_rejected_for_rfdetr(monkeypatch, tmp_path): + from libreyolo.export import tflite as tflite_module + + monkeypatch.setattr(tflite_module, "check_tflite_export_available", lambda: None) + + with pytest.raises(NotImplementedError, match="RF-DETR"): + tflite_module.export_tflite( + str(tmp_path / "model.onnx"), + str(tmp_path / "model.tflite"), + int8=True, + calibration_data=_FakeCalibration([_constant_images(1.0)], num_samples=1), + metadata={"model_family": "rfdetr"}, + ) + + +def test_tflite_int8_export_selects_full_integer_artifact(monkeypatch, tmp_path): + from libreyolo.export import tflite as tflite_module + + captured = _int8_converter( + monkeypatch, + { + "model_float32.tflite": b"fp32", + "model_integer_quant.tflite": b"int8", + "model_full_integer_quant.tflite": b"full-int8", + }, + ) + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"fake onnx") + dst = tmp_path / "model.tflite" + + result = tflite_module.export_tflite( + str(onnx_path), + str(dst), + int8=True, + calibration_data=_FakeCalibration([_constant_images(1.0, 2.0)], 2), + metadata={"model_family": "yolo9", "precision": "int8"}, + ) + + assert result == str(dst) + assert dst.read_bytes() == b"full-int8" + + cmd = captured["cmd"] + # The integer path runs on tf_converter; flatbuffer_direct aborts on any op + # it cannot keep in int8 end to end. + assert cmd[cmd.index("-tb") + 1] == "tf_converter" + assert "-oiqt" in cmd + cind = cmd.index("-cind") + assert cmd[cind + 1] == "images" + assert cmd[cind + 2].endswith(".npy") + # mean 0 / std 1: the calibration batches are already preprocessed, so a + # second normalization inside onnx2tf would calibrate the wrong ranges. + assert cmd[cind + 3 : cind + 5] == ["0.0", "1.0"] + + sidecar = json.loads(Path(str(dst) + ".json").read_text()) + assert sidecar["precision"] == "int8" + + +def test_tflite_int8_export_falls_back_to_integer_quant(monkeypatch, tmp_path): + from libreyolo.export import tflite as tflite_module + + _int8_converter( + monkeypatch, + { + "model_float32.tflite": b"fp32", + "model_integer_quant.tflite": b"int8", + }, + ) + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"fake onnx") + dst = tmp_path / "model.tflite" + + tflite_module.export_tflite( + str(onnx_path), + str(dst), + int8=True, + calibration_data=_FakeCalibration([_constant_images(1.0)], 1), + ) + + assert dst.read_bytes() == b"int8" + + +def test_tflite_int8_export_never_returns_a_float_artifact(monkeypatch, tmp_path): + from libreyolo.export import tflite as tflite_module + + _int8_converter(monkeypatch, {"model_float32.tflite": b"fp32"}) + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"fake onnx") + dst = tmp_path / "model.tflite" + + with pytest.raises(RuntimeError, match="did not produce an INT8"): + tflite_module.export_tflite( + str(onnx_path), + str(dst), + int8=True, + calibration_data=_FakeCalibration([_constant_images(1.0)], 1), + ) + + assert not dst.exists() + + +def test_tflite_exporter_forwards_int8_and_calibration(monkeypatch, tmp_path): + import libreyolo.export.exporter as exporter_module + + _mock_onnx_available(monkeypatch) + + image_dir = tmp_path / "images" + image_dir.mkdir() + rng = np.random.default_rng(0) + for idx in range(3): + cv2.imwrite( + str(image_dir / f"{idx}.jpg"), + rng.integers(0, 256, size=(24, 24, 3), dtype=np.uint8), + ) + data_yaml = tmp_path / "data.yaml" + data_yaml.write_text( + "\n".join( + [ + f"path: {tmp_path.as_posix()}", + "train: images", + "val: images", + "nc: 1", + "names:", + " 0: object", + ] + ), + encoding="utf-8", + ) + + def preprocess(img_rgb, imgsz): + h, w = imgsz if isinstance(imgsz, tuple) else (imgsz, imgsz) + resized = cv2.resize(img_rgb, (w, h)).astype(np.float32) / 255.0 + return np.transpose(resized, (2, 0, 1)), 1.0 + + wrapper = _make_wrapper() + wrapper._get_preprocess_numpy.return_value = preprocess + exporter = TFLiteExporter(wrapper) + output_path = tmp_path / "model.tflite" + captured = {} + + def fake_export_onnx(_nn_model, _dummy, **kwargs): + Path(kwargs["output_path"]).write_bytes(b"onnx") + return kwargs["output_path"] + + def fake_export_tflite(**kwargs): + captured.update(kwargs) + Path(kwargs["output_path"]).write_bytes(b"tflite") + return kwargs["output_path"] + + monkeypatch.setattr(exporter_module, "export_onnx", fake_export_onnx) + monkeypatch.setattr( + "libreyolo.export.tflite.check_tflite_export_available", lambda: None + ) + monkeypatch.setattr("libreyolo.export.tflite.export_tflite", fake_export_tflite) + + exporter( + output_path=str(output_path), + imgsz=16, + simplify=False, + int8=True, + data=str(data_yaml), + ) + + assert captured["int8"] is True + assert captured["half"] is False + assert captured["metadata"]["precision"] == "int8" + assert captured["calibration_data"].num_samples == 3 + assert next(iter(captured["calibration_data"])).shape[1:] == (3, 16, 16) + + def test_tflite_backend_restores_yolonas_channel_first_outputs(): from libreyolo.backends.tflite import TFLiteBackend