def save_model_to_weights_file(weights_file, model):
"""Stash model weights in a dictionary and pickle them to a file. We map
GPU device scoped names to unscoped names (e.g., 'gpu_0/conv1_w' ->
'conv1_w').
"""
logger.info('Saving parameters and momentum to {}'.format(
os.path.abspath(weights_file)))
blobs = {}
# Save all parameters
for param in model.params:
scoped_name = str(param)
unscoped_name = c2_utils.UnscopeName(scoped_name)
if unscoped_name not in blobs:
logger.debug(' {:s} -> {:s}'.format(scoped_name, unscoped_name))
blobs[unscoped_name] = workspace.FetchBlob(scoped_name)
# Save momentum
for param in model.TrainableParams():
scoped_name = str(param) + '_momentum'
unscoped_name = c2_utils.UnscopeName(scoped_name)
if unscoped_name not in blobs:
logger.debug(' {:s} -> {:s}'.format(scoped_name, unscoped_name))
blobs[unscoped_name] = workspace.FetchBlob(scoped_name)
# Save preserved blobs
for scoped_name in workspace.Blobs():
if scoped_name.startswith('__preserve__/'):
unscoped_name = c2_utils.UnscopeName(scoped_name)
if unscoped_name not in blobs:
logger.debug(' {:s} -> {:s} (preserved)'.format(
scoped_name, unscoped_name))
blobs[unscoped_name] = workspace.FetchBlob(scoped_name)
cfg_yaml = yaml.dump(cfg)
save_object(dict(blobs=blobs, cfg=cfg_yaml), weights_file)
def multi_gpu_generate_rpn_on_dataset(weights_file, dataset_name,
_proposal_file_ignored, num_images,
output_dir):
"""Multi-gpu inference on a dataset."""
# Retrieve the test_net binary path
binary_dir = envu.get_runtime_dir()
binary_ext = envu.get_py_bin_ext()
binary = os.path.join(binary_dir, 'test_net' + binary_ext)
assert os.path.exists(binary), 'Binary \'{}\' not found'.format(binary)
# Pass the target dataset via the command line
opts = ['TEST.DATASETS', '("{}",)'.format(dataset_name)]
opts += ['TEST.WEIGHTS', weights_file]
# Run inference in parallel in subprocesses
outputs = subprocess_utils.process_in_parallel('rpn_proposals', num_images,
binary, output_dir, opts)
# Collate the results from each subprocess
boxes, scores, ids = [], [], []
for rpn_data in outputs:
boxes += rpn_data['boxes']
scores += rpn_data['scores']
ids += rpn_data['ids']
rpn_file = os.path.join(output_dir, 'rpn_proposals.pkl')
cfg_yaml = yaml.dump(cfg)
save_object(dict(boxes=boxes, scores=scores, ids=ids, cfg=cfg_yaml),
rpn_file)
logger.info('Wrote RPN proposals to {}'.format(os.path.abspath(rpn_file)))
return boxes, scores, ids, rpn_file
def evaluate_proposal_file(dataset, proposal_file, output_dir):
"""Evaluate box proposal average recall."""
roidb = dataset.get_roidb(gt=True, proposal_file=proposal_file)
results = task_evaluation.evaluate_box_proposals(dataset, roidb)
task_evaluation.log_box_proposal_results(results)
recall_file = os.path.join(output_dir, 'rpn_proposal_recall.pkl')
save_object(results, recall_file)
return results
def _do_detection_eval(json_dataset, res_file, output_dir):
coco_dt = json_dataset.COCO.loadRes(str(res_file))
coco_eval = COCOeval(json_dataset.COCO, coco_dt, 'bbox')
coco_eval.evaluate()
coco_eval.accumulate()
_log_detection_eval_metrics(json_dataset, coco_eval)
eval_file = os.path.join(output_dir, 'detection_results.pkl')
save_object(coco_eval, eval_file)
logger.info('Wrote json eval results to: {}'.format(eval_file))
return coco_eval
def _do_keypoint_eval(json_dataset, res_file, output_dir):
ann_type = 'keypoints'
imgIds = json_dataset.COCO.getImgIds()
imgIds.sort()
coco_dt = json_dataset.COCO.loadRes(res_file)
coco_eval = COCOeval(json_dataset.COCO, coco_dt, ann_type)
coco_eval.params.imgIds = imgIds
coco_eval.evaluate()
coco_eval.accumulate()
eval_file = os.path.join(output_dir, 'keypoint_results.pkl')
save_object(coco_eval, eval_file)
logger.info('Wrote json eval results to: {}'.format(eval_file))
coco_eval.summarize()
return coco_eval
def generate_rpn_on_range(weights_file,
dataset_name,
_proposal_file_ignored,
output_dir,
ind_range=None,
gpu_id=0):
"""Run inference on all images in a dataset or over an index range of images
in a dataset using a single GPU.
"""
assert cfg.MODEL.RPN_ONLY or cfg.MODEL.FASTER_RCNN
roidb, start_ind, end_ind, total_num_images = get_roidb(
dataset_name, ind_range)
logger.info('Output will be saved to: {:s}'.format(
os.path.abspath(output_dir)))
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
nu.initialize_gpu_from_weights_file(
model,
weights_file,
gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
boxes, scores, ids = generate_proposals_on_roidb(
model,
roidb,
start_ind=start_ind,
end_ind=end_ind,
total_num_images=total_num_images,
gpu_id=gpu_id,
)
cfg_yaml = yaml.dump(cfg)
if ind_range is not None:
rpn_name = 'rpn_proposals_range_%s_%s.pkl' % tuple(ind_range)
else:
rpn_name = 'rpn_proposals.pkl'
rpn_file = os.path.join(output_dir, rpn_name)
save_object(dict(boxes=boxes, scores=scores, ids=ids, cfg=cfg_yaml),
rpn_file)
logger.info('Wrote RPN proposals to {}'.format(os.path.abspath(rpn_file)))
return boxes, scores, ids, rpn_file
def _do_python_eval(json_dataset, salt, output_dir='output'):
info = voc_info(json_dataset)
year = info['year']
anno_path = info['anno_path']
image_set_path = info['image_set_path']
devkit_path = info['devkit_path']
cachedir = os.path.join(devkit_path, 'annotations_cache')
aps = []
# The PASCAL VOC metric changed in 2010
use_07_metric = True if int(year) < 2010 else False
logger.info('VOC07 metric? ' + ('Yes' if use_07_metric else 'No'))
if not os.path.isdir(output_dir):
os.mkdir(output_dir)
for _, cls in enumerate(json_dataset.classes):
if cls == '__background__':
continue
filename = _get_voc_results_file_template(json_dataset,
salt).format(cls)
rec, prec, ap = voc_eval(filename,
anno_path,
image_set_path,
cls,
cachedir,
ovthresh=0.5,
use_07_metric=use_07_metric)
aps += [ap]
logger.info('AP for {} = {:.4f}'.format(cls, ap))
res_file = os.path.join(output_dir, cls + '_pr.pkl')
save_object({'rec': rec, 'prec': prec, 'ap': ap}, res_file)
logger.info('Mean AP = {:.4f}'.format(np.mean(aps)))
logger.info('~~~~~~~~')
logger.info('Results:')
for ap in aps:
logger.info('{:.3f}'.format(ap))
logger.info('{:.3f}'.format(np.mean(aps)))
logger.info('~~~~~~~~')
logger.info('')
logger.info('----------------------------------------------------------')
logger.info('Results computed with the **unofficial** Python eval code.')
logger.info('Results should be very close to the official MATLAB code.')
logger.info('Use `./tools/reval.py --matlab ...` for your paper.')
logger.info('-- Thanks, The Management')
logger.info('----------------------------------------------------------')
def multi_gpu_test_net_on_dataset(weights_file, dataset_name, proposal_file,
num_images, output_dir):
"""Multi-gpu inference on a dataset."""
binary_dir = envu.get_runtime_dir()
binary_ext = envu.get_py_bin_ext()
binary = os.path.join(binary_dir, 'test_net' + binary_ext)
assert os.path.exists(binary), 'Binary \'{}\' not found'.format(binary)
# Pass the target dataset and proposal file (if any) via the command line
opts = ['TEST.DATASETS', '("{}",)'.format(dataset_name)]
opts += ['TEST.WEIGHTS', weights_file]
if proposal_file:
opts += ['TEST.PROPOSAL_FILES', '("{}",)'.format(proposal_file)]
# Run inference in parallel in subprocesses
# Outputs will be a list of outputs from each subprocess, where the output
# of each subprocess is the dictionary saved by test_net().
outputs = subprocess_utils.process_in_parallel('detection', num_images,
binary, output_dir, opts)
# Collate the results from each subprocess
all_boxes = [[] for _ in range(cfg.MODEL.NUM_CLASSES)]
all_segms = [[] for _ in range(cfg.MODEL.NUM_CLASSES)]
all_keyps = [[] for _ in range(cfg.MODEL.NUM_CLASSES)]
for det_data in outputs:
all_boxes_batch = det_data['all_boxes']
all_segms_batch = det_data['all_segms']
all_keyps_batch = det_data['all_keyps']
for cls_idx in range(1, cfg.MODEL.NUM_CLASSES):
all_boxes[cls_idx] += all_boxes_batch[cls_idx]
all_segms[cls_idx] += all_segms_batch[cls_idx]
all_keyps[cls_idx] += all_keyps_batch[cls_idx]
det_file = os.path.join(output_dir, 'detections.pkl')
cfg_yaml = yaml.dump(cfg)
save_object(
dict(all_boxes=all_boxes,
all_segms=all_segms,
all_keyps=all_keyps,
cfg=cfg_yaml), det_file)
logger.info('Wrote detections to: {}'.format(os.path.abspath(det_file)))
return all_boxes, all_segms, all_keyps
def test_net(weights_file,
dataset_name,
proposal_file,
output_dir,
ind_range=None,
gpu_id=0):
"""Run inference on all images in a dataset or over an index range of images
in a dataset using a single GPU.
"""
assert not cfg.MODEL.RPN_ONLY, \
'Use rpn_generate to generate proposals from RPN-only models'
roidb, dataset, start_ind, end_ind, total_num_images = get_roidb_and_dataset(
dataset_name, proposal_file, ind_range)
model = initialize_model_from_cfg(weights_file, gpu_id=gpu_id)
num_images = len(roidb)
num_classes = cfg.MODEL.NUM_CLASSES
all_boxes, all_segms, all_keyps = empty_results(num_classes, num_images)
timers = defaultdict(Timer)
for i, entry in enumerate(roidb):
if cfg.TEST.PRECOMPUTED_PROPOSALS:
# The roidb may contain ground-truth rois (for example, if the roidb
# comes from the training or val split). We only want to evaluate
# detection on the *non*-ground-truth rois. We select only the rois
# that have the gt_classes field set to 0, which means there's no
# ground truth.
box_proposals = entry['boxes'][entry['gt_classes'] == 0]
if len(box_proposals) == 0:
continue
else:
# Faster R-CNN type models generate proposals on-the-fly with an
# in-network RPN; 1-stage models don't require proposals.
box_proposals = None
im = cv2.imread(entry['image'])
with c2_utils.NamedCudaScope(gpu_id):
cls_boxes_i, cls_segms_i, cls_keyps_i = im_detect_all(
model, im, box_proposals, timers)
extend_results(i, all_boxes, cls_boxes_i)
if cls_segms_i is not None:
extend_results(i, all_segms, cls_segms_i)
if cls_keyps_i is not None:
extend_results(i, all_keyps, cls_keyps_i)
if i % 10 == 0: # Reduce log file size
ave_total_time = np.sum([t.average_time for t in timers.values()])
eta_seconds = ave_total_time * (num_images - i - 1)
eta = str(datetime.timedelta(seconds=int(eta_seconds)))
det_time = (timers['im_detect_bbox'].average_time +
timers['im_detect_mask'].average_time +
timers['im_detect_keypoints'].average_time)
misc_time = (timers['misc_bbox'].average_time +
timers['misc_mask'].average_time +
timers['misc_keypoints'].average_time)
logger.info(('im_detect: range [{:d}, {:d}] of {:d}: '
'{:d}/{:d} {:.3f}s + {:.3f}s (eta: {})').format(
start_ind + 1, end_ind, total_num_images,
start_ind + i + 1, start_ind + num_images,
det_time, misc_time, eta))
if cfg.VIS:
im_name = os.path.splitext(os.path.basename(entry['image']))[0]
vis_utils.vis_one_image(im[:, :, ::-1],
'{:d}_{:s}'.format(i, im_name),
os.path.join(output_dir, 'vis'),
cls_boxes_i,
segms=cls_segms_i,
keypoints=cls_keyps_i,
thresh=cfg.VIS_TH,
box_alpha=0.8,
dataset=dataset,
show_class=True)
cfg_yaml = yaml.dump(cfg)
if ind_range is not None:
det_name = 'detection_range_%s_%s.pkl' % tuple(ind_range)
else:
det_name = 'detections.pkl'
det_file = os.path.join(output_dir, det_name)
save_object(
dict(all_boxes=all_boxes,
all_segms=all_segms,
all_keyps=all_keyps,
cfg=cfg_yaml), det_file)
logger.info('Wrote detections to: {}'.format(os.path.abspath(det_file)))
return all_boxes, all_segms, all_keyps
