def METHOD_NAME(im: np.ndarray, driver="PNG", **opts) -> bytes:
    if im.dtype != np.uint8:
        raise ValueError("Only support uint8 images on input")
    if im.ndim == 3:
        h, w, nc = im.shape
        bands = np.transpose(im, axes=(2, 0, 1))
    elif im.ndim == 2:
        (h, w), nc = im.shape, 1
        bands = im.reshape(nc, h, w)
    else:
        raise ValueError(f"Expect 2 or 3 dimensional array got: {im.ndim}")
    rio_opts = dict(width=w, height=h, count=nc, driver=driver, dtype="uint8", **opts)
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", rasterio.errors.NotGeoreferencedWarning)
        with rasterio.MemoryFile() as mem:
            with mem.open(**rio_opts) as dst:
                dst.write(bands)
            return mem.read()