/
Detector_YoLo.py
166 lines (131 loc) · 4.76 KB
/
Detector_YoLo.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
160
161
162
163
164
165
166
#########################################
# Detector y discriminador de objetivos #
#########################################
##############################
# You Only Look Once (YOLO) #
##############################
# https: // gilberttanner.com/blog/yolo-object-detection-with-opencv
# https://www.pyimagesearch.com/2020/01/27/yolo-and-tiny-yolo-object-detection-on-the-raspberry-pi-and-movidius-ncs/
import numpy as np
import tensorflow as tf
import cv2 as cv
from tensorflow.python.keras.api import keras
from tensorflow.keras import datasets, layers, models
from imutils.video import FPS
import Utiles
import Config
import Tracker
def genera_predicciones(modelo, layer_names, labels, image, confidence):
height, width = image.shape[:2]
image = cv.UMat(image)
# Create a blob and pass it through the model
blob = cv.dnn.blobFromImage(
image, 1 / 255.0, (416, 416), swapRB=True, crop=False)
modelo.setInput(blob)
outputs = modelo.forward(layer_names)
# Extract bounding boxes, confidences and classIDs
ROIs, predicciones, tipos = genera_ROIs(
outputs, confidence, width, height)
# Apply Non-Max Suppression
# Elimina las detecciones no maximas
threshold = 0.3
idxs = cv.dnn.NMSBoxes(ROIs, predicciones, confidence, threshold)
return ROIs, predicciones, tipos, idxs
def genera_ROIs(outputs, umbral, width, height):
ROIs = []
predicciones = []
tipos = []
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)
confidence = scores[classID]
# Consider only the predictions that are above the confidence threshold
if confidence > umbral:
# 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))
ROIs.append([x, y, int(w), int(h)])
predicciones.append(float(confidence))
tipos.append(classID)
return ROIs, predicciones, tipos
def dibuja_ROIs(ROIs, predicciones, tipos, idxs):
"Auxiliar, para DEBUG, Genera una imagen con varias ROI"
global image
if len(idxs) > 0:
for i in idxs.flatten():
dibuja_ROI(ROIs[i], tipos[i], predicciones[i])
def dibuja_ROI(ROI, indice, confianza):
# extract bounding box coordinates
x, y, w, h = ROI
# color = [int(c) for c in colores[tipos[i]]]
color = colores[indice]
cv.rectangle(image, (x, y), (x + w, y + h), color, 2)
text = "{}: {:.4f}".format(labels[indice], confianza)
cv.putText(image, text, (x, y - 5),
cv.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
def genera_DNN():
"Genera una Red Neuronal Profunda Preentrenada con los parametros de Config"
modelo = []
modelo = cv.dnn.readNetFromDarknet(
Config.DNN.YoLo.archivo_modelo, Config.DNN.YoLo.archivo_pesos)
if(Config.DNN.gpu):
modelo.setPreferableBackend(cv.dnn.DNN_BACKEND_CUDA)
modelo.setPreferableTarget(cv.dnn.DNN_TARGET_CUDA)
return modelo
def genera_labels():
return open(Config.DNN.archivo_labels).read().strip().split('\n')
def genera_colores(labels):
"Genera colores aleatorios segun los labels"
# colores = np.random.randint(0, 255, size=(len(labels), 3), dtype='uint8')
colores = np.random.uniform(0, 255, size=(len(labels), 3))
colores[0] = [84, 7, 220]
return colores
if __name__ == "__main__":
# DEBUG Prueba de las funciones (No se usara, Archivo usado como libreria)
titulo = "Detector YoLo"
dims = (0, 0)
# Config.Fullscreen(titulo)
out = None
# Para capturar la salida
if Config.VidProp.guardar:
from Config import VidProp
out = cv.VideoWriter(f"Salida/{titulo}.avi",
VidProp.fourcc, VidProp.fps, VidProp.resolu)
# print(out.get(2)) # TODO
# Abre el video y almacena las dimesiones
cap = cv.VideoCapture(Config.VidProp.source)
dims = Utiles.dimensiones_video(cap)
# Crea la red neural
modelo = genera_DNN()
# Extrae las capas de YoLo
layer_names = modelo.getLayerNames()
layer_names = [layer_names[i[0] - 1]
for i in modelo.getUnconnectedOutLayers()]
labels = genera_labels()
colores = genera_colores(labels)
fps = FPS().start()
while cap.isOpened():
ret, image = cap.read()
if not ret:
print('Video file finished.')
break
umbral_confianza = 0.5
ROIs, predicciones, tipos, idxs = genera_predicciones(
modelo, layer_names, labels, image, umbral_confianza)
image = cv.UMat(image)
dibuja_ROIs(ROIs, predicciones, tipos, idxs)
if Config.VidProp.show_fps:
Utiles.dibuja_FPS(image, fps)
cv.imshow(titulo, image)
if (cv.waitKey(1) & 0xFF == 27):
break
if Config.VidProp.guardar:
Utiles.guardar(out, image)
cap.release()
if Config.VidProp.guardar:
out.release()