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)