File size: 8,608 Bytes

2991ff3

"""Shared model converter: ONNX -> OpenVINO IR (fp16/int8/CoreML) + validation.

Used by both the runtime model-convert cap (convert.py) and the offline
scripts/build-camstack-models.py. Conversion ALWAYS starts from ONNX via
ov.convert_model (never the Ultralytics OV exporter, which emits unnamed
tensors that break the runtime's named-output postprocess)."""
from __future__ import annotations
import os
from pathlib import Path

# OpenVINO sizes the pool that COMPILES a model to the machine's core count and
# TBB never gives those workers back. A container cgroup counts THREADS, and
# this one has been measured refusing forks against `pids.max=2048` — so a
# conversion running beside the live pools must not spend `nproc` of the budget
# just to validate an IR. Same knob, same env var and same default (4) as
# OV_COMPILATION_NUM_THREADS in
# packages/addon-pipeline/python/inference_pool.py, which carries the
# measurement; this process is short-lived, so the saving here is a spike, not
# the 214-thread accumulation that motivated the bound. 0 disables it.
_OV_COMPILATION_NUM_THREADS: int = 4
try:
    _OV_COMPILATION_NUM_THREADS = int(
        os.environ.get("CAMSTACK_OV_COMPILATION_THREADS", "4")
    )
except ValueError:
    _OV_COMPILATION_NUM_THREADS = 4


def export_openvino(onnx_path: str, out_xml: str, *, fp16: bool) -> None:
    import openvino as ov
    ov_model = ov.convert_model(onnx_path)
    Path(out_xml).parent.mkdir(parents=True, exist_ok=True)
    ov.save_model(ov_model, out_xml, compress_to_fp16=fp16)


def validate_openvino_ir(xml_path: str) -> None:
    """Compile the IR and assert every input/output tensor carries a name.
    This is the exact defect (unnamed tensors) that broke 14/26 catalog IRs."""
    import openvino as ov
    core = ov.Core()
    model = core.read_model(xml_path)
    # Bound the compilation pool (see _OV_COMPILATION_NUM_THREADS). A plugin
    # REJECTS a property it does not know, and this compile is the VALIDATION —
    # it must never fail for a reason that has nothing to do with the IR, so a
    # refused property falls back to the unbounded compile rather than
    # reporting a bad model.
    if _OV_COMPILATION_NUM_THREADS > 0:
        try:
            compiled = core.compile_model(
                model,
                "CPU",
                {"COMPILATION_NUM_THREADS": _OV_COMPILATION_NUM_THREADS},
            )
        except Exception:
            compiled = core.compile_model(model, "CPU")
    else:
        compiled = core.compile_model(model, "CPU")
    for port in list(compiled.inputs) + list(compiled.outputs):
        try:
            name = port.get_any_name()
        except Exception as exc:  # ov raises when a tensor has no names
            raise ValueError(f"IR {xml_path} has an unnamed tensor: {exc}") from exc
        if not name:
            raise ValueError(f"IR {xml_path} has an empty-named tensor")


CALIBRATION_SUFFIXES = {".jpg", ".jpeg", ".png"}


def calibration_images(calib_dir: str, max_samples: int) -> "list[Path]":
    """The first ``max_samples`` calibration images of ``calib_dir``, in path order.

    A flat folder (COCO val2017) is read exactly as before: file-name order. A
    folder is walked recursively, so the operator's own frames work as they
    come. The hub's annotated export (a folder holding ``annotations.json``)
    calibrates on its ``frames/`` only: its ``crops/`` are classifier inputs,
    not what a root detector is fed.
    """
    root = Path(calib_dir)
    if not root.is_dir():
        raise ValueError(f"calibration folder {calib_dir} does not exist")
    base = root / "frames" if (root / "annotations.json").is_file() else root
    images = sorted(
        (p for p in base.rglob("*") if p.is_file() and p.suffix.lower() in CALIBRATION_SUFFIXES),
        key=lambda p: p.relative_to(base).as_posix(),
    )[:max_samples]
    if not images:
        raise ValueError(f"no calibration images in {calib_dir}")
    return images


def letterbox_for_calibration(rgb, imgsz: int):
    """``rgb`` (HxWx3 uint8) → ``[1,3,imgsz,imgsz]`` float32 in [0,1], letterboxed
    as the inference pool does (``inference_pool.letterbox_image``): scale by the
    long side, truncate, CENTRE with grey 114 padding."""
    import cv2
    import numpy as np

    h, w = rgb.shape[:2]
    scale = imgsz / max(h, w)
    nw, nh = int(w * scale), int(h * scale)
    pad_x, pad_y = (imgsz - nw) // 2, (imgsz - nh) // 2
    resized = rgb if (nw, nh) == (w, h) else cv2.resize(rgb, (nw, nh), interpolation=cv2.INTER_LINEAR)
    canvas = np.full((imgsz, imgsz, 3), 114, dtype=np.uint8)
    canvas[pad_y : pad_y + nh, pad_x : pad_x + nw] = resized
    chw = canvas.transpose(2, 0, 1).astype(np.float32) / 255.0
    return chw[np.newaxis, ...]


def quantize_openvino_int8(
    onnx_path: str, out_xml: str, calib_dir: str, imgsz: int, max_samples: int = 300
) -> None:
    import cv2
    import openvino as ov
    import nncf

    images = calibration_images(calib_dir, max_samples)

    def preprocess(path: Path):
        img = cv2.imread(str(path))
        return letterbox_for_calibration(cv2.cvtColor(img, cv2.COLOR_BGR2RGB), imgsz)

    dataset = nncf.Dataset(images, lambda p: preprocess(p))
    ov_model = ov.convert_model(onnx_path)
    quantized = nncf.quantize(ov_model, dataset, subset_size=len(images))
    Path(out_xml).parent.mkdir(parents=True, exist_ok=True)
    ov.save_model(quantized, out_xml)


def export_coreml_from_onnx(
    onnx_path: str,
    out_pkg: str,
    imgsz: int,
    *,
    input_shape: "tuple[int, ...] | None" = None,
    fp32: bool = False,
) -> None:
    """ONNX -> CoreML mlpackage (darwin only).

    Modern coremltools (>=6) dropped the ONNX frontend entirely — ``ct.convert``
    only accepts TensorFlow / PyTorch / MIL sources, so feeding it an ``.onnx``
    raises "Unable to determine the type of the model". We bridge through
    PyTorch: ``onnx2torch`` rebuilds the ONNX graph as an ``nn.Module``, we trace
    it with a representative image input, then coremltools converts the
    TorchScript. Image models are NCHW ``(1, 3, imgsz, imgsz)``.

    ``input_shape`` overrides that square-NCHW assumption for models whose input
    is neither square nor channels-first — the plate recognizer is NHWC
    ``(1, 64, 128, 3)``. Without it the trace runs on a wrongly-shaped tensor and
    either throws deep inside the graph or, worse, produces a package whose
    declared input silently disagrees with the real one. ``imgsz`` stays the
    positional contract for every existing caller.

    ``fp32`` keeps full precision instead of coremltools' default fp16 weights.
    Detectors and embedders absorb fp16 — a box moves a fraction of a pixel, a
    cosine moves 5e-5 — but a RECOGNIZER cannot: its output is an argmax plus a
    confidence threshold, so a small shift in one softmax flips a stored
    character or silences a real read. Measured on the plate recognizer, fp16
    moved the weakest-character probability by up to 0.16 and changed the label
    on a ground-truthed plate. The default stays fp16 so nothing already
    published moves.

    Dynamic-shape models (e.g. YOLO) emit ``Shape`` ops that coremltools' Torch
    converter cannot lower ("Converter is not implemented … 'Shape'"). We first
    run onnx-simplifier with the input shape pinned, which constant-folds those
    Shape ops away, so the traced graph is static and convertible.
    """
    import coremltools as ct
    import onnx
    import torch
    from onnx2torch import convert as onnx_to_torch
    from onnxsim import simplify

    shape = list(input_shape) if input_shape is not None else [1, 3, imgsz, imgsz]
    onnx_model = onnx.load(onnx_path)
    input_name = onnx_model.graph.input[0].name
    simplified, ok = simplify(onnx_model, overwrite_input_shapes={input_name: shape})
    if not ok:
        simplified = onnx_model  # fall back to the original if simplify bailed

    model = onnx_to_torch(simplified).eval()
    example = torch.rand(*shape)
    with torch.no_grad():
        traced = torch.jit.trace(model, example, strict=False)

    mlmodel = ct.convert(
        traced,
        source="pytorch",
        inputs=[ct.TensorType(name="input", shape=example.shape)],
        minimum_deployment_target=ct.target.iOS16,
        **({"compute_precision": ct.precision.FLOAT32} if fp32 else {}),
    )
    Path(out_pkg).parent.mkdir(parents=True, exist_ok=True)
    mlmodel.save(out_pkg)


def validate_coreml(pkg_path: str) -> None:
    import coremltools as ct
    ct.models.MLModel(pkg_path)  # load; raises if malformed