-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathyolo2.py
159 lines (118 loc) · 6.23 KB
/
yolo2.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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
import numpy as np
import argparse
import cv2
import os
import time
def extract_boxes_confidences_classids(outputs, confidence, width, height):
boxes = []
confidences = []
classIDs = []
for output in outputs:
for detection in output:
# Extract the scores, classid, and the confidence of the prediction
scores = detection[5:]
classID = np.argmax(scores)
conf = scores[classID]
# Consider only the predictions that are above the confidence threshold
if conf > confidence:
# Scale the bounding box back to the size of the image
box = detection[0:4] * np.array([width, height, width, height])
centerX, centerY, w, h = box.astype('int')
# Use the center coordinates, width and height to get the coordinates of the top left corner
x = int(centerX - (w / 2))
y = int(centerY - (h / 2))
boxes.append([x, y, int(w), int(h)])
confidences.append(float(conf))
classIDs.append(classID)
return boxes, confidences, classIDs
def draw_bounding_boxes(image, boxes, confidences, classIDs, idxs, colors):
if len(idxs) > 0:
for i in idxs.flatten():
# extract bounding box coordinates
x, y = boxes[i][0], boxes[i][1]
w, h = boxes[i][2], boxes[i][3]
# draw the bounding box and label on the image
color = [int(c) for c in colors[classIDs[i]]]
cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
text = "{}: {:.4f}".format(labels[classIDs[i]], confidences[i])
cv2.putText(image, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
return image
def make_prediction(net, layer_names, labels, image, confidence, threshold):
height, width = image.shape[:2]
# Create a blob and pass it through the model
blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (416, 416), swapRB=True, crop=False)
net.setInput(blob)
outputs = net.forward(layer_names)
# Extract bounding boxes, confidences and classIDs
boxes, confidences, classIDs = extract_boxes_confidences_classids(outputs, confidence, width, height)
# Apply Non-Max Suppression
idxs = cv2.dnn.NMSBoxes(boxes, confidences, confidence, threshold)
return boxes, confidences, classIDs, idxs
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('-w', '--weights', type=str, default='model/yolov3.weights', help='Path to model weights')
parser.add_argument('-cfg', '--config', type=str, default='model/yolov3.cfg', help='Path to configuration file')
parser.add_argument('-l', '--labels', type=str, default='model/coco.names', help='Path to label file')
parser.add_argument('-c', '--confidence', type=float, default=0.5, help='Minimum confidence for a box to be detected.')
parser.add_argument('-t', '--threshold', type=float, default=0.3, help='Threshold for Non-Max Suppression')
parser.add_argument('-u', '--use_gpu', default=False, action='store_true', help='Use GPU (OpenCV must be compiled for GPU). For more info checkout: https://www.pyimagesearch.com/2020/02/03/how-to-use-opencvs-dnn-module-with-nvidia-gpus-cuda-and-cudnn/')
parser.add_argument('-s', '--save', default=False, action='store_true', help='Whether or not the output should be saved')
parser.add_argument('-sh', '--show', default=True, action="store_false", help='Show output')
input_group = parser.add_mutually_exclusive_group()
input_group.add_argument('-i', '--image_path', type=str, default='', help='Path to the image file.')
input_group.add_argument('-v', '--video_path', type=str, default='', help='Path to the video file.')
args = parser.parse_args()
# Get the labels
labels = open(args.labels).read().strip().split('\n')
# Create a list of colors for the labels
colors = np.random.randint(0, 255, size=(len(labels), 3), dtype='uint8')
# Load weights using OpenCV
net = cv2.dnn.readNetFromDarknet(args.config, args.weights)
if args.use_gpu:
print('Using GPU')
net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
if args.save:
print('Creating output directory if it doesn\'t already exist')
os.makedirs('output', exist_ok=True)
# Get the ouput layer names
layer_names = net.getLayerNames()
layer_names = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
if args.image_path != '':
image = cv2.imread(args.image_path)
boxes, confidences, classIDs, idxs = make_prediction(net, layer_names, labels, image, args.confidence, args.threshold)
image = draw_bounding_boxes(image, boxes, confidences, classIDs, idxs, colors)
# show the output image
if args.show:
cv2.imshow('YOLO Object Detection', image)
cv2.waitKey(0)
if args.save:
cv2.imwrite(f'output/{args.image_path.split("/")[-1]}', image)
else:
if args.video_path != '':
cap = cv2.VideoCapture(args.video_path)
else:
cap = cv2.VideoCapture(0)
if args.save:
width = int(cap.get(3))
height = int(cap.get(4))
fps = cap.get(cv2.CAP_PROP_FPS)
name = args.video_path.split("/")[-1] if args.video_path else 'camera.avi'
out = cv2.VideoWriter(f'output/{name}', cv2.VideoWriter_fourcc('M','J','P','G'), fps, (width, height))
while cap.isOpened():
ret, image = cap.read()
if not ret:
print('Video file finished.')
break
boxes, confidences, classIDs, idxs = make_prediction(net, layer_names, labels, image, args.confidence, args.threshold)
image = draw_bounding_boxes(image, boxes, confidences, classIDs, idxs, colors)
if args.show:
cv2.imshow('YOLO Object Detection', image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
if args.save:
out.write(image)
cap.release()
if args.save:
out.release()
cv2.destroyAllWindows()