diff --git a/docs/source/apps.rst b/docs/source/apps.rst index 1a11fc62c6..4b62abd2e3 100644 --- a/docs/source/apps.rst +++ b/docs/source/apps.rst @@ -134,3 +134,8 @@ Applications :members: .. automodule:: monai.apps.detection.transforms.dictionary :members: + +`Box coder` +~~~~~~~~~~~ +.. automodule:: monai.apps.detection.utils.box_coder + :members: diff --git a/monai/apps/detection/utils/__init__.py b/monai/apps/detection/utils/__init__.py new file mode 100644 index 0000000000..1e97f89407 --- /dev/null +++ b/monai/apps/detection/utils/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/monai/apps/detection/utils/box_coder.py b/monai/apps/detection/utils/box_coder.py new file mode 100644 index 0000000000..e8b430af67 --- /dev/null +++ b/monai/apps/detection/utils/box_coder.py @@ -0,0 +1,246 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# ========================================================================= +# Adapted from https://github.com/pytorch/vision/blob/main/torchvision/models/detection/_utils.py +# which has the following license... +# https://github.com/pytorch/vision/blob/main/LICENSE +# +# BSD 3-Clause License + +# Copyright (c) Soumith Chintala 2016, +# All rights reserved. + +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: + +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. + +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. + +# * Neither the name of the copyright holder nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. + +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +""" +This script is modified from torchvision to support N-D images, + +https://github.com/pytorch/vision/blob/main/torchvision/models/detection/_utils.py +""" + +import math +from typing import Sequence, Tuple, Union + +import torch +from torch import Tensor + +from monai.data.box_utils import ( + COMPUTE_DTYPE, + CenterSizeMode, + CornerCornerModeTypeB, + StandardMode, + convert_box_mode, + convert_box_to_standard_mode, + is_valid_box_values, +) +from monai.utils.module import look_up_option + + +def encode_boxes(gt_boxes: Tensor, proposals: Tensor, weights: Tensor) -> Tensor: + """ + Encode a set of proposals with respect to some reference ground truth (gt) boxes. + + Args: + gt_boxes: gt boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + proposals: boxes to be encoded, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + weights: the weights for ``(cx, cy, w, h) or (cx,cy,cz, w,h,d)`` + + Return: + encoded gt, target of box regression that is used to convert proposals into gt_boxes, Nx4 or Nx6 torch tensor. + """ + + if gt_boxes.shape[0] != proposals.shape[0]: + raise ValueError("gt_boxes.shape[0] should be equal to proposals.shape[0].") + spatial_dims = look_up_option(len(weights), [4, 6]) // 2 + + if not is_valid_box_values(gt_boxes): + raise ValueError("gt_boxes is not valid. Please check if it contains empty boxes.") + if not is_valid_box_values(proposals): + raise ValueError("proposals is not valid. Please check if it contains empty boxes.") + + # implementation starts here + ex_cccwhd: Tensor = convert_box_mode(proposals, src_mode=StandardMode, dst_mode=CenterSizeMode) # type: ignore + gt_cccwhd: Tensor = convert_box_mode(gt_boxes, src_mode=StandardMode, dst_mode=CenterSizeMode) # type: ignore + targets_dxyz = ( + weights[:spatial_dims].unsqueeze(0) + * (gt_cccwhd[:, :spatial_dims] - ex_cccwhd[:, :spatial_dims]) + / ex_cccwhd[:, spatial_dims:] + ) + targets_dwhd = weights[spatial_dims:].unsqueeze(0) * torch.log( + gt_cccwhd[:, spatial_dims:] / ex_cccwhd[:, spatial_dims:] + ) + + targets = torch.cat((targets_dxyz, targets_dwhd), dim=1) + # torch.log may cause NaN or Inf + if torch.isnan(targets).any() or torch.isinf(targets).any(): + raise ValueError("targets is NaN or Inf.") + return targets + + +class BoxCoder: + """ + This class encodes and decodes a set of bounding boxes into + the representation used for training the regressors. + + Args: + weights: 4-element tuple or 6-element tuple + boxes_xform_clip: high threshold to prevent sending too large values into torch.exp() + + Example: + .. code-block:: python + + box_coder = BoxCoder(weights=[1., 1., 1., 1., 1., 1.]) + gt_boxes = torch.tensor([[1,2,1,4,5,6],[1,3,2,7,8,9]]) + proposals = gt_boxes + torch.rand(gt_boxes.shape) + rel_gt_boxes = box_coder.encode_single(gt_boxes, proposals) + gt_back = box_coder.decode_single(rel_gt_boxes, proposals) + # We expect gt_back to be equal to gt_boxes + """ + + def __init__(self, weights: Tuple[float], boxes_xform_clip: Union[float, None] = None) -> None: + if boxes_xform_clip is None: + boxes_xform_clip = math.log(1000.0 / 16) + self.spatial_dims = look_up_option(len(weights), [4, 6]) // 2 + self.weights = weights + self.boxes_xform_clip = boxes_xform_clip + + def encode(self, gt_boxes: Sequence[Tensor], proposals: Sequence[Tensor]) -> Tuple[Tensor]: + """ + Encode a set of proposals with respect to some ground truth (gt) boxes. + + Args: + gt_boxes: list of gt boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + proposals: list of boxes to be encoded, each element is Mx4 or Mx6 torch tensor. + The box mode is assumed to be ``StandardMode`` + + Return: + A tuple of encoded gt, target of box regression that is used to + convert proposals into gt_boxes, Nx4 or Nx6 torch tensor. + """ + boxes_per_image = [len(b) for b in gt_boxes] + # concat the lists to do computation + concat_gt_boxes = torch.cat(tuple(gt_boxes), dim=0) + concat_proposals = torch.cat(tuple(proposals), dim=0) + concat_targets = self.encode_single(concat_gt_boxes, concat_proposals) + # split to tuple + targets: Tuple[Tensor] = concat_targets.split(boxes_per_image, 0) + return targets + + def encode_single(self, gt_boxes: Tensor, proposals: Tensor) -> Tensor: + """ + Encode proposals with respect to ground truth (gt) boxes. + + Args: + gt_boxes: gt boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + proposals: boxes to be encoded, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + + Return: + encoded gt, target of box regression that is used to convert proposals into gt_boxes, Nx4 or Nx6 torch tensor. + """ + dtype = gt_boxes.dtype + device = gt_boxes.device + weights = torch.as_tensor(self.weights, dtype=dtype, device=device) + targets = encode_boxes(gt_boxes, proposals, weights) + return targets + + def decode(self, rel_codes: Tensor, reference_boxes: Sequence[Tensor]) -> Tensor: + """ + From a set of original reference_boxes and encoded relative box offsets, + + Args: + rel_codes: encoded boxes, Nx4 or Nx6 torch tensor. + boxes: a list of reference boxes, each element is Mx4 or Mx6 torch tensor. + The box mode is assumed to be ``StandardMode`` + + Return: + decoded boxes, Nx1x4 or Nx1x6 torch tensor. The box mode will be ``StandardMode`` + """ + if not isinstance(reference_boxes, Sequence) or (not isinstance(rel_codes, torch.Tensor)): + raise ValueError("Input arguments wrong type.") + boxes_per_image = [b.size(0) for b in reference_boxes] + # concat the lists to do computation + concat_boxes = torch.cat(tuple(reference_boxes), dim=0) + box_sum = 0 + for val in boxes_per_image: + box_sum += val + if box_sum > 0: + rel_codes = rel_codes.reshape(box_sum, -1) + pred_boxes = self.decode_single(rel_codes, concat_boxes) + if box_sum > 0: + pred_boxes = pred_boxes.reshape(box_sum, -1, 2 * self.spatial_dims) + return pred_boxes + + def decode_single(self, rel_codes: Tensor, reference_boxes: Tensor) -> Tensor: + """ + From a set of original boxes and encoded relative box offsets, + + Args: + rel_codes: encoded boxes, Nx4 or Nx6 torch tensor. + reference_boxes: reference boxes, Nx4 or Nx6 torch tensor. The box mode is assumed to be ``StandardMode`` + + Return: + decoded boxes, Nx4 or Nx6 torch tensor. The box mode will to be ``StandardMode`` + """ + reference_boxes = reference_boxes.to(rel_codes.dtype) + + pred_boxes = [] + boxes_cccwhd = convert_box_mode(reference_boxes, src_mode=StandardMode, dst_mode=CenterSizeMode) + for axis in range(self.spatial_dims): + whd_axis = boxes_cccwhd[:, axis + self.spatial_dims] + ctr_xyz_axis = boxes_cccwhd[:, axis] + dxyz_axis = rel_codes[:, axis] / self.weights[axis] + dwhd_axis = rel_codes[:, self.spatial_dims + axis] / self.weights[axis + self.spatial_dims] + + # Prevent sending too large values into torch.exp() + dwhd_axis = torch.clamp(dwhd_axis.to(COMPUTE_DTYPE), max=self.boxes_xform_clip) + + pred_ctr_xyx_axis = dxyz_axis * whd_axis + ctr_xyz_axis + pred_whd_axis = torch.exp(dwhd_axis) * whd_axis + pred_whd_axis = pred_whd_axis.to(dxyz_axis.dtype) + + # When convert float32 to float16, Inf or Nan may occur + if torch.isnan(pred_whd_axis).any() or torch.isinf(pred_whd_axis).any(): + raise ValueError("pred_whd_axis is NaN or Inf.") + + # Distance from center to box's corner. + c_to_c_whd_axis = ( + torch.tensor(0.5, dtype=pred_ctr_xyx_axis.dtype, device=pred_whd_axis.device) * pred_whd_axis + ) + + pred_boxes.append(pred_ctr_xyx_axis - c_to_c_whd_axis) + pred_boxes.append(pred_ctr_xyx_axis + c_to_c_whd_axis) + + pred_boxes_xxyyzz = torch.stack(pred_boxes, dim=1) + return convert_box_to_standard_mode(pred_boxes_xxyyzz, mode=CornerCornerModeTypeB) # type: ignore diff --git a/tests/test_box_coder.py b/tests/test_box_coder.py new file mode 100644 index 0000000000..86ca7a98c2 --- /dev/null +++ b/tests/test_box_coder.py @@ -0,0 +1,41 @@ +# Copyright (c) MONAI Consortium +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest + +import torch + +from monai.apps.detection.utils.box_coder import BoxCoder +from monai.transforms import CastToType +from tests.utils import assert_allclose + + +class TestBoxTransform(unittest.TestCase): + def test_value(self): + box_coder = BoxCoder(weights=[1, 1, 1, 1, 1, 1]) + test_dtype = [torch.float32, torch.float16] + for dtype in test_dtype: + gt_boxes_0 = torch.rand((10, 3)).abs() + gt_boxes_1 = gt_boxes_0 + torch.rand((10, 3)).abs() + 10 + gt_boxes = torch.cat((gt_boxes_0, gt_boxes_1), dim=1) + gt_boxes = CastToType(dtype=dtype)(gt_boxes) + + proposals_0 = (gt_boxes_0 + torch.rand(gt_boxes_0.shape)).abs() + proposals_1 = proposals_0 + torch.rand(gt_boxes_0.shape).abs() + 10 + proposals = torch.cat((proposals_0, proposals_1), dim=1) + + rel_gt_boxes = box_coder.encode_single(gt_boxes, proposals) + gt_back = box_coder.decode_single(rel_gt_boxes, proposals) + assert_allclose(gt_back, gt_boxes, type_test=True, device_test=True, atol=0.1) + + +if __name__ == "__main__": + unittest.main()