File size: 3,166 Bytes
dab6ba1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
import PIL
import numpy as np
from PIL import Image, ImageColor, ImageDraw
from PIL.Image import Image as PILImage
from pymatting.alpha.estimate_alpha_cf import estimate_alpha_cf
from pymatting.foreground.estimate_foreground_ml import estimate_foreground_ml
from pymatting.util.util import stack_images
from rembg.bg import post_process, naive_cutout, apply_background_color
from scipy.ndimage import binary_erosion


def alpha_matting_cutout(img: PILImage, trimap: np.ndarray) -> PILImage:
    if img.mode == "RGBA" or img.mode == "CMYK":
        img = img.convert("RGB")

    img = np.asarray(img)

    img_normalized = img / 255.0
    trimap_normalized = trimap / 255.0

    alpha = estimate_alpha_cf(img_normalized, trimap_normalized)
    foreground = estimate_foreground_ml(img_normalized, alpha)
    cutout = stack_images(foreground, alpha)

    cutout = np.clip(cutout * 255, 0, 255).astype(np.uint8)
    return Image.fromarray(cutout)


def generate_trimap(
        mask: PILImage,
        foreground_threshold: int,
        background_threshold: int,
        erode_structure_size: int,
) -> np.ndarray:
    mask = np.asarray(mask)

    is_foreground = mask > foreground_threshold
    is_background = mask < background_threshold

    structure = None
    if erode_structure_size > 0:
        structure = np.ones(
            (erode_structure_size, erode_structure_size), dtype=np.uint8
        )

    is_foreground = binary_erosion(is_foreground, structure=structure)
    is_background = binary_erosion(is_background, structure=structure, border_value=1)

    trimap = np.full(mask.shape, dtype=np.uint8, fill_value=128)
    trimap[is_foreground] = 255
    trimap[is_background] = 0

    return trimap


def get_background_dominant_color(img: PILImage, mask: PILImage) -> tuple:
    negative_img = img.copy()
    negative_mask = PIL.ImageOps.invert(mask)
    negative_img.putalpha(negative_mask)
    negative_img = negative_img.resize((1, 1))
    r, g, b, a = negative_img.getpixel((0, 0))
    return r, g, b, 255


def remove(session, img: PILImage, smoot: bool, matting: tuple, color) -> (PILImage, PILImage):
    mask = session.predict(img)[0]

    if smoot:
        mask = PIL.Image.fromarray(post_process(np.array(mask)))

    fg_t, bg_t, erode = matting

    if fg_t > 0 or bg_t > 0 or erode > 0:
        mask = generate_trimap(mask, *matting)
        try:
            cutout = alpha_matting_cutout(img, mask)
            mask = PIL.Image.fromarray(mask)
        except ValueError as err:
            raise err
    else:
        cutout = naive_cutout(img, mask)

    if color is True:
        color = get_background_dominant_color(img, mask)
        cutout = apply_background_color(cutout, color)
    elif isinstance(color, str):
        r, g, b = ImageColor.getcolor(color, "RGB")
        cutout = apply_background_color(cutout, (r, g, b, 255))

    return cutout, mask


def make_label(text, width=600, height=200, color="black") -> PILImage:
    image = Image.new("RGB", (width, height), color)
    draw = ImageDraw.Draw(image)
    text_width, text_height = draw.textsize(text)
    draw.text(((width-text_width)/2, height/2), text)
    return image