View Source YoloFastNMS

A fast Non-Maximum Suppression (NMS) implementation in Rust for YOLO object detection outputs. This library is designed as a companion to the yolo library, specifically optimizing one of the most computationally intensive parts of the YOLO detection pipeline in Elixir/Nx, running much faster (~100x, ~4ms vs ~400ms on my MacBook Air M3)!

The NMS operation is crucial for filtering overlapping bounding boxes in object detection, but is computationally expensive when implemented in Elixir/Nx. This Rust NIF implementation provides significant speed improvements over the pure Elixir version, helping achieve near real-time performance for object detection tasks.

The current version is specifically built around YOLOv8 models trained on the COCO dataset, which outputs tensors in the shape of {84, 8400} (4 bounding box coordinates + 80 class probabilities × 8400 detections). The library efficiently filters these detections based on confidence scores prob_threshold and Intersection over Union (IoU) iou_threshold thresholds. Future development will focus on making the number of classes and detections dynamic, enabling compatibility with different YOLO and custom models.

Installation

def deps do
  [
    {:yolo_fast_nms, "~> 0.1.0"}
  ]
end

Stand-alone usage

{84, 8400} = tensor_nx.shape
iex> YoloFastNMS.run(tensor_nx, 0.4, 0.5)
[
  [cx, cy, w, h, prob, class_idx]
  ...
] 

With YOLO library

model = YOLO.load(...)

YOLO.detect(mat, nms_fun: &YoloFastNMS.run/3)