import numpy as np from PIL import Image, ImageDraw import math import random def RandomBrush( max_tries, s, min_num_vertex = 4, max_num_vertex = 18, mean_angle = 2*math.pi / 5, angle_range = 2*math.pi / 15, min_width = 12, max_width = 48): H, W = s, s average_radius = math.sqrt(H*H+W*W) / 8 mask = Image.new('L', (W, H), 0) for _ in range(np.random.randint(max_tries)): num_vertex = np.random.randint(min_num_vertex, max_num_vertex) angle_min = mean_angle - np.random.uniform(0, angle_range) angle_max = mean_angle + np.random.uniform(0, angle_range) angles = [] vertex = [] for i in range(num_vertex): if i % 2 == 0: angles.append(2*math.pi - np.random.uniform(angle_min, angle_max)) else: angles.append(np.random.uniform(angle_min, angle_max)) h, w = mask.size vertex.append((int(np.random.randint(0, w)), int(np.random.randint(0, h)))) for i in range(num_vertex): r = np.clip( np.random.normal(loc=average_radius, scale=average_radius//2), 0, 2*average_radius) new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w) new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h) vertex.append((int(new_x), int(new_y))) draw = ImageDraw.Draw(mask) width = int(np.random.uniform(min_width, max_width)) draw.line(vertex, fill=1, width=width) for v in vertex: draw.ellipse((v[0] - width//2, v[1] - width//2, v[0] + width//2, v[1] + width//2), fill=1) if np.random.random() > 0.5: mask.transpose(Image.FLIP_LEFT_RIGHT) if np.random.random() > 0.5: mask.transpose(Image.FLIP_TOP_BOTTOM) mask = np.asarray(mask, np.uint8) if np.random.random() > 0.5: mask = np.flip(mask, 0) if np.random.random() > 0.5: mask = np.flip(mask, 1) return mask def RandomMask(s, hole_range=[0,1]): coef = min(hole_range[0] + hole_range[1], 1.0) while True: mask = np.ones((s, s), np.uint8) def Fill(max_size): w, h = np.random.randint(max_size), np.random.randint(max_size) ww, hh = w // 2, h // 2 x, y = np.random.randint(-ww, s - w + ww), np.random.randint(-hh, s - h + hh) mask[max(y, 0): min(y + h, s), max(x, 0): min(x + w, s)] = 0 def MultiFill(max_tries, max_size): for _ in range(np.random.randint(max_tries)): Fill(max_size) MultiFill(int(3 * coef), s // 2) MultiFill(int(2 * coef), s) mask = np.logical_and(mask, 1 - RandomBrush(int(4 * coef), s)) # hole denoted as 0, reserved as 1 hole_ratio = 1 - np.mean(mask) if hole_range is not None and (hole_ratio <= hole_range[0] or hole_ratio >= hole_range[1]): continue return mask[np.newaxis, ...].astype(np.float32) def BatchRandomMask(batch_size, s, hole_range=[0, 1]): return np.stack([RandomMask(s, hole_range=hole_range) for _ in range(batch_size)], axis=0) if __name__ == '__main__': res = 512 # res = 256 cnt = 2000 tot = 0 for i in range(cnt): mask = RandomMask(s=res) tot += mask.mean() print(tot / cnt)