def reshape_image(image, image_height, image_width, num_channels, len_shape):
    """
    Reshape image into input dimensions

    Parameters
    ----------
    image : NumPy array
        image to reshape
    image_height : int
        image height
    image_width : int
        image width
    num_channels : int
        number of channels in image
    len_shape : int
        number of axes in input image

    Returns
    ----------
    image : NumPy array
        reshaped image
    """
    if len_shape == 3:
        image = image.reshape(image_height, image_width, num_channels)
    else:
        image = image.reshape(image_height, image_width)

    return image