scripts/optimize.py

#!/usr/bin/env python3
"""Generate optimized ECSeg ONNX variants from an FP32 baseline.

Every variant PRESERVES the public graph I/O contract (input `images` float32 [1,3,640,640];
outputs labels int64 [1,300], boxes/scores/masks float32) because the app's curated-import
verification checks tensor dtype+shape against a pinned contract. A variant that changed an
output dtype would fail that gate, so we never let optimization touch the boundary tensors.

Subcommands (each writes <out_dir>/<stem>.<variant>.onnx and runs onnx.checker):

  graphopt     onnxslim constant-folding + lossless simplification (runtime-agnostic; no
               ORT-CPU-specific fusions baked in, so it stays portable to ort-web WASM)
  fp16         onnxconverter-common float16 weights, keep_io_types=True (Cast at boundaries),
               with a conservative op_block_list for ops whose fp16 kernels are risky on WASM
  int8-dynamic weight-only dynamic INT8 (per-channel) on MatMul/Gemm
  int8-static  static QDQ INT8 with real calibration images (Conv+MatMul), optional exclude list

Usage:
  python optimize.py graphopt      IN.onnx  OUT_DIR
  python optimize.py fp16          IN.onnx  OUT_DIR
  python optimize.py int8-dynamic  IN.onnx  OUT_DIR
  python optimize.py int8-static   IN.onnx  OUT_DIR --calib-dir DIR [--calib-limit N] [--exclude-mask-head]
"""

from __future__ import annotations

import argparse
import os
import sys

import onnx

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
import ecseg_common as ec  # noqa: E402


def _check(path: str) -> None:
    model = onnx.load(path, load_external_data=False)
    onnx.checker.check_model(model)
    print(f"  onnx.checker OK: {path} ({os.path.getsize(path)/2**20:.2f} MiB)")


def _stem(in_path: str) -> str:
    base = os.path.basename(in_path)
    return base[:-5] if base.endswith(".onnx") else base


# --------------------------------------------------------------------------- graphopt --------

def graphopt(in_path: str, out_dir: str) -> str:
    import onnxslim

    out = os.path.join(out_dir, f"{_stem(in_path)}.graphopt.onnx")
    # onnxslim performs constant folding, dead-node elimination and safe fusions while keeping the
    # graph runtime-agnostic (unlike ORT's offline optimizer, which can bake in CPU-specific layout
    # fusions that don't transfer to ort-web WASM).
    model = onnxslim.slim(in_path)
    onnx.save(model, out)
    _check(out)
    return out


# ------------------------------------------------------------------------------ fp16 ---------

# Ops we keep in fp32 even inside an fp16 graph: boundary-sensitive / sampling / reduction ops
# whose fp16 kernels are either absent or numerically fragile on the ort-web WASM CPU EP.
FP16_OP_BLOCK_LIST = [
    "GridSample",  # deformable mask sampling — precision-critical, fp16 kernel support uncertain
    "Resize",
    "LayerNormalization",
    "Softmax",
    "ReduceSum",
    "ReduceMean",
    "ReduceMax",
    "Erf",
    "TopK",
    "GatherElements",
    "Einsum",
]


def fp16(in_path: str, out_dir: str, hi_range: bool = False) -> str:
    from onnxconverter_common import float16

    suffix = "fp16-hi" if hi_range else "fp16"
    out = os.path.join(out_dir, f"{_stem(in_path)}.{suffix}.onnx")
    model = onnx.load(in_path)
    kwargs = dict(
        keep_io_types=True,          # inputs/outputs stay float32 -> contract preserved
        disable_shape_infer=False,
        op_block_list=FP16_OP_BLOCK_LIST,
    )
    if hi_range:
        # Use fp16's true representable range instead of the library's conservative ±1e4 clamp, so large
        # mask logits are not clipped. min_positive_val = smallest fp16 normal (2**-14).
        kwargs["max_finite_val"] = 65504.0
        kwargs["min_positive_val"] = 5.9604645e-08
    converted = float16.convert_float_to_float16(model, **kwargs)
    onnx.save(converted, out)
    _check(out)
    return out


# ------------------------------------------------------------------------ int8-dynamic -------

def int8_dynamic(in_path: str, out_dir: str) -> str:
    from onnxruntime.quantization import QuantType, quantize_dynamic
    from onnxruntime.quantization.shape_inference import quant_pre_process

    stem = _stem(in_path)
    pre = os.path.join(out_dir, f"{stem}.preproc.onnx")
    quant_pre_process(in_path, pre, skip_symbolic_shape=False)

    out = os.path.join(out_dir, f"{stem}.int8-dynamic.onnx")
    quantize_dynamic(
        pre,
        out,
        weight_type=QuantType.QInt8,
        per_channel=True,
        # Conv dynamic quantization is not well supported; restrict to the matmul/gemm family.
        op_types_to_quantize=["MatMul", "Gemm"],
    )
    os.remove(pre)
    _check(out)
    return out


# ------------------------------------------------------------------------- int8-static -------

class _CalibReader:
    """onnxruntime CalibrationDataReader over app-faithful preprocessed images."""

    def __init__(self, files, input_name="images"):
        self.files = files
        self.input_name = input_name
        self._it = iter(files)

    def get_next(self):
        path = next(self._it, None)
        if path is None:
            return None
        return {self.input_name: ec.preprocess_file(path)}

    def rewind(self):
        self._it = iter(self.files)


# Mask-head node-name substrings to optionally keep in fp32 (the segmentation branch is the most
# quantization-sensitive part). Discovered empirically from the graph; matched case-insensitively.
MASK_HEAD_HINTS = ["mask", "seg", "pixel_decoder", "GridSample"]


def int8_static(
    in_path: str,
    out_dir: str,
    calib_dir,
    calib_limit=None,
    exclude_mask_head=False,
) -> str:
    from onnxruntime.quantization import (
        CalibrationMethod,
        QuantFormat,
        QuantType,
        quantize_static,
    )
    from onnxruntime.quantization.shape_inference import quant_pre_process

    stem = _stem(in_path)
    suffix = "int8-static-selective" if exclude_mask_head else "int8-static"
    pre = os.path.join(out_dir, f"{stem}.preproc.onnx")
    quant_pre_process(in_path, pre, skip_symbolic_shape=False)

    files = ec.list_images(calib_dir if isinstance(calib_dir, list) else [calib_dir], limit=calib_limit)
    if not files:
        raise SystemExit(f"No calibration images found under {calib_dir}")
    print(f"  calibration images: {len(files)}")

    nodes_to_exclude = []
    if exclude_mask_head:
        model = onnx.load(pre, load_external_data=False)
        for node in model.graph.node:
            name = (node.name or "").lower()
            if any(h.lower() in name for h in MASK_HEAD_HINTS):
                nodes_to_exclude.append(node.name)
        print(f"  excluding {len(nodes_to_exclude)} mask-head nodes from quantization")

    out = os.path.join(out_dir, f"{stem}.{suffix}.onnx")
    quantize_static(
        pre,
        out,
        _CalibReader(files),
        quant_format=QuantFormat.QDQ,
        per_channel=True,
        weight_type=QuantType.QInt8,
        activation_type=QuantType.QUInt8,
        calibrate_method=CalibrationMethod.MinMax,
        nodes_to_exclude=nodes_to_exclude,
        extra_options={"AddQDQPairToWeight": False},
    )
    os.remove(pre)
    _check(out)
    return out


# -------------------------------------------------------------------------------- cli --------

def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    sub = ap.add_subparsers(dest="cmd", required=True)

    for name in ("graphopt", "fp16", "int8-dynamic"):
        p = sub.add_parser(name)
        p.add_argument("in_path")
        p.add_argument("out_dir")
        if name == "fp16":
            p.add_argument("--hi-range", action="store_true",
                           help="use fp16's full ±65504 range instead of the default ±1e4 clamp")

    ps = sub.add_parser("int8-static")
    ps.add_argument("in_path")
    ps.add_argument("out_dir")
    ps.add_argument("--calib-dir", action="append", required=True)
    ps.add_argument("--calib-limit", type=int, default=None)
    ps.add_argument("--exclude-mask-head", action="store_true")

    args = ap.parse_args()
    os.makedirs(args.out_dir, exist_ok=True)

    if args.cmd == "graphopt":
        graphopt(args.in_path, args.out_dir)
    elif args.cmd == "fp16":
        fp16(args.in_path, args.out_dir, hi_range=getattr(args, "hi_range", False))
    elif args.cmd == "int8-dynamic":
        int8_dynamic(args.in_path, args.out_dir)
    elif args.cmd == "int8-static":
        int8_static(
            args.in_path, args.out_dir, args.calib_dir, args.calib_limit, args.exclude_mask_head
        )
    return 0


if __name__ == "__main__":
    sys.exit(main())