123456789101112131415161718192021222324 |
- # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/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
|