File size: 5,121 Bytes
1fae98d |
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 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
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,
},
"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,
},
}
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
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"])
MASK_MODES = {
"256train": make_lama_mask,
"256narrow": make_narrow_lama_mask,
"512train": make_512_lama_mask,
"512train-large": make_512_lama_mask_large
}
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)
|