Non-maximum suppression

flat_bug.nms.base_nms_(objects: Any, overlap_fn: Callable, scores: Tensor, collate_fn: Callable = None, overlap_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.

  • overlap_fn (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 overlap_fn 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.

  • overlap_threshold (float, optional) – The overlap (e.g. 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 overlap_fn 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.cluster_overlap_boxes(boxes: ~torch.Tensor, overlap_threshold: float = 0.5, overlap_fn: ~typing.Callable[[~torch.Tensor], ~torch.Tensor] = <function iou_boxes>, time: bool = False) → Tuple[List[Tensor], Tensor]

Cluster boxes via connected components.

Note: Implementation relies on overlap_fn being symmetric (e.g. IoU/IoS).

flat_bug.nms.fancy_nms(objects: Any, overlap_fn: Callable, scores: Tensor, overlap_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 overlap (e.g. IoU) matrix

  3. Create a boolean matrix where overlap > overlap_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.

  • overlap_fn (Callable) – A function that calculates the symmetric overlap (e.g. 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.

  • overlap_threshold (Union[float, int], optional) – The overlap (e.g. 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.ios_boxes(rectangles: Tensor, other_rectangles: Tensor | None = None) → Tensor

Calculates the intersection over smaller (IoS) 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 IoS of the rectangles with themselves is calculated.

Returns:

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

Return type:

out (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.nms_boxes(boxes: Tensor, scores: Tensor, overlap_threshold: float | int = 0.5, overlap_fn: Callable[[Tensor], Tensor] | str | None = None) → Tensor

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

flat_bug.nms.nms_masks(masks: ~torch.Tensor, scores: ~torch.Tensor, overlap_threshold: float = 0.5, return_indices: bool = False, group_first: bool = True, boxes: ~torch.Tensor = None, overlap_fn: ~typing.Callable[[~torch.Tensor], ~torch.Tensor] | str = <torch.jit.torch.jit.ScriptFunction object>, overlap_fn_boxes: ~typing.Callable[[~torch.Tensor], ~torch.Tensor] | str | None = 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.

  • overlap_threshold (float, optional) – The overlap (e.g. 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_masks_(masks: ~torch.Tensor, scores: ~torch.Tensor, overlap_threshold: float = 0.5, overlap_fn: ~typing.Callable[[~torch.Tensor, ~torch.Tensor], ~torch.Tensor] = <torch.jit.torch.jit.ScriptFunction object>) → Tensor

Performs non-maximum suppression on a set of masks.

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.

  • overlap_threshold (float, optional) – The overlap (e.g. IoU) threshold for non-maximum suppression. Defaults to 0.5.

Returns:

A tensor containing the indices of the picked masks.

Return type:

out (torch.Tensor)

flat_bug.nms.nms_polygons(polygons: List[Tensor], scores: Tensor, overlap_threshold: float | int = 0.5, return_indices: bool = False, group_first: bool = True, boxes: Tensor | None = None, overlap_fn: Callable[[List[Tensor], List[Tensor]], Tensor] | str = 'IoU', overlap_fn_boxes: Callable[[Tensor], Tensor] | str | 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.

  • overlap_threshold (float, optional) – The overlap (e.g. 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]])