# @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)
09-25
09-25
09-25