def _resize()

in neuralcompression/models/scale_hyperprior.py [0:0]

• 13 lines of code
• 5 McCabe index (conditional complexity)

def _resize(x: Tensor, target_shape: Sequence[int]) -> Tensor:
    """
    Resizes a tensor to the target shape.

    Given a tensor of shape [batch_size, num_channels, height, width] and
    a [height, width] target_shape, this function resizes the input
    tensor to the taget shape using center cropping if the input is larger
    than the target shape, and using nearest neighbour interpolation if
    the input is smaller than the target shape.

    Args:
        x: the tensor to resize, of shape [batch_size, num_channels,
            height, width].
        target_shape: a [height, width] pair to reshape the input to.

    Returns:
        the resized tensor.
    """

    height, width = x.shape[2:]
    target_height, target_width = target_shape

    if height >= target_height and width >= target_width:
        return tvtF.center_crop(x, target_shape)
    elif height <= target_height and width <= target_width:
        return F.interpolate(x, target_shape, mode="nearest")
    else:
        raise ValueError(
            f"Input tensor (with shape {x.shape}) is larger than the"
            f" target shape of {target_shape} along one height/width axis"
            " and is smaller than the target shape along the other axis."
        )