SHARE
TWEET

Untitled

a guest Aug 17th, 2019 79 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. from __future__ import division, print_function
  2.  
  3. import os
  4.  
  5. import tensorflow as tf
  6. import numpy as np
  7. import argparse
  8. import cv2
  9.  
  10. from utils.misc_utils import parse_anchors, read_class_names
  11. from utils.nms_utils import gpu_nms
  12. from utils.plot_utils import get_color_table, plot_one_box
  13. from utils.data_aug import letterbox_resize
  14.  
  15. from model import yolov3
  16.  
  17. path = '/home/ml/Downloads/old/YOLOv3_TensorFlow/data/my_input/'
  18. for i in os.listdir(path):
  19.     parser = argparse.ArgumentParser(description="YOLO-V3 test single image test procedure.")
  20.     parser.add_argument("--input_image", type=str, default=path+i,
  21.                         help="The path of the input image.")
  22.     parser.add_argument("--anchor_path", type=str, default="/home/ml/Downloads/old/YOLOv3_TensorFlow/data/yolo_anchors.txt",
  23.                         help="The path of the anchor txt file.")
  24.     parser.add_argument("--new_size", nargs='*', type=int, default=[416, 416],
  25.                         help="Resize the input image with `new_size`, size format: [width, height]")
  26.     parser.add_argument("--letterbox_resize", type=lambda x: (str(x).lower() == 'true'), default=True,
  27.                         help="Whether to use the letterbox resize.")
  28.     parser.add_argument("--class_name_path", type=str, default="/home/ml/Downloads/old/YOLOv3_TensorFlow/data/data.names",
  29.                         help="The path of the class names.")
  30.     parser.add_argument("--restore_path", type=str, default="/home/ml/Downloads/old/YOLOv3_TensorFlow/data/darknet_weightsbest_model_Epoch_68_step_621_mAP_0.9167_loss_2.0931_lr_1e-05",
  31.                         help="The path of the weights to restore.")
  32.     args = parser.parse_args()
  33.  
  34.     args.anchors = parse_anchors(args.anchor_path)
  35.     args.classes = read_class_names(args.class_name_path)
  36.     args.num_class = len(args.classes)
  37.  
  38.     color_table = get_color_table(args.num_class)
  39.  
  40.     img_ori = cv2.imread(args.input_image)
  41.     if args.letterbox_resize:
  42.         img, resize_ratio, dw, dh = letterbox_resize(img_ori, args.new_size[0], args.new_size[1])
  43.     else:
  44.         height_ori, width_ori = img_ori.shape[:2]
  45.         img = cv2.resize(img_ori, tuple(args.new_size))
  46.     img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  47.     img = np.asarray(img, np.float32)
  48.     img = img[np.newaxis, :] / 255.
  49.  
  50.     with tf.Session() as sess:
  51.         input_data = tf.placeholder(tf.float32, [1, args.new_size[1], args.new_size[0], 3], name='input_data')
  52.         yolo_model = yolov3(args.num_class, args.anchors)
  53.         with tf.variable_scope('yolov3', reuse=True) as yolov3:
  54.             pred_feature_maps = yolo_model.forward(input_data, False, reuse=False)
  55.         pred_boxes, pred_confs, pred_probs = yolo_model.predict(pred_feature_maps)
  56.  
  57.         pred_scores = pred_confs * pred_probs
  58.  
  59.         boxes, scores, labels = gpu_nms(pred_boxes, pred_scores, args.num_class, max_boxes=200, score_thresh=0.1,
  60.                                         nms_thresh=0.45)
  61.  
  62.         saver = tf.train.Saver()
  63.         saver.restore(sess, args.restore_path)
  64.         boxes_, scores_, labels_ = sess.run([boxes, scores, labels], feed_dict={input_data: img})
  65.  
  66.         # rescale the coordinates to the original image
  67.         if args.letterbox_resize:
  68.             boxes_[:, [0, 2]] = (boxes_[:, [0, 2]] - dw) / resize_ratio
  69.             boxes_[:, [1, 3]] = (boxes_[:, [1, 3]] - dh) / resize_ratio
  70.         else:
  71.             boxes_[:, [0, 2]] *= (width_ori/float(args.new_size[0]))
  72.             boxes_[:, [1, 3]] *= (height_ori/float(args.new_size[1]))
  73.  
  74.         print("box coords:")
  75.         print(boxes_)
  76.         print('*' * 30)
  77.         print("scores:")
  78.         print(scores_)
  79.         print('*' * 30)
  80.         print("labels:")
  81.         print(labels_)
  82.  
  83.         crop_img = img_ori[boxes_[0][1].astype('int'):boxes_[0][3].astype('int'), boxes_[0][0].astype('int'):boxes_[0][2].astype('int')]
  84.         # cv2.imshow("cropped", crop_img)
  85.  
  86.         for j in range(len(boxes_)):
  87.             x0, y0, x1, y1 = boxes_[j]
  88.             plot_one_box(img_ori, [x0, y0, x1, y1], label=args.classes[labels_[j]] + ', {:.2f}%'.format(scores_[j] * 100), color=color_table[labels_[j]])
  89.         # cv2.imshow('Detection result', img_ori)
  90.         cv2.imwrite(f'detection_{i}', img_ori)
  91.         cv2.imwrite(f'crop_{i}', crop_img)
  92.  
  93. cv2.waitKey(0)
  94.      
  95. with tf.variable_scope('yolov3', reuse=True) as yolov3:
  96.             pred_feature_maps = yolo_model.forward(input_data, False, reuse=True)
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top