Exemple #1
0
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)
Exemple #5
0
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
Exemple #8
0
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)
Exemple #9
0
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
Exemple #10
0
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
    )
Exemple #11
0
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
Exemple #12
0
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
Exemple #13
0
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
Exemple #14
0
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)
Exemple #17
0
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
Exemple #18
0
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
Exemple #19
0
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
Exemple #21
0
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
Exemple #22
0
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
Exemple #23
0
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