def process_in_parallel(tag, total_range_size, binary, output_dir, opts=''):
"""Run the specified binary cfg.NUM_GPUS times in parallel, each time as a
subprocess that uses one GPU. The binary must accept the command line
arguments `--range {start} {end}` that specify a data processing range.
"""
# Snapshot the current cfg state in order to pass to the inference
# subprocesses
cfg_file = os.path.join(output_dir, '{}_range_config.yaml'.format(tag))
with open(cfg_file, 'w') as f:
envu.yaml_dump(cfg, stream=f)
subprocess_env = os.environ.copy()
processes = []
subinds = np.array_split(range(total_range_size), cfg.NUM_GPUS)
# Determine GPUs to use
cuda_visible_devices = os.environ.get('CUDA_VISIBLE_DEVICES')
if cuda_visible_devices:
gpu_inds = map(int, cuda_visible_devices.split(','))
assert -1 not in gpu_inds, \
'Hiding GPU indices using the \'-1\' index is not supported'
else:
gpu_inds = reversed(range(cfg.NUM_GPUS))
# Run the binary in cfg.NUM_GPUS subprocesses
for i, gpu_ind in enumerate(gpu_inds):
start = subinds[i][0]
end = subinds[i][-1] + 1
subprocess_env['CUDA_VISIBLE_DEVICES'] = str(gpu_ind)
cmd = '{binary} --range {start} {end} --cfg {cfg_file} NUM_GPUS 1 {opts}'
cmd = cmd.format(binary=shlex_quote(binary),
start=int(start),
end=int(end),
cfg_file=shlex_quote(cfg_file),
opts=' '.join([shlex_quote(opt) for opt in opts]))
logger.info('{} range command {}: {}'.format(tag, i, cmd))
if i == 0:
subprocess_stdout = subprocess.PIPE
else:
filename = os.path.join(
output_dir, '%s_range_%s_%s.stdout' % (tag, start, end))
subprocess_stdout = open(filename, 'w') # NOQA (close below)
p = subprocess.Popen(cmd,
shell=True,
env=subprocess_env,
stdout=subprocess_stdout,
stderr=subprocess.STDOUT,
bufsize=1)
processes.append((i, p, start, end, subprocess_stdout))
print('Sleep 15s for gpu_id ', gpu_ind)
time.sleep(15)
# Log output from inference processes and collate their results
outputs = []
for i, p, start, end, subprocess_stdout in processes:
log_subprocess_output(i, p, output_dir, tag, start, end)
if i > 0:
subprocess_stdout.close()
range_file = os.path.join(output_dir,
'%s_range_%s_%s.pkl' % (tag, start, end))
range_data = load_object(range_file)
outputs.append(range_data)
return outputs
def test_merge_cfg_from_file(self):
with tempfile.NamedTemporaryFile() as f:
envu.yaml_dump(cfg, f)
s = cfg.MODEL.TYPE
cfg.MODEL.TYPE = 'dummy'
assert cfg.MODEL.TYPE != s
core_config.merge_cfg_from_file(f.name)
assert cfg.MODEL.TYPE == s
def test_deprecated_key_from_file(self):
# You should see logger messages like:
# "Deprecated config key (ignoring): MODEL.DILATION"
with tempfile.NamedTemporaryFile() as f:
cfg2 = copy.deepcopy(cfg)
cfg2.MODEL.DILATION = 2
envu.yaml_dump(cfg2, f)
with self.assertRaises(AttributeError):
_ = cfg.MODEL.DILATION # noqa
core_config.merge_cfg_from_file(f.name)
with self.assertRaises(AttributeError):
_ = cfg.MODEL.DILATION # noqa
def test_renamed_key_from_file(self):
# You should see logger messages like:
# "Key EXAMPLE.RENAMED.KEY was renamed to EXAMPLE.KEY;
# please update your config"
with tempfile.NamedTemporaryFile() as f:
cfg2 = copy.deepcopy(cfg)
cfg2.EXAMPLE = AttrDict()
cfg2.EXAMPLE.RENAMED = AttrDict()
cfg2.EXAMPLE.RENAMED.KEY = 'foobar'
envu.yaml_dump(cfg2, f)
with self.assertRaises(AttributeError):
_ = cfg.EXAMPLE.RENAMED.KEY # noqa
with self.assertRaises(KeyError):
core_config.merge_cfg_from_file(f.name)
def main(args):
logger = logging.getLogger(__name__)
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
cfg_orig = load_cfg(envu.yaml_dump(cfg))
im = cv2.imread(args.im_file)
if args.rpn_pkl is not None:
proposal_boxes, _proposal_scores = get_rpn_box_proposals(im, args)
workspace.ResetWorkspace()
else:
proposal_boxes = None
cls_boxes, cls_segms, cls_keyps, cls_bodys = None, None, None, None
for i in range(0, len(args.models_to_run), 2):
pkl = args.models_to_run[i]
yml = args.models_to_run[i + 1]
cfg.immutable(False)
merge_cfg_from_cfg(cfg_orig)
merge_cfg_from_file(yml)
if len(pkl) > 0:
weights_file = pkl
else:
weights_file = cfg.TEST.WEIGHTS
cfg.NUM_GPUS = 1
assert_and_infer_cfg(cache_urls=False)
model = model_engine.initialize_model_from_cfg(weights_file)
with c2_utils.NamedCudaScope(0):
cls_boxes_, cls_segms_, cls_keyps_ , cls_bodys_= \
model_engine.im_detect_all(model, im, proposal_boxes)
cls_boxes = cls_boxes_ if cls_boxes_ is not None else cls_boxes
cls_segms = cls_segms_ if cls_segms_ is not None else cls_segms
cls_keyps = cls_keyps_ if cls_keyps_ is not None else cls_keyps
cls_bodys = cls_bodys_ if cls_bodys_ is not None else cls_bodys
workspace.ResetWorkspace()
out_name = os.path.join(
args.output_dir, '{}'.format(os.path.basename(args.im_file) + '.pdf'))
logger.info('Processing {} -> {}'.format(args.im_file, out_name))
with open('test_vis.pkl', 'w') as f:
pickle.dump(
{
'im': im,
'cls_boxes': np.array(cls_boxes),
'cls_bodys': np.array(cls_bodys)
}, f)
vis_utils.vis_one_image(im[:, :, ::-1],
args.im_file,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
cls_bodys,
dataset=dummy_coco_dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2)
def test_immutability(self):
# Top level immutable
a = AttrDict()
a.foo = 0
a.immutable(True)
with self.assertRaises(AttributeError):
a.foo = 1
a.bar = 1
assert a.is_immutable()
assert a.foo == 0
a.immutable(False)
assert not a.is_immutable()
a.foo = 1
assert a.foo == 1
# Recursively immutable
a.level1 = AttrDict()
a.level1.foo = 0
a.level1.level2 = AttrDict()
a.level1.level2.foo = 0
a.immutable(True)
assert a.is_immutable()
with self.assertRaises(AttributeError):
a.level1.level2.foo = 1
a.level1.bar = 1
assert a.level1.level2.foo == 0
# Serialize immutability state
a.immutable(True)
a2 = core_config.load_cfg(envu.yaml_dump(a))
assert a.is_immutable()
assert a2.is_immutable()
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 = envu.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 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 = envu.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 = envu.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 main(args):
logger = logging.getLogger(__name__)
dummy_nucoco_dataset = dummy_datasets.get_nucoco_dataset()
cfg_orig = load_cfg(envu.yaml_dump(cfg))
## Load image
coco = COCO_PLUS(args.ann_file, args.imgs_dir)
image_id = coco.dataset['images'][args.im_ind]['id']
img_path = os.path.join(args.imgs_dir, coco.imgs[image_id]["file_name"])
im = cv2.imread(img_path)
## Get the proposals for this image
proposals = rrpn_loader(args.rpn_pkl)
proposal_boxes = proposals[image_id]['boxes']
_proposal_scores = proposals[image_id]['scores']
workspace.ResetWorkspace()
## run models
cls_boxes, cls_segms, cls_keyps = None, None, None
for i in range(0, len(args.models_to_run), 2):
pkl = args.models_to_run[i]
yml = args.models_to_run[i + 1]
cfg.immutable(False)
merge_cfg_from_cfg(cfg_orig)
merge_cfg_from_file(yml)
if len(pkl) > 0:
weights_file = pkl
else:
weights_file = cfg.TEST.WEIGHTS
cfg.NUM_GPUS = 1
assert_and_infer_cfg(cache_urls=False)
model = model_engine.initialize_model_from_cfg(weights_file)
with c2_utils.NamedCudaScope(0):
cls_boxes_, cls_segms_, cls_keyps_ = \
model_engine.im_detect_all(model, im, proposal_boxes)
cls_boxes = cls_boxes_ if cls_boxes_ is not None else cls_boxes
cls_segms = cls_segms_ if cls_segms_ is not None else cls_segms
cls_keyps = cls_keyps_ if cls_keyps_ is not None else cls_keyps
workspace.ResetWorkspace()
out_name = os.path.join(
args.output_dir, '{}'.format(os.path.basename(img_path) + '.pdf')
)
logger.info('Processing {} -> {}'.format(img_path, out_name))
vis_utils.vis_one_image(
im[:, :, ::-1],
img_path,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dummy_nucoco_dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2
)
def multi_gpu_test_net_on_dataset(
weights_file, dataset_info, 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.DATASET_INFO.NAME', dataset_info.NAME]
opts += ['TEST.DATASET_INFO.DATA_IM_DIR', dataset_info.DATA_IM_DIR]
opts += ['TEST.DATASET_INFO.DATA_ANN_FN', dataset_info.DATA_ANN_FN]
opts += ['TEST.DATASET_INFO.DATA_RAW_DIR', dataset_info.DATA_RAW_DIR]
opts += ['TEST.DATASET_INFO.DATA_DEVKIT_DIR', dataset_info.DATA_DEVKIT_DIR]
opts += ['TEST.DATASET_INFO.IM_PREFIX', dataset_info.IM_PREFIX]
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 = envu.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 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 = envu.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 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 = envu.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 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 = envu.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 test_merge_cfg_from_cfg(self):
# Test: merge from deepcopy
s = 'dummy0'
cfg2 = copy.deepcopy(cfg)
cfg2.MODEL.TYPE = s
core_config.merge_cfg_from_cfg(cfg2)
assert cfg.MODEL.TYPE == s
# Test: merge from yaml
s = 'dummy1'
cfg2 = core_config.load_cfg(envu.yaml_dump(cfg))
cfg2.MODEL.TYPE = s
core_config.merge_cfg_from_cfg(cfg2)
assert cfg.MODEL.TYPE == s
# Test: merge with a valid key
s = 'dummy2'
cfg2 = AttrDict()
cfg2.MODEL = AttrDict()
cfg2.MODEL.TYPE = s
core_config.merge_cfg_from_cfg(cfg2)
assert cfg.MODEL.TYPE == s
# Test: merge with an invalid key
s = 'dummy3'
cfg2 = AttrDict()
cfg2.FOO = AttrDict()
cfg2.FOO.BAR = s
with self.assertRaises(KeyError):
core_config.merge_cfg_from_cfg(cfg2)
# Test: merge with converted type
cfg2 = AttrDict()
cfg2.TRAIN = AttrDict()
cfg2.TRAIN.SCALES = [1]
core_config.merge_cfg_from_cfg(cfg2)
assert type(cfg.TRAIN.SCALES) is tuple
assert cfg.TRAIN.SCALES[0] == 1
# Test: merge with invalid type
cfg2 = AttrDict()
cfg2.TRAIN = AttrDict()
cfg2.TRAIN.SCALES = 1
with self.assertRaises(ValueError):
core_config.merge_cfg_from_cfg(cfg2)
def save_model_to_weights_file(weights_file, model, cur_iter=None):
"""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)
# Save roidb shuffling:
if 'roidb_state' not in blobs and type(cur_iter) != type(None):
blobs['roidb_state'] = model.roi_data_loader.get_perm_state(
cur_iter + 1) # give iters_done.
else:
logger.info("roidb state not stored")
if cfg.TRAIN.PADA:
if 'weight_db' not in blobs:
blobs['weight_db'] = model.class_weight_db.get_state()
cfg_yaml = envu.yaml_dump(cfg)
save_object(dict(blobs=blobs, cfg=cfg_yaml), weights_file)
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
# )
outputs = subprocess_utils.process_in_parallel(
'feature', num_images, binary, output_dir, opts
)
all_feats = []
for feat_data in outputs:
all_feats.append(feat_data['all_feats'])
# all_feats = np.concatenate(all_feats, axis=0)
all_feats = np.vstack(all_feats)
feat_file = os.path.join(output_dir, 'features.pkl')
cfg_yaml = yaml.dump(cfg)
#save_object(
# dict(
# all_feats=all_feats,
# cfg=cfg_yaml
# ), feat_file
#)
logger.info('Wrote detections to: {}'.format(os.path.abspath(feat_file)))
feat_file = os.path.join(output_dir, 'features.npy')
np.save(feat_file, all_feats)
logger.info('Wrote detections to: {}'.format(os.path.abspath(feat_file)))
return all_feats
# 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 = envu.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_batch(
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, using batch inference
"""
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
)
# roidb = roidb[:500]
# Debug purposes
# roidb = [{'image': im} for im in glob.glob('/staging/leuven/stg_00027/imob/detectron/lib/datasets/data/coco/coco_val2014/*.png')][:500]
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)
ims = []
for i, entry in enumerate(roidb):
if cfg.TEST.PRECOMPUTED_PROPOSALS:
raise NotImplementedError('Precomputed proposals not implemented for batch inference, set TEST.IMS_PER_BATCH to 1')
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(roidb[0]['image'])
im = cv2.imread(entry['image'])
ims.append(im)
if not ((len(ims) == cfg.TEST.IMS_PER_BATCH) or (i == (num_images - 1))):
continue
with c2_utils.NamedCudaScope(gpu_id):
cls_boxes_batch, cls_segms_batch, cls_keyps_batch = im_detect_all_batch(
model, ims, box_proposals, timers
)
for n in range(len(ims)):
local_i = i - len(ims) + n + 1
cls_boxes_i = cls_boxes_batch[n]
cls_segms_i = cls_segms_batch[n] if cls_segms_batch else None
cls_keyps_i = cls_keyps_batch[n] if cls_keyps_batch else None
extend_results(local_i, all_boxes, cls_boxes_i)
if cls_segms_i is not None:
extend_results(local_i, all_segms, cls_segms_i)
if cls_keyps_i is not None:
extend_results(local_i, all_keyps, cls_keyps_i)
if local_i % (10 * cfg.TEST.IMS_PER_BATCH) == 0: # Reduce log file size
ave_total_time = np.sum([t.average_time for t in timers.values()])
eta_seconds = int(ave_total_time * (num_images - local_i - 1) / float(cfg.TEST.IMS_PER_BATCH))
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 + local_i + 1,
start_ind + num_images, det_time, misc_time, eta
)
)
# This will now only show the last image of each batch
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
)
ims = []
cfg_yaml = envu.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
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'
if cfg.TEST.IMS_PER_BATCH != 1:
return test_net_batch(weights_file, dataset_name, proposal_file, output_dir, ind_range, gpu_id)
roidb, dataset, start_ind, end_ind, total_num_images = get_roidb_and_dataset(
dataset_name, proposal_file, ind_range
)
# roidb = roidb[:500]
# roidb = [{'image': im} for im in glob.glob('/staging/leuven/stg_00027/imob/detectron/lib/datasets/data/coco/coco_val2014/*.png')][:500]
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 = envu.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
def test_net(weights_file,
dataset_name,
proposal_file,
output_dir,
ind_range=None,
gpu_id=0,
subset_pointer=None):
"""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'
# determine file name
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)
# load results if already present
if os.path.exists(det_file):
res = load_object(det_file)
all_boxes, all_segms, all_keyps = res['all_boxes'], res[
'all_segms'], res['all_keyps']
else:
roidb, dataset, start_ind, end_ind, total_num_images = get_roidb_and_dataset(
dataset_name, proposal_file, ind_range)
if subset_pointer is not None:
voc_subset = subset_pointer.subset
this_sub = voc_subset[:len(roidb)]
# subset_pointer.subset = voc_subset[len(roidb):]
# filter roidb:
roidb = [roi for taking, roi in zip(this_sub, roidb) if taking]
total_num_images = len(roidb)
end_ind = total_num_images
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)
if cfg.TEST.COLLECT_ALL:
all_feats = []
all_class_weights = np.empty(shape=(num_images, num_classes),
dtype=np.float32)
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, sum_softmax, topk_feats = im_detect_all(
model,
im,
box_proposals,
timers,
return_feats=cfg.TEST.COLLECT_ALL)
# print('nfeats:', topk_feats.shape[0])
# print(topk_feats)
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 cfg.TEST.COLLECT_ALL:
all_class_weights[i] = sum_softmax
all_feats.append(
topk_feats
) # will accumulate about 9 Gb of feats on COCO train set (118K imgs)
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 = envu.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)))
if cfg.TEST.COLLECT_ALL:
save_object(all_class_weights,
os.path.join(output_dir, 'class_weights.pkl'))
save_object(all_feats,
os.path.join(output_dir, 'feature_vectors.pkl'))
logger.info(
'Wrote class weights and feature vectors to output folder')
return all_boxes, all_segms, all_keyps
def process_in_parallel(
tag, total_range_size, binary, output_dir, opts=''
):
"""Run the specified binary cfg.NUM_GPUS times in parallel, each time as a
subprocess that uses one GPU. The binary must accept the command line
arguments `--range {start} {end}` that specify a data processing range.
"""
# Snapshot the current cfg state in order to pass to the inference
# subprocesses
cfg_file = os.path.join(output_dir, '{}_range_config.yaml'.format(tag))
with open(cfg_file, 'w') as f:
envu.yaml_dump(cfg, stream=f)
subprocess_env = os.environ.copy()
processes = []
subinds = np.array_split(range(total_range_size), cfg.NUM_GPUS)
# Determine GPUs to use
cuda_visible_devices = os.environ.get('CUDA_VISIBLE_DEVICES')
if cuda_visible_devices:
gpu_inds = map(int, cuda_visible_devices.split(','))
assert -1 not in gpu_inds, \
'Hiding GPU indices using the \'-1\' index is not supported'
else:
gpu_inds = range(cfg.NUM_GPUS)
# Run the binary in cfg.NUM_GPUS subprocesses
for i, gpu_ind in enumerate(gpu_inds):
start = subinds[i][0]
end = subinds[i][-1] + 1
subprocess_env['CUDA_VISIBLE_DEVICES'] = str(gpu_ind)
cmd = '{binary} --range {start} {end} --cfg {cfg_file} NUM_GPUS 1 {opts}'
cmd = cmd.format(
binary=shlex_quote(binary),
start=int(start),
end=int(end),
cfg_file=shlex_quote(cfg_file),
opts=' '.join([shlex_quote(opt) for opt in opts])
)
logger.info('{} range command {}: {}'.format(tag, i, cmd))
if i == 0:
subprocess_stdout = subprocess.PIPE
else:
filename = os.path.join(
output_dir, '%s_range_%s_%s.stdout' % (tag, start, end)
)
subprocess_stdout = open(filename, 'w') # NOQA (close below)
p = subprocess.Popen(
cmd,
shell=True,
env=subprocess_env,
stdout=subprocess_stdout,
stderr=subprocess.STDOUT,
bufsize=1
)
processes.append((i, p, start, end, subprocess_stdout))
# Log output from inference processes and collate their results
outputs = []
for i, p, start, end, subprocess_stdout in processes:
log_subprocess_output(i, p, output_dir, tag, start, end)
if i > 0:
subprocess_stdout.close()
range_file = os.path.join(
output_dir, '%s_range_%s_%s.pkl' % (tag, start, end)
)
range_data = load_object(range_file)
outputs.append(range_data)
return outputs
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 = envu.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
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, all_track = empty_results(
num_classes, num_images)
box_proposals_prev = None
im_prev = None
fpn_res_sum_prev = None
cls_boxes_prev = None
cls_segms_prev = None
cls_keyps_prev = None
boxes_prev = None
im_scale_prev = None
color_inds_prev = None
colors = vis_utils.distinct_colors(20)
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):
if cfg.MODEL.TRACKING_ON:
# Single image inference on the first frame
if i == 0:
im_prev = im
box_proposals_prev = box_proposals
cls_boxes_list, cls_segms_list, cls_keyps_list, track_mat_i, extras = multi_im_detect_all(
model, [im], [box_proposals], timers, tracking=False)
im_scale_list, boxes_list, fpn_res_sum_list = extras
cls_boxes_i = cls_boxes_list[0]
cls_segms_i = cls_segms_list[0]
cls_keyps_i = cls_keyps_list[0]
im_scale_prev = im_scale_list[0]
boxes_prev = boxes_list[0]
fpn_res_sum_prev = fpn_res_sum_list[0]
# Pairwise image inference on the second frame
elif i == 1:
cls_boxes_list, cls_segms_list, cls_keyps_list, track_mat_i, extras = multi_im_detect_all(
model, [im_prev, im],
[box_proposals_prev, box_proposals], timers)
boxes_list, im_scale_list, fpn_res_sum_list = extras
cls_boxes_i = cls_boxes_list[1]
cls_segms_i = cls_segms_list[1]
cls_keyps_i = cls_keyps_list[1]
boxes_prev = boxes_list[1]
im_scale_prev = im_scale_list[1]
fpn_res_sum_prev = fpn_res_sum_list[1]
# Sequential image inference after the second frame
else:
cls_boxes_i, cls_segms_i, cls_keyps_i, track_mat_i, extras = im_detect_all_seq(
model, im, box_proposals,
(cls_boxes_i, fpn_res_sum_prev, boxes_prev,
im_scale_prev), timers)
boxes_prev, im_scale_prev, fpn_res_sum_prev = extras
else:
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 track_mat_i is not None:
all_track[i] = track_mat_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 +
timers['im_detect_tracking'].average_time)
misc_time = (timers['misc_bbox'].average_time +
timers['misc_mask'].average_time +
timers['misc_keypoints'].average_time +
timers['misc_tracking'].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]
if MODEL.TRACKING_ON:
# Pairwise visualization for the first image pair
if i == 1:
_, _, _, _, color_inds_prev = vis_utils.vis_image_pair_opencv(
im_list,
cls_boxes_list,
cls_segms_list,
cls_keyps_list,
track_mat_i,
dataset=dataset,
show_track=True,
show_box=True,
)
# Sequential visualization after the first image pair
else:
_, _, _, color_inds_prev = vis_utils.vis_image_pair_opencv(
[None, im], [cls_boxes_prev, cls_boxes],
[cls_segms_prev, cls_segms],
[cls_keyps_prev, cls_keyps],
cls_track,
dataset=dataset,
show_track=True,
show_box=True,
color_inds_list=[color_inds_prev, None])
else:
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)
im_prev = im
box_proposals_prev = box_proposals
cfg_yaml = envu.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,
all_track=all_track,
cfg=cfg_yaml), det_file)
logger.info('Wrote detections to: {}'.format(os.path.abspath(det_file)))
return all_boxes, all_segms, all_keyps, all_track