def test_txt_encoder_coco(config):
db = composites_coco(config, 'train', '2017')
all_tables = AllCategoriesTables(db)
all_tables.build_nntables_for_all_categories(True)
sequence_db = sequence_loader(db, all_tables)
net = TextEncoder(db)
loader = DataLoader(sequence_db,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.num_workers)
for cnt, batched in enumerate(loader):
input_inds = batched['word_inds'].long()
input_lens = batched['word_lens'].long()
print('Checking the output shapes')
out = net(input_inds, input_lens)
out_embs, out_rfts, out_msks, out_hids = out
print(out_rfts.size(), out_embs.size(), out_msks.size())
if isinstance(out_hids, tuple):
print(out_hids[0].size())
else:
print(out_hids.size())
print('m: ', out_msks[-1])
print('Checking the embedding')
embeded = net.embedding(input_inds)
v1 = embeded[0, 0]
idx = input_inds[0, 0].data.item()
v2 = db.lang_vocab.vectors[idx]
diff = v2 - v1
print('Diff: (should be zero)', torch.sum(diff.abs_()))
break
def test_coco_decoder(config):
db = composites_coco(config, 'train', '2017')
all_tables = AllCategoriesTables(db)
all_tables.build_nntables_for_all_categories(True)
sequence_db = sequence_loader(db, all_tables)
text_encoder = TextEncoder(db)
img_encoder = VolumeEncoder(config)
what_decoder = WhatDecoder(config)
where_decoder = WhereDecoder(config)
print('txt_encoder', get_n_params(text_encoder))
print('img_encoder', get_n_params(img_encoder))
print('what_decoder', get_n_params(what_decoder))
print('where_decoder', get_n_params(where_decoder))
loader = DataLoader(sequence_db,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.num_workers)
for cnt, batched in enumerate(loader):
word_inds = batched['word_inds'].long()
word_lens = batched['word_lens'].long()
bg_imgs = batched['background'].float()
encoder_states = text_encoder(word_inds, word_lens)
bg_feats = img_encoder(bg_imgs)
prev_bgfs = bg_feats[:, 0].unsqueeze(1)
what_outs = what_decoder((prev_bgfs, None, None), encoder_states)
obj_logits, rnn_feats_2d, nxt_hids_2d, prev_bgfs, att_ctx, att_wei = what_outs
print('------------------------------------------')
print('obj_logits', obj_logits.size())
print('rnn_feats_2d', rnn_feats_2d.size())
print('nxt_hids_2d', nxt_hids_2d.size())
print('prev_bgfs', prev_bgfs.size())
print('att_ctx', att_ctx.size())
print('att_wei', att_wei.size())
print('------------------------------------------')
_, obj_inds = torch.max(obj_logits + 1.0, dim=-1)
curr_fgfs = indices2onehots(obj_inds.cpu().data,
config.output_vocab_size)
# curr_fgfs = curr_fgfs.unsqueeze(1)
if config.cuda:
curr_fgfs = curr_fgfs.cuda()
where_outs = where_decoder(
(rnn_feats_2d, curr_fgfs, prev_bgfs, att_ctx), encoder_states)
coord_logits, attri_logits, patch_vectors, where_ctx, where_wei = where_outs
print('coord_logits ', coord_logits.size())
print('attri_logits ', attri_logits.size())
print('patch_vectors', patch_vectors.size())
# print('att_ctx', where_ctx.size())
# print('att_wei', where_wei.size())
break
def validate_epoch(self, val_db, epoch):
if self.cfg.cuda and self.cfg.parallel:
net = self.net.module
else:
net = self.net
all_losses, all_accuracies = [], []
for G in range(5):
val_db.cfg.sent_group = G
seq_db = sequence_loader(val_db)
val_loader = DataLoader(seq_db,
batch_size=self.cfg.batch_size, shuffle=False,
num_workers=self.cfg.num_workers, pin_memory=True)
for cnt, batched in enumerate(val_loader):
##################################################################
## Batched data
##################################################################
input_inds, input_lens, fg_onehots, bg_imgs, \
fg_imgs, neg_imgs, fg_resnets, neg_resnets,\
gt_inds, gt_msks, patch_inds = \
self.batch_data(batched)
##################################################################
## Validate one step
##################################################################
self.net.eval()
with torch.no_grad():
inputs = (input_inds, input_lens, bg_imgs, fg_onehots, fg_imgs, neg_imgs, fg_resnets, neg_resnets)
inf_outs, _, pos_vecs, neg_vecs = self.net(inputs)
pred_loss, embed_loss, attn_loss, pred_accu = self.evaluate(inf_outs, pos_vecs, neg_vecs, gt_inds, gt_msks)
loss = pred_loss + embed_loss + attn_loss
all_losses.append(np.array(
[
loss.cpu().data.item(),
pred_loss.cpu().data.item(),
embed_loss.cpu().data.item(),
attn_loss.cpu().data.item()
]))
all_accuracies.append(pred_accu.cpu().data.numpy())
print(epoch, G, cnt)
all_losses = np.stack(all_losses, 0)
all_accuracies = np.stack(all_accuracies, 0)
return all_losses, all_accuracies
def test_vol_encoder(config):
db = composites_coco(config, 'train', '2017')
all_tables = AllCategoriesTables(db)
all_tables.build_nntables_for_all_categories(True)
sequence_db = sequence_loader(db, all_tables)
img_encoder = VolumeEncoder(config)
print(get_n_params(img_encoder))
# print(img_encoder)
loader = DataLoader(sequence_db,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.num_workers)
for cnt, batched in enumerate(loader):
x = batched['background'].float()
y = img_encoder(x)
print('y.size()', y.size())
break
def test_shape_encoder(config):
db = coco(config, 'train', '2017')
all_tables = AllCategoriesTables(db)
all_tables.build_nntables_for_all_categories(True)
sequence_db = sequence_loader(db, all_tables)
img_encoder = ShapeEncoder(config)
print(get_n_params(img_encoder))
loader = DataLoader(sequence_db,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.num_workers)
for cnt, batched in enumerate(loader):
x = batched['foreground'].float()
y = batched['foreground_resnets'].float()
y = img_encoder(x, y)
print('y.size()', y.size())
print('y max', torch.max(y))
print('y min', torch.min(y))
print('y norm', torch.norm(y, dim=-1)[0, 0])
break
def test_puzzle_model(config):
output_dir = osp.join(config.model_dir, 'test_puzzle_model')
maybe_create(output_dir)
plt.switch_backend('agg')
db = composites_coco(config, 'train', '2017')
all_tables = AllCategoriesTables(db)
all_tables.build_nntables_for_all_categories(True)
sequence_db = sequence_loader(db, all_tables)
loader = DataLoader(sequence_db,
batch_size=config.batch_size,
shuffle=False,
num_workers=config.num_workers)
net = PuzzleModel(db)
net.eval()
for cnt, batched in enumerate(loader):
word_inds = batched['word_inds'].long()
word_lens = batched['word_lens'].long()
bg_images = batched['background'].float()
fg_images = batched['foreground'].float()
neg_images = batched['negative'].float()
fg_resnets = batched['foreground_resnets'].float()
neg_resnets = batched['negative_resnets'].float()
fg_inds = batched['fg_inds'].long()
gt_inds = batched['out_inds'].long()
gt_msks = batched['out_msks'].float()
fg_onehots = indices2onehots(fg_inds, config.output_vocab_size)
inf_outs, _, positive_feats, negative_feats = net(
(word_inds, word_lens, bg_images, fg_onehots, fg_images,
neg_images, fg_resnets, neg_resnets))
obj_logits, coord_logits, attri_logits, patch_vectors, enc_msks, what_wei, where_wei = inf_outs
print('teacher forcing')
print('obj_logits ', obj_logits.size())
print('coord_logits ', coord_logits.size())
print('attri_logits ', attri_logits.size())
print('patch_vectors ', patch_vectors.size())
print('patch_vectors max:', torch.max(patch_vectors))
print('patch_vectors min:', torch.min(patch_vectors))
print('patch_vectors norm:',
torch.norm(patch_vectors, dim=-2)[0, 0, 0])
print('positive_feats ', positive_feats.size())
print('negative_feats ', negative_feats.size())
if config.what_attn:
print('what_att_logits ', what_wei.size())
if config.where_attn > 0:
print('where_att_logits ', where_wei.size())
print('----------------------')
_, pred_vecs = net.collect_logits_and_vectors(inf_outs, gt_inds)
print('pred_vecs', pred_vecs.size())
print('*******************')
# # inf_outs, env = net.inference(word_inds, word_lens, -1, 0.0, 0, gt_inds, gt_vecs)
# inf_outs, env = net.inference(word_inds, word_lens, -1, 2.0, 0, None, None, all_tables)
# obj_logits, coord_logits, attri_logits, patch_vectors, enc_msks, what_wei, where_wei = inf_outs
# print('scheduled sampling')
# print('obj_logits ', obj_logits.size())
# print('coord_logits ', coord_logits.size())
# print('attri_logits ', attri_logits.size())
# print('patch_vectors ', patch_vectors.size())
# if config.what_attn:
# print('what_att_logits ', what_wei.size())
# if config.where_attn > 0:
# print('where_att_logits ', where_wei.size())
# print('----------------------')
# sequences = env.batch_redraw(True)
# for i in range(len(sequences)):
# sequence = sequences[i]
# image_idx = batched['image_index'][i]
# name = '%03d_'%i + str(image_idx).zfill(12)
# out_path = osp.join(output_dir, name+'.png')
# color = cv2.imread(batched['color_path'][i], cv2.IMREAD_COLOR)
# color, _, _ = create_squared_image(color)
# fig = plt.figure(figsize=(32, 16))
# plt.suptitle(batched['sentence'][i], fontsize=30)
# for j in range(min(len(sequence), 14)):
# plt.subplot(3, 5, j+1)
# partially_completed_img = clamp_array(sequence[j][:,:,-3:], 0, 255).astype(np.uint8)
# partially_completed_img = partially_completed_img[:,:,::-1]
# plt.imshow(partially_completed_img)
# plt.axis('off')
# plt.subplot(3, 5, 15)
# plt.imshow(color[:,:,::-1])
# plt.axis('off')
# fig.savefig(out_path, bbox_inches='tight')
# plt.close(fig)
break
def sample_for_eval(self, test_db, nn_table=None):
##############################################################
# Output prefix
##############################################################
# gt_dir = osp.join(self.cfg.model_dir, 'gt')
# frame_dir = osp.join(self.cfg.model_dir, 'proposal_images')
# noice_dir = osp.join(self.cfg.model_dir, 'proposal_noices')
# label_dir = osp.join(self.cfg.model_dir, 'proposal_labels')
# mask_dir = osp.join(self.cfg.model_dir, 'proposal_masks')
# info_dir = osp.join(self.cfg.model_dir, 'proposal_info')
main_dir = 'puzzle_results'
maybe_create(main_dir)
gt_dir = osp.join(main_dir, 'gt')
frame_dir = osp.join(main_dir, 'proposal_images')
noice_dir = osp.join(main_dir, 'proposal_noices')
label_dir = osp.join(main_dir, 'proposal_labels')
mask_dir = osp.join(main_dir, 'proposal_masks')
info_dir = osp.join(main_dir, 'proposal_info')
maybe_create(gt_dir)
maybe_create(frame_dir)
maybe_create(noice_dir)
maybe_create(label_dir)
maybe_create(mask_dir)
maybe_create(info_dir)
seq_db = sequence_loader(test_db)
##############################################################
# Main loop
##############################################################
if self.cfg.cuda and self.cfg.parallel:
net = self.net.module
else:
net = self.net
# start_ind = 0
# end_ind = len(seq_db)
start_ind = self.cfg.seed * 1250
end_ind = (self.cfg.seed+1) * 1250
# start_ind = 35490
# end_ind = len(seq_db)
for i in range(start_ind, end_ind):
entry = seq_db[i]
gt_scene = test_db.scenedb[i]
image_index = gt_scene['image_index']
image_path = test_db.color_path_from_index(image_index)
name = osp.splitext(osp.basename(image_path))[0]
gt_path = osp.join(gt_dir, osp.basename(image_path))
# save gt
shutil.copy2(image_path, gt_path)
##############################################################
# Inputs
##############################################################
input_inds_np = np.array(entry['word_inds'])
input_lens_np = np.array(entry['word_lens'])
input_inds = torch.from_numpy(input_inds_np).long().unsqueeze(0)
input_lens = torch.from_numpy(input_lens_np).long().unsqueeze(0)
if self.cfg.cuda:
input_inds = input_inds.cuda()
input_lens = input_lens.cuda()
##############################################################
# Inference
##############################################################
self.net.eval()
with torch.no_grad():
inf_outs, env = net.inference(input_inds, input_lens, -1, 1.0, 0, None, None, nn_table)
frame, noice, mask, label, env_info = env.batch_redraw(return_sequence=False)
frame = frame[0][0]; noice = noice[0][0]; mask = mask[0][0]; label = label[0][0]
env_info = env_info[0]
frame_path = osp.join(frame_dir, name+'.jpg')
noice_path = osp.join(noice_dir, name+'.jpg')
mask_path = osp.join(mask_dir, name+'.png')
label_path = osp.join(label_dir, name+'.png')
info_path = osp.join(info_dir, name+'.json')
if self.cfg.use_color_volume:
frame, _ = heuristic_collage(frame, 83)
noice, _ = heuristic_collage(noice, 83)
else:
frame = frame[:,:,-3:]
noice = noice[:,:,-3:]
cv2.imwrite(frame_path, clamp_array(frame, 0, 255).astype(np.uint8))
cv2.imwrite(noice_path, clamp_array(noice, 0, 255).astype(np.uint8))
cv2.imwrite(mask_path, clamp_array(255*mask, 0, 255))
cv2.imwrite(label_path, label)
# info
pred_info = {}
pred_info['width'] = env_info['width']
pred_info['height'] = env_info['height']
pred_info['clses'] = env_info['clses'].tolist()
pred_info['boxes'] = [x.tolist() for x in env_info['boxes']]
current_patches = env_info['patches']
current_image_indices = []
current_instance_inds = []
for j in range(len(pred_info['clses'])):
current_image_indices.append(current_patches[j]['image_index'])
current_instance_inds.append(current_patches[j]['instance_ind'])
pred_info['image_indices'] = current_image_indices
pred_info['instance_inds'] = current_instance_inds
with open(info_path, 'w') as fp:
json.dump(pred_info, fp, indent=4, sort_keys=True)
print('sampling: %d, %s'%(i, name))
def sample_for_vis(self, epoch, test_db, N, random_or_not=False, nn_table=None):
##############################################################
# Output prefix
##############################################################
output_dir = osp.join(self.cfg.model_dir, '%03d'%epoch, 'vis')
maybe_create(output_dir)
seq_db = sequence_loader(test_db)
##############################################################
# Main loop
##############################################################
plt.switch_backend('agg')
if random_or_not:
indices = np.random.permutation(range(len(test_db.scenedb)))
else:
indices = range(len(test_db.scenedb))
indices = indices[:min(N, len(test_db.scenedb))]
if self.cfg.cuda and self.cfg.parallel:
net = self.net.module
else:
net = self.net
for i in indices:
entry = seq_db[i]
gt_scene = test_db.scenedb[i]
image_index = gt_scene['image_index']
image_path = test_db.color_path_from_index(image_index)
gt_img = cv2.imread(image_path, cv2.IMREAD_COLOR)
gt_img, _, _ = create_squared_image(gt_img)
gt_img = cv2.resize(gt_img, (self.cfg.input_image_size[0], self.cfg.input_image_size[1]))
##############################################################
# Inputs
##############################################################
input_inds_np = np.array(entry['word_inds'])
input_lens_np = np.array(entry['word_lens'])
input_inds = torch.from_numpy(input_inds_np).long().unsqueeze(0)
input_lens = torch.from_numpy(input_lens_np).long().unsqueeze(0)
if self.cfg.cuda:
input_inds = input_inds.cuda()
input_lens = input_lens.cuda()
##############################################################
# Inference
##############################################################
self.net.eval()
with torch.no_grad():
inf_outs, env = net.inference(input_inds, input_lens, -1, 1.0, 0, None, None, nn_table)
frames, _, _, _, _ = env.batch_redraw(return_sequence=True)
frames = frames[0]
_, _, _, _, _, what_wei, where_wei = inf_outs
if self.cfg.what_attn:
what_attn_words = self.decode_attention(
input_inds_np, input_lens_np, what_wei.squeeze(0))
if self.cfg.where_attn > 0:
where_attn_words = self.decode_attention(
input_inds_np, input_lens_np, where_wei.squeeze(0))
##############################################################
# Draw
##############################################################
fig = plt.figure(figsize=(32, 32))
plt.suptitle(entry['sentence'], fontsize=40)
for j in range(len(frames)):
subtitle = ''
if self.cfg.what_attn:
subtitle = subtitle + ' '.join(what_attn_words[j])
if self.cfg.where_attn > 0:
subtitle = subtitle + '\n' + ' '.join(where_attn_words[j])
plt.subplot(4, 4, j+1)
plt.title(subtitle, fontsize=30)
if self.cfg.use_color_volume:
vis_img, _ = heuristic_collage(frames[j], 83)
else:
vis_img = frames[j][:,:,-3:]
vis_img = clamp_array(vis_img[ :, :, ::-1], 0, 255).astype(np.uint8)
plt.imshow(vis_img)
plt.axis('off')
plt.subplot(4, 4, 16)
plt.imshow(gt_img[:, :, ::-1])
plt.axis('off')
name = osp.splitext(osp.basename(image_path))[0]
out_path = osp.join(output_dir, name+'.png')
fig.savefig(out_path, bbox_inches='tight')
plt.close(fig)
print('sampling: %d, %d'%(epoch, i))
def train_epoch(self, train_db, epoch):
if self.cfg.cuda and self.cfg.parallel:
net = self.net.module
else:
net = self.net
train_db.cfg.sent_group = -1
seq_db = sequence_loader(train_db)
train_loader = DataLoader(seq_db,
batch_size=self.cfg.batch_size, shuffle=True,
num_workers=self.cfg.num_workers, pin_memory=True)
all_losses, all_accuracies = [], []
for cnt, batched in enumerate(train_loader):
##################################################################
## Batched data
##################################################################
input_inds, input_lens, fg_onehots, bg_imgs, \
fg_imgs, neg_imgs, fg_resnets, neg_resnets,\
gt_inds, gt_msks, patch_inds = \
self.batch_data(batched)
##################################################################
## Train one step
##################################################################
self.net.train()
self.net.zero_grad()
inputs = (input_inds, input_lens, bg_imgs, fg_onehots, fg_imgs, neg_imgs, fg_resnets, neg_resnets)
inf_outs, _, pos_vecs, neg_vecs = self.net(inputs)
pred_loss, embed_loss, attn_loss, pred_accu = self.evaluate(inf_outs, pos_vecs, neg_vecs, gt_inds, gt_msks)
loss = pred_loss + embed_loss + attn_loss
loss.backward()
self.optimizer.step()
##################################################################
## Collect info
##################################################################
all_losses.append(np.array(
[
loss.cpu().data.item(),
pred_loss.cpu().data.item(),
embed_loss.cpu().data.item(),
attn_loss.cpu().data.item()
]))
all_accuracies.append(pred_accu.cpu().data.numpy())
##################################################################
## Print info
##################################################################
if cnt % self.cfg.log_per_steps == 0:
print('Epoch %03d, iter %07d:'%(epoch, cnt))
tmp_losses = np.stack(all_losses, 0)
tmp_accuracies = np.stack(all_accuracies, 0)
print('losses: ', np.mean(tmp_losses[:,0]), np.mean(tmp_losses[:,1]), np.mean(tmp_losses[:,2]), np.mean(tmp_losses[:,3]))
print('accuracies: ', np.mean(tmp_accuracies[:,0]), np.mean(tmp_accuracies[:,1]), np.mean(tmp_accuracies[:,2]), np.mean(tmp_accuracies[:,3]))
print('-------------------------')
all_losses = np.stack(all_losses, 0)
all_accuracies = np.stack(all_accuracies, 0)
return all_losses, all_accuracies
def test_coco_dataloader(config):
db = composites_coco(config, 'train')
all_tables = AllCategoriesTables(db)
all_tables.build_nntables_for_all_categories(True)
sequence_db = sequence_loader(db, all_tables)
output_dir = osp.join(config.model_dir, 'test_coco_dataloader')
maybe_create(output_dir)
loader = DataLoader(sequence_db,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.num_workers)
start = time()
for cnt, batched in enumerate(loader):
x = batched['background'].float()
y = batched['foreground'].float()
z = batched['negative'].float()
x = sequence_color_volumn_preprocess(x, len(db.classes))
y = sequence_onehot_volumn_preprocess(y, len(db.classes))
z = sequence_onehot_volumn_preprocess(z, len(db.classes))
# cv2.imwrite('mask0.png', y[0,2,:,:,-4].cpu().data.numpy())
# cv2.imwrite('mask1.png', y[1,2,:,:,-4].cpu().data.numpy())
# cv2.imwrite('mask2.png', y[2,2,:,:,-4].cpu().data.numpy())
# cv2.imwrite('mask3.png', y[3,2,:,:,-4].cpu().data.numpy())
# cv2.imwrite('label0.png', y[0,2,:,:,3].cpu().data.numpy())
# cv2.imwrite('label1.png', y[1,2,:,:,3].cpu().data.numpy())
# cv2.imwrite('label2.png', y[2,2,:,:,3].cpu().data.numpy())
# cv2.imwrite('label3.png', y[3,2,:,:,3].cpu().data.numpy())
# cv2.imwrite('color0.png', y[0,2,:,:,-3:].cpu().data.numpy())
# cv2.imwrite('color1.png', y[1,2,:,:,-3:].cpu().data.numpy())
# cv2.imwrite('color2.png', y[2,2,:,:,-3:].cpu().data.numpy())
# cv2.imwrite('color3.png', y[3,2,:,:,-3:].cpu().data.numpy())
# cv2.imwrite('bg0.png', x[0,3,:,:,9:12].cpu().data.numpy())
# cv2.imwrite('bg1.png', x[1,3,:,:,9:12].cpu().data.numpy())
# cv2.imwrite('bg2.png', x[2,3,:,:,9:12].cpu().data.numpy())
# cv2.imwrite('bg3.png', x[3,3,:,:,9:12].cpu().data.numpy())
x = (x - 128.0).permute(0, 1, 4, 2, 3)
y = (y - 128.0).permute(0, 1, 4, 2, 3)
z = (z - 128.0).permute(0, 1, 4, 2, 3)
print('background', x.size())
print('foreground', y.size())
print('negative', z.size())
print('word_inds', batched['word_inds'].size())
print('word_lens', batched['word_lens'].size())
print('fg_inds', batched['fg_inds'].size())
print('patch_inds', batched['patch_inds'].size())
print('out_inds', batched['out_inds'].size())
print('out_msks', batched['out_msks'].size())
print('foreground_resnets', batched['foreground_resnets'].size())
print('negative_resnets', batched['negative_resnets'].size())
print('foreground_resnets', batched['foreground_resnets'][0, 0])
print('negative_resnets', batched['negative_resnets'][0, 0])
print('out_msks', batched['out_msks'][0])
print('patch_inds', batched['patch_inds'][0])
plt.switch_backend('agg')
bg_images = x
fg_images = y
neg_images = z
bsize, ssize, n, h, w = bg_images.size()
bg_images = bg_images.view(bsize * ssize, n, h, w)
bg_images = tensors_to_vols(bg_images)
bg_images = bg_images.reshape(bsize, ssize, h, w, n)
bsize, ssize, n, h, w = fg_images.size()
fg_images = fg_images.view(bsize * ssize, n, h, w)
fg_images = tensors_to_vols(fg_images)
fg_images = fg_images.reshape(bsize, ssize, h, w, n)
bsize, ssize, n, h, w = neg_images.size()
neg_images = neg_images.view(bsize * ssize, n, h, w)
neg_images = tensors_to_vols(neg_images)
neg_images = neg_images.reshape(bsize, ssize, h, w, n)
for i in range(bsize):
bg_seq = bg_images[i]
fg_seq = fg_images[i]
neg_seq = neg_images[i]
image_idx = batched['image_index'][i]
fg_inds = batched['fg_inds'][i]
name = '%03d_' % i + str(image_idx).zfill(12)
out_path = osp.join(output_dir, name + '.png')
color = cv2.imread(batched['image_path'][i], cv2.IMREAD_COLOR)
color, _, _ = create_squared_image(color)
fig = plt.figure(figsize=(48, 32))
plt.suptitle(batched['sentence'][i], fontsize=30)
for j in range(min(len(bg_seq), 15)):
bg, _ = heuristic_collage(bg_seq[j], 83)
bg_mask = 255 * np.ones((bg.shape[1], bg.shape[0]))
row, col = np.where(np.sum(np.absolute(bg), -1) == 0)
bg_mask[row, col] = 0
# bg = bg_seq[j][:,:,-3:]
# bg_mask = bg_seq[j][:,:,-4]
bg_mask = np.repeat(bg_mask[..., None], 3, -1)
fg_color = fg_seq[j][:, :, -3:]
# fg_mask = fg_seq[j][:,:,fg_inds[j+1]]
fg_mask = fg_seq[j][:, :, -4]
neg_color = neg_seq[j][:, :, -3:]
# neg_mask = neg_seq[j][:,:,fg_inds[j+1]]
neg_mask = neg_seq[j][:, :, -4]
color_pair = np.concatenate((fg_color, neg_color), 1)
mask_pair = np.concatenate((fg_mask, neg_mask), 1)
mask_pair = np.repeat(mask_pair[..., None], 3, -1)
patch = np.concatenate((color_pair, mask_pair), 0)
patch = cv2.resize(patch, (bg.shape[1], bg.shape[0]))
partially_completed_img = np.concatenate((bg, bg_mask, patch),
1)
partially_completed_img = clamp_array(partially_completed_img,
0, 255).astype(np.uint8)
partially_completed_img = partially_completed_img[:, :, ::-1]
plt.subplot(4, 4, j + 1)
plt.imshow(partially_completed_img)
plt.axis('off')
plt.subplot(4, 4, 16)
plt.imshow(color[:, :, ::-1])
plt.axis('off')
fig.savefig(out_path, bbox_inches='tight')
plt.close(fig)
if cnt == 3:
break
print("Time", time() - start)