import paddlehub as hub
import cv2
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt 
import matplotlib.image as mpimg 

# 待抠图照片
image = 'ce.jpg'
# 背景
background = 'ce10.jpg'
# 合成照片名称
save_image = 'save_image.jpg'

human_seg = hub.Module(name="deeplabv3p_xception65_humanseg")
result = human_seg.segmentation(images=[cv2.imread(image)],visualization=True)
#print(result[0]['save_path'])
# or
# result = human_seg.segmentation(paths=['/PATH/TO/IMAGE'])
# 预测结果展示


test_img_path = result[0]['save_path']
#print(test_img_path)

#cv2.imwrite('mei',img)

def blend_images(fore_image, base_image):
    """
    将抠出的人物图像换背景
    fore_image: 前景图片，抠出的人物图片
    base_image: 背景图片
    """
    # 读入图片
    base_image = Image.open(base_image).convert('RGB')
    fore_image = Image.open(fore_image).resize(base_image.size)

    # 图片加权合成
    scope_map = np.array(fore_image)[:,:,-1] / 255
    scope_map = scope_map[:,:,np.newaxis]
    scope_map = np.repeat(scope_map, repeats=3, axis=2)
    res_image = np.multiply(scope_map, np.array(fore_image)[:,:,:3]) + np.multiply((1-scope_map), np.array(base_image))
    
    #保存图片
    res_image = Image.fromarray(np.uint8(res_image))
    res_image.save('mei.jpg')

blend_images(test_img_path, background)

# 展示合成图片
plt.figure(figsize=(10,10))
img = mpimg.imread(save_image)
plt.imshow(img) 
plt.axis('off') 
plt.show()