Function pixelunshuffle

monai/networks/utils.py:415–454

Apply pixel unshuffle to the tensor `x` with spatial dimensions `spatial_dims` and scaling factor `scale_factor`. Inverse operation of pixelshuffle. See: Shi et al., 2016, "Real-Time Single Image and Video Super-Resolution Using an Efficient Sub-Pixel Convolutional Neural Network."

(x: torch.Tensor, spatial_dims: int, scale_factor: int)

def pixelunshuffle(x: torch.Tensor, spatial_dims: int, scale_factor: int) -> torch.Tensor:
    """
    Apply pixel unshuffle to the tensor `x` with spatial dimensions `spatial_dims` and scaling factor `scale_factor`.
    Inverse operation of pixelshuffle.

    See: Shi et al., 2016, "Real-Time Single Image and Video Super-Resolution
    Using an Efficient Sub-Pixel Convolutional Neural Network."

    See: Aitken et al., 2017, "Checkerboard artifact free sub-pixel convolution".

    Args:
        x: Input tensor with shape BCHW[D]
        spatial_dims: number of spatial dimensions, typically 2 or 3 for 2D or 3D
        scale_factor: factor to reduce the spatial dimensions by, must be >=1

    Returns:
        Unshuffled version of `x` with shape (B, C*(r**d), H/r, W/r) for 2D
        or (B, C*(r**d), D/r, H/r, W/r) for 3D, where r is the scale_factor
        and d is spatial_dims.

    Raises:
        ValueError: When spatial dimensions are not divisible by scale_factor
    """
    dim, factor = spatial_dims, scale_factor
    input_size = list(x.size())
    batch_size, channels = input_size[:2]
    scale_factor_mult = factor**dim
    new_channels = channels * scale_factor_mult

    if any(d % factor != 0 for d in input_size[2:]):
        raise ValueError(
            f"All spatial dimensions must be divisible by factor {factor}. " f", spatial shape is: {input_size[2:]}"
        )
    output_size = [batch_size, new_channels] + [d // factor for d in input_size[2:]]
    reshaped_size = [batch_size, channels] + sum([[d // factor, factor] for d in input_size[2:]], [])

    permute_indices = [0, 1] + [(2 * i + 3) for i in range(spatial_dims)] + [(2 * i + 2) for i in range(spatial_dims)]
    x = x.reshape(reshaped_size).permute(permute_indices)
    x = x.reshape(output_size)
    return x


@contextmanager