| """Resize transformation.""" |
|
|
| from __future__ import annotations |
|
|
| import math |
|
|
| import numpy as np |
| import torch |
| from vis4d.common.typing import NDArrayF32, NDArrayI64 |
| from vis4d.data.const import CommonKeys as K |
| from vis4d.data.transforms.base import Transform |
| from vis4d.data.transforms.resize import ResizeParam, resize_tensor |
|
|
|
|
| @Transform(K.images, ["transforms.resize", K.input_hw]) |
| class GenResizeParameters: |
| """Generate the parameters for a resize operation.""" |
|
|
| def __init__( |
| self, shape: tuple[int, int], scales: tuple[float, float] | float = 1.0 |
| ) -> None: |
| """Create a new instance of the class.""" |
| self.shape = shape |
| self.scales = scales |
|
|
| def __call__( |
| self, images: list[NDArrayF32] |
| ) -> tuple[list[ResizeParam], list[tuple[int, int]]]: |
| """Compute the parameters and put them in the data dict.""" |
| if isinstance(self.scales, float): |
| random_scale = self.scales |
| else: |
| random_scale = np.random.uniform(self.scales[0], self.scales[1]) |
|
|
| shape = ( |
| math.ceil(self.shape[0] * random_scale - 0.5), |
| math.ceil(self.shape[1] * random_scale - 0.5), |
| ) |
|
|
| output_ratio = shape[1] / shape[0] |
|
|
| image = images[0] |
|
|
| input_h, input_w = (image.shape[1], image.shape[2]) |
| input_ratio = input_w / input_h |
|
|
| if output_ratio > input_ratio: |
| scale = shape[0] / input_h |
| else: |
| scale = shape[1] / input_w |
|
|
| target_shape = ( |
| math.ceil(input_h * scale - 0.5), |
| math.ceil(input_w * scale - 0.5), |
| ) |
|
|
| scale_factor = (target_shape[0] / input_h, target_shape[1] / input_w) |
|
|
| resize_params = [ |
| ResizeParam(target_shape=target_shape, scale_factor=scale_factor) |
| ] * len(images) |
| target_shapes = [target_shape] * len(images) |
|
|
| return resize_params, target_shapes |
|
|
|
|
| @Transform( |
| [K.panoptic_masks, "transforms.resize.target_shape"], K.panoptic_masks |
| ) |
| class ResizePanopticMasks: |
| """Resize panoptic segmentation masks.""" |
|
|
| def __call__( |
| self, |
| masks_list: list[NDArrayI64], |
| target_shape_list: list[tuple[int, int]], |
| ) -> list[NDArrayI64]: |
| """Resize masks.""" |
| for i, (masks, target_shape) in enumerate( |
| zip(masks_list, target_shape_list) |
| ): |
| masks_ = torch.from_numpy(masks) |
| masks_ = ( |
| resize_tensor( |
| masks_.float().unsqueeze(0).unsqueeze(0), |
| target_shape, |
| interpolation="nearest", |
| ) |
| .type(masks_.dtype) |
| .squeeze(0) |
| .squeeze(0) |
| ) |
| masks_list[i] = masks_.numpy() |
| return masks_list |
|
|
|
|
| @Transform([K.boxes3d, "transforms.resize.scale_factor"], K.boxes3d) |
| class ResizeBoxes3D: |
| """Resize list of 2D bounding boxes.""" |
|
|
| def __call__( |
| self, |
| boxes_list: list[NDArrayF32], |
| scale_factors: list[tuple[float, float]], |
| ) -> list[NDArrayF32]: |
| """Resize 2D bounding boxes. |
| |
| Args: |
| boxes_list: (list[NDArrayF32]): The bounding boxes to be resized. |
| scale_factors (list[tuple[float, float]]): scaling factors. |
| |
| Returns: |
| list[NDArrayF32]: Resized bounding boxes according to parameters in |
| resize. |
| """ |
| for i, (boxes, scale_factor) in enumerate( |
| zip(boxes_list, scale_factors) |
| ): |
| boxes[:, 2] /= scale_factor[0] |
| boxes_list[i] = boxes |
| return boxes_list |
|
|
