Spaces:
Running on Zero
Running on Zero
| """Pad transformation.""" | |
| from __future__ import annotations | |
| import numpy as np | |
| import torch | |
| import torch.nn.functional as F | |
| from mapdet3d.common.typing import NDArrayF32, NDArrayUI8 | |
| from mapdet3d.data.const import CommonKeys as K | |
| from .base import Transform | |
| class PadImages: | |
| """Pad batch of images at the bottom right.""" | |
| def __init__( | |
| self, | |
| stride: int = 32, | |
| mode: str = "constant", | |
| value: float = 0.0, | |
| update_input_hw: bool = False, | |
| shape: tuple[int, int] | None = None, | |
| pad2square: bool = False, | |
| ) -> None: | |
| """Creates an instance of PadImage. | |
| Args: | |
| stride (int, optional): Chooses padding size so that the input will | |
| be divisible by stride. Defaults to 32. | |
| mode (str, optional): Padding mode. One of constant, reflect, | |
| replicate or circular. Defaults to "constant". | |
| value (float, optional): Value for constant padding. | |
| Defaults to 0.0. | |
| shape (tuple[int, int], optional): Shape of the padded image | |
| (H, W). Defaults to None. | |
| pad2square (bool, optional): Pad to square. Defaults to False. | |
| """ | |
| if pad2square: | |
| assert ( | |
| shape is None | |
| ), "Cannot specify shape when pad2square is True." | |
| self.stride = stride | |
| self.mode = mode | |
| self.value = value | |
| self.update_input_hw = update_input_hw | |
| self.shape = shape | |
| self.pad2square = pad2square | |
| def __call__( | |
| self, images: list[NDArrayF32], input_hw: list[tuple[int, int]] | |
| ) -> tuple[list[NDArrayF32], list[tuple[int, int]]]: | |
| """Pad images to consistent size.""" | |
| heights = [im.shape[1] for im in images] | |
| widths = [im.shape[2] for im in images] | |
| max_hw = _get_max_shape( | |
| heights, widths, self.stride, self.shape, self.pad2square | |
| ) | |
| # generate params for torch pad | |
| for i, (image, h, w) in enumerate(zip(images, heights, widths)): | |
| pad_param = (0, max_hw[1] - w, 0, max_hw[0] - h) | |
| image_ = torch.from_numpy(image).permute(0, 3, 1, 2) | |
| image_ = F.pad( # pylint: disable=not-callable | |
| image_, pad_param, self.mode, self.value | |
| ) | |
| images[i] = image_.permute(0, 2, 3, 1).numpy() | |
| if self.update_input_hw: | |
| input_hw[i] = image_.shape[2], image_.shape[3] | |
| return images, input_hw | |
| class PadDepthMaps: | |
| """Pad batch of depth maps at the bottom right.""" | |
| def __init__( | |
| self, | |
| stride: int = 32, | |
| mode: str = "constant", | |
| value: int = 0, | |
| shape: tuple[int, int] | None = None, | |
| pad2square: bool = False, | |
| ) -> None: | |
| """Creates an instance. | |
| Args: | |
| stride (int, optional): Chooses padding size so that the input will | |
| be divisible by stride. Defaults to 32. | |
| mode (str, optional): Padding mode. One of constant, reflect, | |
| replicate or circular. Defaults to "constant". | |
| value (float, optional): Value for constant padding. | |
| Defaults to 0.0. | |
| shape (tuple[int, int], optional): Shape of the padded image | |
| (H, W). Defaults to None. | |
| pad2square (bool, optional): Pad to square. Defaults to False. | |
| """ | |
| if pad2square: | |
| assert ( | |
| shape is None | |
| ), "Cannot specify shape when pad2square is True." | |
| self.stride = stride | |
| self.mode = mode | |
| self.value = value | |
| self.shape = shape | |
| self.pad2square = pad2square | |
| def __call__(self, depth_maps: list[NDArrayF32]) -> list[NDArrayF32]: | |
| """Pad images to consistent size.""" | |
| heights = [depth.shape[0] for depth in depth_maps] | |
| widths = [depth.shape[1] for depth in depth_maps] | |
| max_hw = _get_max_shape( | |
| heights, widths, self.stride, self.shape, self.pad2square | |
| ) | |
| # generate params for torch pad | |
| for i, (depth, h, w) in enumerate(zip(depth_maps, heights, widths)): | |
| pad_param = ((0, max_hw[0] - h), (0, max_hw[1] - w)) | |
| depth_maps[i] = np.pad( # type: ignore | |
| depth, pad_param, mode=self.mode, constant_values=self.value | |
| ) | |
| return depth_maps | |
| class PadSegMasks: | |
| """Pad batch of segmentation masks at the bottom right.""" | |
| def __init__( | |
| self, | |
| stride: int = 32, | |
| mode: str = "constant", | |
| value: int = 255, | |
| shape: tuple[int, int] | None = None, | |
| pad2square: bool = False, | |
| ) -> None: | |
| """Creates an instance of PadSegMasks. | |
| Args: | |
| stride (int, optional): Chooses padding size so that the input will | |
| be divisible by stride. Defaults to 32. | |
| mode (str, optional): Padding mode. One of constant, reflect, | |
| replicate or circular. Defaults to "constant". | |
| value (float, optional): Value for constant padding. | |
| Defaults to 0.0. | |
| shape (tuple[int, int], optional): Shape of the padded image | |
| (H, W). Defaults to None. | |
| pad2square (bool, optional): Pad to square. Defaults to False. | |
| """ | |
| if pad2square: | |
| assert ( | |
| shape is None | |
| ), "Cannot specify shape when pad2square is True." | |
| self.stride = stride | |
| self.mode = mode | |
| self.value = value | |
| self.shape = shape | |
| self.pad2square = pad2square | |
| def __call__(self, masks: list[NDArrayUI8]) -> list[NDArrayUI8]: | |
| """Pad images to consistent size.""" | |
| heights = [mask.shape[0] for mask in masks] | |
| widths = [mask.shape[1] for mask in masks] | |
| max_hw = _get_max_shape( | |
| heights, widths, self.stride, self.shape, self.pad2square | |
| ) | |
| # generate params for torch pad | |
| for i, (mask, h, w) in enumerate(zip(masks, heights, widths)): | |
| pad_param = ((0, max_hw[0] - h), (0, max_hw[1] - w)) | |
| masks[i] = np.pad( # type: ignore | |
| mask, pad_param, mode=self.mode, constant_values=self.value | |
| ) | |
| return masks | |
| def _get_max_shape( | |
| heights: list[int], | |
| widths: list[int], | |
| stride: int, | |
| shape: tuple[int, int] | None, | |
| pad2square: bool, | |
| ) -> tuple[int, int]: | |
| """Get max shape for padding. | |
| Args: | |
| stride (int): Chooses padding size so that the input will be divisible | |
| by stride. | |
| shape (tuple[int, int], optional): Shape of the padded image (H, W). | |
| Defaults to None. | |
| heights (list[int]): List of heights of input. | |
| widths (list[int]): List of widths of input. | |
| pad2square (bool): Pad to square. | |
| Returns: | |
| tuple[int, int]: Max shape for padding. | |
| """ | |
| if pad2square: | |
| max_size = max(heights + widths) | |
| max_hw = (max_size, max_size) | |
| elif shape is not None: | |
| max_hw = shape | |
| else: | |
| max_hw = max(heights), max(widths) | |
| return tuple(_make_divisible(x, stride) for x in max_hw) | |
| def _make_divisible(x: int, stride: int) -> int: | |
| """Ensure divisibility by stride.""" | |
| return (x + (stride - 1)) // stride * stride | |