-
Notifications
You must be signed in to change notification settings - Fork 15
/
video_handler.py
68 lines (57 loc) · 2.57 KB
/
video_handler.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
from ball_detect_track import *
from plot_tools import *
from feet_detect import *
import skvideo.io
TOPCUT = 320
FLANN_INDEX_KDTREE = 1
index_params = dict(algorithm=FLANN_INDEX_KDTREE, trees=5)
search_params = dict(checks=50)
flann = cv2.FlannBasedMatcher(index_params, search_params)
class VideoHandler:
def __init__(self, pano, video, ball_detector, feet_detector, map_2d):
self.M1 = np.load("Rectify1.npy")
self.sift = cv2.xfeatures2d.SIFT_create()
self.pano = pano
self.video = video
self.kp1, self.des1 = self.sift.compute(pano, self.sift.detect(pano))
self.feet_detector = feet_detector
self.ball_detector = ball_detector
self.map_2d = map_2d
def run_detectors(self):
writer = skvideo.io.FFmpegWriter("demo.mp4")
time_index = 0
while self.video.isOpened():
ok, frame = self.video.read()
if not ok:
break
else:
if 0 <= time_index < 200:
print("\r Computing DEMO: " + str(int(100 * time_index / 200)) + "%",
flush=True, end='')
frame = frame[TOPCUT:, :]
M = self.get_homography(frame, self.des1, self.kp1)
frame, self.map_2d, map_2d_text = self.feet_detector.get_players_pos(M, self.M1, frame, time_index,
self.map_2d)
frame, ball_map_2d = self.ball_detector.ball_tracker(M, self.M1, frame, self.map_2d.copy())
vis = np.vstack((frame, cv2.resize(map_2d_text, (frame.shape[1], frame.shape[1] // 2))))
# cv2.imshow("Tracking", vis)
writer.writeFrame(cv2.cvtColor(vis, cv2.COLOR_BGR2RGB))
k = cv2.waitKey(1) & 0xff
if k == 27:
break
time_index += 1
self.video.release()
writer.close()
cv2.destroyAllWindows()
def get_homography(self, frame, des1, kp1):
kp2 = self.sift.detect(frame)
kp2, des2 = self.sift.compute(frame, kp2)
matches = flann.knnMatch(des1, des2, k=2)
good = []
for m, n in matches:
if m.distance < 0.7 * n.distance:
good.append(m)
src_pts = np.float32([kp1[m.queryIdx].pt for m in good]).reshape(-1, 1, 2)
dst_pts = np.float32([kp2[m.trainIdx].pt for m in good]).reshape(-1, 1, 2)
M, mask = cv2.findHomography(dst_pts, src_pts, cv2.RANSAC, 5.0)
return M