import numpy as np import random from PIL import Image, ImageDraw def get_video_masks_by_moving_random_stroke( video_len, imageWidth=320, imageHeight=180, nStroke=5, nVertexBound=[10, 30], maxHeadSpeed=15, maxHeadAcceleration=(15, 0.5), brushWidthBound=(5, 20), boarderGap=None, nMovePointRatio=0.5, maxPiontMove=10, maxLineAcceleration=5, maxInitSpeed=5 ): ''' Get video masks by random strokes which move randomly between each frame, including the whole stroke and its control points Parameters ---------- imageWidth: Image width imageHeight: Image height nStroke: Number of drawed lines nVertexBound: Lower/upper bound of number of control points for each line maxHeadSpeed: Max head speed when creating control points maxHeadAcceleration: Max acceleration applying on the current head point ( a head point and its velosity decides the next point) brushWidthBound (min, max): Bound of width for each stroke boarderGap: The minimum gap between image boarder and drawed lines nMovePointRatio: The ratio of control points to move for next frames maxPiontMove: The magnitude of movement for control points for next frames maxLineAcceleration: The magnitude of acceleration for the whole line Examples ---------- object_like_setting = { "nVertexBound": [5, 20], "maxHeadSpeed": 15, "maxHeadAcceleration": (15, 3.14), "brushWidthBound": (30, 50), "nMovePointRatio": 0.5, "maxPiontMove": 10, "maxLineAcceleration": (5, 0.5), "boarderGap": 20, "maxInitSpeed": 10, } rand_curve_setting = { "nVertexBound": [10, 30], "maxHeadSpeed": 20, "maxHeadAcceleration": (15, 0.5), "brushWidthBound": (3, 10), "nMovePointRatio": 0.5, "maxPiontMove": 3, "maxLineAcceleration": (5, 0.5), "boarderGap": 20, "maxInitSpeed": 6 } get_video_masks_by_moving_random_stroke(video_len=5, nStroke=3, **object_like_setting) ''' assert(video_len >= 1) # Initilize a set of control points to draw the first mask mask = Image.new(mode='1', size=(imageWidth, imageHeight), color=1) control_points_set = [] for i in range(nStroke): brushWidth = np.random.randint(brushWidthBound[0], brushWidthBound[1]) Xs, Ys, velocity = get_random_stroke_control_points( imageWidth=imageWidth, imageHeight=imageHeight, nVertexBound=nVertexBound, maxHeadSpeed=maxHeadSpeed, maxHeadAcceleration=maxHeadAcceleration, boarderGap=boarderGap, maxInitSpeed=maxInitSpeed ) control_points_set.append((Xs, Ys, velocity, brushWidth)) draw_mask_by_control_points(mask, Xs, Ys, brushWidth, fill=0) # Generate the following masks by randomly move strokes and their control points masks = [mask] for i in range(video_len - 1): mask = Image.new(mode='1', size=(imageWidth, imageHeight), color=1) for j in range(len(control_points_set)): Xs, Ys, velocity, brushWidth = control_points_set[j] new_Xs, new_Ys = random_move_control_points( Xs, Ys, velocity, nMovePointRatio, maxPiontMove, maxLineAcceleration, boarderGap ) control_points_set[j] = (new_Xs, new_Ys, velocity, brushWidth) for Xs, Ys, velocity, brushWidth in control_points_set: draw_mask_by_control_points(mask, Xs, Ys, brushWidth, fill=0) masks.append(mask) return masks def random_accelerate(velocity, maxAcceleration, dist='uniform'): speed, angle = velocity d_speed, d_angle = maxAcceleration if dist == 'uniform': speed += np.random.uniform(-d_speed, d_speed) angle += np.random.uniform(-d_angle, d_angle) elif dist == 'guassian': speed += np.random.normal(0, d_speed / 2) angle += np.random.normal(0, d_angle / 2) else: raise NotImplementedError(f'Distribution type {dist} is not supported.') return (speed, angle) def random_move_control_points(Xs, Ys, lineVelocity, nMovePointRatio, maxPiontMove, maxLineAcceleration, boarderGap=15): new_Xs = Xs.copy() new_Ys = Ys.copy() # move the whole line and accelerate speed, angle = lineVelocity new_Xs += int(speed * np.cos(angle)) new_Ys += int(speed * np.sin(angle)) lineVelocity = random_accelerate(lineVelocity, maxLineAcceleration, dist='guassian') # choose points to move chosen = np.arange(len(Xs)) np.random.shuffle(chosen) chosen = chosen[:int(len(Xs) * nMovePointRatio)] for i in chosen: new_Xs[i] += np.random.randint(-maxPiontMove, maxPiontMove) new_Ys[i] += np.random.randint(-maxPiontMove, maxPiontMove) return new_Xs, new_Ys def get_random_stroke_control_points( imageWidth, imageHeight, nVertexBound=(10, 30), maxHeadSpeed=10, maxHeadAcceleration=(5, 0.5), boarderGap=20, maxInitSpeed=10 ): ''' Implementation the free-form training masks generating algorithm proposed by JIAHUI YU et al. in "Free-Form Image Inpainting with Gated Convolution" ''' startX = np.random.randint(imageWidth) startY = np.random.randint(imageHeight) Xs = [startX] Ys = [startY] numVertex = np.random.randint(nVertexBound[0], nVertexBound[1]) angle = np.random.uniform(0, 2 * np.pi) speed = np.random.uniform(0, maxHeadSpeed) for i in range(numVertex): speed, angle = random_accelerate((speed, angle), maxHeadAcceleration) speed = np.clip(speed, 0, maxHeadSpeed) nextX = startX + speed * np.sin(angle) nextY = startY + speed * np.cos(angle) if boarderGap is not None: nextX = np.clip(nextX, boarderGap, imageWidth - boarderGap) nextY = np.clip(nextY, boarderGap, imageHeight - boarderGap) startX, startY = nextX, nextY Xs.append(nextX) Ys.append(nextY) velocity = get_random_velocity(maxInitSpeed, dist='guassian') return np.array(Xs), np.array(Ys), velocity def get_random_velocity(max_speed, dist='uniform'): if dist == 'uniform': speed = np.random.uniform(max_speed) elif dist == 'guassian': speed = np.abs(np.random.normal(0, max_speed / 2)) else: raise NotImplementedError(f'Distribution type {dist} is not supported.') angle = np.random.uniform(0, 2 * np.pi) return (speed, angle) def draw_mask_by_control_points(mask, Xs, Ys, brushWidth, fill=255): radius = brushWidth // 2 - 1 for i in range(1, len(Xs)): draw = ImageDraw.Draw(mask) startX, startY = Xs[i - 1], Ys[i - 1] nextX, nextY = Xs[i], Ys[i] draw.line((startX, startY) + (nextX, nextY), fill=fill, width=brushWidth) for x, y in zip(Xs, Ys): draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill=fill) return mask # modified from https://github.com/naoto0804/pytorch-inpainting-with-partial-conv/blob/master/generate_data.py def get_random_walk_mask(imageWidth=320, imageHeight=180, length=None): action_list = [[0, 1], [0, -1], [1, 0], [-1, 0]] canvas = np.zeros((imageHeight, imageWidth)).astype("i") if length is None: length = imageWidth * imageHeight x = random.randint(0, imageHeight - 1) y = random.randint(0, imageWidth - 1) x_list = [] y_list = [] for i in range(length): r = random.randint(0, len(action_list) - 1) x = np.clip(x + action_list[r][0], a_min=0, a_max=imageHeight - 1) y = np.clip(y + action_list[r][1], a_min=0, a_max=imageWidth - 1) x_list.append(x) y_list.append(y) canvas[np.array(x_list), np.array(y_list)] = 1 return Image.fromarray(canvas * 255).convert('1') def get_masked_ratio(mask): """ Calculate the masked ratio. mask: Expected a binary PIL image, where 0 and 1 represent masked(invalid) and valid pixel values. """ hist = mask.histogram() return hist[0] / np.prod(mask.size)