--- a +++ b/utils/mask_generator.py @@ -0,0 +1,122 @@ +from random import randint + +import cv2 as cv +import matplotlib.pyplot as plt +import numpy as np +from scipy.special import binom + +bernstein = lambda n, k, t: binom(n, k) * t ** k * (1. - t) ** (n - k) + + +def bezier(points, num=200): + N = len(points) + t = np.linspace(0, 1, num=num) + curve = np.zeros((num, 2)) + for i in range(N): + curve += np.outer(bernstein(N - 1, i, t), points[i]) + return curve + + +class Segment(): + def __init__(self, p1, p2, angle1, angle2, **kw): + self.p1 = p1; + self.p2 = p2 + self.angle1 = angle1; + self.angle2 = angle2 + self.numpoints = kw.get("numpoints", 100) + r = kw.get("r", 0.3) + d = np.sqrt(np.sum((self.p2 - self.p1) ** 2)) + self.r = r * d + self.p = np.zeros((4, 2)) + self.p[0, :] = self.p1[:] + self.p[3, :] = self.p2[:] + self.calc_intermediate_points(self.r) + + def calc_intermediate_points(self, r): + self.p[1, :] = self.p1 + np.array([self.r * np.cos(self.angle1), + self.r * np.sin(self.angle1)]) + self.p[2, :] = self.p2 + np.array([self.r * np.cos(self.angle2 + np.pi), + self.r * np.sin(self.angle2 + np.pi)]) + self.curve = bezier(self.p, self.numpoints) + + +def get_curve(points, **kw): + segments = [] + for i in range(len(points) - 1): + seg = Segment(points[i, :2], points[i + 1, :2], points[i, 2], points[i + 1, 2], **kw) + segments.append(seg) + curve = np.concatenate([s.curve for s in segments]) + return segments, curve + + +def ccw_sort(p): + d = p - np.mean(p, axis=0) + s = np.arctan2(d[:, 0], d[:, 1]) + return p[np.argsort(s), :] + + +def get_bezier_curve(a, rad=0.2, edgy=0): + """ given an array of points *a*, create a curve through + those points. + *rad* is a number between 0 and 1 to steer the distance of + control points. + *edgy* is a parameter which controls how "edgy" the curve is, + edgy=0 is smoothest.""" + p = np.arctan(edgy) / np.pi + .5 + a = ccw_sort(a) + a = np.append(a, np.atleast_2d(a[0, :]), axis=0) + d = np.diff(a, axis=0) + ang = np.arctan2(d[:, 1], d[:, 0]) + f = lambda ang: (ang >= 0) * ang + (ang < 0) * (ang + 2 * np.pi) + ang = f(ang) + ang1 = ang + ang2 = np.roll(ang, 1) + ang = p * ang1 + (1 - p) * ang2 + (np.abs(ang2 - ang1) > np.pi) * np.pi + ang = np.append(ang, [ang[0]]) + a = np.append(a, np.atleast_2d(ang).T, axis=1) + s, c = get_curve(a, r=rad, method="var") + x, y = c.T + return x, y, a + + +def get_random_points(n=5, scale=0.8, mindst=None, rec=0): + """ create n random points in the unit square, which are *mindst* + apart, then scale them.""" + + mindst = mindst or .7 / n + a = np.random.rand(n, 2) + d = np.sqrt(np.sum(np.diff(ccw_sort(a), axis=0), axis=1) ** 2) + if np.all(d >= mindst) or rec >= 200: + return a * scale + else: + return get_random_points(n=n, scale=scale, mindst=mindst, rec=rec + 1) + + +# ======================================= +# developer: Vajria +# ======================================= + +def generate_a_mask(rad=0.25, edgy=0.05, imsize=512, location=None, scale=None): + # for c in np.array([[0,0], [0,1], [1,0], [1,1]]): + if location is None: + location = [randint(0, np.floor(imsize / 1.5)), + randint(0, np.floor(imsize / 1.5))] # location region to generate a mask + else: + location = location + a = get_random_points(n=5, scale=128) + location + x, y, _ = get_bezier_curve(a, rad=rad, edgy=edgy) + + c = np.vstack((x, y)).T # make a (x, y) coordinate paires + + c = c.astype(int) # covert into integer + img = np.zeros((imsize, imsize, 3), np.uint8) # make a image + pts = c.reshape((-1, 1, 2)) + img = cv.fillPoly(img, [pts], (255, 255, 255)) # (255,255, 255) = color (White) + + return img[:, :, 0] / 255 + + +if __name__ == '__main__': + mask = generate_a_mask() + plt.imshow(mask) + plt.show()