def _composite_fg(self, fg, alpha, idx):
if np.random.rand() < 0.5:
idx2 = np.random.randint(self.fg_num_unique) + idx
idx2 = idx2 % self.fg_num_unique
fg2 = get_raw("fg", idx2)
alpha2 = get_raw("a", idx2)
alpha2 = np.reshape(alpha2, (alpha2.shape[0], alpha2.shape[1]))
alpha2 = alpha2.astype(np.float32) / 255.0
h, w = alpha.shape
fg2 = cv.resize(fg2, (w, h), interpolation=maybe_random_interp(cv.INTER_NEAREST))
alpha2 = cv.resize(alpha2, (w, h), interpolation=maybe_random_interp(cv.INTER_NEAREST))
alpha_tmp = 1 - (1 - alpha) * (1 - alpha2)
if np.any(alpha_tmp < 1):
fg = fg.astype(np.float32) * alpha[:,:,None] + fg2.astype(np.float32) * (1 - alpha[:,:,None])
# The overlap of two 50% transparency should be 25%
alpha = alpha_tmp
fg = fg.astype(np.uint8)
if np.random.rand() < 0.25:
fg = cv.resize(fg, (640, 640), interpolation=maybe_random_interp(cv.INTER_NEAREST))
alpha = cv.resize(alpha, (640, 640), interpolation=maybe_random_interp(cv.INTER_NEAREST))
return fg, alpha
def _composite_fg(self, img, alpha, fg, bg, idx):
if np.random.rand() < 0.5:
idx2 = np.random.randint(self.fg_num) + idx
img2, alpha2, fg2, bg2 = process(idx2 % self.fg_num, 0)
h, w = alpha.shape
fg2 = cv.resize(fg2, (w, h),
interpolation=maybe_random_interp(
cv.INTER_NEAREST))
alpha2 = cv.resize(alpha2, (w, h),
interpolation=maybe_random_interp(
cv.INTER_NEAREST))
alpha_tmp = 1 - (1 - alpha / 255.0) * (1 - alpha2 / 255.0)
if np.any(alpha_tmp < 1):
img, alpha, _, _ = composite4(fg, fg2, alpha, w, h)
alpha = alpha_tmp * 255.0
img = img.astype(np.uint8)
img, alpha, fg, bg = composite4(img, bg, alpha, w, h)
if np.random.rand() < 0.25:
img = cv.resize(img, (640, 640),
interpolation=maybe_random_interp(
cv.INTER_NEAREST))
alpha = cv.resize(alpha, (640, 640),
interpolation=maybe_random_interp(
cv.INTER_NEAREST))
return img, alpha
def _composite_fg(img, alpha, fg, bg, idx):
idx2 = 3
img2, alpha2, fg2, bg2 = process(idx2, 1)
h, w = alpha.shape
fg2 = cv.resize(fg2, (w, h),
interpolation=maybe_random_interp(cv.INTER_NEAREST))
alpha2 = cv.resize(alpha2, (w, h),
interpolation=maybe_random_interp(cv.INTER_NEAREST))
cv.imshow("a", fg[:, :, ::-1])
cv.waitKey(0)
cv.destroyAllWindows()
cv.imshow("a", fg2[:, :, ::-1])
cv.waitKey(0)
cv.destroyAllWindows()
alpha_tmp = 1 - (1 - alpha) * (1 - alpha2)
if np.any(alpha_tmp < 1):
fg = fg.astype(np.float32) * alpha[:, :, None] + fg2.astype(
np.float32) * (1 - alpha[:, :, None])
# The overlap of two 50% transparency should be 25%
# alpha = alpha_tmp
fg = fg.astype(np.uint8)
cv.imshow("a", fg[:, :, ::-1])
cv.waitKey(0)
cv.destroyAllWindows()
img, alpha, fg, bg = composite4(fg, bg, alpha, w, h)
if np.random.rand() < 0.25:
fg = cv.resize(fg, (640, 640),
interpolation=maybe_random_interp(cv.INTER_NEAREST))
alpha = cv.resize(alpha, (640, 640),
interpolation=maybe_random_interp(cv.INTER_NEAREST))
return img, alpha
def _composite_fg(alpha, fg, idx):
idx2 = 15
alpha2 = get_raw("a", idx2)
alpha2 = np.reshape(alpha2, (alpha2.shape[0], alpha2.shape[1]))
fg2 = get_raw("fg", idx2)
cv.imshow("fg2", fg2[:, :, ::-1].astype(np.uint8))
cv.waitKey(0)
cv.destroyAllWindows()
h, w = alpha.shape
fg2 = cv.resize(fg2, (w, h),
interpolation=maybe_random_interp(cv.INTER_NEAREST))
alpha2 = cv.resize(alpha2, (w, h),
interpolation=maybe_random_interp(cv.INTER_NEAREST))
alpha_tmp = 1 - (1 - alpha / 255.0) * (1 - alpha2 / 255.0)
if np.any(alpha_tmp < 1):
img, alpha, _, _ = composite4(fg, fg2, alpha, w, h)
alpha = alpha_tmp * 255.0
return img, alpha
def __getitem__(self, i):
fcount = self.fgs[i]
if i % args.batch_size == 0:
self.current_index = fcount
alpha = get_raw("a", fcount)
alpha = np.reshape(alpha, (alpha.shape[0], alpha.shape[1]))
fg = get_raw("fg", fcount)
self.current_fg = fg
self.current_alpha = alpha
self.is_resize = True if np.random.rand() < 0.25 else False
else:
fg = self.current_fg
alpha = self.current_alpha
bcount = np.random.randint(num_bgs)
img, _, _, bg = process(fcount, bcount)
if self.is_resize:
interpolation = maybe_random_interp(cv.INTER_NEAREST)
img = cv.resize(img, (640, 640), interpolation=interpolation)
# fg = cv.resize(fg, (640, 640), interpolation=interpolation)
alpha = cv.resize(alpha, (640, 640), interpolation=interpolation)
# bg = cv.resize(bg, (640, 640), interpolation=interpolation)
# crop size 320:640:480 = 1:1:1
different_sizes = [(320, 320), (480, 480), (640, 640)]
crop_size = random.choice(different_sizes)
trimap = gen_trimap(alpha)
x, y = random_choice(trimap, crop_size)
img = safe_crop(img, x, y, crop_size)
alpha = safe_crop(alpha, x, y, crop_size)
trimap = gen_trimap(alpha)
# Flip array left to right randomly (prob=1:1)
if np.random.random_sample() > 0.5:
img = np.fliplr(img)
trimap = np.fliplr(trimap)
alpha = np.fliplr(alpha)
x = torch.zeros((4, im_size, im_size), dtype=torch.float)
img = transforms.ToPILImage()(img)
img = self.transformer(img)
x[0:3, :, :] = img
x[3, :, :] = torch.from_numpy(trimap.copy() / 255.)
y = np.empty((2, im_size, im_size), dtype=np.float32)
y[0, :, :] = alpha / 255.
mask = np.equal(trimap, 128).astype(np.float32)
y[1, :, :] = mask
return x, y
def __call__(self, sample):
fg = sample['fg']
alpha = sample['alpha']
# fg, alpha = sample['fg'], sample['alpha']
rows, cols, ch = fg.shape
if np.maximum(rows, cols) < 1024:
params = self.get_params((0, 0), self.translate, self.scale, self.shear, self.flip, fg.size)
else:
params = self.get_params(self.degrees, self.translate, self.scale, self.shear, self.flip, fg.size)
center = (cols * 0.5 + 0.5, rows * 0.5 + 0.5)
M = self._get_inverse_affine_matrix(center, *params)
M = np.array(M).reshape((2, 3))
fg = cv.warpAffine(fg, M, (cols, rows),
flags=maybe_random_interp(cv.INTER_NEAREST) + cv.WARP_INVERSE_MAP)
alpha = cv.warpAffine(alpha, M, (cols, rows),
flags=maybe_random_interp(cv.INTER_NEAREST) + cv.WARP_INVERSE_MAP)
sample['fg'], sample['alpha'] = fg, alpha
return sample
def __call__(self, sample):
crop_size = sample['size']
self.output_size = crop_size
self.margin = self.output_size[0] // 2
fg, alpha, trimap = sample['fg'], sample['alpha'], sample['trimap']
bg = sample['bg']
h, w = trimap.shape
bg = cv.resize(bg, (w, h), interpolation=maybe_random_interp(cv.INTER_CUBIC))
if w < self.output_size[0]+1 or h < self.output_size[1]+1:
ratio = 1.1*self.output_size[0]/h if h < w else 1.1*self.output_size[1]/w
# self.logger.warning("Size of {} is {}.".format(name, (h, w)))
while h < self.output_size[0]+1 or w < self.output_size[1]+1:
fg = cv.resize(fg, (int(w*ratio), int(h*ratio)), interpolation=maybe_random_interp(cv.INTER_NEAREST))
alpha = cv.resize(alpha, (int(w*ratio), int(h*ratio)),
interpolation=maybe_random_interp(cv.INTER_NEAREST))
trimap = cv.resize(trimap, (int(w*ratio), int(h*ratio)), interpolation=cv.INTER_NEAREST)
bg = cv.resize(bg, (int(w*ratio), int(h*ratio)), interpolation=maybe_random_interp(cv.INTER_CUBIC))
h, w = trimap.shape
small_trimap = cv.resize(trimap, (w//4, h//4), interpolation=cv.INTER_NEAREST)
unknown_list = list(zip(*np.where(small_trimap[self.margin//4:(h-self.margin)//4,
self.margin//4:(w-self.margin)//4] == 128)))
unknown_num = len(unknown_list)
if len(unknown_list) < 10:
# self.logger.warning("{} does not have enough unknown area for crop.".format(name))
left_top = (np.random.randint(0, h-self.output_size[0]+1), np.random.randint(0, w-self.output_size[1]+1))
else:
idx = np.random.randint(unknown_num)
left_top = (unknown_list[idx][0]*4, unknown_list[idx][1]*4)
fg_crop = fg[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1],:]
alpha_crop = alpha[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1]]
bg_crop = bg[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1],:]
trimap_crop = trimap[left_top[0]:left_top[0]+self.output_size[0], left_top[1]:left_top[1]+self.output_size[1]]
if len(np.where(trimap==128)[0]) == 0:
self.logger.error("Does not have enough unknown area for crop. Resized to target size."
"left_top: {}".format(left_top))
fg_crop = cv.resize(fg, self.output_size[::-1], interpolation=maybe_random_interp(cv.INTER_NEAREST))
alpha_crop = cv.resize(alpha, self.output_size[::-1], interpolation=maybe_random_interp(cv.INTER_NEAREST))
trimap_crop = cv.resize(trimap, self.output_size[::-1], interpolation=cv.INTER_NEAREST)
bg_crop = cv.resize(bg, self.output_size[::-1], interpolation=maybe_random_interp(cv.INTER_CUBIC))
# cv.imwrite('../tmp/tmp.jpg', fg.astype(np.uint8))
# cv.imwrite('../tmp/tmp.png', (alpha*255).astype(np.uint8))
# cv.imwrite('../tmp/tmp2.png', trimap.astype(np.uint8))
# raise ValueError("{} does not have enough unknown area for crop.".format(name))
sample['fg'], sample['alpha'], sample['trimap'] = fg_crop.copy(), alpha_crop.copy(), trimap_crop.copy()
sample['bg'] = bg_crop.copy()
return sample
np.random.shuffle(names)
split_index = math.ceil(num_fgs - num_fgs * valid_ratio)
names_train = names[:split_index]
print(len(names_train))
names_valid = names[split_index:]
return names_train, names_valid
if __name__ == "__main__":
img, alpha, fg, bg = process(2, 19)
h, w = alpha.shape
cv.imshow("fg", fg[:, :, ::-1].astype(np.uint8))
cv.waitKey(0)
cv.destroyAllWindows()
img, alpha = _composite_fg(alpha, img, 2)
cv.imshow("fg", fg[:, :, ::-1].astype(np.uint8))
cv.waitKey(0)
cv.destroyAllWindows()
cv.imshow("combination", img[:, :, ::-1].astype(np.uint8))
cv.waitKey(0)
cv.destroyAllWindows()
img, alpha, fg, bg = composite4(img, bg, alpha, w, h)
interpolation = maybe_random_interp(cv.INTER_NEAREST)
img = cv.resize(img, (640, 640), interpolation=interpolation)
fg = cv.resize(fg, (640, 640), interpolation=interpolation)
alpha = cv.resize(alpha, (640, 640), interpolation=interpolation)
bg = cv.resize(bg, (640, 640), interpolation=interpolation)
cv.imshow("with background", img[:, :, ::-1].astype(np.uint8))
cv.waitKey(0)
cv.destroyAllWindows()