3637516412
Co-authored-by: UltralyticsAssistant <web@ultralytics.com> Co-authored-by: Laughing <61612323+Laughing-q@users.noreply.github.com> Co-authored-by: Laughing-q <1185102784@qq.com> Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
25 lines
715 B
Python
25 lines
715 B
Python
# Ultralytics YOLO 🚀, AGPL-3.0 license
|
|
|
|
|
|
def adjust_bboxes_to_image_border(boxes, image_shape, threshold=20):
|
|
"""
|
|
Adjust bounding boxes to stick to image border if they are within a certain threshold.
|
|
|
|
Args:
|
|
boxes (torch.Tensor): (n, 4)
|
|
image_shape (tuple): (height, width)
|
|
threshold (int): pixel threshold
|
|
|
|
Returns:
|
|
adjusted_boxes (torch.Tensor): adjusted bounding boxes
|
|
"""
|
|
|
|
# Image dimensions
|
|
h, w = image_shape
|
|
|
|
# Adjust boxes
|
|
boxes[boxes[:, 0] < threshold, 0] = 0 # x1
|
|
boxes[boxes[:, 1] < threshold, 1] = 0 # y1
|
|
boxes[boxes[:, 2] > w - threshold, 2] = w # x2
|
|
boxes[boxes[:, 3] > h - threshold, 3] = h # y2
|
|
return boxes
|