| from absl import app, flags, logging |
| from absl.flags import FLAGS |
| import cv2 |
| import os |
| import shutil |
| import numpy as np |
| import tensorflow as tf |
| from core.yolov4 import filter_boxes |
| from tensorflow.python.saved_model import tag_constants |
| import core.utils as utils |
| from core.config import cfg |
|
|
| flags.DEFINE_string('weights', './checkpoints/yolov4-416', |
| 'path to weights file') |
| flags.DEFINE_string('framework', 'tf', 'select model type in (tf, tflite, trt)' |
| 'path to weights file') |
| flags.DEFINE_string('model', 'yolov4', 'yolov3 or yolov4') |
| flags.DEFINE_boolean('tiny', False, 'yolov3 or yolov3-tiny') |
| flags.DEFINE_integer('size', 416, 'resize images to') |
| flags.DEFINE_string('annotation_path', "./data/dataset/val2017.txt", 'annotation path') |
| flags.DEFINE_string('write_image_path', "./data/detection/", 'write image path') |
| flags.DEFINE_float('iou', 0.5, 'iou threshold') |
| flags.DEFINE_float('score', 0.25, 'score threshold') |
|
|
| def main(_argv): |
| INPUT_SIZE = FLAGS.size |
| STRIDES, ANCHORS, NUM_CLASS, XYSCALE = utils.load_config(FLAGS) |
| CLASSES = utils.read_class_names(cfg.YOLO.CLASSES) |
|
|
| predicted_dir_path = './mAP/predicted' |
| ground_truth_dir_path = './mAP/ground-truth' |
| if os.path.exists(predicted_dir_path): shutil.rmtree(predicted_dir_path) |
| if os.path.exists(ground_truth_dir_path): shutil.rmtree(ground_truth_dir_path) |
| if os.path.exists(cfg.TEST.DECTECTED_IMAGE_PATH): shutil.rmtree(cfg.TEST.DECTECTED_IMAGE_PATH) |
|
|
| os.mkdir(predicted_dir_path) |
| os.mkdir(ground_truth_dir_path) |
| os.mkdir(cfg.TEST.DECTECTED_IMAGE_PATH) |
|
|
| |
| if FLAGS.framework == 'tflite': |
| interpreter = tf.lite.Interpreter(model_path=FLAGS.weights) |
| interpreter.allocate_tensors() |
| input_details = interpreter.get_input_details() |
| output_details = interpreter.get_output_details() |
| print(input_details) |
| print(output_details) |
| else: |
| saved_model_loaded = tf.saved_model.load(FLAGS.weights, tags=[tag_constants.SERVING]) |
| infer = saved_model_loaded.signatures['serving_default'] |
|
|
| num_lines = sum(1 for line in open(FLAGS.annotation_path)) |
| with open(cfg.TEST.ANNOT_PATH, 'r') as annotation_file: |
| for num, line in enumerate(annotation_file): |
| annotation = line.strip().split() |
| image_path = annotation[0] |
| image_name = image_path.split('/')[-1] |
| image = cv2.imread(image_path) |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
| bbox_data_gt = np.array([list(map(int, box.split(','))) for box in annotation[1:]]) |
|
|
| if len(bbox_data_gt) == 0: |
| bboxes_gt = [] |
| classes_gt = [] |
| else: |
| bboxes_gt, classes_gt = bbox_data_gt[:, :4], bbox_data_gt[:, 4] |
| ground_truth_path = os.path.join(ground_truth_dir_path, str(num) + '.txt') |
|
|
| print('=> ground truth of %s:' % image_name) |
| num_bbox_gt = len(bboxes_gt) |
| with open(ground_truth_path, 'w') as f: |
| for i in range(num_bbox_gt): |
| class_name = CLASSES[classes_gt[i]] |
| xmin, ymin, xmax, ymax = list(map(str, bboxes_gt[i])) |
| bbox_mess = ' '.join([class_name, xmin, ymin, xmax, ymax]) + '\n' |
| f.write(bbox_mess) |
| print('\t' + str(bbox_mess).strip()) |
| print('=> predict result of %s:' % image_name) |
| predict_result_path = os.path.join(predicted_dir_path, str(num) + '.txt') |
| |
| image_size = image.shape[:2] |
| |
| image_data = cv2.resize(np.copy(image), (INPUT_SIZE, INPUT_SIZE)) |
| image_data = image_data / 255. |
| image_data = image_data[np.newaxis, ...].astype(np.float32) |
|
|
| if FLAGS.framework == 'tflite': |
| interpreter.set_tensor(input_details[0]['index'], image_data) |
| interpreter.invoke() |
| pred = [interpreter.get_tensor(output_details[i]['index']) for i in range(len(output_details))] |
| if FLAGS.model == 'yolov4' and FLAGS.tiny == True: |
| boxes, pred_conf = filter_boxes(pred[1], pred[0], score_threshold=0.25) |
| else: |
| boxes, pred_conf = filter_boxes(pred[0], pred[1], score_threshold=0.25) |
| else: |
| batch_data = tf.constant(image_data) |
| pred_bbox = infer(batch_data) |
| for key, value in pred_bbox.items(): |
| boxes = value[:, :, 0:4] |
| pred_conf = value[:, :, 4:] |
|
|
| boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression( |
| boxes=tf.reshape(boxes, (tf.shape(boxes)[0], -1, 1, 4)), |
| scores=tf.reshape( |
| pred_conf, (tf.shape(pred_conf)[0], -1, tf.shape(pred_conf)[-1])), |
| max_output_size_per_class=50, |
| max_total_size=50, |
| iou_threshold=FLAGS.iou, |
| score_threshold=FLAGS.score |
| ) |
| boxes, scores, classes, valid_detections = [boxes.numpy(), scores.numpy(), classes.numpy(), valid_detections.numpy()] |
|
|
| |
| |
| |
|
|
| with open(predict_result_path, 'w') as f: |
| image_h, image_w, _ = image.shape |
| for i in range(valid_detections[0]): |
| if int(classes[0][i]) < 0 or int(classes[0][i]) > NUM_CLASS: continue |
| coor = boxes[0][i] |
| coor[0] = int(coor[0] * image_h) |
| coor[2] = int(coor[2] * image_h) |
| coor[1] = int(coor[1] * image_w) |
| coor[3] = int(coor[3] * image_w) |
|
|
| score = scores[0][i] |
| class_ind = int(classes[0][i]) |
| class_name = CLASSES[class_ind] |
| score = '%.4f' % score |
| ymin, xmin, ymax, xmax = list(map(str, coor)) |
| bbox_mess = ' '.join([class_name, score, xmin, ymin, xmax, ymax]) + '\n' |
| f.write(bbox_mess) |
| print('\t' + str(bbox_mess).strip()) |
| print(num, num_lines) |
|
|
| if __name__ == '__main__': |
| try: |
| app.run(main) |
| except SystemExit: |
| pass |
|
|
|
|
|
|
