def get_roidb(imdb_name):
imdb = get_imdb(imdb_name)
print('Loaded dataset `{:s}` for training'.format(imdb.name))
imdb.set_proposal_method(cfg.TRAIN.PROPOSAL_METHOD)
print('Set proposal method: {:s}'.format(cfg.TRAIN.PROPOSAL_METHOD))
roidb = get_training_roidb(imdb)
return roidb
def combined_roidb(imdb_names):
"""
Combine multiple roidbs
"""
def get_roidb(imdb_name):
imdb = get_imdb(imdb_name)
print('Loaded dataset `{:s}` for training'.format(imdb.name))
imdb.set_proposal_method(cfg.TRAIN.PROPOSAL_METHOD)
print('Set proposal method: {:s}'.format(cfg.TRAIN.PROPOSAL_METHOD))
roidb = get_training_roidb(imdb)
return roidb
roidbs = [get_roidb(s) for s in imdb_names.split('+')]
roidb = roidbs[0]
if len(roidbs) > 1:
for r in roidbs[1:]:
roidb.extend(r)
tmp = get_imdb(imdb_names.split('+')[1])
imdb = datasets.imdb.imdb(imdb_names, tmp.classes)
else:
imdb = get_imdb(imdb_names)
return imdb, roidb
def frcnn_demo(args):
if args['cfg_file']:
cfg_from_file(args['cfg_file'])
if args['set_cfgs']:
cfg_from_list(args['set_cfgs'])
#cfg.TEST.HAS_RPN = True # Use RPN for proposals
print('Using config:')
pprint.pprint(cfg)
imdb = get_imdb(args['imdb_name'])
# load network
if args['net'] == 'vgg16':
net = vgg16()
elif args['net'] == 'res50':
net = resnetv1(num_layers=50)
elif args['net'] == 'res101':
net = resnetv1(num_layers=101)
elif args['net'] == 'res152':
net = resnetv1(num_layers=152)
else:
raise NotImplementedError
# load model
net.create_architecture(imdb.num_classes,
tag='default',
anchor_scales=cfg.ANCHOR_SCALES,
anchor_ratios=cfg.ANCHOR_RATIOS)
net.load_state_dict(torch.load(args['model']))
net.eval()
net.cuda()
print('Loaded network {:s}'.format(args['model']))
for im_name in args['im_names']:
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print('Demo for data/demo/{}'.format(im_name))
demo(net,
im_name,
imdb,
args['output_dir'],
thresh=args['score_thresh'])
def frcnn_test(imdbtest_name,
network,
model,
output_dir,
score_thresh,
max_per_image,
write_images=False):
"""
args = {#'imdb_name':imdb_name,
'imdbtest_name':imdbtest_name,
'net':network,
'cfg_file':None,
'set_cfgs':None,
'tag':tag,
'comp_mode':comp,
'max_per_image':max_per_image,
'output_dir':output_dir,
'model':model}
"""
#if cfg_file:
# cfg_from_file(cfg_file)
#if set_cfgs:
# cfg_from_list(set_cfgs)
print('Using config:')
pprint.pprint(cfg)
#filename = os.path.splitext(os.path.basename(args['model']))[0]
imdb_test = get_imdb(imdbtest_name)
#imdb_test.competition_mode(args['comp_mode'])
# load network
if network == 'vgg16':
net = vgg16()
elif network == 'res50':
net = resnetv1(num_layers=50)
elif network == 'res101':
net = resnetv1(num_layers=101)
elif network == 'res152':
net = resnetv1(num_layers=152)
else:
raise NotImplementedError
# load model
net.create_architecture(imdb_test.num_classes,
tag='default',
anchor_scales=cfg.ANCHOR_SCALES,
anchor_ratios=cfg.ANCHOR_RATIOS)
net.eval()
# def count_parameters(model):
# return sum(p.numel() for p in model.parameters() if p.requires_grad)
# print(count_parameters(net))
# exit()
net.cuda()
print(('Loading model check point from {:s}').format(model))
net.load_state_dict(torch.load(model))
print('Loaded.')
all_boxes = test_net(net,
imdb_test,
output_dir,
max_per_image=max_per_image,
thresh=score_thresh)
if write_images:
num_images = len(imdb_test.image_index)
for i in range(num_images):
im_path = imdb_test.image_path_at(i)
im = cv2.imread(im_path)
im = im[:, :, (2, 1, 0)]
fig, ax = plt.subplots(1, 1)
ax.imshow(im, aspect='equal')
for t_i in all_boxes[1][i]:
ax.add_patch(
plt.Rectangle((t_i[0], t_i[1]),
t_i[2] - t_i[0],
t_i[3] - t_i[1],
fill=False,
linewidth=1.0))
im_name = osp.basename(im_path)
im_output = osp.join(output_dir, im_name)
plt.axis('off')
plt.tight_layout()
plt.draw()
plt.savefig(im_output)
plt.close()