def run_model(args, checkpoint, output_dir, loader=None): vocab = checkpoint['model_kwargs']['vocab'] model = build_model(args, checkpoint) if loader is None: loader = build_loader(args, checkpoint) img_dir = makedir(output_dir, 'images') graph_dir = makedir(output_dir, 'graphs', args.save_graphs) gt_img_dir = makedir(output_dir, 'images_gt', args.save_gt_imgs) data_path = os.path.join(output_dir, 'data.pt') data = { 'vocab': vocab, 'objs': [], 'masks_pred': [], 'boxes_pred': [], 'masks_gt': [], 'boxes_gt': [], 'filenames': [], } if args.model == 'bm_FH': FH_dir_train = args.FH_dir_train FH_dir_val = args.FH_dir_val if args.model == 'bm_FHrec': FH_dir_train = args.FHrec_dir_train FH_dir_val = args.FHrec_dir_val if args.model == 'bm_FH64': FH_dir_train = args.FH64_dir_train FH_dir_val = args.FH64_dir_val if args.model == 'bm_FHrec64': FH_dir_train = args.FHrec64_dir_train FH_dir_val = args.FHrec64_dir_val FH_objs_train, FH_edges_train, IDs_train = torch.load(FH_dir_train) FH_objs_val, FH_edges_val, IDs_val = torch.load(FH_dir_val) IDs_train = torch.tensor(IDs_train) IDs_val = torch.tensor(IDs_val) if args.which_data == 'train': IDs = IDs_train FH_objs = FH_objs_train FH_edges = FH_edges_train else: IDs = IDs_val FH_objs = FH_objs_val FH_edges = FH_edges_val img_idx = 0 for batch in loader: masks = None if len(batch) == 6: imgs, objs, boxes, triples, obj_to_img, triple_to_img = [x.cuda() for x in batch] elif len(batch) == 7: imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img = [x.cuda() for x in batch] imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img, imgs_ids = [x.cuda() for x in batch] # get FH by images within a batch fh_obj, fh_edge = [],[] for i in range(imgs_ids.shape[0]): idd = ((IDs == imgs_ids[i].item()).nonzero()) fh_obj_i = FH_objs[idd] fh_obj.append(fh_obj_i) fh_edge_i = FH_edges[idd] fh_edge.append(fh_edge_i) fh_obj = torch.cat(fh_obj) fh_edge = torch.cat(fh_edge) imgs_gt = imagenet_deprocess_batch(imgs) boxes_gt = None masks_gt = None if args.use_gt_boxes: boxes_gt = boxes if args.use_gt_masks: masks_gt = masks # Run the model with predicted masks model_out = model(objs, triples, fh_obj, fh_edge, obj_to_img, boxes_gt=boxes_gt, masks_gt=masks_gt) # model_out = model(objs, triples, obj_to_img, # boxes_gt=boxes_gt, masks_gt=masks_gt) imgs_pred, boxes_pred, masks_pred, _ = model_out imgs_pred = imagenet_deprocess_batch(imgs_pred) obj_data = [objs, boxes_pred, masks_pred] _, obj_data = split_graph_batch(triples, obj_data, obj_to_img, triple_to_img) objs, boxes_pred, masks_pred = obj_data obj_data_gt = [boxes.data] if masks is not None: obj_data_gt.append(masks.data) triples, obj_data_gt = split_graph_batch(triples, obj_data_gt, obj_to_img, triple_to_img) boxes_gt, masks_gt = obj_data_gt[0], None if masks is not None: masks_gt = obj_data_gt[1] for i in range(imgs_pred.size(0)): img_filename = '%04d.png' % img_idx if args.save_gt_imgs: img_gt = imgs_gt[i].numpy().transpose(1, 2, 0) img_gt_path = os.path.join(gt_img_dir, img_filename) imsave(img_gt_path, img_gt) img_pred = imgs_pred[i] img_pred_np = imgs_pred[i].numpy().transpose(1, 2, 0) img_path = os.path.join(img_dir, img_filename) imsave(img_path, img_pred_np) data['objs'].append(objs[i].cpu().clone()) data['masks_pred'].append(masks_pred[i].cpu().clone()) data['boxes_pred'].append(boxes_pred[i].cpu().clone()) data['boxes_gt'].append(boxes_gt[i].cpu().clone()) data['filenames'].append(img_filename) cur_masks_gt = None if masks_gt is not None: cur_masks_gt = masks_gt[i].cpu().clone() data['masks_gt'].append(cur_masks_gt) if args.save_graphs: graph_img = draw_scene_graph(vocab, objs[i], triples[i]) graph_path = os.path.join(graph_dir, img_filename) imsave(graph_path, graph_img) img_idx += 1 torch.save(data, data_path) print('Saved %d images' % img_idx)
def run_model(args, checkpoint, output_dir, fn, loader=None): vocab = checkpoint['model_kwargs']['vocab'] print(vocab.keys()) print(vocab['pred_name_to_idx']) dic_pred = vocab[ 'pred_name_to_idx'] #{'inside': 5, 'left of': 1, '__in_image__': 0, 'right of': 2, 'below': 4, 'above': 3, 'surrounding': 6} model = build_model(args, checkpoint) if loader is None: loader = build_loader(args, checkpoint) data = { 'vocab': vocab, 'objs': [], 'masks_pred': [], 'boxes_pred': [], 'masks_gt': [], 'boxes_gt': [], 'filenames': [], } which_data = args.which_data makedir(output_dir, which_data) FN_path = os.path.join(output_dir, args.which_data, args.model + '_FN_%s.csv' % fn) f = open(FN_path, 'w') # f = open('err_bm_FH_%s.csv'%fn, 'w') with f: fieldnames = [ 'imageID', 'left of', 'right of', 'above', 'below', 'inside', 'surrounding' ] writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() class2idx = { "left of": 0, "right of": 1, "above": 2, "below": 3, "inside": 4, "surrounding": 5 } idx2class = {v: k for k, v in class2idx.items()} # save_dir = makedir(save_dir,fn) count_edge_gt = [] count_edge_pre = [] img_idx = 0 ibatch = 0 for batch in loader: ibatch += 1 masks = None # if len(batch) == 6: # imgs, objs, boxes, triples, obj_to_img, triple_to_img = [x.cuda() for x in batch] # elif len(batch) == 7: # imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img = [x.cuda() for x in batch] imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img, imgs_ids = [ x.cuda() for x in batch ] imgs_gt = imagenet_deprocess_batch(imgs) boxes_gt = None masks_gt = None if args.use_gt_boxes: boxes_gt = boxes if args.use_gt_masks: masks_gt = masks # Run the model with predicted masks model_out = model(objs, triples, obj_to_img, boxes_gt=boxes_gt, masks_gt=masks_gt) boxes_pred, masks_pred = model_out obj_data = [objs, boxes_pred, masks_pred] _, obj_data = split_graph_batch(triples, obj_data, obj_to_img, triple_to_img) objs, boxes_pred, masks_pred = obj_data obj_data_gt = [boxes.data] if masks is not None: obj_data_gt.append(masks.data) triples, obj_data_gt = split_graph_batch(triples, obj_data_gt, obj_to_img, triple_to_img) boxes_gt, masks_gt = obj_data_gt[0], None if masks is not None: masks_gt = obj_data_gt[1] for i in range(imgs_gt.size(0)): # for edges triples_i = triples[i] img_id = imgs_ids[i].item() i_edge_gt, i_edge_pre = [], [] for k in range(triples_i.shape[0]): if (triples_i[k][1] != 0): idx_s, idx_o = triples_i[k][0], triples_i[k][2] bbxs_of_img = boxes_gt[i] masks_of_img = masks_gt[i] box_s, box_o = bbxs_of_img[idx_s], bbxs_of_img[idx_o] mask_s, mask_o = masks_of_img[idx_s], masks_of_img[idx_o] edge_gt = get_relationship(box_s, box_o, mask_s, mask_o) count_edge_gt.append(edge_gt) i_edge_gt.append(edge_gt) bbxs_of_img = boxes_pred[i] masks_of_img = masks_pred[i] box_s, box_o = bbxs_of_img[idx_s], bbxs_of_img[idx_o] mask_s, mask_o = masks_of_img[idx_s], masks_of_img[idx_o] mask_s, mask_o = torch.round(mask_s).type( torch.long), torch.round(mask_o).type(torch.long) edge_pre = get_relationship(box_s, box_o, mask_s, mask_o) count_edge_pre.append(edge_pre) i_edge_pre.append(edge_pre) edges_items = list(set(i_edge_gt + i_edge_pre)) edges_items = [idx2class[x - 1] for x in edges_items] dictOfclass = { i: edges_items[i] for i in range(0, len(edges_items)) } cmi = confusion_matrix(i_edge_pre, i_edge_gt) # axis y predicted axis x true total = cmi.sum(axis=1) numacc = [cmi[i][i] for i in range(cmi.shape[0])] numFP = total - numacc cmi = cmi / cmi.sum(axis=0) acci = [cmi[i][i] for i in range(cmi.shape[0])] rowFP, rowAcc = {'imageID': img_id}, {'imageID': img_id} for q in range(len(dictOfclass)): rowFP.update({dictOfclass[q]: numFP[q]}) rowAcc.update({dictOfclass[q]: acci[q]}) img_idx += 1 FN_path = os.path.join(output_dir, args.which_data, args.model + '_FN_%s.csv' % fn) f = open(FN_path, 'a') with f: # 'imageID', 'left of', 'right of', 'above', 'below', 'inside', 'surounding' writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writerow(rowFP) print('%d images' % img_idx) # break print('gt', len(count_edge_gt)) print('pre', len(count_edge_pre)) cm = confusion_matrix(count_edge_pre, count_edge_gt) cm = cm / cm.sum(axis=0) confusion_matrix_df = pd.DataFrame( cm) #.rename(columns=idx2class, index=idx2class) print(confusion_matrix_df) label = {'a': '5%', 'b': '10%', 'c': '20%', 'd': '50%', 'e': '100%'} acc = [ confusion_matrix_df[i][i] for i in range(confusion_matrix_df.shape[0]) ] print('acc', acc) edgenames = [ 'left of', 'right of', 'above', 'below', 'inside', 'surrounding' ] accTotal = {'model': label[fn]} for q in range(0, len(acc)): accTotal.update({edgenames[q]: acc[q]}) err_path = os.path.join(output_dir, args.which_data, args.model + '_acc.csv') fil = open(err_path, 'a') with fil: fieldnames = [ 'model', 'left of', 'right of', 'above', 'below', 'inside', 'surrounding' ] writer = csv.DictWriter(fil, fieldnames=fieldnames) writer.writerow(accTotal) print('over')
def run_model(args, checkpoint, output_dir, fn, loader=None): vocab = checkpoint['model_kwargs']['vocab'] print(vocab.keys()) print(vocab['pred_name_to_idx']) dic_pred = vocab[ 'pred_name_to_idx'] #{'inside': 5, 'left of': 1, '__in_image__': 0, 'right of': 2, 'below': 4, 'above': 3, 'surrounding': 6} model = build_model(args, checkpoint) if loader is None: loader = build_loader(args, checkpoint) data = { 'vocab': vocab, 'objs': [], 'masks_pred': [], 'boxes_pred': [], 'masks_gt': [], 'boxes_gt': [], 'filenames': [], } which_data = args.which_data # save_dir = makedir(output_dir, args.which_data) FN_path = os.path.join(output_dir, args.which_data, args.model + '_FN_%s.csv' % fn) f = open(FN_path, 'w') # f = open('err_bm_FH_%s.csv'%fn, 'w') with f: fieldnames = [ 'imageID', 'left of', 'right of', 'above', 'below', 'inside', 'surrounding' ] writer = csv.DictWriter(f, fieldnames=fieldnames) writer.writeheader() if args.model == 'bm_FH': FH_dir_train = args.FH_dir_train FH_dir_val = args.FH_dir_val if args.model == 'bm_FHrec': FH_dir_train = args.FHrec_dir_train FH_dir_val = args.FHrec_dir_val FH_objs_train, FH_edges_train, IDs_train = torch.load( FH_dir_train) #torch.load('dataFH/train_FH.npy') FH_objs_val, FH_edges_val, IDs_val = torch.load( FH_dir_val) #torch.load('dataFH/val_FH.npy') IDs_train = torch.tensor(IDs_train) IDs_val = torch.tensor(IDs_val) if args.which_data == 'train': IDs = IDs_train FH_objs = FH_objs_train FH_edges = FH_edges_train else: IDs = IDs_val FH_objs = FH_objs_val FH_edges = FH_edges_val class2idx = { "left of": 0, "right of": 1, "above": 2, "below": 3, "inside": 4, "surrounding": 5 } idx2class = {v: k for k, v in class2idx.items()} count_edge_gt = [] count_edge_pre = [] img_idx = 0 ibatch = 0 for batch in loader: ibatch += 1 masks = None imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img, imgs_ids = [ x.cuda() for x in batch ] # get FH by images within a batch fh_obj, fh_edge = [], [] for i in range(imgs_ids.shape[0]): idd = ((IDs == imgs_ids[i].item()).nonzero()) fh_obj_i = FH_objs[idd] fh_obj.append(fh_obj_i) fh_edge_i = FH_edges[idd] fh_edge.append(fh_edge_i) fh_obj = torch.cat(fh_obj) fh_edge = torch.cat(fh_edge) imgs_gt = imagenet_deprocess_batch(imgs) boxes_gt = None masks_gt = None if args.use_gt_boxes: boxes_gt = boxes if args.use_gt_masks: masks_gt = masks # Run the model with predicted masks model_out = model(objs, triples, fh_obj, fh_edge, obj_to_img, boxes_gt=boxes_gt, masks_gt=masks_gt) boxes_pred, masks_pred = model_out obj_data = [objs, boxes_pred, masks_pred] _, obj_data = split_graph_batch(triples, obj_data, obj_to_img, triple_to_img) objs, boxes_pred, masks_pred = obj_data obj_data_gt = [boxes.data] if masks is not None: obj_data_gt.append(masks.data) triples, obj_data_gt = split_graph_batch(triples, obj_data_gt, obj_to_img, triple_to_img) boxes_gt, masks_gt = obj_data_gt[0], None if masks is not None: masks_gt = obj_data_gt[1] for i in range(imgs_gt.size(0)): # for edges triples_i = triples[i] img_id = imgs_ids[i].item() i_edge_gt, i_edge_pre = [], [] for k in range(triples_i.shape[0]): if (triples_i[k][1] != 0): idx_s, idx_o = triples_i[k][0], triples_i[k][2] bbxs_of_img = boxes_gt[i] masks_of_img = masks_gt[i] box_s, box_o = bbxs_of_img[idx_s], bbxs_of_img[idx_o] mask_s, mask_o = masks_of_img[idx_s], masks_of_img[idx_o] edge_gt = get_relationship(box_s, box_o, mask_s, mask_o) count_edge_gt.append(edge_gt) i_edge_gt.append(edge_gt) bbxs_of_img = boxes_pred[i] masks_of_img = masks_pred[i] box_s, box_o = bbxs_of_img[idx_s], bbxs_of_img[idx_o] mask_s, mask_o = masks_of_img[idx_s], masks_of_img[idx_o] mask_s, mask_o = torch.round(mask_s).type( torch.long), torch.round(mask_o).type(torch.long) edge_pre = get_relationship(box_s, box_o, mask_s, mask_o) count_edge_pre.append(edge_pre) i_edge_pre.append(edge_pre) edges_items = list(set(i_edge_gt + i_edge_pre)) # print(i_edge_pre, i_edge_gt,edges_items) edges_items = [idx2class[x - 1] for x in edges_items] # print(edges_items) dictOfclass = { i: edges_items[i] for i in range(0, len(edges_items)) } # edges_items = {k:idx2class[edges_items[k]-1] for k in range(len(edges_items))} # print(edges_items) cmi = confusion_matrix(i_edge_pre, i_edge_gt) # axis y predicted axis x true # confusion_matrix_df = pd.DataFrame(cmi).rename(columns=dictOfclass, index = dictOfclass)#idx2class, index=idx2class) # print(confusion_matrix_df) # FP # print(cmi) total = cmi.sum(axis=1) numacc = [cmi[i][i] for i in range(cmi.shape[0])] numFP = total - numacc # acc cmi = cmi / cmi.sum(axis=0) # print(cmi) acci = [cmi[i][i] for i in range(cmi.shape[0])] rowFP, rowAcc = {'imageID': img_id}, {'imageID': img_id} for q in range(len(dictOfclass)): rowFP.update({dictOfclass[q]: numFP[q]}) rowAcc.update({dictOfclass[q]: acci[q]}) # print(rowi) img_idx += 1 FN_path = os.path.join(output_dir, args.which_data, args.model + '_FN_%s.csv' % fn) f = open(FN_path, 'a') # f = open('err_bm_FH_%s.csv'%fn, 'a') with f: # 'imageID', 'left of', 'right of', 'above', 'below', 'inside', 'surounding' writer = csv.DictWriter(f, fieldnames=fieldnames) # writer = csv.writer(f) writer.writerow(rowFP) # writer.writerow(rowAcc) # writer.writerow([img_id]+acci) print('%d images' % img_idx) # break print('gt', len(count_edge_gt)) print('pre', len(count_edge_pre)) cm = confusion_matrix(count_edge_pre, count_edge_gt) cm = cm / cm.sum(axis=0) confusion_matrix_df = pd.DataFrame( cm) #.rename(columns=idx2class, index=idx2class) print(confusion_matrix_df) label = {'a': '5%', 'b': '10%', 'c': '20%', 'd': '50%', 'e': '100%'} acc = [ confusion_matrix_df[i][i] for i in range(confusion_matrix_df.shape[0]) ] print('acc', acc) # np.savetxt('scores.csv', acc, delimiter=',', fmt='%s') # fn = 'a' # filename = 'accuracy_' + args.model + '_' + fn + '_' + '.txt' # err_path = os.path.join(output_dir,args.which_data, filename) # np.savetxt(err_path, acc, delimiter=',', fmt='%s')s # np.savetxt(err_path, [p for p in zip(class2idx.keys(), acc)], delimiter=',', fmt='%s') # print(['bm_FH_'+fn]+acc) # save in total edgenames = [ 'left of', 'right of', 'above', 'below', 'inside', 'surrounding' ] accTotal = {'model': label[fn]} for q in range(0, len(acc)): accTotal.update({edgenames[q]: acc[q]}) err_path = os.path.join(output_dir, args.which_data, args.model + '_acc.csv') fil = open(err_path, 'a') with fil: fieldnames = [ 'model', 'left of', 'right of', 'above', 'below', 'inside', 'surrounding' ] writer = csv.DictWriter(fil, fieldnames=fieldnames) writer.writerow(accTotal) print('over')
def run_model(args, checkpoint, output_dir, fn, loader=None): vocab = checkpoint['model_kwargs']['vocab'] print(vocab.keys()) print(vocab['pred_name_to_idx']) dic_pred = vocab[ 'pred_name_to_idx'] #{'inside': 5, 'left of': 1, '__in_image__': 0, 'right of': 2, 'below': 4, 'above': 3, 'surrounding': 6} model = build_model(args, checkpoint) if loader is None: loader = build_loader(args, checkpoint) data = { 'vocab': vocab, 'objs': [], 'masks_pred': [], 'boxes_pred': [], 'masks_gt': [], 'boxes_gt': [], 'filenames': [], } which_data = args.which_data save_dir = makedir(output_dir, which_data) FH_objs_train, FH_edges_train, IDs_train = torch.load( args.FH_dir_train) #torch.load('dataFH/train_FH.npy') FH_objs_val, FH_edges_val, IDs_val = torch.load( args.FH_dir_val) #torch.load('dataFH/val_FH.npy') IDs_train = torch.tensor(IDs_train) IDs_val = torch.tensor(IDs_val) if args.which_data == 'train': IDs = IDs_train FH_objs = FH_objs_train FH_edges = FH_edges_train else: IDs = IDs_val FH_objs = FH_objs_val FH_edges = FH_edges_val count_edge_gt = [] count_edge_pre = [] img_idx = 0 ibatch = 0 for batch in loader: ibatch += 1 masks = None imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img, imgs_ids = [ x.cuda() for x in batch ] # get FH by images within a batch fh_obj, fh_edge = [], [] for i in range(imgs_ids.shape[0]): idd = ((IDs == imgs_ids[i].item()).nonzero()) fh_obj_i = FH_objs[idd] fh_obj.append(fh_obj_i) fh_edge_i = FH_edges[idd] fh_edge.append(fh_edge_i) fh_obj = torch.cat(fh_obj) fh_edge = torch.cat(fh_edge) imgs_gt = imagenet_deprocess_batch(imgs) boxes_gt = None masks_gt = None if args.use_gt_boxes: boxes_gt = boxes if args.use_gt_masks: masks_gt = masks # Run the model with predicted masks model_out = model(objs, triples, fh_obj, fh_edge, obj_to_img, boxes_gt=boxes_gt, masks_gt=masks_gt) boxes_pred, masks_pred = model_out obj_data = [objs, boxes_pred, masks_pred] _, obj_data = split_graph_batch(triples, obj_data, obj_to_img, triple_to_img) objs, boxes_pred, masks_pred = obj_data obj_data_gt = [boxes.data] if masks is not None: obj_data_gt.append(masks.data) triples, obj_data_gt = split_graph_batch(triples, obj_data_gt, obj_to_img, triple_to_img) boxes_gt, masks_gt = obj_data_gt[0], None if masks is not None: masks_gt = obj_data_gt[1] for i in range(imgs_gt.size(0)): # for edges triples_i = triples[i] for k in range(triples_i.shape[0]): if (triples_i[k][1] != 0): idx_s, idx_o = triples_i[k][0], triples_i[k][2] bbxs_of_img = boxes_gt[i] masks_of_img = masks_gt[i] box_s, box_o = bbxs_of_img[idx_s], bbxs_of_img[idx_o] mask_s, mask_o = masks_of_img[idx_s], masks_of_img[idx_o] edge_gt = get_relationship(box_s, box_o, mask_s, mask_o) count_edge_gt.append(edge_gt) # print('gt:', triples_i[k][1].item(), edge_gt) bbxs_of_img = boxes_pred[i] masks_of_img = masks_pred[i] box_s, box_o = bbxs_of_img[idx_s], bbxs_of_img[idx_o] mask_s, mask_o = masks_of_img[idx_s], masks_of_img[idx_o] mask_s, mask_o = torch.round(mask_s).type( torch.long), torch.round(mask_o).type(torch.long) edge_pre = get_relationship(box_s, box_o, mask_s, mask_o) count_edge_pre.append(edge_pre) img_idx += 1 print('%d images' % img_idx) class2idx = { "left of": 0, "right of": 1, "above": 2, "below": 3, "inside": 4, "surrounding": 5 } idx2class = {v: k for k, v in class2idx.items()} # break print('gt', len(count_edge_gt)) print('pre', len(count_edge_pre)) cm = confusion_matrix(count_edge_pre, count_edge_gt) # y, x cm = cm / cm.sum(axis=0) confusion_matrix_df = pd.DataFrame(cm).rename(columns=idx2class, index=idx2class) label = {'a': '5%', 'b': '10%', 'c': '20%', 'd': '50%', 'e': '100%'} ax = sns.heatmap(confusion_matrix_df, annot=True, cmap='Blues_r', vmin=0, vmax=1) title = 'M1_bm_FH_' + args.which_data + '_' + label[fn] ax.set(title=title, ylabel='Predicted label', xlabel='True label') fig = ax.get_figure() filename = 'CM1_bm_FH_' + fn + '_' + args.which_data + '.png' CM_path = os.path.join(output_dir, args.which_data, filename) fig.savefig(CM_path) fig.clf() print('over')
def run_model(args, checkpoint, output_dir, fn, loader=None): vocab = checkpoint['model_kwargs']['vocab'] model = build_model(args, checkpoint) if loader is None: loader = build_loader(args, checkpoint) data = { 'vocab': vocab, 'objs': [], 'masks_pred': [], 'boxes_pred': [], 'masks_gt': [], 'boxes_gt': [], 'filenames': [], } which_data = args.which_data save_dir = makedir(output_dir, which_data) # save_dir = makedir(save_dir,fn) img_idx = 0 ibatch = 0 for batch in loader: ibatch += 1 masks = None if len(batch) == 6: imgs, objs, boxes, triples, obj_to_img, triple_to_img = [ x.cuda() for x in batch ] elif len(batch) == 7: imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img = [ x.cuda() for x in batch ] # imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img, imgs_ids = [x.cuda() for x in batch] # imgs_print = imagenet_deprocess_batch(imgs) # grid = torchvision.utils.make_grid(imgs_print) # writer.add_image('img/real', grid, ibatch-1) imgs_gt = imagenet_deprocess_batch(imgs) boxes_gt = None masks_gt = None if args.use_gt_boxes: boxes_gt = boxes if args.use_gt_masks: masks_gt = masks # Run the model with predicted masks model_out = model(objs, triples, obj_to_img, boxes_gt=boxes_gt, masks_gt=masks_gt) # boxes_pred, masks_pred = model_out imgs_pred, boxes_pred, masks_pred, _ = model_out obj_data = [objs, boxes_pred, masks_pred] _, obj_data = split_graph_batch(triples, obj_data, obj_to_img, triple_to_img) objs, boxes_pred, masks_pred = obj_data obj_data_gt = [boxes.data] if masks is not None: obj_data_gt.append(masks.data) triples, obj_data_gt = split_graph_batch(triples, obj_data_gt, obj_to_img, triple_to_img) boxes_gt, masks_gt = obj_data_gt[0], None if masks is not None: masks_gt = obj_data_gt[1] imgs_bbx = torch.zeros(imgs_gt.size(), dtype=torch.uint8) imgs_bbx_pre = torch.zeros(imgs_gt.size(), dtype=torch.uint8) white_bbx_gt, white_bbx_gtb = torch.zeros( imgs_gt.size(), dtype=torch.uint8), torch.zeros(imgs_gt.size(), dtype=torch.uint8) white_bbx_pre, white_bbx_preb = torch.zeros( imgs_gt.size(), dtype=torch.uint8), torch.zeros(imgs_gt.size(), dtype=torch.uint8) for i in range(imgs_gt.size(0)): black_gt = np.zeros([args.image_size[0], args.image_size[1], 3]) black_gtb = np.zeros([args.image_size[0], args.image_size[1], 3]) img = np.copy(imgs_gt[i].numpy().transpose(1, 2, 0)) layer = np.zeros(list(args.image_size)) masks_of_img = masks_gt[i] bbxs_of_img = boxes_gt[i] num_of_objs = bbxs_of_img.size(0) for j in range(num_of_objs - 1): # color = tuple(np.random.randint(256, size=3)) color = colors[j % len(colors)] mask = masks_of_img[j].cpu().clone().numpy() mask = np.round(mask) bbx = (bbxs_of_img[j].cpu().numpy() * args.image_size[0]).astype(int) bbx = np.clip(bbx, 0, args.image_size[0] - 1) wbbx = bbx[2] - bbx[0] hbbx = bbx[3] - bbx[1] if not wbbx > 0: wbbx = 1 print('gt', wbbx, hbbx) if not hbbx > 0: hbbx = 1 print('gt', wbbx, hbbx) maskPIL = Image.fromarray(mask.astype(np.uint8)) maskPIL = maskPIL.resize((wbbx, hbbx), resample=Image.BILINEAR) layer[bbx[1]:bbx[3], bbx[0]:bbx[2]] = np.array(maskPIL) img = apply_mask(img, layer, color) masked_imgPIL = Image.fromarray(img.astype(np.uint8)) draw = ImageDraw.Draw(masked_imgPIL) draw.rectangle(bbx.tolist(), width=1, outline=color) img = np.array(masked_imgPIL) black_gt = apply_mask(black_gt, layer, color) masked_blackPIL = Image.fromarray(black_gt.astype(np.uint8)) draw2 = ImageDraw.Draw(masked_blackPIL) draw2.rectangle(bbx.tolist(), width=1, outline=color) black_gt = np.array(masked_blackPIL) blackPIL = Image.fromarray(black_gtb.astype(np.uint8)) draw2b = ImageDraw.Draw(blackPIL) draw2b.rectangle(bbx.tolist(), width=1, outline=color) black_gtb = np.array(blackPIL) imgs_bbx[i] = torchvision.transforms.ToTensor()( masked_imgPIL) * 255 white_bbx_gt[i] = torchvision.transforms.ToTensor()( masked_blackPIL) * 255 white_bbx_gtb[i] = torchvision.transforms.ToTensor()( blackPIL) * 255 black_gt = np.zeros([args.image_size[0], args.image_size[1], 3]) black_gtb = np.zeros([args.image_size[0], args.image_size[1], 3]) img = np.copy(imgs_gt[i].numpy().transpose(1, 2, 0)) layer = np.zeros(list(args.image_size)) bbxs_of_img = boxes_pred[i] masks_of_img = masks_pred[i] num_of_objs = bbxs_of_img.size(0) for j in range(num_of_objs - 1): color = colors[j % len(colors)] mask = masks_of_img[j].cpu().clone().numpy() mask = np.round(mask) bbx = (bbxs_of_img[j].cpu().numpy() * args.image_size[0]).astype(int) bbx = np.clip(bbx, 0, args.image_size[0] - 1) wbbx = bbx[2] - bbx[0] hbbx = bbx[3] - bbx[1] if not wbbx > 0: wbbx = 1 print('pred', wbbx, hbbx) if not hbbx > 0: hbbx = 1 print('pred', wbbx, hbbx) maskPIL = Image.fromarray(mask.astype(np.uint8)) maskPIL = maskPIL.resize((wbbx, hbbx), resample=Image.BILINEAR) # print('wwbx,hbbx:',wbbx, hbbx, maskPIL2.size, bbx) layer[bbx[1]:bbx[3], bbx[0]:bbx[2]] = np.array(maskPIL) img = apply_mask(img, layer, color) masked_imgPIL = Image.fromarray(img.astype(np.uint8)) draw = ImageDraw.Draw(masked_imgPIL) draw.rectangle(bbx.tolist(), width=1, outline=color) img = np.array(masked_imgPIL) black_gt = apply_mask(black_gt, layer, color) masked_blackPIL = Image.fromarray(black_gt.astype(np.uint8)) draw2 = ImageDraw.Draw(masked_blackPIL) draw2.rectangle(bbx.tolist(), width=1, outline=color) black_gt = np.array(masked_blackPIL) blackPIL = Image.fromarray(black_gtb.astype(np.uint8)) draw2b = ImageDraw.Draw(blackPIL) draw2b.rectangle(bbx.tolist(), width=1, outline=color) black_gtb = np.array(blackPIL) imgs_bbx_pre[i] = torchvision.transforms.ToTensor()( masked_imgPIL) * 255 white_bbx_pre[i] = torchvision.transforms.ToTensor()( masked_blackPIL) * 255 white_bbx_preb[i] = torchvision.transforms.ToTensor()( blackPIL) * 255 img_idx += 1 imgs_orig = imagenet_deprocess_batch(imgs) grid1 = torchvision.utils.make_grid(imgs_orig) toSave = grid1 # GT # imgs_grid_GT = imgs_bbx.byte() # grid2 = torchvision.utils.make_grid(imgs_grid_GT) # toSave = torch.cat((grid1,grid2),1) white_grid_GT = white_bbx_gt.byte() grid3 = torchvision.utils.make_grid(white_grid_GT) toSave = torch.cat((toSave, grid3), 1) white_grid_GTb = white_bbx_gtb.byte() grid3b = torchvision.utils.make_grid(white_grid_GTb) toSave = torch.cat((toSave, grid3b), 1) # PRE imgs_pred = imagenet_deprocess_batch(imgs_pred) gridx = torchvision.utils.make_grid(imgs_pred) toSave = torch.cat((toSave, gridx), 1) # imgs_grid_pre = imgs_bbx_pre.byte() # grid4 = torchvision.utils.make_grid(imgs_grid_pre) # toSave = torch.cat((toSave, grid4),1) white_grid_pre = white_bbx_pre.byte() grid5 = torchvision.utils.make_grid(white_grid_pre) toSave = torch.cat((toSave, grid5), 1) white_grid_preb = white_bbx_preb.byte() grid5b = torchvision.utils.make_grid(white_grid_preb) toSave = torch.cat((toSave, grid5b), 1) toSavePIL = torchvision.transforms.ToPILImage()(toSave) save_dir = 'output' fn = 'M1re' grids_path = os.path.join(save_dir, '%d' % img_idx + fn + '.png') # grids_path = os.path.join(save_dir, '%d'%img_id + fn + '.png') toSavePIL.save(grids_path) print('Saved %d images' % img_idx)
def run_model(args, checkpoint, output_dir, loader=None): vocab = checkpoint['model_kwargs']['vocab'] model = build_model(args, checkpoint) if loader is None: loader = build_loader(args, checkpoint) img_dir = makedir(output_dir, 'images') graph_dir = makedir(output_dir, 'graphs', args.save_graphs) gt_img_dir = makedir(output_dir, 'images_gt', args.save_gt_imgs) data_path = os.path.join(output_dir, 'data.pt') data = { 'vocab': vocab, 'objs': [], 'masks_pred': [], 'boxes_pred': [], 'masks_gt': [], 'boxes_gt': [], 'filenames': [], } img_idx = 0 for batch in loader: masks = None if len(batch) == 6: imgs, objs, boxes, triples, obj_to_img, triple_to_img = [ x.cuda() for x in batch ] elif len(batch) == 7: imgs, objs, boxes, masks, triples, obj_to_img, triple_to_img = [ x.cuda() for x in batch ] imgs_gt = imagenet_deprocess_batch(imgs) boxes_gt = None masks_gt = None if args.use_gt_boxes: boxes_gt = boxes if args.use_gt_masks: masks_gt = masks # Run the model with predicted masks model_out = model(objs, triples, obj_to_img, boxes_gt=boxes_gt, masks_gt=masks_gt) imgs_pred, boxes_pred, masks_pred, _ = model_out imgs_pred = imagenet_deprocess_batch(imgs_pred) obj_data = [objs, boxes_pred, masks_pred] _, obj_data = split_graph_batch(triples, obj_data, obj_to_img, triple_to_img) objs, boxes_pred, masks_pred = obj_data obj_data_gt = [boxes.data] if masks is not None: obj_data_gt.append(masks.data) triples, obj_data_gt = split_graph_batch(triples, obj_data_gt, obj_to_img, triple_to_img) boxes_gt, masks_gt = obj_data_gt[0], None if masks is not None: masks_gt = obj_data_gt[1] for i in range(imgs_pred.size(0)): img_filename = '%04d.png' % img_idx if args.save_gt_imgs: img_gt = imgs_gt[i].numpy().transpose(1, 2, 0) img_gt_path = os.path.join(gt_img_dir, img_filename) imsave(img_gt_path, img_gt) img_pred = imgs_pred[i] img_pred_np = imgs_pred[i].numpy().transpose(1, 2, 0) img_path = os.path.join(img_dir, img_filename) imsave(img_path, img_pred_np) data['objs'].append(objs[i].cpu().clone()) data['masks_pred'].append(masks_pred[i].cpu().clone()) data['boxes_pred'].append(boxes_pred[i].cpu().clone()) data['boxes_gt'].append(boxes_gt[i].cpu().clone()) data['filenames'].append(img_filename) cur_masks_gt = None if masks_gt is not None: cur_masks_gt = masks_gt[i].cpu().clone() data['masks_gt'].append(cur_masks_gt) if args.save_graphs: graph_img = draw_scene_graph(vocab, objs[i], triples[i]) graph_path = os.path.join(graph_dir, img_filename) imsave(graph_path, graph_img) img_idx += 1 torch.save(data, data_path) print('Saved %d images' % img_idx)