Acquiring ONNX Models - digitalcarp/RawNNX GitHub Wiki

Environment Setup

Some models require exporting the ONNX models using Python scripts. The instructions assume you have Python and uv installed. To install uv, see the official installation instructions.

The instructions use an environment variable $EXPORT_SCRIPTS_DIR to refer to scripts/onnx_export/ in the git repo. This directory provides export scripts and uv environment setup configurations. You can export the environment variable to the appropriate path. Alternatively, you can simply call the scripts using the appropriate absolute file path.

To begin, create an empty directory where you can setup the Python environment and run commands. On Linux, this involves running the following:

mkdir rawnnx_export
cd rawnnx_export
cp $EXPORT_SCRIPTS_DIR/pyproject.toml
uv sync

Exporting ONNX models requires you have checkpoints (e.g. pre-trained model weights) already available for the model. It is assumed these are all placed in a directory pointed to by $CHECKPOINT_DIR. Alternatively, you can provide absolute file paths to the checkpoints instead.

Model Specific Instructions

BiRefNet

Download official release of the ONNX model from here. The model used for testing is BiRefNet_lite-general-2K-epoch_232.onnx.

Depth Anything

Download from link. The dynamic versions of the models is the ones we want.

face-parsing PyTorch

The numbers in the model output correspond to areas of the parsed face. Download the pre-trained checkpoint linked here.

git clone https://github.com/zllrunning/face-parsing.PyTorch
cd face-parsing.PyTorch
cp $EXPORT_SCRIPTS_DIR/face-parsing.PyTorch/* .
uv run export_onnx_model.py --checkpoint $CHECKPOINT_DIR/checkpoint.pth --output bisenet_resnet18_face_parsing.onnx

High Quality Segment Anything

Download checkpoints from the official GitHub page here.

Encoder Model

git clone https://github.com/SysCV/sam-hq.git
cd sam-hq
cp $EXPORT_SCRIPTS_DIR/sam-hq/* .
# Light HQ-SAM (vit_tiny)
uv run export_encoder_onnx_model.py \
    --checkpoint $CHECKPOINT_DIR/sam_hq_vit_tiny.pth \
    --model-type vit_tiny \
    --output sam_hq_vit_t_encoder.onnx
# vit_b
uv run export_encoder_onnx_model.py \
    --checkpoint $CHECKPOINT_DIR/sam_hq_vit_b.pth \
    --model-type vit_b \
    --output sam_hq_vit_b_encoder.onnx

Decoder Model

Use the provided ONNX export script as shown here.

You'll need to make some modifications to the script:

  • Line 38: help="In ['default', 'vit_h', 'vit_l', 'vit_b', 'vit_tiny']. Which type of SAM model to export.",
    • Add vit_tiny as a valid input for --model-type
    • If the model is vit_tiny, the interm_embeddings randn should be changed to randn(1, 1, ...)
  • Line 142: encoder_embed_dim_dict = {"vit_b":768,"vit_l":1024,"vit_h":1280,"vit_tiny":160}
git clone https://github.com/SysCV/sam-hq.git
cd sam-hq

uv add --optional sam-hq .

# You may want to use the ONNX export script in this repo and not in sam-hq...
# Light HQ-SAM (vit_tiny)
uv run $EXPORT_SCRIPTS_DIR/export_onnx_model.py \
    --checkpoint $CHECKPOINT_DIR/sam_hq_vit_tiny.pth \
    --model-type vit_tiny \
    --output sam_hq_vit_t_decoder.onnx
# vit_b
uv run $EXPORT_SCRIPTS_DIR/export_onnx_model.py \
    --checkpoint $CHECKPOINT_DIR/sam_hq_vit_b.pth \
    --model-type vit_b \
    --output sam_hq_vit_b_decoder.onnx

uv remove --optional sam-hq segment-anything

MobileSAM

Download the checkpoint files from the link in the official GitHub here.

Encoder Model

The helpful Acly repo provides export scripts for the encoder. Place them in the root directory of the MobileSAM git repository and call the script.

git clone https://github.com/ChaoningZhang/MobileSAM.git
cd MobileSAM

curl --output export_image_encoder.py https://huggingface.co/Acly/MobileSAM/raw/main/mobile_sam_encoder_onnx/export_image_encoder.py
curl --output onnx_image_encoder.py https://huggingface.co/Acly/MobileSAM/raw/main/mobile_sam_encoder_onnx/onnx_image_encoder.py
curl --output mobile_sam.pt https://github.com/ChaoningZhang/MobileSAM/raw/refs/heads/master/weights/mobile_sam.pt

uv run export_image_encoder.py \
    --checkpoint $CHECKPOINT_DIR/mobile_sam.pt \
    --model-type vit_t \
    --output mobile_sam_encoder.onnx

Decoder Model

Download the checkpoint files from the link in the official GitHub here. Afterwards, use the provided ONNX export script as shown here.

git clone https://github.com/ChaoningZhang/MobileSAM.git
cd MobileSAM

uv add --optional mobile-sam .

uv run $EXPORT_SCRIPTS_DIR/export_onnx_model.py \
    --checkpoint $CHECKPOINT_DIR/mobile_sam.pt \
    --model-type vit_t \
    --output mobile_sam_decoder.onnx

uv remove --optional mobile-sam mobile-sam

MODNet

The official repo is has ONNX export scripts here. In the README, there is a link to a direct download to the ONNX version of the official image matting model.

U2-Net

Download the pre-trained models from the official GitHub here. There are also links to specialized pre-trained models in the updates section. The scenarios the specialized models are trained for include:

  • Portraits
  • Humans
  • Clothes
git clone https://github.com/NathanUA/U-2-Net.git
cd U-2-Net

cp $EXPORT_SCRIPTS_DIR/u2net/* .
# You can use replace u2netp with u2net.
uv run export_onnx_model.py \
    --checkpoint=$CHECKPOINT_DIR/u2netp.pth \
    --model-type=u2netp \
    --output=u2netp.onnx

YuNet

Download models from here.