IDphotoChangeBackgroundColor

# @Time : 2020/10/30 下午2:49 
# @Author : wgiplash
# @File : IDphotoChangeBackgroundColor.py 
# @Software: PyCharm


import cv2 as cv
import numpy as np
import os

def processing(image_path, output_path):
    image = cv.imread(image_path)
    h, w, ch = image.shape
    data = image.reshape((-1, 3))
    data = np.float32(data)

    criteria = (cv.TERM_CRITERIA_EPS + cv.TERM_CRITERIA_MAX_ITER, 10, 1.0)
    num_clusters = 4
    _, label, _ = cv.kmeans(data, num_clusters, None, criteria, num_clusters, cv.KMEANS_RANDOM_CENTERS)

    indx = label[0][0]
    mask = np.ones((h, w), dtype=np.uint8) * 255
    label = np.reshape(label, (h, w))

    mask[label == indx] = 0
    Kernel = cv.getStructuringElement(cv.MORPH_RECT, (3, 3))
    cv.erode(mask, Kernel, mask)
    mask = cv.GaussianBlur(mask, (5, 5), 0)

    background_w = np.ones(image.shape, dtype=np.float) * 255
    red = np.array([0, 0, 255])
    background_r = np.tile(red, (h, w, 1))
    blue = np.array([255, 0, 0])
    background_b = np.tile(blue, (h, w, 1))

    alpha = mask.astype(np.float32) / 255
    foreground = alpha[..., None] * image
    background = 1 - alpha[..., None]

    image_white = foreground + background * background_w
    image_red = foreground + background * background_r
    image_blue = foreground + background * background_b

    white_path = os.path.join(output_path, 'white.jpg')
    red_path = os.path.join(output_path, 'red.jpg')
    blue_path = os.path.join(output_path, 'blue.jpg')

    cv.imwrite(white_path, image_white.astype(np.uint8))
    cv.imwrite(red_path, image_red.astype(np.uint8))
    cv.imwrite(blue_path, image_blue.astype(np.uint8))


if __name__ == '__main__':

    input_image = "/data_1/timg.jpeg"
    reuslt_path = "./"
    processing(input_image, reuslt_path)