diff --git a/README.md b/README.md index 95cafca..cfdc408 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,10 @@ 2.[卷积神经网络模型人像分割](https://github.com/itainf/aiphoto/wiki/%E5%8D%B7%E7%A7%AF%E6%A8%A1%E5%9E%8B%E4%BA%BA%E5%83%8F%E5%88%86%E5%89%B2) +3.[利用PyMatting替换背景颜色]() + +4.[利用dlib人脸检测裁剪照片]() + ### 更新记录 2020年7月4日更新 diff --git a/__init__.py b/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/bj.png b/bj.png new file mode 100644 index 0000000..ff82f68 Binary files /dev/null and b/bj.png differ diff --git a/img/meinv.jpg b/img/meinv.jpg index 901c080..9db6b10 100644 Binary files a/img/meinv.jpg and b/img/meinv.jpg differ diff --git a/img/meinv_alpha.png b/img/meinv_alpha.png new file mode 100644 index 0000000..68cd5a8 Binary files /dev/null and b/img/meinv_alpha.png differ diff --git a/img/meinv_alpha_resize.png b/img/meinv_alpha_resize.png new file mode 100644 index 0000000..a045667 Binary files /dev/null and b/img/meinv_alpha_resize.png differ diff --git a/img/meinv_id_grid.png b/img/meinv_id_grid.png new file mode 100644 index 0000000..1b04e3d Binary files /dev/null and b/img/meinv_id_grid.png differ diff --git a/img/meinv_trimap_resize.png b/img/meinv_trimap_resize.png new file mode 100644 index 0000000..6b69f05 Binary files /dev/null and b/img/meinv_trimap_resize.png differ diff --git a/img/trimap/meinv_org_trimap.png b/img/trimap/meinv_org_trimap.png deleted file mode 100644 index bc0dc63..0000000 Binary files a/img/trimap/meinv_org_trimap.png and /dev/null differ diff --git a/img/trimap/meinv_resize_trimap.png b/img/trimap/meinv_resize_trimap.png deleted file mode 100644 index b80bf23..0000000 Binary files a/img/trimap/meinv_resize_trimap.png and /dev/null differ diff --git a/main.py b/main.py new file mode 100644 index 0000000..28a71a1 --- /dev/null +++ b/main.py @@ -0,0 +1,25 @@ + +from u_2_net import my_u2net_test +from to_background import to_background +from to_background import to_standard_trimap +import numpy as np +from PIL import Image +if __name__ == "__main__": + org_img = "..\\aiphoto\\img\\meinv.jpg" + alpha_img = "..\\aiphoto\\img\\meinv_alpha.png" + alpha_resize_img = "..\\aiphoto\\img\\meinv_alpha_resize.png" + # # + # 通过u_2_net 获取 alpha + my_u2net_test.test_seg_trimap(org_img, alpha_img, alpha_resize_img) + # + # # 通过alpha 获取 trimap + trimap = "..\\aiphoto\\img\\meinv_trimap_resize.png" + to_standard_trimap.to_standard_trimap(alpha_resize_img, trimap) + # + # # 证件照添加纯色背景 + id_image = "..\\aiphoto\\img\\meinv_id_grid.png" + # to_background.to_background(org_img, trimap, id_image, "blue") + to_background.to_background_grid(org_img, trimap, id_image) + # image = Image.open(id_image) + # data = image.getdata() + # np.savetxt("data6.txt", data,fmt='%d',delimiter=',') diff --git a/to_background/__pycache__/to_background.cpython-37.pyc b/to_background/__pycache__/to_background.cpython-37.pyc new file mode 100644 index 0000000..84e1a98 Binary files /dev/null and b/to_background/__pycache__/to_background.cpython-37.pyc differ diff --git a/to_background/__pycache__/to_standard_trimap.cpython-37.pyc b/to_background/__pycache__/to_standard_trimap.cpython-37.pyc new file mode 100644 index 0000000..da2d021 Binary files /dev/null and b/to_background/__pycache__/to_standard_trimap.cpython-37.pyc differ diff --git a/to_background/to_background.py b/to_background/to_background.py new file mode 100644 index 0000000..953ce6b --- /dev/null +++ b/to_background/to_background.py @@ -0,0 +1,65 @@ +from pymatting import * +from PIL import Image + +colour_dict = { + "white": (255, 255, 255), + "red": (255, 0, 0), + "blue": (67, 142, 219) +} + + +def to_background(org, resize_trimap, id_image, colour): + """ + org:原始图片 + resize_trimap:trimap + id_image:新图片 + colour: 背景颜色 + """ + scale = 1.0 + image = load_image(org, "RGB", scale, "box") + trimap = load_image(resize_trimap, "GRAY", scale, "nearest") + im = Image.open(org) + # estimate alpha from image and trimap + alpha = estimate_alpha_cf(image, trimap) + + new_background = Image.new('RGB', im.size, colour_dict[colour]) + new_background.save("bj.png") + # load new background + new_background = load_image("bj.png", "RGB", scale, "box") + + + # estimate foreground from image and alpha + foreground, background = estimate_foreground_ml(image, alpha, return_background=True) + + # blend foreground with background and alpha + new_image = blend(foreground, new_background, alpha) + save_image(id_image, new_image) + + +def to_background_grid(org, resize_trimap, id_image): + """ + org:原始图片 + resize_trimap:trimap + id_image:新图片 + colour: 背景颜色 + """ + scale = 1.0 + image = load_image(org, "RGB", scale, "box") + trimap = load_image(resize_trimap, "GRAY", scale, "nearest") + im = Image.open(org) + # estimate alpha from image and trimap + alpha = estimate_alpha_cf(image, trimap) + + # estimate foreground from image and alpha + foreground, background = estimate_foreground_ml(image, alpha, return_background=True) + images = [image] + for k,v in colour_dict.items(): + new_background = Image.new('RGB', im.size, v) + new_background.save("bj.png") + new_background = load_image("bj.png", "RGB", scale, "box") + new_image = blend(foreground, new_background, alpha) + images.append(new_image) + + grid = make_grid(images) + save_image(id_image, grid) + diff --git a/to_background/to_standard_trimap.py b/to_background/to_standard_trimap.py new file mode 100644 index 0000000..253361f --- /dev/null +++ b/to_background/to_standard_trimap.py @@ -0,0 +1,39 @@ +from PIL import Image + + +def to_standard_trimap(alpha, trimap): + # Alpha图生成 trimap + print(alpha) + image = Image.open(alpha) + print(image) + # image = image.convert("P") + # image_file.save('meinv_resize_trimap.png') + sp = image.size + width = sp[0] + height = sp[1] + + for yh in range(height): + for xw in range(width): + dot = (xw, yh) + color_d_arr = image.getpixel(dot) + color_d=color_d_arr[0] + + if 0 < color_d <= 60: + image.putpixel(dot, (0,0,0)) + if 60 < color_d <= 200: + image.putpixel(dot, (128,128,128)) + if 200 < color_d <= 255: + image.putpixel(dot, (255,255,255)) + + image.save(trimap) + + + + + +# to_standard_trimap("..\\img\\trimap\\meinv_resize_trimap.png", "meinv_resize_bz_trimap.png") +# +# +# image = Image.open("meinv_resize_bz_trimap.png") +# data = image.getdata() +# np.savetxt("data4.txt", data,fmt='%d',delimiter=',') diff --git a/u_2_net/__pycache__/my_u2net_test.cpython-37.pyc b/u_2_net/__pycache__/my_u2net_test.cpython-37.pyc new file mode 100644 index 0000000..5302c4f Binary files /dev/null and b/u_2_net/__pycache__/my_u2net_test.cpython-37.pyc differ diff --git a/u_2_net/my_u2net_test.py b/u_2_net/my_u2net_test.py index e3d7b1d..e2b7015 100644 --- a/u_2_net/my_u2net_test.py +++ b/u_2_net/my_u2net_test.py @@ -44,7 +44,7 @@ def preprocess(image): def pre_net(): # 采用n2net 模型数据 model_name = 'u2net' - model_dir = 'saved_models/'+ model_name + '/' + model_name + '.pth' + model_dir = '..\\aiphoto\\u_2_net\\saved_models/'+ model_name + '/' + model_name + '.pth' print("...load U2NET---173.6 MB") net = U2NET(3,1) # 指定cpu @@ -75,8 +75,8 @@ def get_im(pred): return im -def test_seg_trimap(org,org_trimap,resize_trimap): - # 将原始图片转换成trimap +def test_seg_trimap(org,alpha,alpha_resize): + # 将原始图片转换成 Alpha图 # org:原始图片 # org_trimap: # resize_trimap: 调整尺寸的trimap @@ -91,13 +91,13 @@ def test_seg_trimap(org,org_trimap,resize_trimap): pred = normPRED(pred) # 将数据转换成图片 im = get_im(pred) - im.save(org_trimap) + im.save(alpha) sp = image.size # 根据原始图片调整尺寸 imo = im.resize((sp[0], sp[1]), resample=Image.BILINEAR) - imo.save(resize_trimap) + imo.save(alpha_resize) -if __name__ == "__main__": - test_seg_trimap("..\\img\\meinv.jpg","..\\img\\trimap\\meinv_org_trimap.png","..\\img\\trimap\\meinv_resize_trimap.png") - #pil_wait_blue() +# if __name__ == "__main__": +# test_seg_trimap("..\\img\\meinv.jpg","..\\img\\trimap\\meinv_alpha.png","..\\img\\trimap\\meinv_alpha_resize.png") +# #pil_wait_blue()