def puzzle_model_inference(config):
traindb = composites_coco(config, 'train', '2017')
valdb = composites_coco(config, 'test', '2017')
auxdb = composites_coco(config, 'aux', '2017')
trainer = PuzzleTrainer(traindb)
t0 = time()
patch_dir_name = 'patch_feature_' + 'with_bg' if config.use_patch_background else 'without_bg'
if not osp.exists(osp.join(traindb.root_dir, patch_dir_name)):
trainer.dump_shape_vectors(traindb)
all_tables = AllCategoriesTables(traindb)
all_tables.build_nntables_for_all_categories(True)
print("NN completes (time %.2fs)" % (time() - t0))
t0 = time()
if config.for_visualization:
trainer.sample_for_vis(0,
testdb,
len(valdb.scenedb),
nn_table=all_tables)
trainer.sample_for_vis(0,
auxdb,
len(auxdb.scenedb),
nn_table=all_tables)
else:
trainer.sample_for_eval(testdb, nn_table=all_tables)
trainer.sample_for_eval(auxdb, nn_table=all_tables)
print("Sampling completes (time %.2fs)" % (time() - t0))
def test_txt_encoder_coco(config):
transformer = image_normalize('background')
db = coco(config, 'train', transform=transformer)
pca_table = AllCategoriesTables(db)
pca_table.run_PCAs_and_build_nntables_in_feature_space()
net = TextEncoder(db)
loader = DataLoader(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_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_nntable(config):
db = coco(config, 'test')
output_dir = osp.join(config.model_dir, 'test_nntable')
maybe_create(output_dir)
t0 = time()
all_tables = AllCategoriesTables(db)
all_tables.build_nntables_for_all_categories(False)
print("NN completes (time %.2fs)" % (time() - t0))
def __init__(self, db, batch_size=None, nn_table=None):
self.db = db
self.cfg = db.cfg
self.batch_size = batch_size if batch_size is not None else self.cfg.batch_size
if nn_table is None:
self.nn_table = AllCategoriesTables(db)
self.nn_table.build_nntables_for_all_categories()
else:
self.nn_table = nn_table
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 embedding_model_inference_preparation(config):
traindb = coco(config, 'train')
valdb = coco(config, 'val')
trainer = EmbeddingTrainer(traindb)
t0 = time()
trainer.dump_shape_vectors(traindb)
print("Dump shape vectors completes (time %.2fs)" % (time() - t0))
t0 = time()
all_tables = AllCategoriesTables(traindb)
all_tables.build_nntables_for_all_categories(False)
print("NN completes (time %.2fs)" % (time() - t0))
def train_scene_model(config):
transformer = image_normalize('background')
traindb = coco(config, 'train', transform=transformer)
valdb = coco(config, 'val', transform=transformer)
testdb = coco(config, 'test', transform=transformer)
pca_table = AllCategoriesTables(traindb)
pca_table.run_PCAs_and_build_nntables_in_feature_space()
trainer = SupervisedTrainer(traindb)
# we use the official validation set as test set
trainer.train(traindb, testdb, valdb)
def embedding_model_inference(config):
traindb = coco(config, 'train')
valdb = coco(config, 'val')
trainer = EmbeddingTrainer(traindb)
t0 = time()
all_tables = AllCategoriesTables(traindb)
all_tables.build_nntables_for_all_categories(True)
print("NN completes (time %.2fs)" % (time() - t0))
t0 = time()
trainer.sample(0, valdb, 50, random_or_not=False)
print("Sampling completes (time %.2fs)" % (time() - t0))
def puzzle_model_inference_preparation(config):
traindb = coco(config, 'train', '2017')
testdb = coco(config, 'test', '2017')
trainer = PuzzleTrainer(traindb)
t0 = time()
trainer.dump_shape_vectors(traindb)
trainer.dump_shape_vectors(testdb)
print("Dump shape vectors completes (time %.2fs)" % (time() - t0))
t0 = time()
all_tables = AllCategoriesTables(traindb)
all_tables.build_nntables_for_all_categories(False)
print("NN completes (time %.2fs)" % (time() - t0))
def composites_demo(config):
traindb = composites_coco(config, 'train', '2017')
trainer = PuzzleTrainer(traindb)
t0 = time()
patch_dir_name = 'patch_feature_' + 'with_bg' if config.use_patch_background else 'without_bg'
if not osp.exists(osp.join(traindb.root_dir, patch_dir_name)):
trainer.dump_shape_vectors(traindb)
all_tables = AllCategoriesTables(traindb)
all_tables.build_nntables_for_all_categories(True)
print("NN completes (time %.2fs)" % (time() - t0))
t0 = time()
input_sentences = json_load('examples/composites_samples.json')
trainer.sample_demo(input_sentences, all_tables)
print("Sampling completes (time %.2fs)" % (time() - t0))
def overfit_scene_model(config):
config.log_per_steps = 1
transformer = image_normalize('background')
traindb = coco(config, 'train', transform=transformer)
valdb = coco(config, 'val', transform=transformer)
testdb = coco(config, 'test', transform=transformer)
traindb.scenedb = traindb.scenedb[:config.batch_size]
valdb.scenedb = valdb.scenedb[:config.batch_size]
testdb.scenedb = testdb.scenedb[:config.batch_size]
print('build pca table')
pca_table = AllCategoriesTables(traindb)
pca_table.run_PCAs_and_build_nntables_in_feature_space()
print('create trainer')
trainer = SupervisedTrainer(traindb)
print('start training')
# trainer.train(traindb, traindb, traindb)
trainer.train(traindb, valdb, testdb)
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 puzzle_model_inference(config):
traindb = coco(config, 'train', '2017')
valdb = coco(config, 'val', '2017')
auxdb = coco(config, 'aux', '2017')
trainer = PuzzleTrainer(traindb)
t0 = time()
all_tables = AllCategoriesTables(traindb)
all_tables.build_nntables_for_all_categories(True)
print("NN completes (time %.2fs)" % (time() - t0))
t0 = time()
if config.for_visualization:
trainer.sample_for_vis(0,
valdb,
len(valdb.scenedb),
nn_table=all_tables)
trainer.sample_for_vis(0,
auxdb,
len(auxdb.scenedb),
nn_table=all_tables)
else:
trainer.sample_for_eval(valdb, nn_table=all_tables)
# trainer.sample_for_eval(auxdb, nn_table=all_tables)
print("Sampling completes (time %.2fs)" % (time() - t0))
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_step_by_step(config):
db = coco(config, 'train', '2017')
output_dir = osp.join(config.model_dir, 'test_step_by_step')
maybe_create(output_dir)
all_tables = AllCategoriesTables(db)
all_tables.build_nntables_for_all_categories(True)
seq_db = sequence_loader(db, all_tables)
env = simulator(db, config.batch_size, all_tables)
env.reset()
loader = DataLoader(seq_db,
batch_size=config.batch_size,
shuffle=True,
num_workers=config.num_workers)
for cnt, batched in enumerate(loader):
out_inds = batched['out_inds'].long().numpy()
out_vecs = batched['out_vecs'].float().numpy()
sequences = []
for i in range(out_inds.shape[1]):
frames = env.batch_render_to_pytorch(out_inds[:, i], out_vecs[:,
i])
sequences.append(frames)
sequences = torch.stack(sequences, dim=1)
# sequences = [tensors_to_vols(x) for x in sequences]
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['image_path'][i], cv2.IMREAD_COLOR)
color, _, _ = create_squared_image(color)
fig = plt.figure(figsize=(32, 32))
plt.suptitle(batched['sentence'][i], fontsize=30)
for j in range(min(len(sequence), 14)):
plt.subplot(4, 4, j + 1)
seq_np = sequence[j].cpu().data.numpy()
if config.use_color_volume:
partially_completed_img, _ = heuristic_collage(seq_np, 83)
else:
partially_completed_img = seq_np[:, :, -3:]
partially_completed_img = clamp_array(partially_completed_img,
0, 255).astype(np.uint8)
partially_completed_img = partially_completed_img[:, :, ::-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)
break
def test_scene_model(config):
output_dir = osp.join(config.model_dir, 'test_scene_model')
maybe_create(output_dir)
plt.switch_backend('agg')
transformer = image_normalize('background')
db = coco(config, 'train', transform=transformer)
pca_table = AllCategoriesTables(db)
pca_table.run_PCAs_and_build_nntables_in_feature_space()
loader = DataLoader(db,
batch_size=config.batch_size, shuffle=False,
num_workers=config.num_workers)
net = SceneModel(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_inds = batched['fg_inds'].long()
gt_inds = batched['out_inds'].long()
gt_vecs = batched['out_vecs'].float()
gt_msks = batched['out_msks'].float()
fg_onehots = indices2onehots(fg_inds, config.output_vocab_size)
# inf_outs, _ = net((word_inds, word_lens, bg_images, fg_onehots))
# obj_logits, coord_logits, attri_logits, pca_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('pca_vectors ', pca_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('----------------------')
inf_outs, env = net.inference(word_inds, word_lens, -1, 0, 0, gt_inds, gt_vecs)
# inf_outs, env = net.inference(word_inds, word_lens, -1, 2.0, 0, None, None)
obj_logits, coord_logits, attri_logits, pca_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('pca_vectors ', pca_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], 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 train(self, train_db, val_db, test_db):
##################################################################
## LOG
##################################################################
logz.configure_output_dir(self.cfg.model_dir)
logz.save_config(self.cfg)
##################################################################
## NN table
##################################################################
if self.cfg.use_hard_mining:
self.train_tables = AllCategoriesTables(train_db)
self.val_tables = AllCategoriesTables(val_db)
self.train_tables.build_nntables_for_all_categories(True)
self.val_tables.build_nntables_for_all_categories(True)
##################################################################
## Main loop
##################################################################
start = time()
min_val_loss = 100000000
for epoch in range(self.epoch, self.cfg.n_epochs):
##################################################################
## Training
##################################################################
torch.cuda.empty_cache()
train_loss, train_accu = self.train_epoch(train_db, epoch)
##################################################################
## Validation
##################################################################
torch.cuda.empty_cache()
val_loss, val_accu = self.validate_epoch(val_db, epoch)
##################################################################
## Logging
##################################################################
# update optim scheduler
current_val_loss = np.mean(val_loss[:, 0])
# self.optimizer.update(current_val_loss, epoch)
logz.log_tabular("Time", time() - start)
logz.log_tabular("Iteration", epoch)
logz.log_tabular("AverageLoss", np.mean(train_loss[:, 0]))
logz.log_tabular("AveragePredLoss", np.mean(train_loss[:, 1]))
logz.log_tabular("AverageEmbedLoss", np.mean(train_loss[:, 2]))
logz.log_tabular("AverageAttnLoss", np.mean(train_loss[:, 3]))
logz.log_tabular("AverageObjAccu", np.mean(train_accu[:, 0]))
logz.log_tabular("AverageCoordAccu", np.mean(train_accu[:, 1]))
logz.log_tabular("AverageScaleAccu", np.mean(train_accu[:, 2]))
logz.log_tabular("AverageRatioAccu", np.mean(train_accu[:, 3]))
logz.log_tabular("ValAverageLoss", np.mean(val_loss[:, 0]))
logz.log_tabular("ValAveragePredLoss", np.mean(val_loss[:, 1]))
logz.log_tabular("ValAverageEmbedLoss", np.mean(val_loss[:, 2]))
logz.log_tabular("ValAverageAttnLoss", np.mean(val_loss[:, 3]))
logz.log_tabular("ValAverageObjAccu", np.mean(val_accu[:, 0]))
logz.log_tabular("ValAverageCoordAccu", np.mean(val_accu[:, 1]))
logz.log_tabular("ValAverageScaleAccu", np.mean(val_accu[:, 2]))
logz.log_tabular("ValAverageRatioAccu", np.mean(val_accu[:, 3]))
logz.dump_tabular()
##################################################################
## Checkpoint
##################################################################
if self.cfg.use_hard_mining:
if (epoch + 1) % 3 == 0:
torch.cuda.empty_cache()
t0 = time()
self.dump_shape_vectors(train_db)
torch.cuda.empty_cache()
self.dump_shape_vectors(val_db)
print("Dump shape vectors completes (time %.2fs)" %
(time() - t0))
torch.cuda.empty_cache()
t0 = time()
self.train_tables.build_nntables_for_all_categories(False)
self.val_tables.build_nntables_for_all_categories(False)
print("NN completes (time %.2fs)" % (time() - t0))
self.save_checkpoint(epoch)
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 generate_simulated_scenes(config, split, year):
db = coco(config, split, year)
data_dir = osp.join(config.data_dir, 'coco')
if (split == 'test') or (split == 'aux'):
images_dir = osp.join(data_dir, 'crn_images', 'train' + year)
noices_dir = osp.join(data_dir, 'crn_noices', 'train' + year)
labels_dir = osp.join(data_dir, 'crn_labels', 'train' + year)
masks_dir = osp.join(data_dir, 'crn_masks', 'train' + year)
else:
images_dir = osp.join(data_dir, 'crn_images', split + year)
noices_dir = osp.join(data_dir, 'crn_noices', split + year)
labels_dir = osp.join(data_dir, 'crn_labels', split + year)
masks_dir = osp.join(data_dir, 'crn_masks', split + year)
maybe_create(images_dir)
maybe_create(noices_dir)
maybe_create(labels_dir)
maybe_create(masks_dir)
traindb = coco(config, 'train', '2017')
nn_tables = AllCategoriesTables(traindb)
nn_tables.build_nntables_for_all_categories(True)
# start_ind = 0
# end_ind = len(db.scenedb)
start_ind = 25000 + 14000 * config.seed
end_ind = 25000 + 14000 * (config.seed + 1)
patches_per_class = traindb.patches_per_class
color_transfer_threshold = 0.8
for i in range(start_ind, end_ind):
entry = db.scenedb[i]
width = entry['width']
height = entry['height']
xywhs = entry['boxes']
masks = entry['masks']
clses = entry['clses']
image_index = entry['image_index']
instance_inds = entry['instance_inds']
full_mask = np.zeros((height, width), dtype=np.float32)
full_label = np.zeros((height, width), dtype=np.float32)
full_image = np.zeros((height, width, 3), dtype=np.float32)
full_noice = np.zeros((height, width, 3), dtype=np.float32)
original_image = cv2.imread(db.color_path_from_index(image_index),
cv2.IMREAD_COLOR)
for j in range(len(masks)):
src_img = original_image.astype(np.float32).copy()
xywh = xywhs[j]
mask = masks[j]
cls_idx = clses[j]
instance_ind = instance_inds[j]
embed_path = db.patch_path_from_indices(
image_index, instance_ind, 'patch_feature', 'pkl',
config.use_patch_background)
with open(embed_path, 'rb') as fid:
query_vector = pickle.load(fid)
n_samples = min(
100, len(patches_per_class[cls_idx])
) #min(config.n_nntable_trees, len(patches_per_class[cls_idx]))
candidate_patches = nn_tables.retrieve(cls_idx, query_vector,
n_samples)
candidate_patches = [
x for x in candidate_patches
if x['instance_ind'] != instance_ind
]
assert (len(candidate_patches) > 1)
# candidate_instance_ind = instance_ind
# candidate_patch = None
# while (candidate_instance_ind == instance_ind):
# cid = np.random.randint(0, len(candidate_patches))
# candidate_patch = candidate_patches[cid]
# candidate_instance_ind = candidate_patch['instance_ind']
candidate_patch = find_closest_patch(db, traindb, image_index,
instance_ind,
candidate_patches)
# stenciling
src_mask = COCOmask.decode(mask)
dst_mask = COCOmask.decode(candidate_patch['mask'])
src_xyxy = xywh_to_xyxy(xywh, width, height)
dst_xyxy = xywh_to_xyxy(candidate_patch['box'],
candidate_patch['width'],
candidate_patch['height'])
dst_mask = dst_mask[dst_xyxy[1]:(dst_xyxy[3] + 1),
dst_xyxy[0]:(dst_xyxy[2] + 1)]
dst_mask = cv2.resize(
dst_mask,
(src_xyxy[2] - src_xyxy[0] + 1, src_xyxy[3] - src_xyxy[1] + 1),
interpolation=cv2.INTER_NEAREST)
src_mask[src_xyxy[1]:(src_xyxy[3]+1), src_xyxy[0]:(src_xyxy[2]+1)] = \
np.minimum(dst_mask, src_mask[src_xyxy[1]:(src_xyxy[3]+1), src_xyxy[0]:(src_xyxy[2]+1)])
# color transfer
if random.random() > color_transfer_threshold:
candidate_index = candidate_patch['image_index']
candidate_image = cv2.imread(
traindb.color_path_from_index(candidate_index),
cv2.IMREAD_COLOR).astype(np.float32)
candidate_cropped = candidate_image[dst_xyxy[1]:(dst_xyxy[3] +
1),
dst_xyxy[0]:(dst_xyxy[2] +
1)]
candidate_cropped = cv2.resize(candidate_cropped,
(src_xyxy[2] - src_xyxy[0] + 1,
src_xyxy[3] - src_xyxy[1] + 1),
interpolation=cv2.INTER_CUBIC)
original_cropped = src_img[src_xyxy[1]:(src_xyxy[3] + 1),
src_xyxy[0]:(src_xyxy[2] + 1)]
transfer_cropped = Monge_Kantorovitch_color_transfer(
original_cropped, candidate_cropped)
src_img[src_xyxy[1]:(src_xyxy[3] + 1),
src_xyxy[0]:(src_xyxy[2] + 1)] = transfer_cropped
# im1 = cv2.resize(full_image, (128, 128))
# im2 = cv2.resize(src_img[src_xyxy[1]:(src_xyxy[3]+1), src_xyxy[0]:(src_xyxy[2]+1), :], (128, 128))
# # im2 = cv2.resize(np.repeat(255*src_mask[...,None], 3, -1), (128, 128))
# im3 = cv2.resize(candidate_image, (128, 128))
# im4 = cv2.resize(candidate_cropped, (128, 128))
# im = np.concatenate((im1, im2, im3, im4), 1)
# cv2.imwrite("%03d_%03d.png"%(i, j), im)
full_image = compose(full_image, src_img, src_mask)
# boundary elision
radius = int(0.05 * min(width, height))
if np.amin(src_mask) > 0:
src_mask[0, :] = 0
src_mask[-1, :] = 0
src_mask[:, 0] = 0
src_mask[:, -1] = 0
sobelx = cv2.Sobel(src_mask, cv2.CV_64F, 1, 0, ksize=3)
sobely = cv2.Sobel(src_mask, cv2.CV_64F, 0, 1, ksize=3)
sobel = np.abs(sobelx) + np.abs(sobely)
edge = np.zeros_like(sobel)
edge[sobel > 0.9] = 1.0
morp_kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,
(radius, radius))
edge = cv2.dilate(edge, morp_kernel, iterations=1)
row, col = np.where(edge > 0)
n_edge_pixels = len(row)
pixel_indices = np.random.permutation(range(n_edge_pixels))
pixel_indices = pixel_indices[:(n_edge_pixels // 2)]
row = row[pixel_indices]
col = col[pixel_indices]
src_img[row, col, :] = 255
full_mask = np.maximum(full_mask, src_mask)
full_label[src_mask > 0] = cls_idx
full_noice = compose(full_noice, src_img, src_mask)
# im1 = cv2.resize(full_image, (128, 128))
# im2 = cv2.resize(src_img[src_xyxy[1]:(src_xyxy[3]+1), src_xyxy[0]:(src_xyxy[2]+1), :], (128, 128))
# im3 = cv2.resize(candidate_image, (128, 128))
# im4 = cv2.resize(candidate_cropped, (128, 128))
# im = np.concatenate((im1, im2, im3, im4), 1)
# cv2.imwrite("%03d_%03d.png"%(i, j), im)
output_name = str(image_index).zfill(12)
output_path = osp.join(images_dir, output_name + '.jpg')
cv2.imwrite(output_path,
clamp_array(full_image, 0, 255).astype(np.uint8))
output_path = osp.join(noices_dir, output_name + '.jpg')
cv2.imwrite(output_path,
clamp_array(full_noice, 0, 255).astype(np.uint8))
output_path = osp.join(masks_dir, output_name + '.png')
cv2.imwrite(output_path,
clamp_array(255 * full_mask, 0, 255).astype(np.uint8))
output_path = osp.join(labels_dir, output_name + '.png')
cv2.imwrite(output_path, full_label.astype(np.uint8))
print(i, image_index)
def test_coco_dataloader(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)
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_onehot_volumn_preprocess(x, len(db.classes))
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)