def METHOD_NAME(image: np.ndarray, n_pixels: int):
    """Sample the image to improve runtime, expected image format H,W,C."""
    flat_image = image.reshape((-1, image.shape[-1]))
    if flat_image.shape[0] > n_pixels:
        pixel_idxs = np.random.choice(flat_image.shape[0], n_pixels)
    else:
        pixel_idxs = np.arange(flat_image.shape[0])
    sampled_image = flat_image[pixel_idxs, np.newaxis, :]
    return sampled_image