-
Notifications
You must be signed in to change notification settings - Fork 3
/
yolodetect.py
113 lines (94 loc) · 4.2 KB
/
yolodetect.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
from shapely.geometry import Point
from shapely.geometry.polygon import Polygon
import cv2
import numpy as np
from telegram_utils import send_telegram
import datetime
import threading
def isInside(points, centroid):
polygon = Polygon(points)
centroid = Point(centroid)
print(polygon.contains(centroid))
return polygon.contains(centroid)
class YoloDetect():
def __init__(self, detect_class="person", frame_width=1280, frame_height=720):
# Parameters
self.classnames_file = "model/classnames.txt"
self.weights_file = "model/yolov4-tiny.weights"
self.config_file = "model/yolov4-tiny.cfg"
self.conf_threshold = 0.5
self.nms_threshold = 0.4
self.detect_class = detect_class
self.frame_width = frame_width
self.frame_height = frame_height
self.scale = 1 / 255
self.model = cv2.dnn.readNet(self.weights_file, self.config_file)
self.classes = None
self.output_layers = None
self.read_class_file()
self.get_output_layers()
self.last_alert = None
self.alert_telegram_each = 15 # seconds
def read_class_file(self):
with open(self.classnames_file, 'r') as f:
self.classes = [line.strip() for line in f.readlines()]
def get_output_layers(self):
net = cv2.dnn.readNet(self.weights_file, self.config_file)
layer_names = self.model.getLayerNames()
# old
self.output_layers = [layer_names[i - 1] for i in net.getUnconnectedOutLayers()]
# new
# self.output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
def draw_prediction(self, img, class_id, x, y, x_plus_w, y_plus_h, points):
label = str(self.classes[class_id])
color = (0, 255, 0)
cv2.rectangle(img, (x, y), (x_plus_w, y_plus_h), color, 2)
cv2.putText(img, label, (x - 10, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
# Tinh toan centroid
centroid = ((x + x_plus_w) // 2, (y + y_plus_h) // 2)
cv2.circle(img, centroid, 5, (color), -1)
if isInside(points, centroid):
img = self.alert(img)
return isInside(points, centroid)
def alert(self, img):
cv2.putText(img, "Warning !!!", (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
# New thread to send telegram after 15 seconds
if (self.last_alert is None) or (
(datetime.datetime.utcnow() - self.last_alert).total_seconds() > self.alert_telegram_each):
self.last_alert = datetime.datetime.utcnow()
cv2.imwrite("alert.png", cv2.resize(img, dsize=(1280,720), fx=0.2, fy=0.2))
thread = threading.Thread(target=send_telegram)
thread.start()
return img
def detect(self, frame, points):
blob = cv2.dnn.blobFromImage(frame, self.scale, (416, 416), (0, 0, 0), True, crop=False)
self.model.setInput(blob)
outs = self.model.forward(self.output_layers)
# Loc cac object trong khung hinh
class_ids = []
confidences = []
boxes = []
for out in outs:
for detection in out:
scores = detection[5:]
class_id = np.argmax(scores)
confidence = scores[class_id]
if (confidence >= self.conf_threshold) and (self.classes[class_id] == self.detect_class):
center_x = int(detection[0] * self.frame_width)
center_y = int(detection[1] * self.frame_height)
w = int(detection[2] * self.frame_width)
h = int(detection[3] * self.frame_height)
x = center_x - w / 2
y = center_y - h / 2
class_ids.append(class_id)
confidences.append(float(confidence))
boxes.append([x, y, w, h])
indices = cv2.dnn.NMSBoxes(boxes, confidences, self.conf_threshold, self.nms_threshold)
for i in indices:
box = boxes[i]
x = box[0]
y = box[1]
w = box[2]
h = box[3]
self.draw_prediction(frame, class_ids[i], round(x), round(y), round(x + w), round(y + h), points)
return frame