def __call__(
self,
img: np.ndarray,
mode: Optional[Union[NumpyPadMode, str]] = None) -> np.ndarray:
"""
Args:
img: data to be transformed, assuming `img` is channel-first
and padding doesn't apply to the channel dim.
mode: {``"constant"``, ``"edge"``, ``"linear_ramp"``, ``"maximum"``, ``"mean"``,
``"median"``, ``"minimum"``, ``"reflect"``, ``"symmetric"``, ``"wrap"``, ``"empty"``}
One of the listed string values or a user supplied function. Defaults to ``self.mode``.
See also: https://numpy.org/doc/1.18/reference/generated/numpy.pad.html
"""
new_size = compute_divisible_spatial_size(spatial_shape=img.shape[1:],
k=self.k)
spatial_pad = SpatialPad(spatial_size=new_size,
method=Method.SYMMETRIC,
mode=mode or self.mode)
return spatial_pad(img)
def compute_bounding_box(self, img: np.ndarray):
"""
Compute the start points and end points of bounding box to crop.
And adjust bounding box coords to be divisible by `k`.
"""
box_start, box_end = generate_spatial_bounding_box(
img, self.select_fn, self.channel_indices, self.margin)
box_start_ = np.asarray(box_start, dtype=np.int16)
box_end_ = np.asarray(box_end, dtype=np.int16)
orig_spatial_size = box_end_ - box_start_
# make the spatial size divisible by `k`
spatial_size = np.asarray(
compute_divisible_spatial_size(spatial_shape=orig_spatial_size,
k=self.k_divisible))
# update box_start and box_end
box_start_ = box_start_ - np.floor_divide(
np.asarray(spatial_size) - orig_spatial_size, 2)
box_end_ = box_start_ + spatial_size
return box_start_, box_end_
def pad_images(
input_images: Union[List[Tensor], Tensor],
spatial_dims: int,
size_divisible: Union[int, Sequence[int]],
mode: Union[PytorchPadMode, str] = PytorchPadMode.CONSTANT,
**kwargs,
) -> Tuple[Tensor, List[List[int]]]:
"""
Pad the input images, so that the output spatial sizes are divisible by `size_divisible`.
It pads them at the end to create a (B, C, H, W) or (B, C, H, W, D) Tensor.
Padded size (H, W) or (H, W, D) is divisible by size_divisible.
Default padding uses constant padding with value 0.0
Args:
input_images: It can be 1) a tensor sized (B, C, H, W) or (B, C, H, W, D),
or 2) a list of image tensors, each image i may have different size (C, H_i, W_i) or (C, H_i, W_i, D_i).
spatial_dims: number of spatial dimensions of the images, 2D or 3D.
size_divisible: int or Sequence[int], is the expected pattern on the input image shape.
If an int, the same `size_divisible` will be applied to all the input spatial dimensions.
mode: available modes for PyTorch Tensor: {``"constant"``, ``"reflect"``, ``"replicate"``, ``"circular"``}.
One of the listed string values or a user supplied function. Defaults to ``"constant"``.
See also: https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html
kwargs: other arguments for `torch.pad` function.
Return:
- images, a (B, C, H, W) or (B, C, H, W, D) Tensor
- image_sizes, the original spatial size of each image
"""
size_divisible = ensure_tuple_rep(size_divisible, spatial_dims)
# If input_images: Tensor
if isinstance(input_images, Tensor):
orig_size = list(input_images.shape[-spatial_dims:])
new_size = compute_divisible_spatial_size(spatial_shape=orig_size,
k=size_divisible)
all_pad_width = [(0, max(sp_i - orig_size[i], 0))
for i, sp_i in enumerate(new_size)]
pt_pad_width = [
val for sublist in all_pad_width for val in sublist[::-1]
][::-1]
if max(pt_pad_width) == 0:
# if there is no need to pad
return input_images, [orig_size] * input_images.shape[0]
mode_: str = convert_pad_mode(dst=input_images, mode=mode)
return F.pad(input_images, pt_pad_width, mode=mode_,
**kwargs), [orig_size] * input_images.shape[0]
# If input_images: List[Tensor])
image_sizes = [img.shape[-spatial_dims:] for img in input_images]
in_channels = input_images[0].shape[0]
dtype = input_images[0].dtype
device = input_images[0].device
# compute max_spatial_size
image_sizes_t = torch.tensor(image_sizes)
max_spatial_size_t, _ = torch.max(image_sizes_t, dim=0)
if len(max_spatial_size_t) != spatial_dims or len(
size_divisible) != spatial_dims:
raise ValueError(
" Require len(max_spatial_size_t) == spatial_dims ==len(size_divisible)."
)
max_spatial_size = compute_divisible_spatial_size(
spatial_shape=list(max_spatial_size_t), k=size_divisible)
# allocate memory for the padded images
images = torch.zeros([len(image_sizes), in_channels] + max_spatial_size,
dtype=dtype,
device=device)
# Use `SpatialPad` to match sizes, padding in the end will not affect boxes
padder = SpatialPad(spatial_size=max_spatial_size,
method="end",
mode=mode,
**kwargs)
for idx, img in enumerate(input_images):
images[idx, ...] = padder(img)
return images, [list(ss) for ss in image_sizes]