-
Notifications
You must be signed in to change notification settings - Fork 0
/
backup_stream_processor.py
executable file
·351 lines (319 loc) · 13.7 KB
/
backup_stream_processor.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
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
import picamera
import picamera.array
import threading
import cv2
import numpy
from enum import Enum
import time
from core import I2C_Lidar
class State(Enum):
LEARNING = 1
ORIENTING = 2
HUNTING = 3
FINISHED = 4
# Image stream processing thread
class StreamProcessor(threading.Thread):
def __init__(self, core_module, camera):
self.core_module = core_module
self.state = State.LEARNING # Default state
self.reversing = False
super(StreamProcessor, self).__init__()
self.camera = camera
self.stream = picamera.array.PiRGBArray(self.camera)
self.event = threading.Event()
self.terminated = False
self.start()
self.begin = 0
self.min_distance = 80.0
self.back_off_distance = 400
# Auto drive settings
self.autoMaxPower = 1.0 # Maximum output in automatic mode
self.autoMinPower = 0.4 # Minimum output in automatic mode
self.autoMinArea = 100 # Smallest target to move towards
# full image
# self.autoMaxArea = 55000 # Largest target to move towards
# Cropped Image
self.autoMaxArea = 40000 # Largest target to move towards
# Target size at which we use the maximum allowed output
self.autoFullSpeedArea = 20000
# Colour order to visit
self.challengecolours = ['red', 'blue', 'yellow', 'green']
self.arenacolours = [] # Order we've detected in the arena, clockwise
self.colourindex = 0
self.colour = self.challengecolours[self.colourindex]
self.lookingatcolour = ''
# Camera settings
self.imageWidth = 320 # Camera image width
self.imageHeight = 240 # Camera image height
self.imageCentreX = self.imageWidth / 2.0
self.imageCentreY = self.imageHeight / 2.0
self.tickInt = 0
def run(self):
# This method runs in a separate thread
while not self.terminated and self.state != State.FINISHED:
# Wait for an image to be written to the stream
if self.event.wait(1):
try:
# Read the image and do some processing on it
self.stream.seek(0)
self.ProcessImage(self.stream.array)
finally:
# Reset the stream and event
self.stream.seek(0)
self.stream.truncate()
self.event.clear()
print('stream_processor terminated')
# Image processing function
def ProcessImage(self, image):
# View the original image seen by the camera.
# Crop the image down to just the bit with the arena in
image = image[100:240, 0:320]
debug = False
if debug:
cv2.imshow('original', image)
cv2.waitKey(0)
# Blur the image
# image = cv2.medianBlur(image, 5)
# if debug:
# cv2.imshow('blur', image)
# cv2.waitKey
# Convert the image from 'BGR' to HSV colour space
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
# if debug:
# cv2.imshow('cvtColour', image)
# cv2.waitKey(0)
if self.colour == "red":
imrange1 = cv2.inRange(
image,
numpy.array((160, 140, 210)),
numpy.array((180, 255, 255))
)
imrange2 = cv2.inRange(
image,
numpy.array((0, 140, 210)),
numpy.array((5, 255, 255))
)
imrange = cv2.bitwise_xor(imrange1, imrange2)
elif self.colour == 'yellow':
imrange = cv2.inRange(
image,
# numpy.array((15, 85, 64)),
# numpy.array((35, 255, 255))
numpy.array((20, 150, 200)),
numpy.array((35, 255, 255))
)
elif self.colour == "green":
imrange = cv2.inRange(
image,
# numpy.array((50, 96, 64)),
# numpy.array((85, 255, 255))
numpy.array((35, 70, 210)),
numpy.array((55, 255, 255))
)
elif self.colour == 'blue':
imrange = cv2.inRange(
image,
numpy.array((85, 150, 200)),
numpy.array((115, 255, 255))
)
# Blur the mask, not the image
imrange = cv2.medianBlur(imrange, 5)
if debug:
cv2.imshow('imrange', imrange)
cv2.waitKey()
# Find the contours
contourimage, contours, hierarchy = cv2.findContours(
imrange,
cv2.RETR_LIST,
cv2.CHAIN_APPROX_SIMPLE
)
# if debug:
# cv2.imshow('contour', contourimage)
# cv2.waitKey(0)
# Go through each contour
squareness = -1
x = -1
y = -1
area = 0
for (idx, contour) in enumerate(contours):
x, y, w, h = cv2.boundingRect(contour)
cx = x + (w / 2)
cy = y + (h / 2)
area = w * h
contourarea = cv2.contourArea(contour)
extent = float(contourarea) / area
aspect = float(w) / h
cont_squareness = (1.0 / aspect if aspect > 1 else aspect)
cont_squareness *= (1.0 / extent if extent > 1.0 else extent)
if (cont_squareness > squareness):
if (debug):
print("New squarest: %f" % squareness)
print(" extent " + str(extent))
print(" aspect = " + str(aspect))
print(" area = " + str(contourarea))
squareness = cont_squareness
if area > 0:
ball = [cx, cy, area]
else:
ball = None
# Set drives or report ball status
self.SetSpeedFromBall(ball)
def SetSpeedFromBall(self, ball):
""" Make decisions about what we're doing and
set the motor speed from the ball position """
# Tuning constants
backoff = -0.6 # how fast to back out of the corner
seek = 1.0 # how fast to turn when we can't see a ball
# how fast we may turn a wheel backwards when a ball is in sight
hunt_reverse = -0.2
driveLeft = 0.0
driveRight = 0.0
if ball:
x = ball[0]
# y = ball[1]
area = ball[2]
# If we're learning, just seeing a colour is enough
if self.state == State.LEARNING:
# We've seen a ball of a colour - is it what we want?
if not (self.colour in self.arenacolours):
# It's a new colour
self.arenacolours.append(self.colour)
self.lookingatcolour = self.colour
# Have we found all four colours?
if len(self.arenacolours) == 4:
print('Lets remember these for next time')
f = open('arenacolours.txt', 'w')
f.write("{0}\n{1}\n{2}\n{3}".format(
*self.arenacolours))
f.close()
print(
('I found all the colours, now '
'm looking at {0} hunting a {1}'
).format(
self.lookingatcolour,
self.challengecolours[0]
)
)
self.colourindex = 0
self.colour = self.challengecolours[0]
self.state = State.HUNTING
# time.sleep(2)
elif self.state == State.ORIENTING and self.lookingatcolour == '':
# If we can see a colour, set lookingatcolour and go hunting
print(('Im looking at a {0} ball,'
' lets hunt a {1} one').format(
self.colour,
self.challengecolours[0])
)
self.lookingatcolour = self.colour
self.colourindex = 0
self.colour = self.challengecolours[0]
self.state = State.HUNTING
# time.sleep(2)
elif self.state == State.HUNTING and self.reversing is False:
d_front = self.core_module.get_distance(I2C_Lidar.LIDAR_FRONT)
if area < self.autoMinArea:
print('Too small / far')
driveLeft = self.autoMinPower
driveRight = self.autoMinPower
elif d_front <= self.min_distance:
print('Close enough')
# Remember we're looking at the current colour
self.lookingatcolour = self.colour
self.colourindex = self.colourindex + 1
if (self.colourindex >= len(self.challengecolours)):
print('Donezo!')
self.state = State.FINISHED
self.core_module.set_neutral(braked=True)
else:
self.colour = self.challengecolours[self.colourindex]
print('Now looking for %s ball' % (self.colour))
driveLeft = backoff
driveRight = backoff
self.reversing = True
else:
if area < self.autoFullSpeedArea:
speed = 1.0
else:
speed = 1.0 / (area / self.autoFullSpeedArea)
speed *= self.autoMaxPower - self.autoMinPower
speed += self.autoMinPower
direction = (self.imageCentreX - x) / self.imageCentreX
direction = direction * 3
if direction > 0.0:
# Turn right
print('Turn right for %s' % self.colour)
driveLeft = speed
driveRight = speed * (1.0 - direction)
if driveRight < hunt_reverse:
driveRight = hunt_reverse
else:
# Turn left
print('Turn left for %s' % self.colour)
driveLeft = speed * (1.0 + direction)
driveRight = speed
if driveLeft < hunt_reverse:
driveLeft = hunt_reverse
elif self.reversing is False:
# Figure out which direction to seek from arenacolours
if (self.state == State.HUNTING and
(self.arenacolours.index(self.colour) == self.arenacolours.index(self.lookingatcolour) - 1 or
self.arenacolours.index(self.colour) == self.arenacolours.index(self.lookingatcolour) + 3)):
# colour we want is left of looking-at-colour, turn leftwards
print('No {0} ball, {0} is left of {1}, turn left'.format(
self.colour,
self.lookingatcolour)
)
driveLeft = 0 - seek
driveRight = seek
else:
print('No {0} ball, turn right'.format(self.colour))
# turn right like we normally do
driveLeft = seek
driveRight = 0 - seek
if self.reversing is True:
# Drive backwards until front distance > 600mm
d_front = self.core_module.get_distance(I2C_Lidar.LIDAR_FRONT)
print("*** Reversing {} ***".format(d_front))
if d_front > self.back_off_distance:
self.reversing = False
driveLeft = backoff
driveRight = backoff
if self.tickInt == 0:
asciiTick = "| "
elif self.tickInt == 1:
asciiTick = " | "
elif self.tickInt == 2:
asciiTick = " | "
else:
asciiTick = " |"
# DriveLeft * DriveRight if one is negative and the other
# positigve the result will be negative else positive.
if (self.tickInt % 4) == 1 or ((self.tickInt % 2) == 1 and (driveLeft * driveRight < 0)):
# Blip motors, except if reversing
if (self.reversing is False):
driveLeft = 0.01
driveRight = 0.01 # not zero as that turns on brakes
self.tickInt = self.tickInt + 1 if self.tickInt < 3 else 0
# If we're figuring out what colour we're looking at, cycle through
# colours on each tick; otherwise focus on the coloru we're hunting
if self.state == State.LEARNING or self.state == State.ORIENTING:
self.colour = self.challengecolours[self.tickInt]
print('{0} ({1}) {2:4.1f}, {3:4.1f} - {4}, x {5}, {6} > {7}'.format(
self.state,
asciiTick,
driveLeft,
driveRight,
self.arenacolours,
ball[0] if ball else 0,
self.lookingatcolour,
self.colour)
)
if self.state == State.FINISHED:
self.core_module.throttle(-1, -1)
time.sleep(0.1)
self.core_module.set_neutral(braked=True)
else:
self.core_module.throttle(driveLeft * 100, driveRight * 100)
# if (driveLeft == backoff):
# time.sleep(0.8)