batched_nms: Batched Non-maximum Suppression (NMS)

In torchvision: Models, Datasets and Transformations for Images

Batched Non-maximum Suppression (NMS)

Description

Performs non-maximum suppression in a batched fashion. Each index value correspond to a category, and NMS will not be applied between elements of different categories.

Usage

batched_nms(boxes, scores, idxs, iou_threshold)

Arguments

boxes	(Tensor[N, 4]): boxes where NMS will be performed. They are expected to be in (x_{min}, y_{min}, x_{max}, y_{max}) format with 0 \leq x_{min} < x_{max} and 0 \leq y_{min} < y_{max}.
scores	(Tensor[N]): scores for each one of the boxes
idxs	(Tensor[N]): indices of the categories for each one of the boxes.
iou_threshold	(float): discards all overlapping boxes with IoU > iou_threshold

Value

keep (Tensor): int64 tensor with the indices of the elements that have been kept by NMS, sorted in decreasing order of scores

torchvision documentation built on Aug. 8, 2025, 7:27 p.m.