# ECSeg ONNX optimization pipeline

Reproducible tooling that answers one question: **should AnnotateIt replace the FP32 ECSeg-S / ECSeg-M
checkpoints with an optimized variant?** It downloads the pinned FP32 baselines, builds candidate
variants (lossless graph-opt, FP16, dynamic INT8, static INT8/QDQ), checks numerical equivalence against
FP32, and — decisively — measures every variant in the **target runtime** (onnxruntime-web 1.24.3, CPU/WASM)
via a real browser, not Python.

The full write-up lives in [`docs/model-optimization/ecseg-quantization-report.md`](../../docs/model-optimization/ecseg-quantization-report.md).

> **Nothing here is an app dependency.** The Python tools run in a throwaway venv (see
> `requirements.txt`); the browser benchmark uses the repo's own `node_modules/onnxruntime-web`. No
> generated `.onnx` is committed — they are large and rebuildable.

## Layout

| File | Purpose |
|---|---|
| `inspect_onnx.py` | opset / node & initializer counts / weight dtypes / external-data / I/O contract / op histogram |
| `optimize.py` | build variants: `graphopt`, `fp16`, `int8-dynamic`, `int8-static` (all preserve the public I/O contract) |
| `ecseg_common.py` | app-faithful preprocessing (stretch-640 bilinear + ImageNet) + edgecrafter-seg post-processing |
| `correctness.py` | candidate-vs-FP32 raw-output error + end-to-end instance/box/mask-IoU on identical inputs |
| `aggregate_report.py` | consolidate report JSONs → `combined.json` + CSVs + markdown tables |
| `bench/harness.{html,mjs}` | in-browser ORT harness, configured exactly like the app's ECSeg session |
| `bench/bench_browser.mjs` | Playwright driver: serves harness+ORT+models to Chrome/WebKit, measures latency |

## Prerequisites

```bash
# Python 3.12 (onnxruntime has no 3.14 wheels); create the isolated venv:
python3.12 -m venv .venv && source .venv/bin/activate
pip install -r scripts/model-optimization/requirements.txt

# Node deps (playwright, onnxruntime-web) are already in the repo's node_modules.
# The browser benchmark drives your system Google Chrome (Playwright channel 'chrome').
```

## Reproduce

Set a work dir for artifacts (kept OUT of git):

```bash
export WORK=/tmp/ecseg-opt && mkdir -p $WORK/{baseline,variants,bench-inputs,coco/val2017,reports}
```

### 1. Baseline (download + verify pinned SHA-256)

```bash
curl -sL -o $WORK/baseline/ecseg-s.onnx \
  https://huggingface.co/AnnotateIt/edgecrafter-ecseg-s-onnx/resolve/a71d0ee5e27b1083d40cd5ced647139a7ab8d533/model.onnx
curl -sL -o $WORK/baseline/ecseg-m.onnx \
  https://huggingface.co/AnnotateIt/edgecrafter-ecseg-m-onnx/resolve/56812fa474791c8db90740141e359bc767e0dc22/model.onnx
shasum -a 256 $WORK/baseline/ecseg-s.onnx  # fd30a61cee9b798259e89e5b58cbe57a7c65da2b3067f33ccf02b923402522da
shasum -a 256 $WORK/baseline/ecseg-m.onnx  # ded8147572ecec27fe2b384cfac5e5cc38738c9b3be773d5ccbee3c261820474

python scripts/model-optimization/inspect_onnx.py $WORK/baseline/ecseg-s.onnx --json $WORK/reports/inspect-ecseg-s-fp32.json
```

### 2. Calibration / evaluation images

The report calibrated static INT8 on a 50-image COCO val2017 subset and evaluated agreement on the repo's
held-out `datasets/` images (disjoint from calibration). To refetch the COCO subset, use the id list the
report was produced with (`$WORK/coco/ids.txt`) and pull from `http://images.cocodataset.org/val2017/<id>.jpg`.

### 3. Build variants (`<model>` = ecseg-s or ecseg-m)

```bash
V=$WORK/variants
python scripts/model-optimization/optimize.py graphopt     $WORK/baseline/<model>.onnx $V
python scripts/model-optimization/optimize.py fp16         $WORK/baseline/<model>.onnx $V
python scripts/model-optimization/optimize.py int8-dynamic $WORK/baseline/<model>.onnx $V
python scripts/model-optimization/optimize.py int8-static  $WORK/baseline/<model>.onnx $V --calib-dir $WORK/coco/val2017
python scripts/model-optimization/optimize.py int8-static  $WORK/baseline/<model>.onnx $V --calib-dir $WORK/coco/val2017 --exclude-mask-head
cp $WORK/baseline/<model>.onnx $V/<model>.fp32.onnx
```

### 4. Numerical equivalence vs FP32

```bash
python scripts/model-optimization/correctness.py \
  --baseline $WORK/baseline/ecseg-s.onnx \
  --candidate $WORK/variants/ecseg-s.fp16.onnx \
  --images datasets --json $WORK/reports/corr-ecseg-s-fp16.json
```

### 5. Browser latency in the target runtime (the decisive measurement)

```bash
# Export byte-identical preprocessed inputs (so browser feeds match the Python correctness run):
python - <<'PY'
import sys, numpy as np; sys.path.insert(0,'scripts/model-optimization'); import ecseg_common as ec, os
W=os.environ['WORK']
for n in ['000000000139','000000000785','000000002149']:
    ec.preprocess_file(f'{W}/coco/val2017/{n}.jpg').astype(np.float32).tofile(f'{W}/bench-inputs/{n}.f32')
PY

# wasm-mt (cross-origin isolated, threaded pool) must run HEADED — headless Chrome will not bring up
# ORT's WASM worker pool (session.create hangs). wasm-st can run headless (BENCH_HEADLESS=1).
node scripts/model-optimization/bench/bench_browser.mjs \
  --models-dir $WORK/variants --input-dir $WORK/bench-inputs \
  --models ecseg-s.fp32.onnx,ecseg-s.fp16.onnx,ecseg-s.int8-dynamic.onnx,ecseg-s.int8-static.onnx,ecseg-s.int8-static-selective.onnx,ecseg-s.graphopt.onnx \
  --configs wasm-mt,wasm-st --engine chrome --repeats 15 --warmups 2 \
  --out $WORK/reports/results-s.json

# Optional second engine (Safari/WebKit) for cross-browser confirmation:
#   --engine webkit
```

### 6. Consolidate

```bash
python scripts/model-optimization/aggregate_report.py --reports-dir $WORK/reports --variants-dir $WORK/variants
```

## Runtime facts the tooling encodes (from the app source)

- **Storage is uncompressed** — the offline model store keeps the raw ONNX bytes in IndexedDB
  (`offline-model-store.ts`). gzip/Brotli on the wire does **not** shrink the stored artifact; disk bytes ≈
  IndexedDB bytes. That is why the size tables use on-disk `.onnx` size.
- **The I/O contract is verified strictly at curated import** (`custom-model-curated-flow.ts`
  `verifyTensorAgainstSpec`): tensor **dtype and shape** must match `images` f32[1,3,640,640] →
  `labels` i64[1,300], `boxes`/`scores` f32, `masks` f32[1,300,160,160]. Every variant preserves this
  (FP16 keeps IO in float32 via boundary casts; INT8 keeps IO float).
- **ECSeg is pinned to the CPU/WASM EP** (`custom-checkpoints.ts`) — WebGPU fails this graph on
  ort-web 1.24.3 (the graph uses `GridSample`). The benchmark therefore measures CPU/WASM only.
- **Session options** mirrored by the harness: `executionProviders:['cpu']`,
  `graphOptimizationLevel:'all'`, `logSeverityLevel:3`, `executionMode` sequential (1 thread) / parallel
  (pool), `env.wasm.simd=true`.