Spaces:
Sleeping
Sleeping
File size: 7,265 Bytes
c429825 |
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 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 |
import numpy as np
from PIL import Image, ImageDraw
import math
import random
import torch
#import tensorflow as tf
np.random.seed(10)
def random_sq_bbox(img, mask_shape, image_size=256, margin=(16, 16)):
"""Generate a random sqaure mask for inpainting
"""
B, H, W, C = img.shape
h, w = mask_shape
margin_height, margin_width = margin
maxt = image_size - margin_height - h
maxl = image_size - margin_width - w
# bb
t = np.random.randint(margin_height, maxt)
l = np.random.randint(margin_width, maxl)
# make mask
mask = torch.ones([B, C, H, W], device=img.device)
mask[..., t:t+h, l:l+w] = 0
mask = 1 - mask
#Fixed mid box
#mask[..., t:t+h, l:l+w] = 0
return mask, t, t+h, l, l+w
def RandomBrush(
max_tries,
s,
min_num_vertex = 4,
max_num_vertex = 18,
mean_angle = 2*math.pi / 5,
angle_range = 2*math.pi / 15,
min_width = 12,
max_width = 48):
H, W = s, s
average_radius = math.sqrt(H*H+W*W) / 8
mask = Image.new('L', (W, H), 0)
for _ in range(np.random.randint(max_tries)):
num_vertex = np.random.randint(min_num_vertex, max_num_vertex)
angle_min = mean_angle - np.random.uniform(0, angle_range)
angle_max = mean_angle + np.random.uniform(0, angle_range)
angles = []
vertex = []
for i in range(num_vertex):
if i % 2 == 0:
angles.append(2*math.pi - np.random.uniform(angle_min, angle_max))
else:
angles.append(np.random.uniform(angle_min, angle_max))
h, w = mask.size
vertex.append((int(np.random.randint(0, w)), int(np.random.randint(0, h))))
for i in range(num_vertex):
r = np.clip(
np.random.normal(loc=average_radius, scale=average_radius//2),
0, 2*average_radius)
new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w)
new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h)
vertex.append((int(new_x), int(new_y)))
draw = ImageDraw.Draw(mask)
width = int(np.random.uniform(min_width, max_width))
draw.line(vertex, fill=1, width=width)
for v in vertex:
draw.ellipse((v[0] - width//2,
v[1] - width//2,
v[0] + width//2,
v[1] + width//2),
fill=1)
if np.random.random() > 0.5:
mask.transpose(Image.FLIP_LEFT_RIGHT)
if np.random.random() > 0.5:
mask.transpose(Image.FLIP_TOP_BOTTOM)
mask = np.asarray(mask, np.uint8)
if np.random.random() > 0.5:
mask = np.flip(mask, 0)
if np.random.random() > 0.5:
mask = np.flip(mask, 1)
return mask
def RandomMask(s, hole_range=[0,1]):
coef = min(hole_range[0] + hole_range[1], 1.0)
while True:
mask = np.ones((s, s), np.uint8)
def Fill(max_size):
w, h = np.random.randint(max_size), np.random.randint(max_size)
ww, hh = w // 2, h // 2
x, y = np.random.randint(-ww, s - w + ww), np.random.randint(-hh, s - h + hh)
mask[max(y, 0): min(y + h, s), max(x, 0): min(x + w, s)] = 0
def MultiFill(max_tries, max_size):
for _ in range(np.random.randint(max_tries)):
Fill(max_size)
MultiFill(int(10 * coef), s // 2)
MultiFill(int(5 * coef), s)
##comment the following line for lower masking ratios
#mask = np.logical_and(mask, 1 - RandomBrush(int(20 * coef), s))
hole_ratio = 1 - np.mean(mask)
if hole_range is not None and (hole_ratio <= hole_range[0] or hole_ratio >= hole_range[1]):
continue
return mask[np.newaxis, ...].astype(np.float32)
def BatchRandomMask(batch_size, s, hole_range=[0, 1]):
return np.stack([RandomMask(s, hole_range=hole_range) for _ in range(batch_size)], axis = 0)
def random_rotation(shape):
cutoff = 100 #was 30
(n , channels, p, q) = shape
mask = np.zeros((n,p,q))
for i in range(n):
angle = np.random.choice(360, 1)
mask_one = np.ones((p+cutoff,q+cutoff))
mask_one[int((p+cutoff)/2):,:] = 0
im = Image.fromarray(mask_one)
im = im.rotate(angle)
left = (p+cutoff - p)/2
top = (q+cutoff - q)/2
right = (p+cutoff + p)/2
bottom = (q+cutoff + q)/2
# Crop the center of the image
im = im.crop((left, top, right, bottom))
mask[i] = np.array(im)
#mask = np.repeat(mask.reshape([n,1,p,q]), channels, axis=1)
mask = mask.reshape([n,1,p,q])
return mask
class mask_generator:
def __init__(self, mask_type, mask_len_range=None, mask_prob_range=None,
image_size=256, margin=(16, 16)):
"""
(mask_len_range): given in (min, max) tuple.
Specifies the range of box size in each dimension
(mask_prob_range): for the case of random masking,
specify the probability of individual pixels being masked
"""
assert mask_type in ['box', 'random', 'half', 'extreme']
self.mask_type = mask_type
self.mask_len_range = mask_len_range
self.mask_prob_range = mask_prob_range
self.image_size = image_size
self.margin = margin
def _retrieve_box(self, img):
l, h = self.mask_len_range
l, h = int(l), int(h)
mask_h = np.random.randint(l, h)
mask_w = np.random.randint(l, h)
mask, t, tl, w, wh = random_sq_bbox(img,
mask_shape=(mask_h, mask_w),
image_size=self.image_size,
margin=self.margin)
return mask, t, tl, w, wh
def generate_center_mask(self, shape):
assert len(shape) == 2
assert shape[1] % 2 == 0
center = shape[0] // 2
center_size = shape[0] // 4
half_resol = center_size // 2 # for now
ret = torch.zeros(shape, dtype=torch.float32)
ret[
center - half_resol: center + half_resol,
center - half_resol: center + half_resol,
] = 1
ret = ret.unsqueeze(0).unsqueeze(0)
return ret
def __call__(self, img):
if self.mask_type == 'random':
mask = BatchRandomMask(1, self.image_size, hole_range=self.mask_prob_range) #self._retrieve_random(img)
return mask
elif self.mask_type == "half":
mask = random_rotation((1, 3, self.image_size, self.image_size))
elif self.mask_type == 'box':
#mask, t, th, w, wl = self._retrieve_box(img)
mask = self.generate_center_mask((self.image_size,self.image_size)) # self._retrieve_box(img)
return mask #.permute(0,3,1,2)
elif self.mask_type == 'extreme':
mask, t, th, w, wl = self._retrieve_box(img)
mask = 1. - mask
return mask
'''
def tf_mask_generator(s, tf_hole_range):
def random_mask_generator(hole_range):
while True:
yield RandomMask(s, hole_range=hole_range)
return tf.data.Dataset.from_generator(random_mask_generator, tf.float32, tf.TensorShape([1, s, s]), (tf_hole_range,))
''' |