-
Notifications
You must be signed in to change notification settings - Fork 13
/
predict.py
executable file
·84 lines (65 loc) · 2.59 KB
/
predict.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
"""
Predict one image (for testing)
"""
import os
import cv2
import yaml
import argparse
import numpy as np
from keras.layers import Input
from yolov2.zoo import yolov2_darknet19
from yolov2.utils.painter import draw_boxes
from yolov2.utils.parser import parse_label_map
def _main_():
parser = argparse.ArgumentParser(description="Detect objects in an image")
parser.add_argument('--path', type=str, default='./test_imgs/person.jpg',
help="Path to image file")
parser.add_argument('--weights', type=str, default='yolo-coco.weights',
help="Path to pre-trained weight file")
parser.add_argument('--output_dir', type=str, default=None,
help="Output Directory")
parser.add_argument('--iou', type=float, default=0.5,
help="Intersection over Union (IoU) value")
parser.add_argument('--threshold', type=float, default=0.5,
help="Score Threshold value (minimum accuracy)")
# ############
# Parse Config
# ############
args = parser.parse_args()
with open('config.yml', 'r') as stream:
config = yaml.load(stream)
anchors = np.array(config['anchors'])
label_dict = parse_label_map(config['label_map'])
# ###################
# Define Keras Model
# ###################
model_cfg = config['model']
inputs = Input(shape=(None, None, 3))
model = yolov2_darknet19(inputs,
is_training=True,
anchors=anchors,
num_classes=model_cfg['num_classes'],
iou=args.iou,
scores_threshold=args.threshold,
max_boxes=100)
model.load_weights(args.weights)
model.summary()
# #####################
# Make one prediction #
# #####################
image = cv2.imread(args.path)
image = cv2.resize(image, (320, 320))
bboxes, scores, classes = model.predict_on_batch(np.expand_dims(image, axis=0))
# Convert idx to label
classes = [label_dict[idx] for idx in classes]
# #################
# Display Result #
# #################
h, w, _ = image.shape
if args.output_dir is not None:
# Scale relative coordinates into actual coordinates
bboxes = [box * np.array([h, w, h, w]) for box in bboxes]
result = draw_boxes(image, bboxes, classes, scores)
cv2.imwrite(os.path.join(args.output_dir, args.path.split('/')[-1].split('.')[0] + '_result.jpg'), result)
if __name__ == "__main__":
_main_()