-
Notifications
You must be signed in to change notification settings - Fork 15
/
main.py
99 lines (76 loc) · 3.43 KB
/
main.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
import os
from matplotlib import pyplot as plt
from ball_detect_track import BallDetectTrack
from player import Player
from rectify_court import *
from video_handler import *
def get_frames(video_path, central_frame, mod):
frames = []
cap = cv2.VideoCapture(video_path)
index = 0
while cap.isOpened():
ret, frame = cap.read()
if (index % mod) == 0:
frames.append(frame[TOPCUT:, :])
if not ret or frame is None:
cap.release()
print("Released Video Resource")
break
if cv2.waitKey(20) == ord('q'): break
index += 1
cap.release()
cv2.destroyAllWindows()
print(f"Number of frames : {len(frames)}")
plt.title(f"Centrale {frames[central_frame].shape}")
plt.imshow(frames[central_frame])
plt.show()
return frames
#####################################################################
if __name__ == '__main__':
# COURT REAL SIZES
# 28m horizontal lines
# 15m vertical lines
# loading already computed panoramas
if os.path.exists('resources/pano.png'):
pano = cv2.imread("resources/pano.png")
else:
central_frame = 36
frames = get_frames('resources/Short4Mosaicing.mp4', central_frame, mod=3)
frames_flipped = [cv2.flip(frames[i], 1) for i in range(central_frame)]
current_mosaic1 = collage(frames[central_frame:], direction=1)
current_mosaic2 = collage(frames_flipped, direction=-1)
pano = collage([cv2.flip(current_mosaic2, 1)[:, :-10], current_mosaic1])
cv2.imwrite("resources/pano.png", pano)
if os.path.exists('resources/pano_enhanced.png'):
pano_enhanced = cv2.imread("resources/pano_enhanced.png")
plt_plot(pano, "Panorama")
else:
pano_enhanced = pano
for file in os.listdir("resources/snapshots/"):
frame = cv2.imread("resources/snapshots/" + file)[TOPCUT:]
pano_enhanced = add_frame(frame, pano, pano_enhanced, plot=False)
cv2.imwrite("resources/pano_enhanced.png", pano_enhanced)
###################################
pano_enhanced = np.vstack((pano_enhanced,
np.zeros((100, pano_enhanced.shape[1], pano_enhanced.shape[2]), dtype=pano.dtype)))
img = binarize_erode_dilate(pano_enhanced, plot=False)
simplified_court, corners = (rectangularize_court(img, plot=False))
simplified_court = 255 - np.uint8(simplified_court)
plt_plot(simplified_court, "Corner Detection", cmap="gray", additional_points=corners)
rectified = rectify(pano_enhanced, corners, plot=False)
# correspondences map-pano
map = cv2.imread("resources/2d_map.png")
scale = rectified.shape[0] / map.shape[0]
map = cv2.resize(map, (int(scale * map.shape[1]), int(scale * map.shape[0])))
resized = cv2.resize(rectified, (map.shape[1], map.shape[0]))
map = cv2.resize(map, (rectified.shape[1], rectified.shape[0]))
video = cv2.VideoCapture("resources/Short4Mosaicing.mp4")
players = []
for i in range(1, 6):
players.append(Player(i, 'green', hsv2bgr(COLORS['green'][2])))
players.append(Player(i, 'white', hsv2bgr(COLORS['white'][2])))
players.append(Player(0, 'referee', hsv2bgr(COLORS['referee'][2])))
feet_detector = FeetDetector(players)
ball_detect_track = BallDetectTrack(players)
video_handler = VideoHandler(pano_enhanced, video, ball_detect_track, feet_detector, map)
video_handler.run_detectors()