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)