def remove_alpha()

in shap_e/util/image_util.py [0:0]

• 17 lines of code
• 6 McCabe index (conditional complexity)

def remove_alpha(img: Image.Image, mode: str = "random") -> Image.Image:
    """
    No op if the image doesn't have an alpha channel.

    :param: mode: Defaults to "random" but has an option to use a "black" or
        "white" background

    :return: image with alpha removed
    """
    img_arr = np.asarray(img)
    if img_arr.shape[2] == 4:
        # Add bg to get rid of alpha channel
        if mode == "random":
            height, width = img_arr.shape[:2]
            bg = Image.fromarray(
                random.choice([_black_bg, _gray_bg, _checker_bg, _noise_bg])(height, width)
            )
            bg.paste(img, mask=img)
            img = bg
        elif mode == "black" or mode == "white":
            img_arr = img_arr.astype(float)
            rgb, alpha = img_arr[:, :, :3], img_arr[:, :, -1:] / 255
            background = np.zeros((1, 1, 3)) if mode == "black" else np.full((1, 1, 3), 255)
            rgb = rgb * alpha + background * (1 - alpha)
            img = Image.fromarray(np.round(rgb).astype(np.uint8))
    return img