Spaces:
Running on Zero
Running on Zero
File size: 7,474 Bytes
0122a25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 | """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
@Transform(
[K.images, K.input_hw],
[K.images, K.input_hw],
)
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
@Transform(K.depth_maps, K.depth_maps)
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
@Transform(K.seg_masks, K.seg_masks)
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
|