TinyYOLOv2 → SONNX C export: object detection tutorial

This notebook takes the pretrained TinyYOLOv2 (ONNX model zoo, VOC, 20 classes) and runs it through the aidge_export_sonnx pipeline:

Requirements: aidge_core, aidge_onnx, aidge_export_sonnx, a C++ toolchain with make, plus onnx, numpy, Pillow, matplotlib (and optionally onnxruntime for the sanity check at the end).

1. Setup

import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

import numpy as np
import onnx
from onnx import TensorProto, helper
from PIL import Image
import matplotlib.pyplot as plt

import aidge_core
import aidge_onnx
import aidge_export_sonnx

from aidge_export_sonnx.examples.tinyyolo.sonnx_export_tiny_yolo import (
    BLOCKS,
    ChainBuilder,
)
from aidge_export_sonnx.examples.tinyyolo.run_tinyyolo_inference import (
    load_and_preprocess,
    parse_sonnx_results,
    decode_yolo,
    draw_boxes,
    VOC_CLASSES,
    GRID_SIZE,
)

All generated artifacts (downloaded model, .npy weights, rebuilt ONNX graph, exported C++ project) go to a temporary folder. We chdir into it so that relative output paths used by the export API never land in the repository.

WORKDIR = Path(tempfile.mkdtemp(prefix="tinyyolo_tutorial_"))
os.chdir(WORKDIR)
print("Working directory:", WORKDIR)

2. Get the pretrained weights

The model and its weights are obtained by running the extract_tinyyolo_weights.py script shipped with the package. It does two things:

1. downloads tinyyolov2-8.onnx from the ONNX model zoo (if the file is not already present);

2. extracts every initializer as a .npy file.

We point both the model file and the output folder to the temporary working directory, so nothing is written to the repository.

model_file = WORKDIR / "tinyyolov2-8.onnx"
weights_dir = WORKDIR / "weights"

subprocess.run(
    [
        sys.executable,
        "-m",
        "aidge_export_sonnx.examples.tinyyolo.extract_tinyyolo_weights",
        str(model_file),
        "--out",
        str(weights_dir),
    ],
    check=True,
)

print(f"{len(list(weights_dir.glob('*.npy')))} weight files in {weights_dir}")

3. The network definition

The network is described as a list of ConvBnLReluMaxPoolConfig, one per block. Adding or changing a layer means editing one entry of that list. Here is a summary of what the package defines:

print(
    f"{'block':8s} {'out_ch':>6s} {'kernel':>7s} {'bn':>3s} {'lrelu':>5s} {'pool (k/s)':>12s}"
)
for b in BLOCKS:
    pool = f"{b.pool_kernel}/{b.pool_strides}" if b.pool_enabled else "-"
    print(
        f"{b.name:8s} {b.out_channels:6d} {str(b.kernel_hw):>7s} "
        f"{'Y' if b.bn_enabled else '-':>3s} {'Y' if b.leaky_enabled else '-':>5s} {pool:>12s}"
    )

4. Load and preprocess an image

The image is resized to 416×416 and scaled to [0, 1]. We fetch the classic Darknet dog.jpg as a sample; set IMAGE_PATH to your own file to use something else.

import urllib.request

IMAGE_PATH = None  # e.g. "/path/to/my_image.jpg"

if IMAGE_PATH is None:
    IMAGE_PATH = str(WORKDIR / "sample.jpg")
    urllib.request.urlretrieve(
        "https://raw.githubusercontent.com/pjreddie/darknet/master/data/dog.jpg",
        IMAGE_PATH,
    )

x, orig_size = load_and_preprocess(IMAGE_PATH)
print("Original size:", orig_size, "| network input:", x.shape, x.dtype)
plt.imshow(Image.open(IMAGE_PATH))
plt.axis("off")
plt.title("Input image")
plt.show()

5. Build the ONNX graph

ChainBuilder walks over BLOCKS, loads each weight file from weights_dir, computes the SAME_UPPER paddings and appends the Conv, BatchNormalization, LeakyRelu and MaxPool nodes. Everything is float64, which is the datatype used for the export.

input_shape = list(x.shape)
X = helper.make_tensor_value_info("X", TensorProto.DOUBLE, input_shape)

builder = ChainBuilder(weights_dir, input_name="X", input_shape=input_shape)
builder.add_blocks(BLOCKS)
model = builder.build_model(X)
onnx.checker.check_model(model)

onnx_path = WORKDIR / "tinyyolo_chain.onnx"
onnx.save(model, onnx_path)
print(
    f"{len(model.graph.node)} nodes, saved to {onnx_path} ({onnx_path.stat().st_size / 1e6:.0f} MB)"
)

6. Import into aidge, compile with the SONNX backend, schedule

aidge_model = aidge_onnx.load_onnx(str(onnx_path))
print("Node types:", sorted({n.type() for n in aidge_model.get_nodes()}))

aidge_model.compile(
    backend="export_sonnx",
    datatype=aidge_core.dtype.float64,
    dims=[input_shape],
)

scheduler = aidge_core.SequentialScheduler(aidge_model)
scheduler.generate_scheduling()

7. Export to C++, build and run

The export embeds the input tensor in the generated project. make builds bin/run_main, whose standard output contains the raw network output between SONNX_RESULTS_BEGIN / SONNX_RESULTS_END markers.

import time

export_folder = "sonnx_export_tinyyolo"
aidge_export_sonnx.export(
    export_folder_name=export_folder,
    scheduler=scheduler,
    inputs_tensor=[aidge_core.Tensor(x, backend="cpu")],
)

build = subprocess.run(
    ["make", "-C", export_folder, "clean", "all"], capture_output=True, text=True
)
if build.returncode != 0:
    print(build.stdout, build.stderr)
    raise RuntimeError("make failed")

binary = Path(export_folder) / "bin" / "run_main"
t0 = time.perf_counter()
run = subprocess.run(
    [str(binary.resolve())], capture_output=True, text=True, cwd=export_folder
)
elapsed = time.perf_counter() - t0
if run.returncode != 0:
    print(run.stdout, run.stderr)
    raise RuntimeError(f"run_main exited with code {run.returncode}")
print(f"Inference time: {elapsed * 1000:.1f} ms")

8. Decode the output and draw the detections

The raw tensor has shape (1, 125, 13, 13) = 5 anchors × (4 box coordinates + 1 objectness + 20 classes). decode_yolo applies sigmoid to x, y and objectness, softmax to the class scores, exp on w/h scaled by the anchor priors, filters by confidence and applies non-maximum suppression.

values = parse_sonnx_results(run.stdout)
assert values is not None, "No SONNX_RESULTS block found in the binary's output"

output = np.array(values, dtype=np.float64).reshape(
    1, BLOCKS[-1].out_channels, GRID_SIZE, GRID_SIZE
)
print("Raw output shape:", output.shape)

boxes = decode_yolo(output, conf_thresh=0.3, nms_thresh=0.4)
for x1, y1, x2, y2, conf, cls in boxes:
    print(
        f"{VOC_CLASSES[cls]:>12s}  conf={conf:.3f}  box=({x1:.3f}, {y1:.3f}, {x2:.3f}, {y2:.3f})"
    )

result_path = WORKDIR / "result.jpg"
draw_boxes(IMAGE_PATH, boxes, str(result_path), orig_size)
plt.figure(figsize=(8, 8))
plt.imshow(Image.open(result_path))
plt.axis("off")
plt.show()

9. (Optional) Sanity check against ONNX Runtime

We run the original zoo model with ONNX Runtime and compare its output with the SONNX result. The zoo model embeds its own image scaler, so it takes raw 0..255 pixels.

try:
    import onnxruntime as ort
except ImportError:
    print("onnxruntime not installed, skipping this check")
else:
    sess = ort.InferenceSession(str(model_file), providers=["CPUExecutionProvider"])
    inp = sess.get_inputs()[0]
    ref = sess.run(None, {inp.name: (x * 255.0).astype(np.float32)})[0]
    print("max |SONNX - ONNX Runtime| =", float(np.abs(output - ref).max()))
    np.testing.assert_allclose(output, ref, rtol=1e-3, atol=1e-3)
    print("Outputs match.")

10. Cleanup

os.chdir(Path.home())
shutil.rmtree(WORKDIR, ignore_errors=True)
print("Temporary files removed.")