def upfirdn_2d()

in jcm/models/up_or_down_sampling.py [0:0]

• 52 lines of code
• 8 McCabe index (conditional complexity)

def upfirdn_2d(x, k, upx, upy, downx, downy, padx0, padx1, pady0, pady1):
    """Pad, upsample, FIR filter, and downsample a batch of 2D images.

    Accepts a batch of 2D images of the shape `[majorDim, inH, inW, minorDim]`
    and performs the following operations for each image, batched across
    `majorDim` and `minorDim`:
    1. Pad the image with zeros by the specified number of pixels on each side
       (`padx0`, `padx1`, `pady0`, `pady1`). Specifying a negative value
       corresponds to cropping the image.
    2. Upsample the image by inserting the zeros after each pixel (`upx`,
    `upy`).
    3. Convolve the image with the specified 2D FIR filter (`k`), shrinking the
       image so that the footprint of all output pixels lies within the input
       image.
    4. Downsample the image by throwing away pixels (`downx`, `downy`).
    This sequence of operations bears close resemblance to
    scipy.signal.upfirdn().
    The fused op is considerably more efficient than performing the same
    calculation
    using standard TensorFlow ops. It supports gradients of arbitrary order.
    Args:
        x:      Input tensor of the shape `[majorDim, inH, inW, minorDim]`.
        k:      2D FIR filter of the shape `[firH, firW]`.
        upx:    Integer upsampling factor along the X-axis (default: 1).
        upy:    Integer upsampling factor along the Y-axis (default: 1).
        downx:  Integer downsampling factor along the X-axis (default: 1).
        downy:  Integer downsampling factor along the Y-axis (default: 1).
        padx0:  Number of pixels to pad on the left side (default: 0).
        padx1:  Number of pixels to pad on the right side (default: 0).
        pady0:  Number of pixels to pad on the top side (default: 0).
        pady1:  Number of pixels to pad on the bottom side (default: 0).
        impl:   Name of the implementation to use. Can be `"ref"` or `"cuda"`
          (default).

    Returns:
        Tensor of the shape `[majorDim, outH, outW, minorDim]`, and same
        datatype as `x`.
    """
    k = jnp.asarray(k, dtype=np.float32)
    assert len(x.shape) == 4
    inH = x.shape[1]
    inW = x.shape[2]
    minorDim = x.shape[3]
    kernelH, kernelW = k.shape
    assert inW >= 1 and inH >= 1
    assert kernelW >= 1 and kernelH >= 1
    assert isinstance(upx, int) and isinstance(upy, int)
    assert isinstance(downx, int) and isinstance(downy, int)
    assert isinstance(padx0, int) and isinstance(padx1, int)
    assert isinstance(pady0, int) and isinstance(pady1, int)

    # Upsample (insert zeros).
    x = jnp.reshape(x, (-1, inH, 1, inW, 1, minorDim))
    x = jnp.pad(x, [[0, 0], [0, 0], [0, upy - 1], [0, 0], [0, upx - 1], [0, 0]])
    x = jnp.reshape(x, [-1, inH * upy, inW * upx, minorDim])

    # Pad (crop if negative).
    x = jnp.pad(
        x,
        [
            [0, 0],
            [max(pady0, 0), max(pady1, 0)],
            [max(padx0, 0), max(padx1, 0)],
            [0, 0],
        ],
    )
    x = x[
        :,
        max(-pady0, 0) : x.shape[1] - max(-pady1, 0),
        max(-padx0, 0) : x.shape[2] - max(-padx1, 0),
        :,
    ]

    # Convolve with filter.
    x = jnp.transpose(x, [0, 3, 1, 2])
    x = jnp.reshape(x, [-1, 1, inH * upy + pady0 + pady1, inW * upx + padx0 + padx1])
    w = jnp.array(k[::-1, ::-1, None, None], dtype=x.dtype)
    x = jax.lax.conv_general_dilated(
        x,
        w,
        window_strides=(1, 1),
        padding="VALID",
        dimension_numbers=("NCHW", "HWIO", "NCHW"),
    )

    x = jnp.reshape(
        x,
        [
            -1,
            minorDim,
            inH * upy + pady0 + pady1 - kernelH + 1,
            inW * upx + padx0 + padx1 - kernelW + 1,
        ],
    )
    x = jnp.transpose(x, [0, 2, 3, 1])

    # Downsample (throw away pixels).
    return x[:, ::downy, ::downx, :]