/
predict.py
105 lines (77 loc) · 2.95 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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
#! /usr/bin/env python
import argparse
import os
import cv2
import numpy as np
from preprocessing import parse_annotation
from utils import draw_boxes
from frontend import YOLO
import json
import time
os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"]="0" #"0" use cpu
#Simple function to predict bounding boxes using YOLO
def _main_():
config_path = 'config_cargo_door.json'
weights_path = '/media/ubuntu/hdd/tensorflow_data/YOLO/CargoDoor/full_yolo_cargo_door.h5'
image_path = '/media/ubuntu/DANFOSS/airport_images.tar.gz/015.jpg'
with open(config_path) as config_buffer:
config = json.load(config_buffer)
###############################
# Make the model
###############################
yolo = YOLO(architecture = config['model']['architecture'],
input_size = config['model']['input_size'],
labels = config['model']['labels'],
max_box_per_image = config['model']['max_box_per_image'],
anchors = config['model']['anchors'])
###############################
# Load trained weights
###############################
print(weights_path)
yolo.load_weights(weights_path)
###############################
# Predict bounding boxes
###############################
cv2.namedWindow("Detection", cv2.WINDOW_NORMAL)
cv2.startWindowThread() #to make sure we can close it later on
if image_path[-4:] == '.mp4':
video_reader = cv2.VideoCapture(image_path)
nb_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
frame_h = int(video_reader.get(cv2.CAP_PROP_FRAME_HEIGHT))
frame_w = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
for i in range(nb_frames):
_, image = video_reader.read()
boxes = yolo.predict(image)
image = draw_boxes(image, boxes, config['model']['labels'])
print len(boxes), 'box(es) found'
#display frame
image_np = np.uint8(image)
cv2.imshow("Detection", image_np)
k = cv2.waitKey(0) & 0xEFFFFF
if k == 27:
print("You Pressed Escape")
break
video_reader.release()
else:
image = cv2.imread(image_path)
#Predict and time it
t0 = time.time()
boxes = yolo.predict(image)
t1 = time.time()
total = t1-t0
#overlay boxes
image = draw_boxes(image, boxes, config['model']['labels'])
#feedback
print len(boxes), 'box(es) found'
print 'Prediciton took %f seconds'%(total)
#display frame
cv2.imshow("Detection", image)
k = cv2.waitKey(0) & 0xEFFFFF
if k == 27:
print("You Pressed Escape")
cv2.destroyAllWindows()
for i in range (1,5):
cv2.waitKey(1)
if __name__ == '__main__':
_main_()