20200709 利用PyMatting替换背景颜色

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日更新

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=',')

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)
+

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=',')

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()