66 lines
1.9 KiB
Python
66 lines
1.9 KiB
Python
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)
|
||
|