-
Notifications
You must be signed in to change notification settings - Fork 15
/
Copy pathtest_single_image.py
102 lines (79 loc) · 3.56 KB
/
test_single_image.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
####
# edicted by Huangdebo
# test the model using ckpt file
# ***
from __future__ import division, print_function
import tensorflow as tf
import numpy as np
import argparse
import cv2
from utils.misc_utils import parse_anchors, read_class_names
from utils.nms_utils import gpu_nms
from utils.plot_utils import get_color_table, plot_one_box
from model.yolov3 import yolov3
from model.yolov3_tiny import yolov3_tiny
net_name = 'yolov3_tiny'
body_name = 'darknet19'
data_name = 'hardhat'
ckpt_name = 'model-step_10000_loss_0.452858_lr_0.0006925339'
parser = argparse.ArgumentParser(description="%s test single image test procedure."%net_name)
parser.add_argument("--input_image", type=str, default="./data/demo_data/00005.jpg",
help="The path of the input image.")
parser.add_argument("--anchor_path", type=str, default="./data/tiny_yolo_anchors.txt",
help="The path of the anchor txt file.")
parser.add_argument("--new_size", nargs='*', type=int, default=[416, 416],
help="Resize the input image with `new_size`, size format: [width, height]")
parser.add_argument("--class_name_path", type=str, default="./data/data.names",
help="The path of the class names.")
parser.add_argument("--restore_path", type=str, default="./checkpoint/%s_%s/%s"%(net_name,data_name,ckpt_name),
help="The path of the weights to restore.")
args = parser.parse_args()
args.anchors = parse_anchors(args.anchor_path)
args.classes = read_class_names(args.class_name_path)
args.num_class = len(args.classes)
color_table = get_color_table(args.num_class)
img_ori = cv2.imread(args.input_image)
height_ori, width_ori = img_ori.shape[:2]
img = cv2.resize(img_ori, tuple(args.new_size))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.asarray(img, np.float32)
img = img[np.newaxis, :] / 255.
with tf.Session() as sess:
input_data = tf.placeholder(tf.float32, [1, args.new_size[1], args.new_size[0], 3], name='input_data')
# yolo_model = yolov3(args.num_class, args.anchors)
yolo_model = yolov3_tiny(args.num_class, args.anchors)
with tf.variable_scope(net_name):
pred_feature_maps = yolo_model.forward(input_data, False)
pred_boxes, pred_confs, pred_probs = yolo_model.predict(pred_feature_maps)
pred_scores = pred_confs * pred_probs
boxes, scores, labels = gpu_nms(pred_boxes, pred_scores, args.num_class, max_boxes=200, score_thresh=0.4, iou_thresh=0.5)
saver = tf.train.Saver()
saver.restore(sess, args.restore_path)
boxes_, scores_, labels_ = sess.run([boxes, scores, labels], feed_dict={input_data: img})
# rescale the coordinates to the original image
boxes_[:, 0] *= (width_ori/float(args.new_size[0]))
boxes_[:, 2] *= (width_ori/float(args.new_size[0]))
boxes_[:, 1] *= (height_ori/float(args.new_size[1]))
boxes_[:, 3] *= (height_ori/float(args.new_size[1]))
print("box coords:")
print(boxes_)
print('*' * 30)
print("scores:")
print(scores_)
print('*' * 30)
print("labels:")
print(labels_)
class_labels=[0,1,2,3,4]
count=0
for lb in labels_:
for cl in class_labels:
if lb==cl:
count+=1
print("Worker added")
print("Number of workers detected in image:",count)
for i in range(len(boxes_)):
x0, y0, x1, y1 = boxes_[i]
plot_one_box(img_ori, [x0, y0, x1, y1], label=args.classes[labels_[i]], color=color_table[labels_[i]])
cv2.imshow('Detection result', img_ori)
cv2.imwrite('detection_result.jpg', img_ori)
cv2.waitKey(0)