| | |
| |
|
| | 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 |
| |
|
