utils.py 742 B

123456789101112131415161718192021222324
  1. # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
  2. def adjust_bboxes_to_image_border(boxes, image_shape, threshold=20):
  3. """
  4. Adjust bounding boxes to stick to image border if they are within a certain threshold.
  5. Args:
  6. boxes (torch.Tensor): (n, 4)
  7. image_shape (tuple): (height, width)
  8. threshold (int): pixel threshold
  9. Returns:
  10. adjusted_boxes (torch.Tensor): adjusted bounding boxes
  11. """
  12. # Image dimensions
  13. h, w = image_shape
  14. # Adjust boxes
  15. boxes[boxes[:, 0] < threshold, 0] = 0 # x1
  16. boxes[boxes[:, 1] < threshold, 1] = 0 # y1
  17. boxes[boxes[:, 2] > w - threshold, 2] = w # x2
  18. boxes[boxes[:, 3] > h - threshold, 3] = h # y2
  19. return boxes