class YolowNCS(object): def __init__(self, model_name_path, input_size, labels, num_requests=2): self.model = model_name_path + '.xml' self.weights = model_name_path + '.bin' self.labels = labels self.input_size = input_size self.imer = Imager(self.input_size, self.labels) if not os.path.exists(self.model) or not os.path.exists(self.weights): raise ValueError( 'model files {} does not exist.'.format(model_name_path)) self.plugin = IEPlugin(device='MYRIAD') log.info('Loading network files:\n\t{}\n\t{}'.format( self.model, self.weights)) self.net = IENetwork(model=self.model, weights=self.weights) log.info('Preparing inputs') self.input_blob = next(iter(self.net.inputs)) self.net.batch_size = 1 log.info('Loading model to the plugin') self.current_request_id = 0 self.next_request_id = 1 self.num_requests = num_requests self.exec_net = self.plugin.load(network=self.net, num_requests=self.num_requests) def set_input(self, images): if type(images) == str: self.imer.imset_from_path(images) else: self.imer.imset(images) def predict(self, confidence_threshold=.5, iou_threshold=.4, async_mode=False): self.out_list = list() in_frames = self.imer.ncs_preprocess() for in_frame, orig_frame in zip(in_frames, self.imer.ims): objects = list() origin_im_size = orig_frame.shape[:-1] input_dict = {self.input_blob: in_frame} request_handle = self.exec_net.requests[self.current_request_id] if async_mode: next_request_id = self.current_request_id + 1 if next_request_id == self.num_requests: next_request_id = 0 else: next_request_id = self.current_request_id self.exec_net.start_async(request_id=next_request_id, inputs=input_dict) if async_mode: self.current_request_id = next_request_id request_handle.wait() pred_dict = request_handle.outputs for layer_name, out_blob in pred_dict.items(): params = self.net.layers[layer_name].params layer_params = YoloV3Params(params, out_blob.shape[2]) objects += parse_yolo_region(out_blob, in_frame.shape[2:], origin_im_size, layer_params, confidence_threshold) for i in range(len(objects)): if objects[i]['confidence'] == 0: continue for j in range(i + 1, len(objects)): if intersection_over_union(objects[i], objects[j]) > iou_threshold: objects[j]['confidence'] = 0 objects = [ obj for obj in objects if obj['confidence'] >= confidence_threshold ] self.out_list.append(objects) out_frames = self.imer.ncs_visualise_preds(self.out_list) return out_frames def get_output(self): try: return self.out_list except: raise ValueError( 'output does not exist, YolowNCS.predict() method must be called prior to this method.' )