def SW_add_scalars2(self, main_tag, tag_scalar_dict, global_step=None): """Adds many scalar data to summary. Args: tag (string): Data identifier main_tag (string): The parent name for the tags tag_scalar_dict (dict): Key-value pair storing the tag and corresponding values global_step (int): Global step value to record Examples:: writer.add_scalars('run_14h',{'xsinx':i*np.sin(i/r), 'xcosx':i*np.cos(i/r), 'arctanx': numsteps*np.arctan(i/r)}, i) # This function adds three values to the same scalar plot with the tag # 'run_14h' in TensorBoard's scalar section. """ timestamp = time.time() fw_logdir = self.file_writer.get_logdir() for tag, scalar_value in tag_scalar_dict.items(): fw_tag = fw_logdir + "/" + tag #fw_tag_full = fw_logdir + "/" + main_tag + "/" + tag if fw_tag in self.all_writers.keys(): fw = self.all_writers[fw_tag] else: fw = tensorboardX.FileWriter(logdir=fw_tag) self.all_writers[fw_tag] = fw fw.add_summary(tensorboardX.summary.scalar(main_tag, scalar_value), global_step)
def add_tsb_image_last(entry, data, path=".", subfolder="", **kwargs): """Handler that append the last image of the data to a tensorboard event file. :param string entry: Name of the log entry :param Dict data: Data should be a numpy array / torch tensor of shape [3, W, H] :param string subfolder: Subfolder in which put the data :param string path: Root path. Set by DataLogger if used as handler. """ tsb_dir = os.path.join(path, subfolder) os.makedirs(tsb_dir, exist_ok=True) if tsb_writer.file_writer.get_logdir() != tsb_dir: tsb_writer.file_writer = tensorboardX.FileWriter(tsb_dir) last_time = max(data.keys()) value = data[last_time] tsb_writer.add_image(entry, value, last_time)
def add_tsb_scalars_last(entry, data, labels=None, path=".", subfolder="", **kwargs): """Handler that appends the last item of the data dictionary to a tensorboard event file. :param string entry: Name of the log entry. :param Dict data: Data should be numpy arrays of size [n] with constant n. :param List[string] labels: Labels to use for the lines :param string subfolder: Subfolder in which put the data :param string path: Root path. Set by DataLogger if used as handler. """ tsb_dir = os.path.join(path, subfolder) os.makedirs(tsb_dir, exist_ok=True) if tsb_writer.file_writer.get_logdir() != tsb_dir: tsb_writer.file_writer = tensorboardX.FileWriter(tsb_dir) last_time = max(data.keys()) value = data[last_time] if labels is None: labels = [str(a) for a in range(len(value))] scalars_dict = {labels[i]: value[i] for i in range(len(value))} tsb_writer.add_scalars(entry, scalars_dict, last_time)
def __init__(self, model_dir): self.summary_writer = tb.FileWriter(model_dir)
def main(argv): (opts, args) = parser.parse_args(argv) if 'estimate' in opts.mode: mode_idx = int(opts.mode[-1]) global colorPlatte, bones, Evaluation if 'nyu' in opts.config: colorPlatte = utils.util.nyuColorIdx bones = utils.util.nyuBones Evaluation = NYUHandposeEvaluation elif 'icvl' in opts.config: colorPlatte = utils.util.icvlColorIdx bones = utils.util.icvlBones Evaluation = ICVLHandposeEvaluation # Load experiment setting assert isinstance(opts, object) config = NetConfig(opts.config) batch_size = config.hyperparameters['batch_size'] if 'estimate' in opts.mode else 1 test_batch_size = batch_size * 32 max_iterations = config.hyperparameters['max_iterations'] frac = opts.frac dataset_a = get_dataset(config.datasets['train_a']) dataset_b = get_dataset(config.datasets['train_b']) dataset_test = get_dataset(config.datasets['test_b']) train_loader_a = get_data_loader(dataset_a, batch_size, shuffle=True) train_loader_b = get_data_loader(dataset_b, batch_size, shuffle=True) test_loader_real = get_data_loader(dataset_test, test_batch_size, shuffle=False) cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters['trainer'] local_dict = locals() exec(cmd,globals(),local_dict) trainer = local_dict['trainer'] di_a = dataset_a.di di_b = dataset_b.di # Check if resume training iterations = 0 if opts.resume == 1: iterations = trainer.resume(config.snapshot_prefix, idx=-1, load_opt=True) for i in range(iterations//1000): trainer.dis_sch.step() trainer.gen_sch.step() trainer.cuda(opts.gpu) print('using %.2f percent of the labeled real data' % frac) try: if 'estimate' in opts.mode and (mode_idx == 3 or mode_idx == 4): trainer.load_vae(config.snapshot_prefix, 2+frac) else: trainer.load_vae(config.snapshot_prefix, frac) except: print('Failed to load the parameters of vae') if 'estimate' in opts.mode: if opts.idx != 0: trainer.resume(config.snapshot_prefix, idx=opts.idx, est=mode_idx==5) if frac > 0. and frac < 1.: dataset_b.set_nmax(frac) #trainer.dis.freeze_layers() ############################################################################################### # Setup logger and repare image outputs train_writer = tensorboardX.FileWriter("%s/%s" % (opts.log,os.path.splitext(os.path.basename(opts.config))[0])) image_directory, snapshot_directory = prepare_snapshot_and_image_folder(config.snapshot_prefix, iterations, config.image_save_iterations) best_err, best_acc = 100., 0. start_time = time.time() for ep in range(0, MAX_EPOCHS): for it, ((images_a, labels_a, com_a, M_a, cube_a, _), (images_b,labels_b, com_b, M_b, cube_b, _)) in \ enumerate(izip(train_loader_a,train_loader_b)): if images_a.size(0) != batch_size or images_b.size(0) != batch_size: continue images_a = Variable(images_a.cuda(opts.gpu)) images_b = Variable(images_b.cuda(opts.gpu)) labels_a = Variable(labels_a.cuda(opts.gpu)) labels_b = Variable(labels_b.cuda(opts.gpu)) com_a = Variable(com_a.cuda(opts.gpu)) com_b = Variable(com_b.cuda(opts.gpu)) trainer.dis.train() if opts.mode == 'pretrain': if (iterations+1) % 1000 == 0: trainer.dis_sch.step() trainer.gen_sch.step() print('lr %.8f' % trainer.dis_sch.get_lr()[0]) trainer.dis_update(images_a, labels_a, images_b, labels_b, com_a, com_b, config.hyperparameters) image_outputs = trainer.gen_update(images_a, labels_a, images_b, labels_b, config.hyperparameters) assembled_images = trainer.assemble_outputs(images_a, images_b, image_outputs) else: if (iterations+1) % 100 == 0: trainer.dis_sch.step() image_outputs = trainer.post_update(images_a, labels_a, images_b, labels_b,com_a,com_b, mode_idx, config.hyperparameters) assembled_images = trainer.assemble_outputs(images_a, images_b, image_outputs) # Dump training stats in log file if (iterations+1) % config.display == 0: elapsed_time = time.time() - start_time write_loss(iterations, max_iterations, trainer, train_writer, elapsed_time) start_time = time.time() if (iterations + 1) % config.image_display_iterations == 0: img_filename = '%s/gen.jpg' % (image_directory) torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1) if (iterations+1) % config.image_save_iterations == 0: if opts.mode == 'pretrain':# and (iterations+1) % (2*config.image_save_iterations) != 0: img_filename = '%s/gen_%08d.jpg' % (image_directory, iterations + 1) torchvision.utils.save_image(assembled_images.data / 2 + 0.5, img_filename, nrow=1) write_html(snapshot_directory + "/index.html", iterations + 1, \ config.image_save_iterations, image_directory) else: trainer.dis.eval() score, maxerr = 0, 0 num_samples = 0 maxJntError = [] meanJntError = 0 img2sav = None gt3D = [] joints = [] joints_imgcord = [] codec = cv2.VideoWriter_fourcc(*'XVID') vid = cv2.VideoWriter(os.path.join(image_directory,'gen.avi'), codec, 25, (128*2,128)) for tit, (test_images_b, test_labels_b, com_b, trans_b, cube_b, fn) in enumerate(test_loader_real): test_images_b = Variable(test_images_b.cuda(opts.gpu)) test_labels_b = Variable(test_labels_b.cuda(opts.gpu)) if mode_idx == 0: pred_pose, pred_post, _ = trainer.dis.regress_a(test_images_b) else: pred_pose, pred_post, _ = trainer.dis.regress_b(test_images_b) if True: pred_pose = trainer.vae.decode(pred_post) n = test_labels_b.size(0) gt_pose = test_labels_b.data.cpu().numpy().reshape((n,-1, 3)) pr_pose = pred_pose.data.cpu().numpy().reshape((n,-1, 3)) if tit < 20: for i in range(0, n, 4): real_img = visPair(di_b, test_images_b[i].data.cpu().numpy(), gt_pose[i].reshape((-1)), \ trans_b[i].numpy(), com_b[i].numpy(), cube_b[i].numpy(), 50.0) est_img = visPair(di_b, test_images_b[i].data.cpu().numpy(), pr_pose[i].reshape((-1)), \ trans_b[i].numpy(), com_b[i].numpy(), cube_b[i].numpy(), 50.0) vid.write(np.hstack((real_img,est_img)).astype('uint8')) both_img = np.vstack((real_img,est_img)) if True and tit < 8: if img2sav is None: img2sav = both_img else: img2sav = np.hstack((img2sav,both_img)) if 'nyu' in opts.config: restrictedJointsEval = np.array([0, 3, 6, 9, 12, 15, 18, 21, 24, 25, 27, 30, 31, 32]) gt_pose = gt_pose[:,restrictedJointsEval] pr_pose = pr_pose[:,restrictedJointsEval] for i in range(n): gt3D.append(gt_pose[i]*(cube_b.numpy()[0]/2.)+ com_b[i].numpy()) joints.append(pr_pose[i]*(cube_b.numpy()[0]/2.)+ com_b[i].numpy()) joints_imgcord.append(di_b.joints3DToImg(pr_pose[i]*(cube_b.numpy()[0]/2.)+ com_b[i].numpy())) score += meanJntError num_samples += test_images_b.size(0) cv2.imwrite(image_directory + '/_test.jpg', img2sav.astype('uint8')) vid.release() hpe = Evaluation(np.array(gt3D), np.array(joints)) mean_err = hpe.getMeanError() over_40 = 100. * hpe.getNumFramesWithinMaxDist(40) / len(gt3D) best_err = np.minimum(best_err, mean_err) best_acc = np.maximum(best_acc, over_40) print("------------ Mean err: {:.4f} ({:.4f}) mm, Max over 40mm: {:.2f} ({:.2f}) %".format(mean_err, best_err, over_40, best_acc)) # Save network weights if (iterations+1) % config.snapshot_save_iterations == 0: if opts.mode == 'pretrain': trainer.save(config.snapshot_prefix, iterations) elif 'estimate' in opts.mode: trainer.save(config.snapshot_prefix+'_est', iterations) iterations += 1 if iterations >= max_iterations: return
optimizer = torch.optim.Adam() else: optimizer = '' raise ValueError() model.cuda() # # DATA LOADER get_loader = get_data_loader(cfg.datasetname) train_data = get_loader() class_names = train_data.dataset.classes print('dataset len: {}'.format(len(train_data.dataset))) tb_dir = os.path.join(cfg.train_dir, cfg.backbone + '_' + cfg.datasetname, time.strftime("%h%d_%H")) writer = tbx.FileWriter(tb_dir) summary_out = [] global_step = 0 timer = Timer() for ep in range(start_epoch, cfg.max_epoch): if ep in cfg.lr_decay_epoches and cfg.solver == 'SGD': lr *= cfg.lr_decay adjust_learning_rate(optimizer, lr) print('adjusting learning rate {:.6f}'.format(lr)) for step, batch in enumerate(train_data): timer.tic() input, anchors_np, im_scale_list, image_ids, gt_boxes_list, rpn_targets, _, _ = batch
def __init__(self, model_dir, sub_dir=""): self.summary_writer = tb.FileWriter(os.path.join(model_dir, sub_dir))
def main(): # config model and lr num_anchors = len(cfg.anchor_ratios) * len(cfg.anchor_scales[0]) * len(cfg.anchor_shift) \ if isinstance(cfg.anchor_scales[0], list) else \ len(cfg.anchor_ratios) * len(cfg.anchor_scales) resnet = resnet50 if cfg.backbone == 'resnet50' else resnet101 detection_model = MaskRCNN if cfg.model_type.lower( ) == 'maskrcnn' else RetinaNet model = detection_model(resnet(pretrained=True, maxpool5=cfg.maxpool5), num_classes=cfg.num_classes, num_anchors=num_anchors, strides=cfg.strides, in_channels=cfg.in_channels, f_keys=cfg.f_keys, num_channels=256, is_training=False, activation=cfg.class_activation) lr = cfg.lr start_epoch = 0 if cfg.restore is not None: meta = load_net(cfg.restore, model) print(meta) if meta[0] >= 0: start_epoch = meta[0] + 1 lr = meta[1] print('Restored from %s, starting from %d epoch, lr:%.6f' % (cfg.restore, start_epoch, lr)) else: raise ValueError('restore is not set') model.cuda() model.eval() class_names = test_data.dataset.classes print('dataset len: {}'.format(len(test_data.dataset))) tb_dir = os.path.join(cfg.train_dir, cfg.backbone + '_' + cfg.datasetname, 'test', time.strftime("%h%d_%H")) writer = tbx.FileWriter(tb_dir) # main loop timer_all = Timer() timer_post = Timer() all_results1 = [] all_results2 = [] all_results_gt = [] for step, batch in enumerate(test_data): timer_all.tic() # NOTE: Targets is in NHWC order!! # input, anchors_np, im_scale_list, image_ids, gt_boxes_list = batch # input = everything2cuda(input) input_t, anchors_np, im_scale_list, image_ids, gt_boxes_list = batch input = everything2cuda(input_t, volatile=True) outs = model(input, gt_boxes_list=None, anchors_np=anchors_np) if cfg.model_type == 'maskrcnn': rpn_logit, rpn_box, rpn_prob, rpn_labels, rpn_bbtargets, rpn_bbwghts, anchors, \ rois, roi_img_ids, rcnn_logit, rcnn_box, rcnn_prob, rcnn_labels, rcnn_bbtargets, rcnn_bbwghts = outs outputs = [ rois, roi_img_ids, rpn_logit, rpn_box, rpn_prob, rcnn_logit, rcnn_box, rcnn_prob, anchors ] targets = [] elif cfg.model_type == 'retinanet': rpn_logit, rpn_box, rpn_prob, _, _, _ = outs outputs = [rpn_logit, rpn_box, rpn_prob] else: raise ValueError('Unknown model type: %s' % cfg.model_type) timer_post.tic() dets_dict = model.get_final_results( outputs, everything2cuda(anchors_np), score_threshold=0.01, max_dets=cfg.max_det_num * cfg.batch_size, overlap_threshold=cfg.overlap_threshold) if 'stage1' in dets_dict: Dets = dets_dict['stage1'] else: raise ValueError('No stage1 results:', dets_dict.keys()) Dets2 = dets_dict['stage2'] if 'stage2' in dets_dict else Dets t3 = timer_post.toc() t = timer_all.toc() formal_res1 = dataset.to_detection_format(copy.deepcopy(Dets), image_ids, im_scale_list) formal_res2 = dataset.to_detection_format(copy.deepcopy(Dets2), image_ids, im_scale_list) all_results1 += formal_res1 all_results2 += formal_res2 Dets_gt = [] for gb in gt_boxes_list: cpy_mask = gb[:, 4] >= 1 gb = gb[cpy_mask] n = cpy_mask.astype(np.int32).sum() res_gt = np.zeros((n, 6)) res_gt[:, :4] = gb[:, :4] res_gt[:, 4] = 1. res_gt[:, 5] = gb[:, 4] Dets_gt.append(res_gt) formal_res_gt = dataset.to_detection_format(Dets_gt, image_ids, im_scale_list) all_results_gt += formal_res_gt if step % cfg.log_image == 0: input_np = everything2numpy(input) summary_out = [] Is = single_shot.draw_detection(input_np, Dets, class_names=class_names) Is = Is.astype(np.uint8) summary_out += log_images(Is, image_ids, step, prefix='Detection/') Is = single_shot.draw_detection(input_np, Dets2, class_names=class_names) Is = Is.astype(np.uint8) summary_out += log_images(Is, image_ids, step, prefix='Detection2/') Imgs = single_shot.draw_gtboxes(input_np, gt_boxes_list, class_names=class_names) Imgs = Imgs.astype(np.uint8) summary_out += log_images(Imgs, image_ids, float(step), prefix='GT') for s in summary_out: writer.add_summary(s, float(step)) if step % cfg.display == 0: print(time.strftime("%H:%M:%S ") + 'Epoch %d iter %d: speed %.3fs (%.3fs)' % (0, step, t, t3) + ' ImageIds: ' + ', '.join(str(s) for s in image_ids), end='\r') res_dict = { 'stage1': all_results1, 'stage2': all_results2, 'gt': all_results_gt } return res_dict
def main(argv): (opts, args) = parser.parse_args(argv) global colorPlatte, bones, Evaluation if 'nyu' in opts.config: colorPlatte = utils.util.nyuColorIdx bones = utils.util.nyuBones Evaluation = NYUHandposeEvaluation elif 'icvl' in opts.config: colorPlatte = utils.util.icvlColorIdx bones = utils.util.icvlBones Evaluation = ICVLHandposeEvaluation # Load experiment setting assert isinstance(opts, object) config = NetConfig(opts.config) batch_size = config.hyperparameters['batch_size_pose'] max_iterations = 200000 #config.hyperparameters['max_iterations'] frac = opts.frac dataset_a = get_dataset(config.datasets['train_a']) dataset_b = get_dataset(config.datasets['train_b']) dataset_test = get_dataset(config.datasets['test_b']) train_loader_a = get_data_loader(dataset_a, batch_size, shuffle=True) train_loader_b = get_data_loader(dataset_b, batch_size, shuffle=True) test_loader_real = get_data_loader(dataset_test, 1, shuffle=True) cmd = "trainer=%s(config.hyperparameters)" % config.hyperparameters[ 'trainer'] local_dict = locals() exec(cmd, globals(), local_dict) trainer = local_dict['trainer'] iterations = 0 trainer.cuda(opts.gpu) dataset_a.pose_only = True dataset_b.pose_only = True if frac > 0. and frac < 1.: dataset_b.set_nmax(frac) di_a = dataset_a.di di_b = dataset_b.di dataset_a.sample_poses() dataset_b.sample_poses() ################################################################### # Setup logger and repare image outputs train_writer = tensorboardX.FileWriter( "%s/%s" % (opts.log, os.path.splitext(os.path.basename(opts.config))[0])) image_directory, snapshot_directory = prepare_snapshot_and_image_folder( config.snapshot_prefix, iterations, config.image_save_iterations) print('using %.2f percent of the labeled real data' % frac) start_time = time.time() for ep in range(0, MAX_EPOCHS): for it, ((labels_a), (labels_b)) in enumerate(izip(train_loader_a, train_loader_b)): if labels_a.size(0) != batch_size or labels_b.size( 0) != batch_size: continue labels_a = Variable(labels_a.cuda(opts.gpu)) labels_b = Variable(labels_b.cuda(opts.gpu)) labels = labels_a if frac > 0.: labels = torch.cat((labels_a, labels_b), 0) if (iterations + 1) % 1000 == 0: trainer.vae_sch.step() recon_pose = trainer.vae_update(labels, config.hyperparameters) # Dump training stats in log file if (iterations + 1) % config.display == 0: elapsed_time = time.time() - start_time write_loss(iterations, max_iterations, trainer, train_writer, elapsed_time) start_time = time.time() if (iterations + 1) % (10 * config.image_save_iterations) == 0: if True: score, maxerr = 0, 0 num_samples = 0 maxJntError = [] img2sav = None gt3D = [] joints = [] for tit, (test_images_b, test_labels_b, com_b, trans_b, cube_b, _) in enumerate(test_loader_real): test_images_b = Variable(test_images_b.cuda(opts.gpu)) test_labels_b = Variable(test_labels_b.cuda(opts.gpu)) pred_pose = trainer.vae.decode( trainer.vae.encode(test_labels_b)[1]) gt3D.append(test_labels_b.data.cpu().numpy().reshape((-1, 3))*(cube_b.numpy()[0]/2.) +\ com_b.numpy()) joints.append(pred_pose.data.cpu().numpy().reshape((-1, 3))*(cube_b.numpy()[0]/2.) +\ com_b.numpy()) if True and tit < 8: real_img = visPair(di_b, test_images_b.data.cpu().numpy(), test_labels_b.data.cpu().numpy(), \ trans_b.numpy(), com_b.numpy(), cube_b.numpy(), 50.0) est_img = visPair(di_b, test_images_b.data.cpu().numpy(), pred_pose.data.cpu().numpy(), \ trans_b.numpy(), com_b.numpy(), cube_b.numpy(), 50.0) if img2sav is None: img2sav = np.vstack((real_img, est_img)) else: img2sav = np.hstack( (img2sav, np.vstack((real_img, est_img)))) num_samples += test_images_b.size(0) cv2.imwrite(image_directory + '/_test.jpg', img2sav.astype('uint8')) #maxerr = Evaluation.plotError(maxJntError, image_directory + '/maxJntError.txt') hpe = Evaluation(np.array(gt3D), np.array(joints)) print("Mean error: {}mm, max error: {}mm".format( hpe.getMeanError(), hpe.getMaxError())) # Save network weights if (iterations + 1) % (4 * config.snapshot_save_iterations) == 0: trainer.save_vae(config.snapshot_prefix, iterations, 2 + frac) iterations += 1 if iterations >= max_iterations: return
def main(): # config model and lr num_anchors = len(cfg.anchor_ratios) * len(cfg.anchor_scales[0]) \ if isinstance(cfg.anchor_scales[0], list) else \ len(cfg.anchor_ratios) * len(cfg.anchor_scales) resnet = resnet50 if cfg.backbone == 'resnet50' else resnet101 detection_model = MaskRCNN if cfg.model_type.lower( ) == 'maskrcnn' else RetinaNet model = detection_model(resnet(pretrained=True), num_classes=cfg.num_classes, num_anchors=num_anchors, strides=cfg.strides, in_channels=cfg.in_channels, f_keys=cfg.f_keys, num_channels=256, is_training=False, activation=cfg.class_activation) lr = cfg.lr start_epoch = 0 if cfg.restore is not None: meta = load_net(cfg.restore, model) print(meta) if meta[0] >= 0: start_epoch = meta[0] + 1 lr = meta[1] print('Restored from %s, starting from %d epoch, lr:%.6f' % (cfg.restore, start_epoch, lr)) else: raise ValueError('restore is not set') model.cuda() model.eval() ANCHORS = np.vstack( [anc.reshape([-1, 4]) for anc in test_data.dataset.ANCHORS]) model.anchors = everything2cuda(ANCHORS.astype(np.float32)) class_names = test_data.dataset.classes print('dataset len: {}'.format(len(test_data.dataset))) tb_dir = os.path.join(cfg.train_dir, cfg.backbone + '_' + cfg.datasetname, 'test', time.strftime("%h%d_%H")) writer = tbx.FileWriter(tb_dir) summary_out = [] # main loop timer_all = Timer() timer_post = Timer() all_results1 = [] all_results2 = [] all_results_gt = [] for step, batch in enumerate(test_data): timer_all.tic() # NOTE: Targets is in NHWC order!! input, image_ids, gt_boxes_list, image_ori = batch input = everything2cuda(input) outs = model(input) timer_post.tic() dets_dict = model.get_final_results( score_threshold=0.05, max_dets=cfg.max_det_num * cfg.batch_size, overlap_threshold=cfg.overlap_threshold) if 'stage1' in dets_dict: Dets = dets_dict['stage1'] else: raise ValueError('No stage1 results:', dets_dict.keys()) Dets2 = dets_dict['stage2'] if 'stage2' in dets_dict else Dets t3 = timer_post.toc() t = timer_all.toc() formal_res1 = dataset.to_detection_format( copy.deepcopy(Dets), image_ids, ori_sizes=[im.shape for im in image_ori]) formal_res2 = dataset.to_detection_format( copy.deepcopy(Dets2), image_ids, ori_sizes=[im.shape for im in image_ori]) all_results1 += formal_res1 all_results2 += formal_res2 if step % cfg.log_image == 0: input_np = everything2numpy(input) summary_out = [] Is = single_shot.draw_detection(input_np, Dets, class_names=class_names) Is = Is.astype(np.uint8) summary_out += log_images(Is, image_ids, step, prefix='Detection/') Is = single_shot.draw_detection(input_np, Dets2, class_names=class_names) Is = Is.astype(np.uint8) summary_out += log_images(Is, image_ids, step, prefix='Detection2/') Imgs = single_shot.draw_gtboxes(input_np, gt_boxes_list, class_names=class_names) Imgs = Imgs.astype(np.uint8) summary_out += log_images(Imgs, image_ids, float(step), prefix='GT') for s in summary_out: writer.add_summary(s, float(step)) if step % cfg.display == 0: print(time.strftime("%H:%M:%S ") + 'Epoch %d iter %d: speed %.3fs (%.3fs)' % (0, step, t, t3) + ' ImageIds: ' + ', '.join(str(s) for s in image_ids), end='\r') res_dict = { 'stage1': all_results1, 'stage2': all_results2, 'gt': all_results_gt } return res_dict