示例#1
0
    def __init__(self, cfg_path, weight_path, name_path, thresh=0.1, nms=0.3, target_dim=416):
        import_yolo()

        DarknetObjectDetector.set_device(0)
        self.target_dim = target_dim
        self.class_names = open(name_path).read().splitlines()
        self.det = DarknetObjectDetector(cfg_path, weight_path, thresh, nms, 0)
        print >> sys.stderr, 'YoloWorker: ready'
示例#2
0
    def __init__(self,
                 cfg_path,
                 weight_path,
                 name_path,
                 thresh=0.1,
                 nms=0.3,
                 target_dim=416,
                 logger=None,
                 *args,
                 **kwargs):
        _import_yolo()

        self.logger = logger or logging

        DarknetObjectDetector.set_device(
            int(os.environ.get("CUDA_VISIBLE_DEVICES", "")))

        self.target_dim = target_dim
        self.class_names = open(name_path).read().splitlines()
        self.det = DarknetObjectDetector(cfg_path, weight_path, thresh, nms, 0)

        self.logger.debug('YoloWorker: ready')
示例#3
0
 def set_device(gpu_id):
     DarknetObjectDetector.set_device(gpu_id)