Source code for wbia.algo.detect.yolo

# -*- coding: utf-8 -*-
"""
Interface to pydarknet yolo object detection.
"""
import logging
import utool as ut
import vtool as vt

(print, rrr, profile) = ut.inject2(__name__, '[yolo]')
logger = logging.getLogger('wbia')

if not ut.get_argflag('--no-pydarknet'):
    try:
        import pydarknet
    except ImportError:
        logger.info(
            'WARNING Failed to import pydarknet. '
            'PyDarknet YOLO detection is unavailable'
        )
        if ut.SUPER_STRICT:
            raise


VERBOSE_DARK = ut.get_argflag('--verbdark') or ut.VERBOSE


# def train_gid_list(ibs, gid_list, trees_path=None, species=None, setup=True,


[docs]def detect_gid_list(ibs, gid_list, downsample=False, **kwargs): """ Args: gid_list (list of int): the list of IBEIS image_rowids that need detection downsample (bool, optional): a flag to indicate if the original image sizes should be used; defaults to True True: ibs.get_image_detectpaths() is used False: ibs.get_image_paths() is used Kwargs (optional): refer to the PyDarknet documentation for configuration settings Args: ibs (wbia.IBEISController): image analysis api gid_list (list of int): the list of IBEIS image_rowids that need detection downsample (bool, optional): a flag to indicate if the original image sizes should be used; defaults to True Kwargs: detector, config_filepath, weights_filepath, verbose Yields: tuple: (gid, gpath, result_list) CommandLine: python -m wbia.algo.detect.yolo detect_gid_list --show Example: >>> # DISABLE_DOCTEST >>> from wbia.algo.detect.yolo import * # NOQA >>> from wbia.core_images import LocalizerConfig >>> import wbia >>> ibs = wbia.opendb(defaultdb='WS_ALL') >>> gid_list = ibs.images()._rowids[0:1] >>> kwargs = config = LocalizerConfig(**{ >>> 'weights_filepath': '/media/raid/work/WS_ALL/localizer_backup/detect.yolo.2.39000.weights', >>> 'config_filepath': '/media/raid/work/WS_ALL/localizer_backup/detect.yolo.2.cfg', >>> }) >>> exec(ut.execstr_dict(config), globals()) >>> #classes_fpath = '/media/raid/work/WS_ALL/localizer_backup/detect.yolo.2.cfg.classes' >>> downsample = False >>> (gid, gpath, result_list) = detect_gid_list(ibs, gid_list, downsample, **config) >>> result = ('(gid, gpath, result_list) = %s' % (ut.repr2((gid, gpath, result_list)),)) >>> print(result) >>> ut.quit_if_noshow() >>> import wbia.plottool as pt >>> ut.show_if_requested() Yields: results (list of dict) """ # Get new gpaths if downsampling if downsample: gpath_list = ibs.get_image_detectpaths(gid_list) neww_list = [vt.open_image_size(gpath)[0] for gpath in gpath_list] oldw_list = [oldw for (oldw, oldh) in ibs.get_image_sizes(gid_list)] downsample_list = [oldw / neww for oldw, neww in zip(oldw_list, neww_list)] orient_list = [1] * len(gid_list) else: gpath_list = ibs.get_image_paths(gid_list) downsample_list = [None] * len(gpath_list) orient_list = ibs.get_image_orientation(gid_list) # Run detection results_iter = detect(gpath_list, **kwargs) # Upscale the results _iter = zip(downsample_list, gid_list, orient_list, results_iter) for downsample, gid, orient, (gpath, result_list) in _iter: # Upscale the results back up to the original image size for result in result_list: if downsample is not None and downsample != 1.0: for key in ['xtl', 'ytl', 'width', 'height']: result[key] = int(result[key] * downsample) bbox = ( result['xtl'], result['ytl'], result['width'], result['height'], ) bbox_list = [bbox] bbox = bbox_list[0] result['xtl'], result['ytl'], result['width'], result['height'] = bbox yield (gid, gpath, result_list)
[docs]def detect( gpath_list, detector=None, config_filepath=None, weights_filepath=None, **kwargs ): """ Args: gpath_list (list of str): the list of image paths that need detection Kwargs (optional): refer to the PyDarknet documentation for configuration settings Returns: iter Example: >>> # DISABLE_DOCTEST >>> from wbia.algo.detect.yolo import * # NOQA >>> from wbia.core_images import LocalizerConfig >>> import wbia >>> ibs = wbia.opendb(defaultdb='WS_ALL') >>> gid_list = ibs.images()._rowids[0:1] >>> gpath_list = ibs.get_image_paths(gid_list) >>> dpath = '/media/raid/work/WS_ALL/localizer_backup/' >>> weights_filepath = join(dpath, 'detect.yolo.2.39000.weights') >>> config_filepath = join(dpath, 'detect.yolo.2.cfg') >>> config = LocalizerConfig( >>> weights_filepath=weights_filepath, >>> config_filepath=config_filepath, >>> ) >>> kwargs = config.asdict() >>> ut.delete_dict_keys(kwargs, ['weights_filepath', 'config_filepath']) >>> ut.delete_dict_keys(kwargs, ['thumbnail_cfg', 'species', 'algo']) """ # Run detection if detector is None: classes_filepath = kwargs.pop('classes_filepath', None) verbose = kwargs.get('verbose', False) detector = pydarknet.Darknet_YOLO_Detector( config_filepath=config_filepath, weights_filepath=weights_filepath, classes_filepath=classes_filepath, verbose=verbose, ) # dark = detector # input_gpath_list = gpath_list results_iter = detector.detect(gpath_list, **kwargs) results_list = list(results_iter) del detector return results_list