finegrain-box-segmenter-ONNX / docs /HOW-CONVERSION-WAS-MADE.md
MarcinEU's picture
Add fp32 ONNX model, card, usage example, comparison samples, and conversion tooling
c07c0dc
|
Raw
History Blame Contribute Delete
16 kB

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

This is a record of how the Hugging Face model 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/ β€” 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:

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.

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. 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 (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/.

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):

  • 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.

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).

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

# 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.