if not hasFrame: break (h_orig, w_orig) = original_frame.shape[:2] log.debug("Preprocessing frame") if original_frame.shape[2] > 1: frame = cv.cvtColor(cv.cvtColor(original_frame, cv.COLOR_BGR2GRAY), cv.COLOR_GRAY2RGB) else: frame = cv.cvtColor(original_frame, cv.COLOR_GRAY2RGB) img_rgb = frame.astype(np.float32) / 255 img_lab = cv.cvtColor(img_rgb, cv.COLOR_RGB2Lab) img_l_rs = cv.resize(img_lab.copy(), (w_in, h_in))[:, :, 0] log.debug("Network inference") res = exec_net.infer(inputs={input_blob: [img_l_rs]}) update_res = (res[output_blob] * color_coeff.transpose()[:, :, np.newaxis, np.newaxis]).sum(1) log.debug("Get results") out = update_res.transpose((1, 2, 0)) out = cv.resize(out, (w_orig, h_orig)) img_lab_out = np.concatenate((img_lab[:, :, 0][:, :, np.newaxis], out), axis=2) img_bgr_out = np.clip(cv.cvtColor(img_lab_out, cv.COLOR_Lab2BGR), 0, 1) if not args.no_show: log.debug("Show results") imshowSize = (640, 480) original_image = cv.resize(original_frame, imshowSize) grayscale_image = cv.resize(frame, imshowSize) colorize_image = (cv.resize(img_bgr_out, imshowSize) * 255).astype(np.uint8)
class Model_YOLO_V3: def __init__(self, model_name, device='CPU', threshold=0.6): self.model_weights = model_name + '.bin' self.model_structure = model_name + '.xml' self.device = device self.threshold = threshold self.nms_threshold = 0.4 self.model = IECore().read_network(model=self.model_structure, weights=self.model_weights) self.input_name = next(iter(self.model.inputs)) self.input_shape = self.model.inputs[self.input_name].shape self.output_names = ['detector/yolo-v3/Conv_6/BiasAdd/YoloRegion', \ 'detector/yolo-v3/Conv_14/BiasAdd/YoloRegion', \ 'detector/yolo-v3/Conv_22/BiasAdd/YoloRegion'] self.anchors = [[116,90, 156,198, 373,326], [30,61, 62,45, 59,119], [10,13, 16,30, 33,23]] self.labels = [] self.load_labels('../models/coco.names') print('device : ', self.device) print('input name : ', self.input_name) print('input shape : ', self.input_shape) print('output names : ', self.output_names) def load_model(self): self.net = IECore().load_network(self.model, self.device) def load_labels(self, file_name): with open(file_name) as f: for id, name in enumerate(f): self.labels.append(name.rstrip('\n')) def predict(self, image): input_img = self.preprocess_input(image) start_time = time.time() output = self.net.infer({self.input_name:input_img}) print('inference time : ', time.time() - start_time) return self.preprocess_output(output, image) def preprocess_input(self, image): p_img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) p_img = cv2.resize(p_img, (self.input_shape[3], self.input_shape[2])) p_img = p_img.transpose((2,0,1)) p_img = p_img.reshape(1, *p_img.shape) return p_img def preprocess_output(self, outputs, image): h, w, c = image.shape _, _, input_h, input_w = self.input_shape num_box = 3 boxes = [] for i in range(num_box): output = np.squeeze(outputs[self.output_names[i]]) anchor = self.anchors[i] _, grid_h, grid_w = output.shape grid_width = input_w / grid_w grid_height = input_h / grid_h output = output.transpose((1,2,0)) output = output.reshape((grid_h, grid_w, num_box, -1)) for ih in range(grid_h): for iw in range(grid_w): for ib in range(num_box): objectness = output[ih][iw][ib][4] if objectness >= self.threshold: tx, ty, tw, th = output[ih][iw][ib][:4] bx = round((iw + self.sigmoid(tx)) * grid_width) * w / input_w by = round((ih + self.sigmoid(ty)) * grid_height) * h / input_h bw = round(anchor[2*ib+0] * np.exp(tw)) * w / input_w bh = round(anchor[2*ib+1] * np.exp(th)) * h / input_h classes = output[ih][iw][ib][5:] boxes.append([bx,by,bh,bw, classes]) for c in range(len(self.labels)): sorted_idx = np.argsort([-box[4][c] for box in boxes]) for i in range(len(sorted_idx)): idx_i = sorted_idx[i] if boxes[idx_i][4][c] != 0: for j in range(i+1, len(sorted_idx)): idx_j = sorted_idx[j] if self.iou(boxes[idx_i], boxes[idx_j]) >= self.nms_threshold: boxes[idx_j][4][c] = 0 ret = [] for box in boxes: for i in range(len(self.labels)): if box[4][i] > 0.5: x0 = int(max(box[0]-box[3]/2, 0)) y0 = int(max(box[1]-box[2]/2, 0)) x1 = int(min(box[0]+box[3]/2, w-1)) y1 = int(min(box[1]+box[2]/2, h-1)) ret.append([x0, y0, x1, y1, self.labels[np.argmax(box[4])]]) return ret def sigmoid(self, x): return 1. / (1. + np.exp(-x)) def iou(self, box1, box2): box1_xmin = box1[0] - box1[3]/2 box1_ymin = box1[1] - box1[2]/2 box1_xmax = box1[0] + box1[3]/2 box1_ymax = box1[1] + box1[2]/2 box2_xmin = box2[0] - box2[3]/2 box2_ymin = box2[1] - box2[2]/2 box2_xmax = box2[0] + box2[3]/2 box2_ymax = box2[1] + box2[2]/2 intersect_w = self._interval_overlap([box1_xmin, box1_xmax], [box2_xmin, box2_xmax]) intersect_h = self._interval_overlap([box1_ymin, box1_ymax], [box2_ymin, box2_ymax]) intersect = intersect_w * intersect_h w1, h1 = box1_xmax-box1_xmin, box1_ymax-box1_ymin w2, h2 = box2_xmax-box2_xmin, box2_ymax-box2_ymin union = w1*h1 + w2*h2 - intersect return float(intersect) / union def _interval_overlap(self, interval_a, interval_b): x1, x2 = interval_a x3, x4 = interval_b if x3 < x1: if x4 < x1: return 0 else: return min(x2,x4) - x1 else: if x2 < x3: return 0 else: return min(x2,x4) - x3