# Converting `finegrain/finegrain-box-segmenter` to ONNX (with box prompting)

This is a record of how the Hugging Face model
[`finegrain/finegrain-box-segmenter`](https://huggingface.co/finegrain/finegrain-box-segmenter)
got converted to ONNX, how the bounding-box prompt is preserved, and how the exported model behaved
under `onnxruntime-node` on both CPU and WebGPU.

Anything load-bearing is sourced inline. The longer notes behind this summary sit in
[`../.research/`](../.research) — six files, one per angle I researched.

---

## 1. What the model actually is

- It's **MVANet** (Multi-View Aggregation Network, CVPR 2024), arXiv:2404.07445. Sources: the model
  card, https://huggingface.co/finegrain/finegrain-box-segmenter ; the paper,
  https://arxiv.org/abs/2404.07445
- The code lives in Finegrain's Refiners library, exposed as `refiners.solutions.BoxSegmenter`. The
  network class is `refiners.foundationals.swin.mvanet.MVANet`. License is MIT.
  Sources: https://github.com/finegrain-ai/refiners ; the model card above.
- Backbone is Swin-B (`embed_dim=128`, `depths=[2,2,18,2]`, `num_heads=[4,8,16,32]`,
  `window_size=12`, `patch_size=4`), input 1024×1024. Sources: `refiners .../swin/mvanet/mvanet.py`,
  and the official `qianyu-dlut/MVANet` `SwinB(...)`
  (https://raw.githubusercontent.com/qianyu-dlut/MVANet/main/model/SwinTransformer.py).
- The weights are one file, `model.safetensors` (~189 MB on disk: 601 tensors, ~94.6 M params, all
  F16) at revision `v0.1`. There is no `config.json` and no `preprocessor_config.json`; every bit of
  pre/post-processing is hard-coded in Refiners' Python. Sources: HF API
  `https://huggingface.co/api/models/finegrain/finegrain-box-segmenter/tree/v0.1`; safetensors
  header read directly (see `../.research/05-hf-files-weights.md`).

## 2. The key insight: the box prompt is a crop, not a network input

This is the thing the whole conversion hinges on. `BoxSegmenter.run()` never feeds the box to the
network. It's used only on the PIL image side, as geometry:

```python
def run(self, img, box_prompt=None):
    if box_prompt is None: box_prompt = (0, 0, img.width, img.height)
    box = self.add_margin(box_prompt)     # expand by margin=0.05 (5%)
    cropped = self.crop_pad(img, box)     # crop original to box; black-pad if out of bounds
    prediction = self.predict(cropped)    # the ONLY neural call, on a 1024x1024 crop
    out = Image.new("L", (img.width, img.height))
    out.paste(prediction, box)            # paste mask back at the box
    return out
```
Source (verbatim):
https://github.com/finegrain-ai/refiners/blob/7ca1774b5f8f172708db647a26c3be68858f285a/src/refiners/solutions/box_segmenter.py#L40-L79

The network's first layer is `Conv2d(3 → 128)`, and the safetensors backs that up:
`ComputeShallow.Conv2d.weight = [128, 3, 3, 3]`, so 3 input channels. Source:
`../.research/01-refiners-boxsegmenter.md`, `../.research/05-hf-files-weights.md`.

> ⚠️ **Conflicting-source trap, resolved.** Finegrain's blog
> (https://blog.finegrain.ai/posts/promptable-hd-segmentation/) describes a *5-channel*
> (RGB + box + mask) promptable variant. That one is a later, commercial model. The public v0.1
> weights are the older 3-channel, crop-based MVANet, and the `[128,3,3,3]` first conv plus the crop
> logic above settle it. We convert the public weights, so the ONNX graph is image → logits, and the
> box prompt gets reimplemented as crop/paste in the runtime wrapper.

### Exact pre/post-processing (has to be reproduced around the ONNX model)

From `BoxSegmenter.predict`
(https://github.com/finegrain-ai/refiners/blob/7ca1774b5f8f172708db647a26c3be68858f285a/src/refiners/solutions/box_segmenter.py#L62-L67):

- Pre: crop the (5%-margin-expanded) box, convert to RGB, resize to 1024×1024 bilinear (a plain
  squash, no aspect preservation), `/255` to `[0,1]`, then ImageNet-normalize with
  `mean=[0.485,0.456,0.406] std=[0.229,0.224,0.225]`. Layout NCHW `[1,3,1024,1024]` float32.
- Net: raw logits `[1,1,1024,1024]` out of the final `Conv2d(128→1)`. No sigmoid in the model. (There
  is an internal sigmoid inside MCRM's token-attention map, but the output head is raw.)
- Post: sigmoid, `×255` to uint8, resize back to the crop size (bilinear), paste into a black
  full-image `L` canvas at the box.

## 3. The ONNX I/O contract

| | name | shape | dtype | notes |
|---|---|---|---|---|
| input | `input` | `[1, 3, 1024, 1024]` | float32 | RGB, ImageNet-normalized, static |
| output | `logits` | `[1, 1, 1024, 1024]` | float32 | raw logits (apply sigmoid host-side) |

The static shapes aren't optional. The pyramid/multi-view stages hard-code intermediate sizes
(`Interpolate((32,32))…(512,512)`), Swin asserts a square, window-divisible input, and the SW-MSA
mask is size-locked. Dynamic H/W would be neither meaningful nor exportable here. Source:
`../.research/02-mvanet-architecture.md`, `../.research/04-onnx-export-mechanics.md`.

## 4. The export recipe (what worked, and why the obvious path didn't)

Script: [`../python/export_onnx.py`](../python/export_onnx.py).

```python
seg = BoxSegmenter(device="cpu")          # downloads model.safetensors v0.1
model = seg.model.eval().float()          # raw nn.Module; outputs LOGITS
x = torch.randn(1, 3, 1024, 1024)
torch.onnx.export(model, (x,), "mvanet_box_segmenter.onnx",
                  dynamo=False, opset_version=17,        # legacy TorchScript exporter
                  input_names=["input"], output_names=["logits"],
                  do_constant_folding=True)
```

Two decisions, both settled by trial:

1. Legacy exporter (`dynamo=False`), not dynamo. torch 2.12's default dynamo exporter captured the
   graph fine but fell over in the decomposition pass on Swin's
   `x.transpose(1,2).reshape(B, num_windows, N, C)` (swin_transformer.py:202): the pass lowered
   `reshape` to a strict `aten.view` that can't handle the non-contiguous transposed strides
   (`ValueError: Cannot view a tensor with shape [605,144,4,32] … as (5,121,144,128)`). This is a
   known `torch.export` limitation, not a model bug. The legacy TorchScript tracer records a real
   `Reshape` (copy-on-need) and exports cleanly, which is the fallback the research pointed to for
   Swin. Source: `../.research/04-onnx-export-mechanics.md` (§2, §4e); the dynamo failure is in this
   repo's history.
2. `adaptive_avg_pool2d` → `avg_pool2d` patch (`--patch-pool`). Refiners' `Pool` calls
   `adaptive_avg_pool2d(x, (h//ratio, w//ratio))` with `assert h % ratio == 0`, so each output cell
   covers exactly `ratio` pixels, which makes it numerically identical to
   `avg_pool2d(kernel=ratio, stride=ratio)` — and that exports to a plain ONNX `AveragePool`. Same
   workaround the community used for vanilla MVANet (Kazuhito00). Source: refiners
   `.../mvanet/utils.py` `Pool`; https://github.com/Kazuhito00/MVANet-ONNX-Sample ;
   `../.research/02-mvanet-architecture.md` (§4.2), `../.research/03-existing-onnx.md`.

The rest of the architecture exported fine at opset 17 (legacy): `torch.roll` (Slice+Concat),
`interpolate` with constant sizes (`Resize`), SDPA/attention, window reshape/permute, and the
Refiners `SetContext/UseContext` shallow skip. No `grid_sample`, no `einsum`, no `unfold`. Source:
`../.research/02-mvanet-architecture.md`, `../.research/04-onnx-export-mechanics.md`.

Resulting op set: `Conv, AveragePool, MatMul, Gemm, Resize, Softmax, LayerNormalization, PRelu,
Add, Mul, Concat, Split, Slice, Reshape, Transpose, …`, all standard ONNX, runs on ORT CPU/WebGPU.

## 5. Parity — the ONNX matches PyTorch

Script: [`../python/verify_parity.py`](../python/verify_parity.py). On CPU, fp32:

- Random input, torch vs ORT logits: `max|Δ| = 1.5e-5`, sigmoid `max|Δ| = 3e-6`.
- Cactus end-to-end (refiners' own golden test image), torch vs ORT mask: `MAE = 0.000 / 255`, i.e.
  bit-identical. Both torch and ORT differ from refiners' published `expected_cactus_mask.png` by the
  *same* 5.4/255, so that gap is a CPU-vs-original-environment difference in the reference, not
  something ONNX introduced. Golden path matched (MVANet direct, full image, no box):
  https://raw.githubusercontent.com/finegrain-ai/refiners/main/tests/e2e/test_mvanet.py

## 6. Testing with `onnxruntime-node`, and reimplementing the prompt

Harness: [`../node/segment.mjs`](../node/segment.mjs) (`onnxruntime-node` + `sharp`). It reproduces
`BoxSegmenter.run()` host-side: `addMargin(5%) → cropPad(black-pad if OOB) → resize 1024 →
ImageNet-normalize → session.run → sigmoid → resize-back → paste`.

Because the box is a crop, the same image with different boxes gives you different objects, which is
what proves the prompt localizes:

- `two-objects.png` (mug | bottle): the mug-box returns only the mug, the bottle-box only the bottle.
- `cats.jpg`: the left-cat box gives the left cat, the right-cat box the right cat, each one
  excluding the other.
- `*_nobg_*.avif` (a 4096² amulet on wood): the box cuts the amulet cleanly off the wood background.

A CPU run is about 13 s/crop (heavy Swin-B at batch-5). Output masks and cutouts land in
[`../assets/out/`](../assets/out).

> **Three `sharp` gotchas, found and fixed in the host pipeline** (each checked against the refiners
> PyTorch reference on the same crop — e.g. raw-1024 fg 6% → 17% → 37.3%, matching refiners' 37.1%):
> 1. resize-before-composite: sharp runs `.resize()` *before* `.composite()` no matter what order you
>    call them in, so for an out-of-bounds (black-padded) box the crop landed at original scale in the
>    top-left of the 1024² input and the model only saw a shrunken fragment. Fix: render the composite
>    to a buffer *before* resizing (`cropPad`).
> 2. composite emits 4 channels: `sharp({create:channels:3}).composite(…).raw()` comes back RGBA, so
>    reading it as 3-channel mis-strides into magenta/green garbage (and the model then emits a noisy,
>    low-confidence mask). Fix: `.removeAlpha()` before `.raw()` (`cropPad`).
> 3. 1-channel raw → 3 channels on output: the mask resize-back has to force `.toColourspace('b-w')`,
>    or the paste stride drifts and produces horizontal striping (`pasteMask`).
>
> The lesson: sharp's `.raw()` channel count and its fixed pipeline order are the two traps. Pin the
> channel count, and never resize a still-pending composite. Validated by `python/reference_cats.py`
> (which runs refiners' official `BoxSegmenter` for ground truth).

## 7. WebGPU and the `maxStorageBuffersPerShaderStage` limit

The BiRefNet-family failure mode (a wide `Concat`/`Split` exceeding WebGPU's per-shader
storage-buffer limit) was checked with `onnx-webgpu-cascade`
([`../node/webgpu-analyze.mjs`](../node/webgpu-analyze.mjs)):

- The widest Concat/Split in this model is fan 4 → 5 storage buffers (Swin `StatefulPad/Pad/Concat`).
- `needsCascade = false` at limit 16 (this RTX 5070 Ti) and at limit 8 (standard WebGPU). Only at
  limit 4 (compatibility mode) does the cascade fix become necessary.
- The actual WebGPU EP run in `onnxruntime-node` worked as-is: 3.1 s (≈4× faster than CPU), logits
  bit-identical to CPU. No storage-buffer error.

So unlike BiRefNet (a `Concat` with 1024 inputs), this model doesn't trip the limit on normal
adapters. The cascade surgery (`V:\MCP\onnx-webgpu-cascade`) is only for the rare compat-mode
(limit-4) devices. Source: `../.research/06-ort-node-testing.md`, and that package's README.

## 7b. Model variants and size

| file | precision | size | use |
|---|---|---|---|
| `mvanet_box_segmenter.onnx` | fp32 | 805 MB | reference; CPU ~13 s + WebGPU ~3.1 s; exact parity |
| `mvanet_box_segmenter_fp16.onnx` | fp16 (I/O fp32) | 403 MB | compact; WebGPU ~3.5 s (clean) / CPU ~14 s; logits match fp32 within fp16 |

The fp16 model loads and runs on both EPs. On WebGPU its logits match fp32 to within fp16 precision
(`[-53.9, 11.2]` vs `[-53.5, 11.1]`) and the mask is visually identical. On the CPU EP it also runs,
but emits harmless "can't constant-fold fp16 Sqrt/Add" warnings; fp16 is meant for the GPU.

Why is fp32 805 MB when the weights are only ~377 MB? The legacy exporter with
`do_constant_folding=True` bakes in ~425 MB of constant tensors — 3377 `Constant` nodes, the Swin
SW-MSA attention masks and positional-embedding tables that fold at the static 1024² shape. Diagnosed
with [`../python/diag_size.py`](../python/diag_size.py).

What didn't help: exporting with `do_constant_folding=False` is *larger* (1126 MB — BatchNorm not
folded into Conv, constants not deduped). `onnxslim` (the right dedupe tool) trips protobuf's 2 GB
serialization limit on the fp32 model, but works on the fp16 model (it eliminates all 3377 `Constant`
nodes). So the slimming lever that actually works is fp16 → onnxslim
([`../python/optimize_onnx.py`](../python/optimize_onnx.py)).

`onnxconverter_common.float16` doesn't produce a load-clean model for this graph on its own; four
fixes were needed, each prompted by a concrete ORT load error:

1. Shape inference ON (`disable_shape_infer=False`), so the converter can see branch types. Otherwise
   it leaves a mixed fp16/fp32 `Concat` in `SplitMultiView`.
2. `Resize`/`Upsample` `roi`/`scales` → float32 (ONNX requires these to be float32 even when X is
   fp16); the converter sometimes casts those Constants, so fix them back to fp32 afterward.
3. General type reconcile: forward-propagate per-tensor dtypes yourself (onnx shape-inference
   under-types this graph) and `Cast→fp16` every fp32 input of a float op. That fixes the
   `WindowSDPA/Div` scale and ~134 similar spots ORT would otherwise reject.
4. Clear stale `value_info`: the original export's fp32 type annotations survive conversion and
   conflict with the now-fp16 tensors (ORT validates node outputs against them), so drop `value_info`.

Order matters: convert → onnxslim → reconcile → clear value_info. Reconcile has to be *last*, or
onnxslim eats the reconcile Casts as "redundant". The result is a 403 MB fp16 model that loads and
runs on CPU and WebGPU with masks identical to fp32.

## 8. Reproduce

```powershell
# Python env (uv, Python 3.12)
uv venv --python 3.12.10
uv pip install "git+https://github.com/finegrain-ai/refiners.git" onnx onnxruntime onnxscript pillow numpy huggingface_hub safetensors
$env:PYTHONUTF8='1'                         # else a torch progress emoji crashes cp1250 consoles

# Export (legacy opset 17 + avg_pool patch) and verify
python python/export_onnx.py --exporter legacy --opset 17 --patch-pool
python python/verify_parity.py

# Node test (onnxruntime-node + sharp)
cd node && npm install
node segment.mjs --image ../assets/two-objects.png --box 290,400,665,765 --name mug --box 1050,185,1185,780 --name bottle
node webgpu-analyze.mjs                      # storage-buffer analysis
```

## 9. Limitations (model, not conversion)

Per the model card, v0.1 can struggle with objects touching image edges, transparent/non-salient
matting, or busy scenes (reflections, hard shadows). With the host-pipeline fixes above, the exported
model segments both cats and the ornate amulet cleanly and matches the refiners PyTorch reference
(right-cat raw-1024 fg 37.3% vs 37.1%). Any leftover roughness is an upstream model trait, not a
conversion artifact — the ONNX↔PyTorch parity is exact. Source: model card.

## 10. Source index

The detailed notes, every claim tied to a URL:
- `../.research/01-refiners-boxsegmenter.md` — BoxSegmenter/MVANet source, pre/post, weights.
- `../.research/02-mvanet-architecture.md` — architecture plus the per-op ONNX/WebGPU risk table.
- `../.research/03-existing-onnx.md` — prior art (no existing box-segmenter ONNX; vanilla MVANet only).
- `../.research/04-onnx-export-mechanics.md` — exporter choice, opsets, pitfalls, version matrix.
- `../.research/05-hf-files-weights.md` — HF files, safetensors keys, weight-conversion lineage.
- `../.research/06-ort-node-testing.md` — onnxruntime-node recipe + WebGPU caveat.