You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
67 lines
2.1 KiB
67 lines
2.1 KiB
# Ultralytics YOLO 🚀, AGPL-3.0 license |
|
|
|
import torch |
|
|
|
|
|
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 |
|
|
|
|
|
def bbox_iou(box1, boxes, iou_thres=0.9, image_shape=(640, 640), raw_output=False): |
|
""" |
|
Compute the Intersection-Over-Union of a bounding box with respect to an array of other bounding boxes. |
|
|
|
Args: |
|
box1 (torch.Tensor): (4, ) |
|
boxes (torch.Tensor): (n, 4) |
|
iou_thres (float): IoU threshold |
|
image_shape (tuple): (height, width) |
|
raw_output (bool): If True, return the raw IoU values instead of the indices |
|
|
|
Returns: |
|
high_iou_indices (torch.Tensor): Indices of boxes with IoU > thres |
|
""" |
|
boxes = adjust_bboxes_to_image_border(boxes, image_shape) |
|
# Obtain coordinates for intersections |
|
x1 = torch.max(box1[0], boxes[:, 0]) |
|
y1 = torch.max(box1[1], boxes[:, 1]) |
|
x2 = torch.min(box1[2], boxes[:, 2]) |
|
y2 = torch.min(box1[3], boxes[:, 3]) |
|
|
|
# Compute the area of intersection |
|
intersection = (x2 - x1).clamp(0) * (y2 - y1).clamp(0) |
|
|
|
# Compute the area of both individual boxes |
|
box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1]) |
|
box2_area = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) |
|
|
|
# Compute the area of union |
|
union = box1_area + box2_area - intersection |
|
|
|
# Compute the IoU |
|
iou = intersection / union # Should be shape (n, ) |
|
if raw_output: |
|
return 0 if iou.numel() == 0 else iou |
|
|
|
# return indices of boxes with IoU > thres |
|
return torch.nonzero(iou > iou_thres).flatten()
|
|
|