-
Notifications
You must be signed in to change notification settings - Fork 3
/
FindAndDrawAvgVideo-Multi-Thread.py
241 lines (199 loc) · 7.9 KB
/
FindAndDrawAvgVideo-Multi-Thread.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
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
# import the necessary packages
import numpy as np
import argparse
import imutils
import time
import cv2
import os
import darknet as dn
import pdb
import threading
from ctypes import *
from tqdm import tqdm
def add_padding(img, pad_l, pad_t, pad_r, pad_b):
height, width, colors = img.shape
# Adding padding to the left side.
pad_left = np.zeros([height, pad_l, 3])
img = np.concatenate((pad_left, img), axis=1)
# Adding padding to the top.
pad_up = np.zeros([pad_t, pad_l + width, 3])
img = np.concatenate((pad_up, img), axis=0)
# Adding padding to the right.
pad_right = np.zeros([height + pad_t, pad_r, 3])
img = np.concatenate((img, pad_right), axis=1)
# Adding padding to the bottom
pad_bottom = np.zeros([pad_b, pad_l + width + pad_r, 3])
img = np.concatenate((img, pad_bottom), axis=0)
return img
def frame_detect(frame,net,meta,frames, frame_num, fh ,fw):
outs = dn.detect(net, meta, frame)
# initialize our lists of detected bounding boxes, confidences,
# and class IDs, respectively
boxes = []
confidences = []
classIDs = []
avgx = 0
avgy = 0
players = 0
# ensure at least one detection exists
if len(outs) > 0:
# loop over the indexes we are keeping
for i in range(len(outs)):
if(outs[i][1] > 0.5):
# extract the bounding box coordinates
x = outs[i][2][0]
y = outs[i][2][1]
w = outs[i][2][2]
h = outs[i][2][3]
#x = x-(w/2)
#y = y-(h/2)
#x = int(x*(100/scale_percent))
#y = int(y*(100/scale_percent))
#w = int(w*(100/scale_percent))
#h = int(h*(100/scale_percent))
# draw a bounding box rectangle and label on the frame
#color = [int(c) for c in COLORS[0]]
#cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
#text = "{}: {:.4f}".format('player', outs[i][1])
#cv2.putText(frame, text, (x, y - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
players += 1
avgx = (avgx * (players-1) + x)/players
avgy = (avgy * (players-1) + y)/players
classIDs.append(1)
confidences.append(float(1))
boxes.append([int(avgx), int(avgy), 3, 3])
classIDs.append(2)
confidences.append(float(1))
boxes.append([int(1920/2), int(1080/2), 3, 3])
if len(outs) > 0:
# loop over the indexes we are keeping
for i in range(len(boxes)):
# extract the bounding box coordinates
(x, y) = (boxes[i][0], boxes[i][1])
(w, h) = (boxes[i][2], boxes[i][3])
# draw a bounding box rectangle and label on the frame
color = [int(c) for c in COLORS[0]]
cv2.rectangle(frame, (x, y), (x + w, y + h), color, 2)
if i == 0:
text = "Center of Players"
else:
text = "Center of Frame"
cv2.putText(frame, text, (x, y - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
cv2.line(frame, (x, y), (boxes[i-1][0], boxes[i-1][1]), color,
thickness=1, lineType=8, shift=0)
# write the output frame to disk
################################################################################
img = np.asarray(frame)
pad_l = int(max(-1 * (boxes[0][0] - fw/2), 0))
pad_u = int(max(-1 * (boxes[0][1] - fh/2), 0))
pad_r = int(max(boxes[0][0] - fw/2, 0))
pad_d = int(max(boxes[0][1] - fh/2, 0))
crop_l = int(max(boxes[0][0] - fw/2, 0))
crop_u = int(max(boxes[0][1] - fh/2, 0))
crop_r = int(min(fw, boxes[0][0] + fw/2))
crop_d = int(min(fh, boxes[0][1] + fh/2))
cropped_image = img[crop_u:crop_d, crop_l:crop_r]
im2 = add_padding(cropped_image, pad_l, pad_u, pad_r, pad_d)
#################################################################################
frames[frame_num]= np.uint8(im2)
# construct the argument parse and parse the arguments
ap = argparse.ArgumentParser()
ap.add_argument("-i", "--input", required=True,
help="path to input video")
ap.add_argument("-o", "--output", required=True,
help="path to output video")
ap.add_argument("-y", "--yolo",
help="base path to YOLO directory")
ap.add_argument("-c", "--confidence", type=float, default=0.5,
help="minimum probability to filter weak detections")
ap.add_argument("-t", "--threshold", type=float, default=0.3,
help="threshold when applyong non-maxima suppression")
args = vars(ap.parse_args())
# load the COCO class labels our YOLO model was trained on
labelsPath = os.path.sep.join(["player5-obj.names"])
LABELS = open(labelsPath).read().strip().split("\n")
# initialize a list of colors to represent each possible class label
np.random.seed(42)
COLORS = np.random.randint(0, 255, size=(len(LABELS), 3),dtype="uint8")
# derive the paths to the YOLO weights and model configuration
#weightsPath = os.path.sep.join([args["yolo"], "player5-yolov3_24000.weights"])
#configPath = os.path.sep.join([args["yolo"], "player4-yolov3.cfg"])
# net = dn.load_net(c_char_p('cfg/player5-yolov3.cfg'.encode('utf-8')), c_char_p('backup/player5-yolov3_24000.weights'.encode('utf-8')), 0)
# meta = dn.load_meta("cfg/player5-obj.data".encode('utf-8'))
cores = 2
nets = []
metas = []
for i in range(cores):
nets.append(dn.load_net(c_char_p('cfg/player5-yolov3.cfg'.encode('utf-8')), c_char_p('backup/player5-yolov3_24000.weights'.encode('utf-8')), 0))
metas.append(dn.load_meta("cfg/player5-obj.data".encode('utf-8')))
# load our YOLO object detector trained on COCO dataset (80 classes)
# and determine only the *output* layer names that we need from YOLO
print("[INFO] loading YOLO from disk...")
#net = cv2.dnn.readNetFromDarknet(configPath, weightsPath)
#ln = net.getLayerNames()
#ln = [ln[i[0] - 1] for i in net.getUnconnectedOutLayers()]
# initialize the video stream, pointer to output video file, and
# frame dimensions
vs = cv2.VideoCapture(args["input"])
writer = None
(W, H) = (None, None)
# try to determine the total number of frames in the video file
total = 0
try:
prop = cv2.cv.CV_CAP_PROP_FRAME_COUNT if imutils.is_cv2() \
else cv2.CAP_PROP_FRAME_COUNT
total = int(vs.get(prop))
print("[INFO] {} total frames in video".format(total))
# an error occurred while trying to determine the total
# number of frames in the video file
except:
print("[INFO] could not determine # of frames in video")
print("[INFO] no approx. completion time can be provided")
total = -1
# loop over frames from the video file stream
count = 0
# check if the video writer is None
fw = int(vs.get(cv2.CAP_PROP_FRAME_WIDTH))
fh = int(vs.get(cv2.CAP_PROP_FRAME_HEIGHT))
if writer is None:
# initialize our video writer
fourcc = cv2.VideoWriter_fourcc(*"XVID")
writer = cv2.VideoWriter(args["output"], fourcc, 59.94,(fh, fw), True)
og_start = time.time()
pbar = tqdm(total=total)
frames = [None]*total
frame_num = 0
threads =[]
while cores > 0:
(grabbed, frame) = vs.read()
if not grabbed:
break
print("test-1")
threads.append(threading.Thread(target=frame_detect, args=(frame,nets[frame_num%cores],metas[frame_num%cores],frames, frame_num, fh ,fw)))
print("test-2")
threads[-1].start()
cores -= 1
frame_num += 1
cores = 2
while len(threads) > 0:
threads[0].join()
threads.pop(0)
(grabbed, frame) = vs.read()
pbar.update(1)
if grabbed:
threads.append(threading.Thread(target=frame_detect, args=(frame,nets[frame_num%cores],metas[frame_num%cores],frames, frame_num, fh ,fw)))
threads[-1].start()
frame_num += 1
pbar.close()
og_end = time.time()
elap = (og_end - og_start)
print("[INFO] single frame took {:.4f} seconds".format(elap))
for im2 in frames:
cv2.imshow("object detection", im2)
cv2.waitKey()
writer.write(np.uint8(im2))
# release the file pointers
print("[INFO] cleaning up...")
writer.release()
vs.release()