Non-maximum suppression

flat_bug.nms.base_nms_(objects: Any, iou_fun: Callable, scores: Tensor, collate_fn: Callable = None, iou_threshold: float = 0.5, strict: bool = True, return_indices: bool = False, **kwargs) → Tensor | Tuple[Any, Tensor]

Implements the standard non-maximum suppression algorithm.

Parameters:

  • objects (Any) – An object which can be indexed by a tensor of indices.

  • iou_fun (Callable) – A function which takes an anchor object and a comparison set (not in the Python sense) of (different) objects and returns the IoU of the anchor object with each object in the comparison set as a tensor of shape (1, n). The reason it is not just (n, ) is to allow for implementations of iou_fun functions between two sets, where the IoU is calculated between each pair of objects from distinct sets.

  • scores (torch.Tensor) – A tensor of shape (n, ) containing the “scores” of the objects, this can merely be though of as a priority score, where the higher the score, the higher the priority of the object - it does not have to be a probability/confidence.

  • collate_fn (Callable, optional) – A function which takes a list of objects and returns a single combined object. Defaults to torch.cat if objects is a tensor and lambda x : x if objects is a list, otherwise it has to be specified.

  • iou_threshold (float, optional) – The IoU threshold for non-maximum suppression. Defaults to 0.5.

  • strict (bool, optional) – A flag to indicate whether to perform strict checks on the algorithm. Defaults to True.

  • return_indices (bool, optional) – A flag to indicate whether to return the indices of the picked objects or the objects themselves. Defaults to False. If True, both the picked objects and scores are returned.

  • **kwargs – Additional keyword arguments to be passed to the iou_fun function.

Returns:

  • torch.Tensor: A tensor of shape (m,) containing the indices of the picked objects.

  • Tuple[Any, torch.Tensor]: A tuple where the first element contains the picked objects and the second element is a tensor of their scores.

Return type:

out (Union[torch.Tensor, Tuple[Any, torch.Tensor]])

flat_bug.nms.detect_duplicate_boxes(boxes: Tensor, scores: Tensor, margin: int = 9, return_indices: bool = False) → Tensor | Tuple[Tensor, Tensor]

Duplicate detection algorithm based on the standard non-maximum suppression algorithm.

Algorithm overview:

  • Instead of IoU we use the maximum difference between the sides of the boxes as the metric for determining whether two boxes are duplicates.

  • To make this metric compatible with NMS we negate the metric and the threshold, such that large side difference are very negative and thus below the threshold.

flat_bug.nms.fancy_nms(objects: Any, iou_fun: Callable, scores: Tensor, iou_threshold: float | int = 0.5, return_indices: bool = False) → Tensor | Tuple[Any, Tensor]

This is a ‘fancy’ implementation of non-maximum suppression. It is not as fast as the non-maximum suppression algorithm, nor does it follow the exact same algorithm, but it is more readable and easier to debug.

The algorithm works as follows:

  1. Sort the objects by score (implicitly)

  2. Calculate the IoU matrix

  3. Create a boolean matrix where IoU > iou_threshold

  4. Fold the boolean matrix sequentially (i.e. row_i = row_i + row_i-1 + … + row_0) (The values on the diagonal of the matrix now correspond to the number of higher-priority objects that suppress the current object, including itself)

  5. objects which are suppressed only by themselves are returned.

Parameters:

  • objects (Any) – Any object collection that can be indexed by a tensor, where the first dimension corresponds to the objects.

  • iou_fun (Callable) – A function that calculates the symmetric IoU matrix of a set of objects returned as a torch.Tensor of shape (n, n), where n is the number of objects. The device should match the device of the scores.

  • scores (torch.Tensor) – A tensor of shape (n, ) containing the scores of the objects.

  • iou_threshold (Union[float, int], optional) – The IoU threshold for non-maximum suppression. Defaults to 0.5.

  • return_indices (bool, optional) – A flag to indicate whether to return the indices of the picked objects or the objects themselves. Defaults to False. If True, both the picked objects and scores are returned.

Returns:

  • torch.Tensor: A tensor of shape (m,) containing the indices of the picked objects.

  • Tuple[Any, torch.Tensor]: A tuple where the first element contains the picked objects and the second element is a tensor of their scores.

Return type:

out (Union[torch.Tensor, Tuple[Any, torch.Tensor]])

flat_bug.nms.iou_boxes(rectangles: Tensor, other_rectangles: Tensor | None = None) → Tensor

Calculates the intersection over union (IoU) of a set of rectangles.

Parameters:

  • rectangles (torch.Tensor) – A tensor of shape (n, 4), where n is the number of rectangles and the 4 columns are the x_min, y_min, x_max and y_max coordinates of the rectangles.

  • other_rectangles (Optional[torch.Tensor], optional) – A tensor of shape (m, 4), where m is the number of rectangles and the 4 columns are the x_min, y_min, x_max and y_max coordinates of the rectangles. Defaults to None, in which case the symmetric IoU of the rectangles with themselves is calculated.

Returns:

A tensor of shape (n, n), where n is the number of rectangles, containing the IoU of each rectangle with each other rectangle.

Return type:

out (torch.Tensor)

flat_bug.nms.iou_masks(masks: Tensor, areas: Tensor | None = None, dtype: dtype = torch.float32) → Tensor

Low-memory wrapper for flat-bug.nms.iou_masks_2sets that calculates the IoU of a set of masks with itself, in the symmetric case.

Parameters:

  • masks (torch.Tensor) – A tensor of shape (n, h, w), where n is the number of masks and h and w are the height and width of the masks.

  • areas (Optional[torch.Tensor], optional) – A tensor of shape (n, ) containing the areas of the masks. Defaults to None, in which case the areas are calculated.

  • dtype (torch.dtype, optional) – The data type of the output tensor. Defaults to torch.float32.

Returns:

A tensor of shape (n, n) containing the IoU of each pair of masks.

Return type:

out (torch.Tensor)

flat_bug.nms.iou_polygons(polygons1: List[Tensor], polygons2: List[Tensor] | None = None, dtype: dtype = torch.float32) → Tensor

Calculates the intersection over union (IoU) of a set of polygons.

The IoU is calculated using:

IoU[i,j] = intersection[i, j] / (area1[i] + area2[j] - intersection[i, j])

and then intersections and areas are calculated with the Shapely library.

OBS: Invalid polygons are handled by using the buffer(0) method from Shapely, which ensures that the function does not crash, but the results are not guaranteed to be “correct” for invalid polygons.

Parameters:

  • polygons1 (List[torch.Tensor]) – A list of tensors of shape (n, 2), where n is the number of vertices in the polygon and the 2 columns are the x and y coordinates of the vertices.

  • polygons2 (Optional[List[torch.Tensor]], optional) – A list of tensors of shape (m, 2), where m is the number of vertices in the polygon and the 2 columns are the x and y coordinates of the vertices. Defaults to None, in which case the symmetric IoU of the polygons with themselves is calculated.

Returns:

A tensor of shape (n, m), where n is the number of polygons in polygons1 and m is the number of polygons in polygons2, containing the IoU of each polygon in polygons1 with each polygon in polygons2.

Return type:

out (torch.Tensor)

flat_bug.nms.nms_boxes(boxes: Tensor, scores: Tensor, iou_threshold: float | int = 0.5) → Tensor

Wrapper for torchvision.ops.nms; the standard non-maximum suppression algorithm.

flat_bug.nms.nms_masks(masks: Tensor, scores: Tensor, iou_threshold: float = 0.5, return_indices: bool = False, group_first: bool = True, boxes: Tensor = None) → Tensor | Tuple[Tensor, Tensor]

Efficiently perform non-maximum suppression on a set of boolean masks.

Defaults to a modified two-stage NMS algorithm, that aims to minimize the number of mask intersection calculations needed.

Parameters:

  • masks (torch.Tensor) – A tensor of shape (n, h, w), where n is the number of masks and h and w are the height and width of the masks.

  • scores (torch.Tensor) – A tensor of shape (n, ) containing the “scores” of the masks, this can merely be though of as a priority score, where the higher the score, the higher the priority of the object - it does not have to be a probability/confidence.

  • iou_threshold (float, optional) – The IoU threshold for non-maximum suppression. Defaults to 0.5.

  • return_indices (bool, optional) – A flag to indicate whether to return the indices of the picked objects or the objects themselves. Defaults to False. If True, both the picked objects and scores are returned.

  • group_first (bool, optional) – A flag to indicate whether two use the two-stage NMS method. Defaults to True.

  • boxes (Optional[torch.Tensor], optional) – Bounding boxes for the masks. A tensor of shape (n, 4), where n is the number of masks and the 4 columns are the x_min, y_min, x_max and y_max coordinates of the bounding boxes.

Returns:

  • torch.Tensor: A tensor of shape (m,) containing the indices of the picked objects.

  • Tuple[torch.Tensor, torch.Tensor]: A tuple where the first element contains the picked masks and the second element is a tensor of their scores.

Return type:

out (Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]])

flat_bug.nms.nms_polygons(polygons: List[Tensor], scores: Tensor, iou_threshold: float | int = 0.5, return_indices: bool = False, group_first: bool = True, boxes: Tensor | None = None) → Tensor | Tuple[List[Tensor], Tensor]

Efficiently perform non-maximum suppression on a set of polygons.

Defaults to a modified two-stage NMS algorithm, that aims to minimize the number of polygon intersection calculations needed (very expensive).

Parameters:

  • polygons (List[torch.Tensor]) – A list of tensors of shape (n, 2), where n is the number of vertices in the polygon and the 2 columns are the x and y coordinates of the vertices.

  • scores (torch.Tensor) – A tensor of shape (n, ) containing the “scores” of the polygons, this can merely be though of as a priority score, where the higher the score, the higher the priority of the object - it does not have to be a probability/confidence.

  • iou_threshold (float, optional) – The IoU threshold for non-maximum suppression. Defaults to 0.5.

  • return_indices (bool, optional) – A flag to indicate whether to return the indices of the picked objects or the objects themselves. Defaults to False. If True, both the picked objects and scores are returned.

  • group_first (bool, optional) – A flag to indicate whether two use the two-stage NMS method. Defaults to True (recommended).

  • boxes (Optional[torch.Tensor], optional) – Bounding boxes for the polygons. A tensor of shape (n, 4), where n is the number of polygons and the 4 columns are the x_min, y_min, x_max and y_max coordinates of the bounding boxes.

Returns:

  • torch.Tensor: A tensor of shape (m,) containing the indices of the picked polygons.

  • Tuple[List[torch.Tensor], torch.Tensor]: A tuple where the first element contains the picked polygons and the second element is a tensor of their scores.

Return type:

out (Union[torch.Tensor, Tuple[List[torch.Tensor], torch.Tensor]])

flat_bug.nms.nms_polygons_(polys: List[Tensor], scores: Tensor, iou_threshold: float = 0.5) → Tensor

Performs non-maximum suppression on a set of polygons.

Parameters:

  • polys (List[torch.Tensor]) – A list of tensors of shape (n, 2), where n is the number of vertices in the polygon and the 2 columns are the x and y coordinates of the vertices.

  • scores (torch.Tensor) – A tensor of shape (n, ) containing the scores of the polygons.

  • iou_threshold (float, optional) – The IoU threshold for non-maximum suppression. Defaults to 0.5.

Returns:

A tensor containing the indices of the picked polygons.

Return type:

out (torch.Tensor)