File size: 4,121 Bytes
194b4ef
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
from modules import script_callbacks
import gradio as gr
import numpy as np
from PIL import Image
import cv2


def check_alpha(img:np.array) -> bool:
    return img.shape[-1] == 4 and (img[:, :, 3] < 255).any()

def preprocess_foreground(fg_img:Image, rotation:float, scale:float):
    if not check_alpha(np.asarray(fg_img)):
        print("Foreground Image contains no transparency...")
        return None

    w, h = fg_img.size
    w = int(float(w) * scale)
    h = int(float(h) * scale)

    return fg_img.resize((w, h), resample=Image.BILINEAR).rotate(rotation, resample=Image.BILINEAR, expand=True)

def process_mask(bg: str, fg: str, tc: float, lt: float, ht: float, x:float, y:float, r:float, s:float):
    if (bg is None) or (fg is None):
        return [None, None]

    fg_img = preprocess_foreground(Image.open(fg), r, s)
    if fg_img is None:
        return [None, None]

    bg_img = Image.open(bg)

    width, height = bg_img.size
    new_fg = Image.new('RGBA', (width, height), (0, 0, 0, 0))
    new_fg.paste(fg_img, (x, y), fg_img)

    mask = np.asarray(new_fg)[:, :, 3]
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    edge = cv2.Canny(mask, lt, ht)
    dilate = cv2.dilate(edge, kernel, iterations=tc)

    bg_img.paste(new_fg, (0, 0), new_fg)

    return [
        bg_img,
        dilate.astype("uint8")
    ]


def aaa_ui():

    with gr.Blocks() as aaa_UI:
        with gr.Row():
            background = gr.Image(
                image_mode="RGBA",
                sources="upload",
                type="filepath",
                label="Background Image",
                show_download_button=False,
                interactive=True,
                height=250
            )

            foreground = gr.Image(
                image_mode="RGBA",
                sources="upload",
                type="filepath",
                label="Foreground Image",
                show_download_button=False,
                interactive=True,
                height=250
            )

            img = gr.Image(
                image_mode="RGBA",
                label="Blended Image",
                elem_id="aaa_img",
                interactive=False,
                height=250,
            )

            mask = gr.Image(
                image_mode="L",
                label="Mask",
                elem_id="aaa_mask",
                interactive=False,
                height=250
            )

        with gr.Row():
            with gr.Column():
                thicc = gr.Slider(
                    label="Mask Thickness", minimum=5, maximum=50, step=5, value=10
                )

                low_t = gr.Slider(
                    label="Low Threshold", minimum=0, maximum=255, step=1, value=0
                )

                high_t = gr.Slider(
                    label="High Threshold", minimum=0, maximum=255, step=1, value=100
                )

            with gr.Column():
                offset_x = gr.Slider(
                    label="X Offset", minimum=-2048, maximum=2048, step=1, value=0
                )

                offset_y = gr.Slider(
                    label="Y Offset", minimum=-2048, maximum=2048, step=1, value=0
                )

                with gr.Row():
                    offset_r = gr.Slider(
                        label="Rotation", minimum=-180, maximum=180, step=1, value=0
                    )

                    offset_s = gr.Slider(
                        label="Scale", minimum=0.0, maximum=2.0, step=0.1, value=1.0
                    )

        with gr.Row():
            proc_btn = gr.Button("Process Mask", variant="primary")
            send_btn = gr.Button("Send to Inpaint", variant="primary")

        proc_btn.click(
            process_mask,
            inputs=[background, foreground, thicc, low_t, high_t, offset_x, offset_y, offset_r, offset_s],
            outputs=[img, mask],
        )

        send_btn.click(
            None, None, None,
            _js="() => { aaa_sendImage2InpaintUpload(); }"
        )

    return [(aaa_UI, "AAA", "sd-webui-aaa")]


script_callbacks.on_ui_tabs(aaa_ui)