|
|
|
|
|
import torch |
|
|
|
|
|
class ImageResizeTransform: |
|
""" |
|
Transform that resizes images loaded from a dataset |
|
(BGR data in NCHW channel order, typically uint8) to a format ready to be |
|
consumed by DensePose training (BGR float32 data in NCHW channel order) |
|
""" |
|
|
|
def __init__(self, min_size: int = 800, max_size: int = 1333): |
|
self.min_size = min_size |
|
self.max_size = max_size |
|
|
|
def __call__(self, images: torch.Tensor) -> torch.Tensor: |
|
""" |
|
Args: |
|
images (torch.Tensor): tensor of size [N, 3, H, W] that contains |
|
BGR data (typically in uint8) |
|
Returns: |
|
images (torch.Tensor): tensor of size [N, 3, H1, W1] where |
|
H1 and W1 are chosen to respect the specified min and max sizes |
|
and preserve the original aspect ratio, the data channels |
|
follow BGR order and the data type is `torch.float32` |
|
""" |
|
|
|
images = images.float() |
|
min_size = min(images.shape[-2:]) |
|
max_size = max(images.shape[-2:]) |
|
scale = min(self.min_size / min_size, self.max_size / max_size) |
|
images = torch.nn.functional.interpolate( |
|
images, |
|
scale_factor=scale, |
|
mode="bilinear", |
|
align_corners=False, |
|
) |
|
return images |
|
|
