Matting/utils/color.py

import os
import cv2
import numpy as np


def get_bg(background, img_shape):
    # 1、纯色
    # 2、通道颜色
    # 3、图片
    bg = np.zeros(img_shape)
    if background is None:
        return None
    if os.path.exists(background):
        bg = cv2.imread(background)
        bg = cv2.resize(bg, (img_shape[1], img_shape[0]))
    elif background == 'r':
        bg[:, :, 2] = 255
    elif background == 'g':
        bg[:, :, 1] = 255
    elif background == 'b':
        bg[:, :, 0] = 255
    elif background == 'w':
        bg[:, :, :] = 255
    elif is_color_hex(background):
        r, g, b, _ = hex_to_rgb(background)
        bg[:, :, 2] = r
        bg[:, :, 1] = g
        bg[:, :, 0] = b
    else:
        return None

    return bg


def is_color_hex(color: str):
    size = len(color)
    if color.startswith("#"):
        return size == 7 or size == 9
    return False


def hex_to_rgb(color: str):
    if color.startswith("#"):
        color = color[1:len(color)]
    r = int(color[0:2], 16)
    g = int(color[2:4], 16)
    b = int(color[4:6], 16)
    a = 100
    if len(color) == 8:
        a = int(color[6:8], 16)
    return r, g, b, a