-
Notifications
You must be signed in to change notification settings - Fork 0
/
social_distance_detector.py
116 lines (106 loc) · 4.23 KB
/
social_distance_detector.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
# -*- coding: utf-8 -*-
"""social_distance_detector.ipynb
Automatically generated by Colaboratory.
Original file is located at
https://colab.research.google.com/drive/13tPLS-tqn_xnTUUwJ6ZmcUc0hc1gb1l8
"""
# import the necessary packages
from socialdistancing.detection import detect_people
from scipy.spatial import distance as dist
import numpy as np
import argparse
import imutils
import cv2
import os
MODEL_PATH = "yolo-coco"
MIN_DISTANCE = 50
# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--input", type=str, default="",
help="path to (optional) input video file")
ap.add_argument("-o", "--output", type=str, default="",
help="path to (optional) output video file")
ap.add_argument("-d", "--display", type=int, default=1,
help="whether or not output frame should be displayed")
args = vars(ap.parse_args())
# load the COCO class labels our YOLO model was trained on
labelsPath = os.path.sep.join([MODEL_PATH, "coco.names"])
LABELS = open(labelsPath).read().strip().split("\n")
# derive the paths to the YOLO weights and model configuration
weightsPath = os.path.sep.join([MODEL_PATH, "yolov3.weights"])
configPath = os.path.sep.join([MODEL_PATH, "yolov3.cfg"])
#load our YOLO object detector trained on COCO dataset (80 classes)
print("[INFO] loading YOLO from disk...")
net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)
#we need from YOLO determine only the *output* layer names
ln = net.getLayerNames()
#print(ln)
ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
# initialize the video stream and pointer to output video file
print("[INFO] accessing video stream...")
vs = cv2.VideoCapture(args["input"] if args["input"] else 0)
writer = None
# loop over the frames from the video stream
while True:
#from the file read the next frame
(grabbed, frame) = vs.read()
# if we have reached the end then the frame was not grabbed of the stream
if not grabbed:
break
#detect people and resize the frame
frame = imutils.resize(frame, width=700)
#print(frame.shape)
results = detect_people(frame, net, ln,personIdx=LABELS.index("person"))
#print(results)
#violate the minimum social
# distance initialize the set of indexes
violate = set()
#in order to compute our pairwise distance maps ensure there are *at least* two people detections (required)
if len(results) >= 2:
# fectch centroids from the results and calculate
# Euclidean distances between all pairs of the centroids
centroids = np.array([r[2] for r in results])
D = dist.cdist(centroids, centroids, metric="euclidean")
# loop over the upper triangular of the distance matrix
for i in range(0, D.shape[0]):
for j in range(i + 1, D.shape[1]):
#print(i,j)
if D[i, j] < MIN_DISTANCE:
#the centroid pairs update our violation set with the indexes of
violate.add(i)
violate.add(j)
#print(i,j)
# loop over the results
for (i, (prob, bbox, centroid)) in enumerate(results):
#extract the bounding box and centroid coordinates and initialize the color of the annotatio, then
(startX, startY, endX, endY) = bbox
(cX, cY) = centroid
#print(cX,cY)
color = (0, 255, 0)
#update the color if the index pair exists within the violation set, then
if i in violate:
color = (0, 0, 255)
#print(color.shape)
# draw (1) a bounding box around the person and (2) the
# centroid coordinates of the person,
cv2.rectangle(frame, (startX, startY), (endX, endY), color, 2)
cv2.circle(frame, (cX, cY), 5, color, 1)
#print(i)
#output frame draw the total number of social distancing violations
text = "Social Distancing Violations: {}".format(len(violate))
cv2.putText(frame, text, (10, frame.shape[0] - 25),
cv2.FONT_HERSHEY_SIMPLEX, 0.85, (0, 0, 255), 3)
if args["display"] > 0:
# show the output frame
cv2.imshow("Frame", frame)
key = cv2.waitKey(1) & 0xFF
# if the `q` key was pressed, break from the loop
if key == ord("q"):
break
if args["output"] != "" and writer is None:
# initialize our video writer
fourcc = cv2.VideoWriter_fourcc(*"MJPG")
writer = cv2.VideoWriter(args["output"], fourcc, 25,
(frame.shape[1], frame.shape[0]), True)
if writer is not None:
writer.write(frame)