forked from yastrebksv/TennisCourtDetector
-
Notifications
You must be signed in to change notification settings - Fork 0
/
court_reference.py
126 lines (110 loc) · 5.85 KB
/
court_reference.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
import cv2
import numpy as np
import matplotlib.pyplot as plt
class CourtReference:
"""
Court reference model
"""
def __init__(self):
self.baseline_top = ((286, 561), (1379, 561))
self.baseline_bottom = ((286, 2935), (1379, 2935))
self.net = ((286, 1748), (1379, 1748))
self.left_court_line = ((286, 561), (286, 2935))
self.right_court_line = ((1379, 561), (1379, 2935))
self.left_inner_line = ((423, 561), (423, 2935))
self.right_inner_line = ((1242, 561), (1242, 2935))
self.middle_line = ((832, 1110), (832, 2386))
self.top_inner_line = ((423, 1110), (1242, 1110))
self.bottom_inner_line = ((423, 2386), (1242, 2386))
self.top_extra_part = (832.5, 580)
self.bottom_extra_part = (832.5, 2910)
self.key_points = [*self.baseline_top, *self.baseline_bottom,
*self.left_inner_line, *self.right_inner_line,
*self.top_inner_line, *self.bottom_inner_line,
*self.middle_line]
self.border_points = [*self.baseline_top, *self.baseline_bottom[::-1]]
self.court_conf = {1: [*self.baseline_top, *self.baseline_bottom],
2: [self.left_inner_line[0], self.right_inner_line[0], self.left_inner_line[1],
self.right_inner_line[1]],
3: [self.left_inner_line[0], self.right_court_line[0], self.left_inner_line[1],
self.right_court_line[1]],
4: [self.left_court_line[0], self.right_inner_line[0], self.left_court_line[1],
self.right_inner_line[1]],
5: [*self.top_inner_line, *self.bottom_inner_line],
6: [*self.top_inner_line, self.left_inner_line[1], self.right_inner_line[1]],
7: [self.left_inner_line[0], self.right_inner_line[0], *self.bottom_inner_line],
8: [self.right_inner_line[0], self.right_court_line[0], self.right_inner_line[1],
self.right_court_line[1]],
9: [self.left_court_line[0], self.left_inner_line[0], self.left_court_line[1],
self.left_inner_line[1]],
10: [self.top_inner_line[0], self.middle_line[0], self.bottom_inner_line[0],
self.middle_line[1]],
11: [self.middle_line[0], self.top_inner_line[1], self.middle_line[1],
self.bottom_inner_line[1]],
12: [*self.bottom_inner_line, self.left_inner_line[1], self.right_inner_line[1]]}
self.line_width = 1
self.court_width = 1117
self.court_height = 2408
self.top_bottom_border = 549
self.right_left_border = 274
self.court_total_width = self.court_width + self.right_left_border * 2
self.court_total_height = self.court_height + self.top_bottom_border * 2
# self.court = cv2.cvtColor(cv2.imread('court_configurations/court_reference.png'), cv2.COLOR_BGR2GRAY)
def build_court_reference(self):
"""
Create court reference image using the lines positions
"""
court = np.zeros((self.court_height + 2 * self.top_bottom_border, self.court_width + 2 * self.right_left_border), dtype=np.uint8)
cv2.line(court, *self.baseline_top, 1, self.line_width)
cv2.line(court, *self.baseline_bottom, 1, self.line_width)
cv2.line(court, *self.net, 1, self.line_width)
cv2.line(court, *self.top_inner_line, 1, self.line_width)
cv2.line(court, *self.bottom_inner_line, 1, self.line_width)
cv2.line(court, *self.left_court_line, 1, self.line_width)
cv2.line(court, *self.right_court_line, 1, self.line_width)
cv2.line(court, *self.left_inner_line, 1, self.line_width)
cv2.line(court, *self.right_inner_line, 1, self.line_width)
cv2.line(court, *self.middle_line, 1, self.line_width)
court = cv2.dilate(court, np.ones((5, 5), dtype=np.uint8))
# plt.imsave('court_configurations/court_reference.png', court, cmap='gray')
self.court = court
return court
def get_important_lines(self):
"""
Returns all lines of the court
"""
lines = [*self.baseline_top, *self.baseline_bottom, *self.net, *self.left_court_line, *self.right_court_line,
*self.left_inner_line, *self.right_inner_line, *self.middle_line,
*self.top_inner_line, *self.bottom_inner_line]
return lines
def get_extra_parts(self):
parts = [self.top_extra_part, self.bottom_extra_part]
return parts
def save_all_court_configurations(self):
"""
Create all configurations of 4 points on court reference
"""
for i, conf in self.court_conf.items():
c = cv2.cvtColor(255 - self.court, cv2.COLOR_GRAY2BGR)
for p in conf:
c = cv2.circle(c, p, 15, (0, 0, 255), 30)
cv2.imwrite(f'court_configurations/court_conf_{i}.png', c)
def get_court_mask(self, mask_type=0):
"""
Get mask of the court
"""
mask = np.ones_like(self.court)
if mask_type == 1: # Bottom half court
# mask[:self.net[0][1] - 1000, :] = 0
mask[:self.net[0][1], :] = 0
elif mask_type == 2: # Top half court
mask[self.net[0][1]:, :] = 0
elif mask_type == 3: # court without margins
mask[:self.baseline_top[0][1], :] = 0
mask[self.baseline_bottom[0][1]:, :] = 0
mask[:, :self.left_court_line[0][0]] = 0
mask[:, self.right_court_line[0][0]:] = 0
return mask
if __name__ == '__main__':
c = CourtReference()
c.build_court_reference()