From 0ca8f22c2a299a4dbed6b79fa2904e46c557bda4 Mon Sep 17 00:00:00 2001 From: tflahaul Date: Sun, 6 Jun 2021 14:21:45 +0200 Subject: [PATCH 1/2] box_rescale function to rescale bounding boxes to a new image shape --- docs/source/ops.rst | 1 + torchvision/ops/boxes.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+) 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..ffc039ce9c3 100644 --- a/torchvision/ops/boxes.py +++ b/torchvision/ops/boxes.py @@ -154,6 +154,25 @@ 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]): actual shape of the image. + out_shape: (Tuple[int, int]): desired shape of the image. + + Returns: + Tensor[N, 4]: reshaped boxes + """ + + scaled_boxes = boxes.new(boxes.shape) + scaled_boxes[:, 0::2] = boxes[:, 0::2] * (out_shape[1] / in_shape[1]) + scaled_boxes[:, 1::2] = boxes[:, 1::2] * (out_shape[0] / in_shape[0]) + return scaled_boxes + + def box_convert(boxes: Tensor, in_fmt: str, out_fmt: str) -> Tensor: """ Converts boxes from given in_fmt to out_fmt. From 9284f37c6ae9968152d911ded78bdb57711af439 Mon Sep 17 00:00:00 2001 From: tflahaul Date: Mon, 7 Jun 2021 13:12:02 +0200 Subject: [PATCH 2/2] [skip ci] changed confusing doc & variable name --- torchvision/ops/boxes.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/torchvision/ops/boxes.py b/torchvision/ops/boxes.py index ffc039ce9c3..b46f6e3c9c7 100644 --- a/torchvision/ops/boxes.py +++ b/torchvision/ops/boxes.py @@ -160,17 +160,16 @@ def box_rescale(boxes: Tensor, in_shape: Tuple[int, int], out_shape: Tuple[int, Args: boxes: (Tensor[N, 4]): boxes to rescale, in ``(x1, y1, x2, y2)`` format. - in_shape: (Tuple[int, int]): actual shape of the image. + in_shape: (Tuple[int, int]): current shape of the image. out_shape: (Tuple[int, int]): desired shape of the image. Returns: - Tensor[N, 4]: reshaped boxes + Tensor[N, 4]: rescaled boxes """ - - scaled_boxes = boxes.new(boxes.shape) - scaled_boxes[:, 0::2] = boxes[:, 0::2] * (out_shape[1] / in_shape[1]) - scaled_boxes[:, 1::2] = boxes[:, 1::2] * (out_shape[0] / in_shape[0]) - return scaled_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: