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'
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')
def set_device(gpu_id):
DarknetObjectDetector.set_device(gpu_id)
