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'
class Darknet_ObjectDetector():
#做好加载
def __init__(self, spec, weight):
self._detector = DarknetObjectDetector(spec, weight)
def detect_object(self, pil_image):
start = time.time()
data = np.array(pil_image).transpose([2, 0, 1]).astype(np.uint8)
#numpy.transpose():矩阵转换操作
#numpy.astype() :数据类型转换
#bp::list detect_object(bp::str img_data, int img_width, int img_height, int img_channel) 返回一个装了bbox的列表
rst = self._detector.detect_object(data.tostring(), pil_image.size[0],
pil_image.size[1], 3)
end = time.time()
ret_rst = [DetBBox(x) for x in rst]
return ret_rst, end - start
@staticmethod
def set_device(gpu_id):
DarknetObjectDetector.set_device(gpu_id)
class YoloWorker(BaseWorker):
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 _detection_message(self, yolo_artifact):
return "Yolo Objects Detected: {}".format(yolo_artifact.filepath)
def imread(self, path):
img = Image.open(path).convert('RGB')
img = img.resize((self.target_dim, self.target_dim), Image.BILINEAR)
data = np.array(img).transpose([2, 0, 1]).astype(np.uint8).tostring()
return data, (img.size[0], img.size[1])
def featurize(self, yolo_artifact, obj):
data, size = obj
bboxes = [
DetBBox(x) for x in self.det.detect_object(str(data), size[0],
size[1], 3).content
]
feats = []
for bbox in bboxes:
yolo_artifact.features.append(
YoloFeature(
self.class_names[bbox.cls],
bbox.confidence,
[bbox.top, bbox.bottom, bbox.left, bbox.right],
))
return yolo_artifact
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')
class YoloWorker(BaseWorker):
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 imread(self, path):
if path[:4] == 'http':
with contextlib.closing(urllib.urlopen(path)) as req:
path = cStringIO.StringIO(req.read())
img = Image.open(path).convert('RGB')
img = img.resize((self.target_dim, self.target_dim), Image.BILINEAR)
data = np.array(img).transpose([2,0,1]).astype(np.uint8).tostring()
return data, (img.size[0], img.size[1])
def featurize(self, meta, obj, return_feat=False):
data, size = obj
bboxes = [DetBBox(x) for x in self.det.detect_object(data, size[0], size[1], 3).content]
feats = []
for bbox in bboxes:
class_name = self.class_names[bbox.cls]
if not return_feat:
print '\t'.join(map(str, [
meta,
class_name,
bbox.confidence,
bbox.top,
bbox.bottom,
bbox.left,
bbox.right
]))
sys.stdout.flush()
else:
feats.append({
"class_name" : class_name,
"confidence" : bbox.confidence,
"bbox" : [bbox.top, bbox.bottom, bbox.left, bbox.right],
})
if return_feat:
return meta, feats
def close(self):
print >> sys.stderr, 'YoloWorker: terminating'
pass
class Darknet_ObjectDetector():
def __init__(self, spec, weight, thresh=0.5, nms=0.4, draw=0):
self._detector = DarknetObjectDetector(spec, weight, thresh, nms, draw)
def detect_object(self, data, size):
res = self._detector.detect_object(data, size[0], size[1], 3)
out = [DetBBox(x) for x in res.content], res.load_time, res.pred_time
# if self.py_resize:
# print >> sys.stderr, "!! BBOX is in transformed dimensions -- need to implement fix"
# pass
return out
@staticmethod
def set_device(gpu_id):
DarknetObjectDetector.set_device(gpu_id)
class Darknet_ObjectDetector():
def __init__(self, spec, weight):
self._detector = DarknetObjectDetector(spec, weight)
def detect_object(self, pil_image):
start = time.time()
data = np.array(pil_image).transpose([2,0,1]).astype(np.uint8)
rst = self._detector.detect_object(data.tostring(), pil_image.size[0], pil_image.size[1], 3)
end = time.time()
ret_rst = [DetBBox(x) for x in rst]
return ret_rst, end-start
@staticmethod
def set_device(gpu_id):
DarknetObjectDetector.set_device(gpu_id)
def set_device(gpu_id):
DarknetObjectDetector.set_device(gpu_id)
def __init__(self, spec, weight):
self._detector = DarknetObjectDetector(spec, weight)
def __init__(self, spec, weight, thresh=0.5, nms=0.4, draw=0):
self._detector = DarknetObjectDetector(spec, weight, thresh, nms, draw)
def __init__(self, net_cfg, weight):
self._detector = DarknetObjectDetector(net_cfg, weight)
