View Source YOLO.NMS (YOLO v0.1.1)

Elixir NMS (Non-Maximum Suppression)

Learn more about Non-Maximum Suppression (NMS) at: https://builtin.com/machine-learning/non-maximum-suppression

This implementation applies NMS independently for each output class. The following steps are executed for each class:

  1. Filter out all bounding boxes with maximum class probability below prob_threshold (default: 0.5).
  2. Select the bounding box with the highest prob.
  3. Remove any remaining bounding boxes with an IoU >= iou_threshold (default: 0.5).

Summary

Functions

Filters detections, keeping only those with a probability higher than :prob_threshold.

IoU (Intersection over Union) measures the ratio of the area of overlap between two bounding boxes to the area of their union.

Applies Non-Maximum Suppression (NMS) for each class, discarding bounding boxes with an IoU exceeding the specified iou_threshold.

Functions

filter_predictions(tensor, prob_threshold)

@spec filter_predictions(Nx.Tensor.t(), float()) :: [[float()]]

Filters detections, keeping only those with a probability higher than :prob_threshold.

The input tensor must have the shape {8400, 84} (transposed YOLOv8 output format).

Returns a list of [bbox_cx, bbox_cy, bbox_w, bbox_h, prob, class_idx].

This implementation is inspired by Hans Elias B. Josephsen's talk (see the filter function at 12:06): https://youtu.be/OsxGB6MbA8o?t=726

iou(list1, list2)

IoU (Intersection over Union) measures the ratio of the area of overlap between two bounding boxes to the area of their union.

IoU = A ∩ B / A U B

nms(bboxes, iou_threshold)

Applies Non-Maximum Suppression (NMS) for each class, discarding bounding boxes with an IoU exceeding the specified iou_threshold.

run(tensor, prob_threshold, iou_threshold)

@spec run(Nx.Tensor.t(), float(), float()) :: [[float()]]

Runs both filter_predictions/2 and nms/2.

  1. Filters out detections with a probability below prob_threshold (p < prob_threshold)
  2. Applies Non-Maximum Suppression (NMS) for each class, discarding bounding boxes with an IoU exceeding the specified iou_threshold.

The input tensor must have the shape {8400, 84} (transposed YOLOv8 output format).

Returns a list of [bbox_cx, bbox_cy, bbox_w, bbox_h, prob, class_idx].