diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 6071b083..3200dc9c 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1154,7 +1154,7 @@ class LabelAnnotator(BaseAnnotator): ) if self.smart_positions: - xyxy = spread_out_boxes(xyxy, step=2, max_iterations=len(xyxy) * 20) + xyxy = spread_out_boxes(xyxy) self._draw_labels( scene=scene, @@ -1449,7 +1449,7 @@ class RichLabelAnnotator(BaseAnnotator): ) if self.smart_positions: - xyxy = spread_out_boxes(xyxy, step=2, max_iterations=len(xyxy) * 20) + xyxy = spread_out_boxes(xyxy) self._draw_labels( draw=draw, diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py index ede1ed35..246d3c60 100644 --- a/supervision/detection/utils.py +++ b/supervision/detection/utils.py @@ -1062,37 +1062,52 @@ def get_unit_vector(xy_1: np.ndarray, xy_2: np.ndarray) -> np.ndarray: def spread_out_boxes( - xyxy: np.ndarray, step: int, max_iterations: int = 100 + xyxy: np.ndarray, + max_iterations: int = 100, + force_multiplier: float = 0.03, ) -> np.ndarray: + """ + Spread out boxes that overlap with each other. + + Args: + xyxy: Numpy array of shape (N, 4) where N is the number of boxes. + max_iterations: Maximum number of iterations to run the algorithm for. + force_multiplier: Multiplier to scale the force vectors by. Similar to + learning rate in gradient descent. + """ if len(xyxy) == 0: return xyxy - xyxy_padded = pad_boxes(xyxy, px=step) + xyxy_padded = pad_boxes(xyxy, px=1) for _ in range(max_iterations): + # NxN iou = box_iou_batch(xyxy_padded, xyxy_padded) np.fill_diagonal(iou, 0) - if np.all(iou == 0): break - i, j = np.unravel_index(np.argmax(iou), iou.shape) + overlap_mask = iou > 0 - xyxy_i, xyxy_j = xyxy_padded[i], xyxy_padded[j] - box_intersection = get_box_intersection(xyxy_i, xyxy_j) - assert ( - box_intersection is not None - ), "Since we checked IoU already, boxes should always intersect" + # Nx2 + centers = (xyxy_padded[:, :2] + xyxy_padded[:, 2:]) / 2 - intersection_center = (box_intersection[:2] + box_intersection[2:]) / 2 - xyxy_i_center = (xyxy_i[:2] + xyxy_i[2:]) / 2 - xyxy_j_center = (xyxy_j[:2] + xyxy_j[2:]) / 2 + # NxNx2 + delta_centers = centers[:, np.newaxis, :] - centers[np.newaxis, :, :] + delta_centers *= overlap_mask[:, :, np.newaxis] - unit_vector_i = get_unit_vector(intersection_center, xyxy_i_center) - unit_vector_j = get_unit_vector(intersection_center, xyxy_j_center) + # Nx2 + force_vectors = np.sum(delta_centers, axis=1) + force_vectors *= force_multiplier + force_vectors[(force_vectors > 0) & (force_vectors < 1)] = 1 + force_vectors[(force_vectors < 0) & (force_vectors > -1)] = -1 - xyxy_padded[i, [0, 2]] += int(unit_vector_i[0] * step) - xyxy_padded[i, [1, 3]] += int(unit_vector_i[1] * step) - xyxy_padded[j, [0, 2]] += int(unit_vector_j[0] * step) - xyxy_padded[j, [1, 3]] += int(unit_vector_j[1] * step) + # Reduce motion along primary axis + primary_axis = np.argmax(np.abs(force_vectors), axis=1) + force_vectors[np.arange(len(force_vectors)), primary_axis] /= 2 - return pad_boxes(xyxy_padded, px=-step) + force_vectors = force_vectors.astype(int) + + xyxy_padded[:, [0, 1]] += force_vectors + xyxy_padded[:, [2, 3]] += force_vectors + + return pad_boxes(xyxy_padded, px=-1) diff --git a/supervision/keypoint/annotators.py b/supervision/keypoint/annotators.py index d968786c..fc450a83 100644 --- a/supervision/keypoint/annotators.py +++ b/supervision/keypoint/annotators.py @@ -363,9 +363,7 @@ class VertexLabelAnnotator: xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding) if self.smart_positions: - xyxy_padded = spread_out_boxes( - xyxy_padded, step=2, max_iterations=len(xyxy_padded) * 20 - ) + xyxy_padded = spread_out_boxes(xyxy_padded) xyxy = pad_boxes(xyxy=xyxy_padded, px=-self.text_padding) for text, color, text_color, box, box_padded in zip(