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
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)