How I built a YOLO pipeline with ONNX and TensorRT for Jetson

I built Edge YOLO TensorRT as a portfolio demo of a detector pipeline: fine-tune YOLO, export ONNX, build a TensorRT engine, and measure inference on the machine running it. This is a demonstration project, with simulated inputs in its tests, not a record of production use.

The README results table is empty on purpose. There are no checked-in device measurements to quote. The project gives me commands for collecting those measurements, but it does not establish a Jetson frame rate or an accuracy improvement.

Keep the model path separate from the runtime

The intended sequence is workstation training, workstation ONNX export, then engine building and benchmarking on the target Jetson. The artifacts are best.pt, artifacts/coco128.onnx, and a precision-specific .engine file.

I put the common image handling in Detector. OnnxRuntimeDetector supplies an ONNX session call. TensorRtDetector supplies a TensorRT runner. Both inherit the same preprocessing, decoding, coordinate conversion, and detection records. That keeps the backend comparison easier to inspect because the surrounding Python path is shared.

The implementation expects a single inference output. ONNX Runtime takes the first output, while the TensorRT binding rejects multiple outputs. This is a narrow YOLO pipeline rather than a general model server.

The code around the graph

Prepare the same tensor for either backend

In src/edge_yolo/preprocess.py, the input conversion is explicit:

def prepare_input(image: np.ndarray, input_size: int) -> tuple[np.ndarray, LetterboxMeta]:
    _require_hwc_uint8(image, "image")
    if input_size < 32:
        raise ValueError("input_size must be >= 32")
    canvas, meta = letterbox(image, input_size)
    blob = canvas.astype(np.float32) / 255.0
    blob = np.transpose(blob, (2, 0, 1))[None]
    return np.ascontiguousarray(blob), meta

Images must be RGB uint8 arrays in height, width, channel order. Letterboxing preserves aspect ratio and pads a square canvas with 114. The returned metadata records scale, padding, and original dimensions so the decoder can put boxes back into image coordinates.

The resulting tensor is contiguous float32, normalized to zero through one, with batch and channel dimensions first. A .npy input still has to satisfy the RGB array contract. PNG and JPEG loading uses Pillow through the optional media dependency.

There is a detail I would keep visible when comparing runs: resizing uses OpenCV's linear interpolation when OpenCV is available, otherwise a NumPy nearest-neighbor fallback. Sharing a function does not guarantee identical pixels across environments with different optional dependencies.

Decode once, then undo the padding

This excerpt from src/edge_yolo/detectors/base.py shows the shared prediction path before it constructs detection objects:

    def predict(self, image: np.ndarray) -> list[Detection]:
        blob, meta = prepare_input(image, self.input_size)
        raw = self.infer(blob)
        boxes, scores, class_ids = decode_output(
            raw,
            self.layout,
            self.conf_threshold,
            self.iou_threshold,
        )
        boxes = scale_boxes(boxes, meta)

The rest of the method attaches confidence, class ID, and class name to each box. src/edge_yolo/postprocess.py supports raw YOLOv8 output, YOLOv5 output with objectness, and six-column end-to-end output. For YOLOv5, the score includes objectness multiplied by the selected class score.

Filtering removes nonfinite scores or coordinates, low confidence predictions, and boxes with nonpositive dimensions. Nonmaximum suppression runs separately for each class, so overlapping boxes from different classes can survive. Coordinate scaling subtracts the letterbox padding, divides by the scale, and clips to the original image dimensions.

Automatic layout selection is a shape heuristic. I would pass an explicit layout when working with an unfamiliar export; the CLI exposes --layout yolov5 for the objectness case. Batch sizes above one are rejected by the decoder.

Treat INT8 as a calibrated build

The INT8 branch in src/edge_yolo/export_tensorrt.py makes the fallback policy visible:

        case Precision.INT8:
            return EngineBuildPlan(
                precision,
                enable_fp16=True,
                enable_int8=True,
                requires_calibrator=True,
            )

ONNX export stays FP32 through half=False. Precision is selected later when building the engine. The builder checks the GPU's advertised fast FP16 and INT8 support, parses the ONNX graph, and reports parser errors if parsing fails.

Calibration lives in src/edge_yolo/calibrate.py. It loads frames through the same preprocessing function, groups them into contiguous float batches, and supplies device pointers through IInt8EntropyCalibrator2. A trailing partial batch is dropped. Cache reads and writes use a separate small class.

The CLI's image-based INT8 route requires a calibration directory even when a cache path is supplied, and it loads those images before creating the calibrator. A cache is not a directory-free entry point here. Representative camera frames matter; a blank array can exercise plumbing but cannot support a useful calibration claim.

Run the pipeline

Python 3.10 or newer is the declared requirement. From README.md, the base development install is:

python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"

For workstation training and ONNX inference, the next README command is pip install -e ".[train,onnx,media]". TensorRT comes from the Jetson environment, not a project pip extra.

Run python -m edge_yolo train --config configs/coco128.yaml. That config starts from yolov8n.pt, uses the public 128-image coco128 subset, and requests 20 epochs at image size 640 with batch size 8. These are settings, not achieved results. The first training run downloads the checkpoint and subset. The expected weights path is runs/detect/coco128/weights/best.pt.

The export commands below are copied from README.md. Run the ONNX step after training; choose the engine build appropriate to the target and available calibration frames:

python -m edge_yolo export-onnx --config configs/export.yaml

python -m edge_yolo export-trt \
  --onnx artifacts/coco128.onnx \
  --precision fp16 \
  --output artifacts/coco128_fp16.engine

python -m edge_yolo export-trt \
  --onnx artifacts/coco128.onnx \
  --precision int8 \
  --output artifacts/coco128_int8.engine \
  --calibration-dir data/calibration \
  --calibration-cache artifacts/coco128_int8.cache

configs/export.yaml selects opset 17, simplification, a static export, and a 2048 MiB workspace limit. The documented inference path is batch one. If an ONNX input has dynamic dimensions, the builder creates an optimization profile with identical minimum, optimum, and maximum shapes.

For Jetson setup, I keep the board-specific instructions in docs/jetson.md. That guide targets JetPack 6, L4T R36, and the TensorRT 10 Python API. The Dockerfile defaults to L4T tag r36.2.0; match the image to the board's release. The guide also covers power modes, clocks, and installing the Jetson PyTorch wheel before Ultralytics when exporting on the board.

Measure the whole prediction call

The timing loop in src/edge_yolo/benchmark.py discards warmup calls:

    for index in range(warmup):
        predict(images[index % len(images)])

    samples_ms: list[float] = []
    for index in range(iterations):
        image = images[index % len(images)]
        start = time.perf_counter()
        predict(image)
        samples_ms.append((time.perf_counter() - start) * 1000.0)

The CLI passes detector.predict, so samples include preprocessing, network execution, and postprocessing. Image loading and model construction happen before this loop. The TensorRT runner synchronizes its stream before returning output, so this path waits for execution rather than just timing submission.

This README command measures the engine on the calling machine:

python -m edge_yolo bench \
  --backend tensorrt \
  --model artifacts/coco128_fp16.engine \
  --precision fp16 \
  --input-size 640 \
  --warmup 20 \
  --iterations 200 \
  --json results/device.json

Without --images, it uses seeded synthetic frames. A real-frame directory changes the workload, but this command still measures timing, not detection accuracy. --precision only labels the report; it does not change the engine.

The report includes mean, p50, p95, p99, minimum, maximum, and FPS. FPS is 1000 divided by mean milliseconds for serial prediction calls. I would record power mode and input conditions beside any published row, and keep the generated JSON with it.

What the tests establish

.github/workflows/ci.yml runs Ruff checks, formatting checks, and pytest on Ubuntu with Python 3.11 for pull requests and pushes to main. The README commands are ruff check ., ruff format --check ., and pytest.

Tests check letterbox padding, normalized tensor layout, box rescaling, class-aware suppression, export arguments, calibration batching, and benchmark report serialization. Backend equivalence uses a fake ONNX session and an injected TensorRT runner returning scripted tensors. It verifies shared decoding, not numerical agreement between real engines.

The project also tests that the README results cells stay empty. CI neither downloads training data nor requires a GPU. Passing it cannot certify a TensorRT build or camera-stream behavior.

Limits and the next run

I would next collect target-device measurements and evaluate accuracy before and after precision changes on labeled frames. Neither result is stored here. The current runner allocates and frees CUDA buffers and creates and destroys a stream on every call; reusing those resources is a concrete optimization to investigate after establishing a baseline.

This repo stops at individual images and serial benchmark calls. Camera capture, sustained streaming, and application integration would be separate work.

Source and reproduction

The repository contains the configs, runtime adapters, tests, and Jetson guide used here. I would start with the offline tests, then train and export, then build and measure on the intended board. The empty results table stays empty until there is a device report to support a row.