File size: 5,213 Bytes
814a594 |
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 |
import random
import numpy as np
import torch
from scipy.special import binom
from scipy import ndimage
import matplotlib.pyplot as plt
from matplotlib.backends.backend_agg import FigureCanvasAgg
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
class Polygon:
def __init__(self, cfg, is_train):
self.max_points = cfg['STROKE_SAMPLER']['POLYGON']['MAX_POINTS']
self.eval_points = cfg['STROKE_SAMPLER']['EVAL']['MAX_ITER']
self.is_train = is_train
def get_random_points_from_mask(self, mask, n=3):
h,w = mask.shape
view_mask = mask.reshape(h*w)
non_zero_idx = view_mask.nonzero()[:,0]
selected_idx = torch.randperm(len(non_zero_idx))[:n]
non_zero_idx = non_zero_idx[selected_idx]
y = (non_zero_idx // w)*1.0/(h+1)
x = (non_zero_idx % w)*1.0/(w+1)
return torch.cat((x[:,None],y[:,None]), dim=1).numpy()
def draw(self, mask=None, box=None):
if mask.sum() < 10:
return torch.zeros(mask.shape).bool() # if mask is empty
if not self.is_train:
return self.draw_eval(mask=mask, box=box)
# box: x1,y1,x2,y2
x1,y1,x2,y2 = box.int().unbind()
rad = 0.2
edgy = 0.05
num_points = random.randint(1, min(self.max_points, mask.sum().item()))
a = self.get_random_points_from_mask(mask[y1:y2,x1:x2], n=num_points)
x,y, _ = get_bezier_curve(a,rad=rad, edgy=edgy)
x = x.clip(0.0, 1.0)
y = y.clip(0.0, 1.0)
points = torch.from_numpy(np.concatenate((y[None,]*(y2-y1-1).item(),x[None,]*(x2-x1-1).item()))).int()
canvas = torch.zeros((y2-y1, x2-x1))
canvas[points.long().tolist()] = 1
rand_mask = torch.zeros(mask.shape)
rand_mask[y1:y2,x1:x2] = canvas
return rand_mask.bool()
def draw_eval(self, mask=None, box=None):
# box: x1,y1,x2,y2
x1,y1,x2,y2 = box.int().unbind()
rad = 0.2
edgy = 0.05
num_points = min(self.eval_points, mask.sum().item())
a = self.get_random_points_from_mask(mask[y1:y2,x1:x2], n=num_points)
rand_masks = []
for i in range(len(a)):
x,y, _ = get_bezier_curve(a[:i+1],rad=rad, edgy=edgy)
x = x.clip(0.0, 1.0)
y = y.clip(0.0, 1.0)
points = torch.from_numpy(np.concatenate((y[None,]*(y2-y1-1).item(),x[None,]*(x2-x1-1).item()))).int()
canvas = torch.zeros((y2-y1, x2-x1))
canvas[points.long().tolist()] = 1
rand_mask = torch.zeros(mask.shape)
rand_mask[y1:y2,x1:x2] = canvas
struct = ndimage.generate_binary_structure(2, 2)
rand_mask = torch.from_numpy((ndimage.binary_dilation(rand_mask, structure=struct, iterations=5).astype(rand_mask.numpy().dtype)))
rand_masks += [rand_mask.bool()]
return torch.stack(rand_masks)
def __repr__(self,):
return 'polygon' |