def initialize_model_from_cfg(args, gpu_id=0): """Initialize a model from the global cfg. Loads test-time weights and set to evaluation mode. """ model = model_builder.Generalized_RCNN() model.eval() if args.cuda: model.cuda() if args.load_ckpt: load_name = args.load_ckpt logger.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(model, checkpoint['model']) if args.load_detectron: logger.info("loading detectron weights %s", args.load_detectron) load_detectron_weight(model, args.load_detectron) model = mynn.DataParallel(model, cpu_keywords=['im_info', 'roidb'], minibatch=True) return model
def main(): """Main function""" args = parse_args() cfg_from_file(args.cfg_file) cfg.DANET.REFINEMENT = EasyDict(cfg.DANET.REFINEMENT) cfg.MSRES_MODEL.EXTRA = EasyDict(cfg.MSRES_MODEL.EXTRA) if args.use_opendr: from lib.utils.opendr_render import opendr_render if cfg.DANET.SMPL_MODEL_TYPE == 'male': smpl_model_path = './data/SMPL_data/basicmodel_m_lbs_10_207_0_v1.0.0.pkl' elif cfg.DANET.SMPL_MODEL_TYPE == 'neutral': smpl_model_path = './data/SMPL_data/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl' elif cfg.DANET.SMPL_MODEL_TYPE == 'female': smpl_model_path = './data/SMPL_data/basicModel_f_lbs_10_207_0_v1.0.0.pkl' dr_render = opendr_render(model_path=smpl_model_path) if not os.path.exists(args.out_dir): os.makedirs(args.out_dir) ### Model ### model = DaNet().cuda() ### Load checkpoint if args.load_ckpt: load_name = args.load_ckpt logging.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(model, checkpoint) del checkpoint model = mynn.DataParallel(model, minibatch=False) model.eval() img_path_list = [ os.path.join(args.img_dir, name) for name in os.listdir(args.img_dir) if name.endswith('.jpg') ] for i, path in enumerate(img_path_list): image = Image.open(path).convert('RGB') img_id = path.split('/')[-1][:-4] image_tensor = torchvision.transforms.ToTensor()(image).unsqueeze( 0).cuda() # run inference pred_results = model.module.infer_net(image_tensor) para_pred = pred_results['para'] cam_pred = para_pred[:, 0:3].contiguous() beta_pred = para_pred[:, 3:13].contiguous() Rs_pred = para_pred[:, 13:].contiguous().view(-1, 24, 3, 3) smpl_pts = model.module.iuv2smpl.smpl(beta_pred, Rs=Rs_pred, get_skin=True) kps3ds_pred = smpl_pts['cocoplus'] vert_pred = smpl_pts['verts'] # input image image_np = image_tensor[0].cpu().numpy() image_np = np.transpose(image_np, (1, 2, 0)) ones_np = np.ones(image_np.shape[:2]) * 255 ones_np = ones_np[:, :, None] image_in_rgba = np.concatenate((image_np, ones_np), axis=2) # estimated global IUV global_iuv = iuv_map2img( *pred_results['visualization']['iuv_pred'])[0].cpu().numpy() global_iuv = np.transpose(global_iuv, (1, 2, 0)) global_iuv = resize(global_iuv, image_np.shape[:2]) global_iuv_rgba = np.concatenate((global_iuv, ones_np), axis=2) # estimated patial IUV part_iuv_pred = pred_results['visualization']['part_iuv_pred'][0] p_iuv_vis = [] for i in range(part_iuv_pred.size(0)): p_u_vis, p_v_vis, p_i_vis = [ part_iuv_pred[i, iuv].unsqueeze(0) for iuv in range(3) ] if p_u_vis.size(1) == 25: p_iuv_vis_i = iuv_map2img(p_u_vis.detach(), p_v_vis.detach(), p_i_vis.detach()) else: p_iuv_vis_i = iuv_map2img( p_u_vis.detach(), p_v_vis.detach(), p_i_vis.detach(), ind_mapping=[0] + model.module.img2iuv.dp2smpl_mapping[i]) p_iuv_vis.append(p_iuv_vis_i) part_iuv = torch.cat(p_iuv_vis, dim=0) part_iuv = make_grid(part_iuv, nrow=6, padding=0).cpu().numpy() part_iuv = np.transpose(part_iuv, (1, 2, 0)) part_iuv_rgba = np.concatenate( (part_iuv, np.ones(part_iuv.shape[:2])[:, :, None] * 255), axis=2) # rendered IUV of the predicted SMPL model smpl_projection = model.module.iuv2smpl.make_uv_image( Rs=Rs_pred[0].unsqueeze(0), beta=beta_pred[0].unsqueeze(0), cam=cam_pred[0].unsqueeze(0), add_smpl_joint=True) render_iuv = smpl_projection['render_image'].squeeze(0).cpu().numpy() render_iuv = np.transpose(render_iuv, (1, 2, 0)) render_iuv = resize(render_iuv, image_np.shape[:2]) img_render_iuv = image_np.copy() img_render_iuv[render_iuv > 0] = render_iuv[render_iuv > 0] img_render_iuv_rgba = np.concatenate((img_render_iuv, ones_np), axis=2) img_vis_list = [ image_in_rgba, global_iuv_rgba, part_iuv_rgba, img_render_iuv_rgba ] if args.use_opendr: # visualize the predicted SMPL model using the opendr renderer joint_pelvis = torch.mean(kps3ds_pred[:, [2, 3]], dim=1).unsqueeze(1) vert_centered = vert_pred - joint_pelvis K, _, _ = model.module.iuv2smpl.camera_matrix(cam_pred) _, _, img_smpl, smpl_rgba = dr_render.render( image_tensor[0].cpu().numpy(), K.cpu().numpy(), vert_centered[0].cpu().numpy()) img_smpl_rgba = np.concatenate((img_smpl, ones_np), axis=2) img_vis_list.extend([img_smpl_rgba, smpl_rgba]) img_vis = np.concatenate(img_vis_list, axis=1) imsave(os.path.join(args.out_dir, img_id + '_result.png'), img_vis)
def main(): """Main function""" args = parse_args() print('Called with args:') print(args) if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") if args.cuda or cfg.NUM_GPUS > 0: cfg.CUDA = True else: raise ValueError("Need Cuda device to run !") if args.dataset == "coco2017": cfg.TRAIN.DATASETS = ('coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 81 elif args.dataset == "keypoints_coco2017": cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train', ) cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError("Unexpected args.dataset: {}".format(args.dataset)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) ### Adaptively adjust some configs ### original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH if args.batch_size is None: args.batch_size = original_batch_size cfg.NUM_GPUS = torch.cuda.device_count() assert (args.batch_size % cfg.NUM_GPUS) == 0, \ 'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS) cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS print('Batch size change from {} (in config file) to {}'.format( original_batch_size, args.batch_size)) print('NUM_GPUs: %d, TRAIN.IMS_PER_BATCH: %d' % (cfg.NUM_GPUS, cfg.TRAIN.IMS_PER_BATCH)) if args.num_workers is not None: cfg.DATA_LOADER.NUM_THREADS = args.num_workers print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS) ### Adjust learning based on batch size change linearly old_base_lr = cfg.SOLVER.BASE_LR cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size print('Adjust BASE_LR linearly according to batch size change: {} --> {}'. format(old_base_lr, cfg.SOLVER.BASE_LR)) ### Overwrite some solver settings from command line arguments if args.optimizer is not None: cfg.SOLVER.TYPE = args.optimizer if args.lr is not None: cfg.SOLVER.BASE_LR = args.lr if args.lr_decay_gamma is not None: cfg.SOLVER.GAMMA = args.lr_decay_gamma timers = defaultdict(Timer) ### Dataset ### timers['roidb'].tic() roidb, ratio_list, ratio_index = combined_roidb_for_training( cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES) timers['roidb'].toc() train_size = len(roidb) logger.info('{:d} roidb entries'.format(train_size)) logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time) sampler = MinibatchSampler(ratio_list, ratio_index) dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True) dataloader = torch.utils.data.DataLoader( dataset, batch_size=args.batch_size, sampler=sampler, num_workers=cfg.DATA_LOADER.NUM_THREADS, collate_fn=collate_minibatch) assert_and_infer_cfg() ### Model ### maskRCNN = Generalized_RCNN() if cfg.CUDA: maskRCNN.cuda() ### Optimizer ### bias_params = [] nonbias_params = [] for key, value in dict(maskRCNN.named_parameters()).items(): if value.requires_grad: if 'bias' in key: bias_params.append(value) else: nonbias_params.append(value) params = [{ 'params': nonbias_params, 'lr': cfg.SOLVER.BASE_LR, 'weight_decay': cfg.SOLVER.WEIGHT_DECAY }, { 'params': bias_params, 'lr': cfg.SOLVER.BASE_LR * (cfg.SOLVER.BIAS_DOUBLE_LR + 1), 'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0 }] if cfg.SOLVER.TYPE == "SGD": optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM) elif cfg.SOLVER.TYPE == "Adam": optimizer = torch.optim.Adam(params) ### Load checkpoint if args.load_ckpt: load_name = args.load_ckpt logging.info("loading checkpoint %s", load_name) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) if args.resume: assert checkpoint['iters_per_epoch'] == train_size // args.batch_size, \ "iters_per_epoch should match for resume" # There is a bug in optimizer.load_state_dict on Pytorch 0.3.1. # However it's fixed on master. # optimizer.load_state_dict(checkpoint['optimizer']) misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer']) if checkpoint['step'] == (checkpoint['iters_per_epoch'] - 1): # Resume from end of an epoch args.start_epoch = checkpoint['epoch'] + 1 args.start_iter = 0 else: # Resume from the middle of an epoch. # NOTE: dataloader is not synced with previous state args.start_epoch = checkpoint['epoch'] args.start_iter = checkpoint['step'] + 1 del checkpoint torch.cuda.empty_cache() if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values) logging.info("loading Detectron weights %s", args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) lr = optimizer.param_groups[0][ 'lr'] # lr of non-bias parameters, for commmand line outputs. maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True) ### Training Setups ### args.run_name = misc_utils.get_run_name() output_dir = misc_utils.get_output_dir(args, args.run_name) args.cfg_filename = os.path.basename(args.cfg_file) if not args.no_save: if not os.path.exists(output_dir): os.makedirs(output_dir) blob = {'cfg': yaml.dump(cfg), 'args': args} with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f: pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL) if args.use_tfboard: from tensorboardX import SummaryWriter # Set the Tensorboard logger tblogger = SummaryWriter(output_dir) ### Training Loop ### maskRCNN.train() training_stats = TrainingStats( args, args.disp_interval, tblogger if args.use_tfboard and not args.no_save else None) iters_per_epoch = int(train_size / args.batch_size) # drop last args.iters_per_epoch = iters_per_epoch ckpt_interval_per_epoch = iters_per_epoch // args.ckpt_num_per_epoch try: logger.info('Training starts !') args.step = args.start_iter global_step = iters_per_epoch * args.start_epoch + args.step for args.epoch in range(args.start_epoch, args.start_epoch + args.num_epochs): # ---- Start of epoch ---- # adjust learning rate if args.lr_decay_epochs and args.epoch == args.lr_decay_epochs[ 0] and args.start_iter == 0: args.lr_decay_epochs.pop(0) net_utils.decay_learning_rate(optimizer, lr, cfg.SOLVER.GAMMA) lr *= cfg.SOLVER.GAMMA for args.step, input_data in zip( range(args.start_iter, iters_per_epoch), dataloader): for key in input_data: if key != 'roidb': # roidb is a list of ndarrays with inconsistent length input_data[key] = list(map(Variable, input_data[key])) training_stats.IterTic() net_outputs = maskRCNN(**input_data) training_stats.UpdateIterStats(net_outputs) loss = net_outputs['total_loss'] optimizer.zero_grad() loss.backward() optimizer.step() training_stats.IterToc() if (args.step + 1) % ckpt_interval_per_epoch == 0: net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer) if args.step % args.disp_interval == 0: log_training_stats(training_stats, global_step, lr) global_step += 1 # ---- End of epoch ---- # save checkpoint net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer) # reset starting iter number after first epoch args.start_iter = 0 # ---- Training ends ---- if iters_per_epoch % args.disp_interval != 0: # log last stats at the end log_training_stats(training_stats, global_step, lr) except (RuntimeError, KeyboardInterrupt): logger.info('Save ckpt on exception ...') net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer) logger.info('Save ckpt done.') stack_trace = traceback.format_exc() print(stack_trace) finally: if args.use_tfboard and not args.no_save: tblogger.close()
def main(): args = parse_args() # Configuration cfg_from_file(args.cfg_file) # Adaptively adjust some configs set_cfg_value() print(pprint.pformat(cfg)) # Create model model = MaskRCNN(config=cfg) if cfg.GPU_COUNT: model = model.cuda() # Select weights file to load if args.last: model_path = find_last(model)[1] print('Continue training on {}'.format(model_path)) else: model_path = os.path.join(cfg.ROOT_DIR, cfg.TRAIN.WEIGHTS) print("Loading weights ", model_path) model.load_weights(model_path) # Multi-GPU support if cfg.GPU_COUNT > 1: model = mynn.DataParallel(model, cpu_keywords=['im_info', 'roidb'], minibatch=True) # save printing logs to file now = datetime.datetime.now() save_log_dir = os.path.join( cfg.TRAIN.LOG_DIR, "{}{:%Y%m%dT%H%M}".format(cfg.MODEL.NAME.lower(), now)) # Train or evaluate if args.command == "train": # save config and model fprintf_log(pprint.pformat(cfg), save_log_dir) fprintf_log(model, save_log_dir) # Training dataset. Use the training set and 35K from the # validation set, as as in the Mask RCNN paper. dataset_train = CocoDataset() dataset_train.load_coco(cfg.DATASET.PATH, "train", year=cfg.DATASET.YEAR, auto_download=cfg.DATASET.DOWNLOAD) dataset_train.load_coco(cfg.DATASET.PATH, "valminusminival", year=cfg.DATASET.YEAR, auto_download=cfg.DATASET.DOWNLOAD) dataset_train.prepare() # Validation dataset dataset_val = CocoDataset() dataset_val.load_coco(cfg.DATASET.PATH, "minival", year=cfg.DATASET.YEAR, auto_download=cfg.DATASET.DOWNLOAD) dataset_val.prepare() # *** This training schedule is an example. Update to your needs *** # Training - Stage 1 print("Training network heads") train_model(model, dataset_train, dataset_val, learning_rate=cfg.SOLVER.BASE_LR, epochs=cfg.TRAIN.TRAIN_SCHEDULE[0], layers='heads') # Training - Stage 2 # Finetune layers from ResNet stage 4 and up print("Fine tune Resnet stage 4 and up") train_model(model, dataset_train, dataset_val, learning_rate=cfg.SOLVER.BASE_LR, epochs=cfg.TRAIN.TRAIN_SCHEDULE[1], layers='4+') # Training - Stage 3 # Fine tune all layers print("Fine tune all layers") train_model(model, dataset_train, dataset_val, learning_rate=cfg.SOLVER.BASE_LR / 10, epochs=cfg.TRAIN.TRAIN_SCHEDULE[2], layers='all') elif args.command == "evaluate": # Validation dataset dataset_val = CocoDataset() coco = dataset_val.load_coco(cfg.DATASET.PATH, "minival", year=cfg.DATASET.YEAR, return_coco=True, auto_download=cfg.DATASET.DOWNLOAD) dataset_val.prepare() print("Running COCO evaluation on {} images.".format(cfg.TEST.NUM_IMG)) evaluate_coco(model, dataset_val, coco, "bbox", limit=int(cfg.TEST.NUM_IMG)) evaluate_coco(model, dataset_val, coco, "segm", limit=int(cfg.TEST.NUM_IMG)) else: print("'{}' is not recognized. " "Use 'train' or 'evaluate'".format(args.command))
def main(): """main function""" if not torch.cuda.is_available(): sys.exit("Need a CUDA device to run the code.") args = parse_args() print('Called with args:') print(args) assert args.image_dir or args.images assert bool(args.image_dir) ^ bool(args.images) if args.dataset.startswith("coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = len(dataset.classes) elif args.dataset.startswith("keypoints_coco"): dataset = datasets.get_coco_dataset() cfg.MODEL.NUM_CLASSES = 2 else: raise ValueError('Unexpected dataset name: {}'.format(args.dataset)) print('load cfg from file: {}'.format(args.cfg_file)) cfg_from_file(args.cfg_file) if args.set_cfgs is not None: cfg_from_list(args.set_cfgs) assert bool(args.load_ckpt) ^ bool(args.load_detectron), \ 'Exactly one of --load_ckpt and --load_detectron should be specified.' cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights assert_and_infer_cfg() maskRCNN = Generalized_RCNN() if args.cuda: maskRCNN.cuda() if args.load_ckpt: load_name = args.load_ckpt print("loading checkpoint %s" % (load_name)) checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage) net_utils.load_ckpt(maskRCNN, checkpoint['model']) if args.load_detectron: print("loading detectron weights %s" % args.load_detectron) load_detectron_weight(maskRCNN, args.load_detectron) maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU maskRCNN.eval() if args.image_dir: imglist = misc_utils.get_imagelist_from_dir(args.image_dir) else: imglist = args.images num_images = len(imglist) if not os.path.exists(args.output_dir): os.makedirs(args.output_dir) for i in xrange(num_images): print('img', i) im = cv2.imread(imglist[i]) assert im is not None timers = defaultdict(Timer) cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers) im_name, _ = os.path.splitext(os.path.basename(imglist[i])) vis_utils.vis_one_image( im[:, :, ::-1], # BGR -> RGB for visualization im_name, args.output_dir, cls_boxes, cls_segms, cls_keyps, dataset=dataset, box_alpha=0.3, show_class=True, thresh=0.7, kp_thresh=2) if args.merge_pdfs and num_images > 1: merge_out_path = '{}/results.pdf'.format(args.output_dir) if os.path.exists(merge_out_path): os.remove(merge_out_path) command = "pdfunite {}/*.pdf {}".format(args.output_dir, merge_out_path) subprocess.call(command, shell=True)