stable-diffusion-finetune/ldm/data/inpainting/synthetic_mask.py

167 lines
5 KiB
Python
Raw Normal View History

2022-07-24 11:23:26 +00:00
from PIL import Image, ImageDraw
import numpy as np
settings = {
"256narrow": {
"p_irr": 1,
"min_n_irr": 4,
"max_n_irr": 50,
"max_l_irr": 40,
"max_w_irr": 10,
"min_n_box": None,
"max_n_box": None,
"min_s_box": None,
"max_s_box": None,
"marg": None,
},
"256train": {
"p_irr": 0.5,
"min_n_irr": 1,
"max_n_irr": 5,
"max_l_irr": 200,
"max_w_irr": 100,
"min_n_box": 1,
"max_n_box": 4,
"min_s_box": 30,
"max_s_box": 150,
"marg": 10,
},
"512train": { # TODO: experimental
"p_irr": 0.5,
"min_n_irr": 1,
"max_n_irr": 5,
"max_l_irr": 450,
"max_w_irr": 250,
"min_n_box": 1,
"max_n_box": 4,
"min_s_box": 30,
"max_s_box": 300,
"marg": 10,
},
2022-07-27 22:07:45 +00:00
"512train-large": { # TODO: experimental
"p_irr": 0.5,
"min_n_irr": 1,
"max_n_irr": 5,
"max_l_irr": 450,
"max_w_irr": 400,
"min_n_box": 1,
"max_n_box": 4,
"min_s_box": 75,
"max_s_box": 450,
"marg": 10,
},
2022-07-24 11:23:26 +00:00
}
def gen_segment_mask(mask, start, end, brush_width):
mask = mask > 0
mask = (255 * mask).astype(np.uint8)
mask = Image.fromarray(mask)
draw = ImageDraw.Draw(mask)
draw.line([start, end], fill=255, width=brush_width, joint="curve")
mask = np.array(mask) / 255
return mask
def gen_box_mask(mask, masked):
x_0, y_0, w, h = masked
mask[y_0:y_0 + h, x_0:x_0 + w] = 1
return mask
def gen_round_mask(mask, masked, radius):
x_0, y_0, w, h = masked
xy = [(x_0, y_0), (x_0 + w, y_0 + w)]
mask = mask > 0
mask = (255 * mask).astype(np.uint8)
mask = Image.fromarray(mask)
draw = ImageDraw.Draw(mask)
draw.rounded_rectangle(xy, radius=radius, fill=255)
mask = np.array(mask) / 255
return mask
def gen_large_mask(prng, img_h, img_w,
marg, p_irr, min_n_irr, max_n_irr, max_l_irr, max_w_irr,
min_n_box, max_n_box, min_s_box, max_s_box):
"""
img_h: int, an image height
img_w: int, an image width
marg: int, a margin for a box starting coordinate
p_irr: float, 0 <= p_irr <= 1, a probability of a polygonal chain mask
min_n_irr: int, min number of segments
max_n_irr: int, max number of segments
max_l_irr: max length of a segment in polygonal chain
max_w_irr: max width of a segment in polygonal chain
min_n_box: int, min bound for the number of box primitives
max_n_box: int, max bound for the number of box primitives
min_s_box: int, min length of a box side
max_s_box: int, max length of a box side
"""
mask = np.zeros((img_h, img_w))
uniform = prng.randint
if np.random.uniform(0, 1) < p_irr: # generate polygonal chain
n = uniform(min_n_irr, max_n_irr) # sample number of segments
for _ in range(n):
y = uniform(0, img_h) # sample a starting point
x = uniform(0, img_w)
a = uniform(0, 360) # sample angle
l = uniform(10, max_l_irr) # sample segment length
w = uniform(5, max_w_irr) # sample a segment width
# draw segment starting from (x,y) to (x_,y_) using brush of width w
x_ = x + l * np.sin(a)
y_ = y + l * np.cos(a)
mask = gen_segment_mask(mask, start=(x, y), end=(x_, y_), brush_width=w)
x, y = x_, y_
else: # generate Box masks
n = uniform(min_n_box, max_n_box) # sample number of rectangles
for _ in range(n):
h = uniform(min_s_box, max_s_box) # sample box shape
w = uniform(min_s_box, max_s_box)
x_0 = uniform(marg, img_w - marg - w) # sample upper-left coordinates of box
y_0 = uniform(marg, img_h - marg - h)
if np.random.uniform(0, 1) < 0.5:
mask = gen_box_mask(mask, masked=(x_0, y_0, w, h))
else:
r = uniform(0, 60) # sample radius
mask = gen_round_mask(mask, masked=(x_0, y_0, w, h), radius=r)
return mask
2022-07-27 22:07:45 +00:00
make_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["256train"])
make_narrow_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["256narrow"])
make_512_lama_mask = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["512train"])
make_512_lama_mask_large = lambda prng, h, w: gen_large_mask(prng, h, w, **settings["512train-large"])
2022-07-24 11:23:26 +00:00
2022-07-27 22:07:45 +00:00
MASK_MODES = {
"256train": make_lama_mask,
"256narrow": make_narrow_lama_mask,
"512train": make_512_lama_mask,
"512train-large": make_512_lama_mask_large
}
2022-07-24 11:23:26 +00:00
if __name__ == "__main__":
import sys
out = sys.argv[1]
prng = np.random.RandomState(1)
kwargs = settings["256train"]
mask = gen_large_mask(prng, 256, 256, **kwargs)
mask = (255 * mask).astype(np.uint8)
mask = Image.fromarray(mask)
mask.save(out)