Diff of /utils/mask_generator.py [000000] .. [92cc18]

Switch to side-by-side view

--- 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()