/
visionwrapper.py
138 lines (110 loc) · 4.73 KB
/
visionwrapper.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
from vision.vision import Vision, Camera
import vision.tools as tools
from postprocessing.postprocessing import Postprocessing
from preprocessing.preprocessing import Preprocessing
import cv2
from copy import deepcopy
class VisionWrapper:
"""
Class that handles vision
"""
def __init__(self, pitch, color, our_side, video_port=0):
"""
Entry point for the SDP system.
Params:
[int] pitch 0 - main pitch, 1 - secondary pitch
[string] colour The colour of our teams plates
[string] our_side the side we're on - 'left' or 'right'
[int] video_port port number for the camera
Fields
pitch
camera
calibration
vision
postporcessing
color
side
preprocessing
model_positions
regular_positions
"""
assert pitch in [0, 1]
assert color in ['yellow', 'blue']
assert our_side in ['left', 'right']
self.pitch = pitch
# Set up camera for frames
self.camera = Camera(port=video_port, pitch=pitch)
self.frame = self.camera.get_frame()
center_point = self.camera.get_adjusted_center(self.frame)
# Set up vision
self.calibration = tools.get_colors(pitch)
self.vision = Vision(
pitch=pitch, color=color, our_side=our_side,
frame_shape=self.frame.shape, frame_center=center_point,
calibration=self.calibration)
# Set up preprocessing and postprocessing
self.postprocessing = Postprocessing()
self.preprocessing = Preprocessing()
self.color = color
self.side = our_side
self.frameQueue = []
def update(self):
"""
Gets this frame's positions from the vision system.
"""
self.frame = self.camera.get_frame()
# Apply preprocessing methods toggled in the UI
self.preprocessed = self.preprocessing.run(self.frame, self.preprocessing.options)
self.frame = self.preprocessed['frame']
if 'background_sub' in self.preprocessed:
cv2.imshow('bg sub', self.preprocessed['background_sub'])
# Find object positions
# model_positions have their y coordinate inverted
self.model_positions, self.regular_positions = self.vision.locate(self.frame)
self.model_positions = self.postprocessing.analyze(self.model_positions)
#self.model_positions = self.averagePositions(3, self.model_positions)
def averagePositions(self, frames, positions_in):
"""
:param frames: number of frames to average
:param positions_in: positions for the current frame
:return: averaged positions
"""
validFrames = self.frameQueue.__len__() + 1
positions_out = deepcopy(positions_in)
# Check that the incoming positions have legal values
for obj in positions_out.items():
if (positions_out[obj[0]].velocity is None):
positions_out[obj[0]].velocity = 0
if positions_out[obj[0]].x is None:
positions_out[obj[0]].x = 0
if positions_out[obj[0]].y is None:
positions_out[obj[0]].y = 0
positions_out[obj[0]].angle = positions_in[obj[0]].angle
# Loop over queue
for positions in self.frameQueue:
# Loop over each object in the position dictionary
isFrameValid = True
for obj in positions.items():
# Check if the current object's positions have legal values
if (obj[1].x is None) or (obj[1].y is None) or (obj[1].angle is None) or (obj[1].velocity is None):
isFrameValid = False
else:
positions_out[obj[0]].x += obj[1].x
positions_out[obj[0]].y += obj[1].y
positions_out[obj[0]].velocity += obj[1].velocity
if not isFrameValid and validFrames > 1:
#validFrames -= 1
pass
# Loop over each object in the position dictionary and average the values
for obj in positions_out.items():
positions_out[obj[0]].velocity /= validFrames
positions_out[obj[0]].x /= validFrames
positions_out[obj[0]].y /= validFrames
# If frameQueue is already full then pop the top entry off
if self.frameQueue.__len__() >= frames:
self.frameQueue.pop(0)
# Add our new positions to the end
self.frameQueue.append(positions_in)
return positions_out
def saveCalibrations(self):
tools.save_colors(self.vision.pitch, self.calibration)