def detect_image(self, image): if self.model_image_size != (None, None): assert self.model_image_size[ 0] % 32 == 0, 'Multiples of 32 required' assert self.model_image_size[ 1] % 32 == 0, 'Multiples of 32 required' image_data = preprocess_image(image, self.model_image_size) #origin image shape, in (height, width) format image_shape = tuple(reversed(image.size)) start = time.time() out_boxes, out_classes, out_scores = self.predict( image_data, image_shape) print('Found {} boxes for {}'.format(len(out_boxes), 'img')) end = time.time() print("Inference time: {:.8f}s".format(end - start)) #draw result on input image image_array = np.array(image, dtype='uint8') image_array = draw_boxes(image_array, out_boxes, out_classes, out_scores, self.class_names, self.colors) return Image.fromarray(image_array)
def detect_image(self, image): if self.model_image_size != (None, None): assert self.model_image_size[ 0] % 32 == 0, 'Multiples of 32 required' assert self.model_image_size[ 1] % 32 == 0, 'Multiples of 32 required' image = Image.open(image) image_data = preprocess_image(image, self.model_image_size) image_shape = image.size self.interpreter.set_tensor(self.input_details[0]['index'], image_data) start = time.time() self.interpreter.invoke() output_data = [ self.interpreter.get_tensor(self.output_details[2]['index']), self.interpreter.get_tensor(self.output_details[0]['index']), self.interpreter.get_tensor(self.output_details[1]['index']) ] out_boxes, out_classes, out_scores = yolo3_postprocess_np( output_data, image_shape, self.anchors, len(self.class_names), self.model_image_size, max_boxes=20, confidence=0.35) self.log.info('Found {} boxes for {}'.format(len(out_boxes), 'img')) end = time.time() self.log.info("Inference time: {:.8f}s".format(end - start)) if out_classes is None or len(out_classes) == 0: self.log.warning("No boxes found!") return image_data, None, None order = out_boxes[:, 0].argsort() out_boxes = out_boxes[order] out_classes = out_classes[order] out_scores = out_scores[order] yref_top = min(out_boxes[:, 1]) yref_top_idx = np.argmin(out_boxes[:, 1]) r_number, r_screen = ([] for i in range(2)) for idx, box in enumerate(out_boxes): _, ymin, _, ymax = box if ymin < 195: r_screen.append(out_classes[idx]) else: r_number.append(out_classes[idx]) #r_number = int("".join("{0}".format(n) for n in r_number)) r_number = "".join(str(n) for n in r_number) r_screen = "".join(str(n) for n in r_screen) #draw result on input image image_array = np.array(image, dtype='uint8') image_array = draw_boxes(image_array, out_boxes, out_classes, out_scores, self.class_names, self.colors) return Image.fromarray(image_array), r_number, r_screen
def get_prediction_class_records(model, model_format, annotation_records, anchors, class_names, model_image_size, conf_threshold, save_result): ''' Do the predict with YOLO model on annotation images to get predict class dict predict class dict would contain image_name, coordinary and score, and sorted by score: pred_classes_records = { 'car': [ ['000001.jpg','94,115,203,232',0.98], ['000002.jpg','82,64,154,128',0.93], ... ], ... } ''' if model_format == 'MNN': #MNN inference engine need create session session = model.createSession() # create txt file to save prediction result, with # save format as annotation file but adding score, like: # # path/to/img1.jpg 50,100,150,200,0,0.86 30,50,200,120,3,0.95 # os.makedirs('result', exist_ok=True) result_file = open(os.path.join('result','detection_result.txt'), 'w') pred_classes_records = OrderedDict() pbar = tqdm(total=len(annotation_records), desc='Eval model') for (image_name, gt_records) in annotation_records.items(): image = Image.open(image_name) if image.mode != 'RGB': image = image.convert('RGB') image_array = np.array(image, dtype='uint8') # support of tflite model if model_format == 'TFLITE': pred_boxes, pred_classes, pred_scores = yolo_predict_tflite(model, image, anchors, len(class_names), conf_threshold) # support of MNN model elif model_format == 'MNN': pred_boxes, pred_classes, pred_scores = yolo_predict_mnn(model, session, image, anchors, len(class_names), conf_threshold) # support of TF 1.x frozen pb model elif model_format == 'PB': pred_boxes, pred_classes, pred_scores = yolo_predict_pb(model, image, anchors, len(class_names), model_image_size, conf_threshold) # support of ONNX model elif model_format == 'ONNX': pred_boxes, pred_classes, pred_scores = yolo_predict_onnx(model, image, anchors, len(class_names), conf_threshold) # normal keras h5 model elif model_format == 'H5': pred_boxes, pred_classes, pred_scores = yolo_predict_keras(model, image, anchors, len(class_names), model_image_size, conf_threshold) else: raise ValueError('invalid model format') #print('Found {} boxes for {}'.format(len(pred_boxes), image_name)) pbar.update(1) # save prediction result to txt result_file.write(image_name) for box, cls, score in zip(pred_boxes, pred_classes, pred_scores): xmin, ymin, xmax, ymax = box box_annotation = " %d,%d,%d,%d,%d,%f" % ( xmin, ymin, xmax, ymax, cls, score) result_file.write(box_annotation) result_file.write('\n') result_file.flush() if save_result: gt_boxes, gt_classes, gt_scores = transform_gt_record(gt_records, class_names) result_dir=os.path.join('result','detection') os.makedirs(result_dir, exist_ok=True) colors = get_colors(class_names) image_array = draw_boxes(image_array, gt_boxes, gt_classes, gt_scores, class_names, colors=None, show_score=False) image_array = draw_boxes(image_array, pred_boxes, pred_classes, pred_scores, class_names, colors) image = Image.fromarray(image_array) # here we handle the RGBA image if(len(image.split()) == 4): r, g, b, a = image.split() image = Image.merge("RGB", (r, g, b)) image.save(os.path.join(result_dir, image_name.split(os.path.sep)[-1])) # Nothing detected if pred_boxes is None or len(pred_boxes) == 0: continue for box, cls, score in zip(pred_boxes, pred_classes, pred_scores): pred_class_name = class_names[cls] xmin, ymin, xmax, ymax = box coordinate = "{},{},{},{}".format(xmin, ymin, xmax, ymax) #append or add predict class item record = [os.path.basename(image_name), coordinate, score] if pred_class_name in pred_classes_records: pred_classes_records[pred_class_name].append(record) else: pred_classes_records[pred_class_name] = list([record]) # sort pred_classes_records for each class according to score for pred_class_list in pred_classes_records.values(): pred_class_list.sort(key=lambda ele: ele[2], reverse=True) pbar.close() result_file.close() return pred_classes_records
def get_prediction_class_records(model, model_format, annotation_records, anchors, class_names, model_image_size, conf_threshold, save_result): ''' Do the predict with YOLO model on annotation images to get predict class dict predict class dict would contain image_name, coordinary and score, and sorted by score: pred_classes_records = { 'car': [ ['00001.jpg','94,115,203,232',0.98], ['00002.jpg','82,64,154,128',0.93], ... ], ... } ''' if model_format == 'MNN': #MNN inference engine need create session session = model.createSession() pred_classes_records = {} pbar = tqdm(total=len(annotation_records), desc='Eval model') for (image_name, gt_records) in annotation_records.items(): image = Image.open(image_name) image_array = np.array(image, dtype='uint8') # support of tflite model if model_format == 'TFLITE': pred_boxes, pred_classes, pred_scores = yolo_predict_tflite( model, image, anchors, len(class_names), conf_threshold) # support of MNN model elif model_format == 'MNN': pred_boxes, pred_classes, pred_scores = yolo_predict_mnn( model, session, image, anchors, len(class_names), conf_threshold) # support of TF 1.x frozen pb model elif model_format == 'PB': pred_boxes, pred_classes, pred_scores = yolo_predict_pb( model, image, anchors, len(class_names), model_image_size, conf_threshold) # normal keras h5 model elif model_format == 'H5': pred_boxes, pred_classes, pred_scores = yolo_predict_keras( model, image, anchors, len(class_names), model_image_size, conf_threshold) else: raise ValueError('invalid model format') #print('Found {} boxes for {}'.format(len(pred_boxes), image_name)) pbar.update(1) if save_result: gt_boxes, gt_classes, gt_scores = transform_gt_record( gt_records, class_names) result_dir = os.path.join('result', 'detection') touchdir(result_dir) colors = get_colors(class_names) image_array = draw_boxes(image_array, gt_boxes, gt_classes, gt_scores, class_names, colors=None, show_score=False) image_array = draw_boxes(image_array, pred_boxes, pred_classes, pred_scores, class_names, colors) image = Image.fromarray(image_array) # here we handle the RGBA image if (len(image.split()) == 4): r, g, b, a = image.split() image = Image.merge("RGB", (r, g, b)) image.save( os.path.join(result_dir, image_name.split(os.path.sep)[-1])) # Nothing detected if pred_boxes is None or len(pred_boxes) == 0: continue for box, cls, score in zip(pred_boxes, pred_classes, pred_scores): pred_class_name = class_names[cls] xmin, ymin, xmax, ymax = box coordinate = "{},{},{},{}".format(xmin, ymin, xmax, ymax) #append or add predict class item if pred_class_name in pred_classes_records: pred_classes_records[pred_class_name].append( [image_name, coordinate, score]) else: pred_classes_records[pred_class_name] = list( [[image_name, coordinate, score]]) # sort pred_classes_records for each class according to score for pred_class_list in pred_classes_records.values(): pred_class_list.sort(key=lambda ele: ele[2], reverse=True) pbar.close() return pred_classes_records
def detect_image(self, image, apply_constraints=False): if self.model_image_size != (None, None): assert self.model_image_size[ 0] % 32 == 0, 'Multiples of 32 required' assert self.model_image_size[ 1] % 32 == 0, 'Multiples of 32 required' image_data = preprocess_image(image, self.model_image_size) # prepare origin image shape, (height, width) format image_shape = np.array([image.size[1], image.size[0]]) image_shape = np.expand_dims(image_shape, 0) start = time.time() out_boxes, out_classes, out_scores = self.predict( image_data, image_shape) end = time.time() print('Found {} boxes for {}'.format(len(out_boxes), 'img')) print("Inference time: {:.8f}s".format(end - start)) if (apply_constraints): print('Applying constraints...') idx_stem, = np.where(self.class_names == 'stem') idx_stick, = np.where(self.class_names == 'stick') # Stem constraint if (len(idx_stick) > 0): stems, = np.where(out_classes == idx_stem) if (len(stems) > 1): higher_weighted_dist = -9999 higher_weighted_dist_idx = -1 for i in stems: # Calculating the weighted difference between the stick and the current stem bottom_dist = abs(out_boxes[i][-1] - out_boxes[idx_stick][-1]) dist = min([ abs(out_boxes[idx_stick][0] - out_boxes[i][1]), abs(out_boxes[idx_stick][1] - out_boxes[i][0]) ]) weighted_dist = bottom_dist * ( 1 / (dist + 0.00001)) + bottom_dist * out_scores[i] if (weighted_dist > higher_weighted_dist): higher_weighted_dist = weighted_dist if (higher_weighted_dist_idx >= 0): out_boxes = np.delete(out_boxes, higher_weighted_dist_idx, axis=0) out_classes = np.delete( out_classes, higher_weighted_dist_idx, axis=0) out_scores = np.delete( out_scores, higher_weighted_dist_idx, axis=0) higher_weighted_dist_idx = i #draw result on input image image_array = np.array(image, dtype='uint8') image_array = draw_boxes(image_array, out_boxes, out_classes, out_scores, self.class_names, self.colors) return Image.fromarray(image_array)
elapsed = time.time() - start fps = 1 / elapsed print('Inference time in ms: %.2f' % (elapsed * 1000)) # post-processing image_shape = tuple((frame.shape[0], frame.shape[1])) out_boxes, out_classes, out_scores = yolo3_postprocess_np( detected_boxes, image_shape, anchors, len(labels), (input_height, input_width), max_boxes=10, confidence=conf_threshold, iou_threshold=iou_threshold, elim_grid_sense=True) # modified draw_boxes function to return an openCV formatted image img_bbox = draw_boxes(img, out_boxes, out_classes, out_scores, labels, colors) # draw information overlay onto the frames cv2.putText(img_bbox, 'Inference Running on : {0}'.format(backend_name), (30, 50), font, font_size, color, font_thickness) cv2.putText( img_bbox, 'FPS : {0} | Inference Time : {1}ms'.format( int(fps), round((elapsed * 1000), 2)), (30, 80), font, font_size, color, font_thickness) if input_mode in 'directory': out_file = "detections_{0}.jpg".format(image_id) out_file = os.path.join(result_dir, out_file) cv2.imwrite(out_file, img_bbox) if input_mode in 'image': cv2.imwrite("detections.jpg", img_bbox) print("Output image is saved in detections.jpg")