Spaces:
Build error
Build error
File size: 4,587 Bytes
edebe10 |
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 |
import random
import torch
from .mask_generators import get_mask_by_input_strokes
class Circle:
def __init__(self, cfg, is_train=True):
self.num_stroke = cfg['STROKE_SAMPLER']['CIRCLE']['NUM_STROKES']
self.stroke_preset = cfg['STROKE_SAMPLER']['CIRCLE']['STROKE_PRESET']
self.stroke_prob = cfg['STROKE_SAMPLER']['CIRCLE']['STROKE_PROB']
self.max_eval = cfg['STROKE_SAMPLER']['EVAL']['MAX_ITER']
self.is_train = is_train
@staticmethod
def get_stroke_preset(stroke_preset):
if stroke_preset == 'object_like':
return {
"nVertexBound": [5, 30],
"maxHeadSpeed": 15,
"maxHeadAcceleration": (10, 1.5),
"brushWidthBound": (20, 50),
"nMovePointRatio": 0.5,
"maxPiontMove": 10,
"maxLineAcceleration": (5, 0.5),
"boarderGap": None,
"maxInitSpeed": 10,
}
elif stroke_preset == 'object_like_middle':
return {
"nVertexBound": [5, 15],
"maxHeadSpeed": 8,
"maxHeadAcceleration": (4, 1.5),
"brushWidthBound": (20, 50),
"nMovePointRatio": 0.5,
"maxPiontMove": 5,
"maxLineAcceleration": (5, 0.5),
"boarderGap": None,
"maxInitSpeed": 10,
}
elif stroke_preset == 'object_like_small':
return {
"nVertexBound": [5, 20],
"maxHeadSpeed": 7,
"maxHeadAcceleration": (3.5, 1.5),
"brushWidthBound": (10, 30),
"nMovePointRatio": 0.5,
"maxPiontMove": 5,
"maxLineAcceleration": (3, 0.5),
"boarderGap": None,
"maxInitSpeed": 4,
}
else:
raise NotImplementedError(f'The stroke presetting "{stroke_preset}" does not exist.')
def get_random_points_from_mask(self, mask, n=5):
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
x = (non_zero_idx % w)*1.0
return torch.cat((x[:,None], y[:,None]), dim=1).numpy()
def draw(self, mask=None, box=None):
if mask.sum() < 10: # if mask is nearly empty
return torch.zeros(mask.shape).bool()
if not self.is_train:
return self.draw_eval(mask=mask, box=box)
stroke_preset_name = random.choices(self.stroke_preset, weights=self.stroke_prob, k=1)[0] # select which kind of object to use
preset = Circle.get_stroke_preset(stroke_preset_name)
nStroke = min(random.randint(1, self.num_stroke), mask.sum().item())
h,w = mask.shape
points = self.get_random_points_from_mask(mask, n=nStroke)
rand_mask = get_mask_by_input_strokes(
init_points=points,
imageWidth=w, imageHeight=h, nStroke=min(nStroke, len(points)), **preset)
rand_mask = (~torch.from_numpy(rand_mask)) * mask
return rand_mask
def draw_eval(self, mask=None, box=None):
stroke_preset_name = random.choices(self.stroke_preset, weights=self.stroke_prob, k=1)[0] # select which kind of object to use
preset = Circle.get_stroke_preset(stroke_preset_name)
nStroke = min(self.max_eval, mask.sum().item())
h,w = mask.shape
points = self.get_random_points_from_mask(mask, n=nStroke)
rand_masks = []
for i in range(len(points)):
rand_mask = get_mask_by_input_strokes(
init_points=points[:i+1],
imageWidth=w, imageHeight=h, nStroke=min(nStroke, len(points[:i+1])), **preset)
rand_masks += [(~torch.from_numpy(rand_mask)) * mask]
return torch.stack(rand_masks)
@staticmethod
def draw_by_points(points, mask, h, w):
stroke_preset_name = random.choices(['object_like', 'object_like_middle', 'object_like_small'], weights=[0.33,0.33,0.33], k=1)[0] # select which kind of object to use
preset = Circle.get_stroke_preset(stroke_preset_name)
rand_mask = get_mask_by_input_strokes(
init_points=points,
imageWidth=w, imageHeight=h, nStroke=len(points), **preset)[None,]
rand_masks = (~torch.from_numpy(rand_mask)) * mask
return rand_masks
def __repr__(self,):
return 'circle' |