diff --git a/docs/source/ops.rst b/docs/source/ops.rst index cdebe9721c3..34737d69026 100644 --- a/docs/source/ops.rst +++ b/docs/source/ops.rst @@ -13,6 +13,7 @@ torchvision.ops .. autofunction:: batched_nms .. autofunction:: remove_small_boxes .. autofunction:: clip_boxes_to_image +.. autofunction:: box_rescale .. autofunction:: box_convert .. autofunction:: box_area .. autofunction:: box_iou diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py index c1f176f4da9..b46f6e3c9c7 100644 --- a/torchvision/ops/boxes.py +++ b/torchvision/ops/boxes.py @@ -154,6 +154,24 @@ def clip_boxes_to_image(boxes: Tensor, size: Tuple[int, int]) -> Tensor: return clipped_boxes.reshape(boxes.shape) +def box_rescale(boxes: Tensor, in_shape: Tuple[int, int], out_shape: Tuple[int, int]) -> Tensor: + """ + Rescale boxes to a new image shape. + + Args: + boxes: (Tensor[N, 4]): boxes to rescale, in ``(x1, y1, x2, y2)`` format. + in_shape: (Tuple[int, int]): current shape of the image. + out_shape: (Tuple[int, int]): desired shape of the image. + + Returns: + Tensor[N, 4]: rescaled boxes + """ + rescaled = boxes.new(boxes.shape) + rescaled[:, 0::2] = boxes[:, 0::2] * (out_shape[1] / in_shape[1]) + rescaled[:, 1::2] = boxes[:, 1::2] * (out_shape[0] / in_shape[0]) + return rescaled + + def box_convert(boxes: Tensor, in_fmt: str, out_fmt: str) -> Tensor: """ Converts boxes from given in_fmt to out_fmt.