forked from theos-ai/easy-yolov7
-
Notifications
You must be signed in to change notification settings - Fork 0
/
drone_manual_control.py
222 lines (185 loc) · 7.68 KB
/
drone_manual_control.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
from djitellopy import Tello
from algorithm.object_detector import YOLOv7
from utils.detections import draw
import cv2
import pygame
import numpy as np
import time
import logging
import json
# Yolo settings
WEIGHTS = 'coco.weights'
CLASSES = 'coco.yaml'
DEVICE = 'cpu'
# Speed of the drone
S = 60
# Frames per second of the pygame window display
# A low number also results in input lag, as input information is processed once per frame.
FPS = 120
# Tello video url
URL = "udp://0.0.0.0:11111"
class FrontEnd(object):
""" Maintains the Tello display and moves it through the keyboard keys.
Press escape key to quit.
The controls are:
- T: Takeoff
- L: Land
- Arrow keys: Forward, backward, left and right.
- A and D: Counter clockwise and clockwise rotations (yaw)
- W and S: Up and down.
- P: Turn mission pad detection on/off
- 0/1/2: Set mission pad detection direction
- Y: Turn yolo detection on/off
"""
def __init__(self):
# Init pygame
pygame.init()
# Creat pygame window
pygame.display.set_caption(URL)
self.screen = pygame.display.set_mode([960, 720])
# Init yolo
self.yolov7 = YOLOv7()
self.yolov7.load(WEIGHTS, classes=CLASSES, device=DEVICE)
# Init Tello object that interacts with the Tello drone
self.tello = Tello()
# Drone velocities between -100~100
self.for_back_velocity = 0
self.left_right_velocity = 0
self.up_down_velocity = 0
self.yaw_velocity = 0
self.speed = 10
self.send_rc_control = False
self.pad_detection = False
self.pad_direction = 0
self.yolo_detection = False
# create update timer
pygame.time.set_timer(pygame.USEREVENT + 1, 1000 // FPS)
Tello.LOGGER.setLevel(logging.WARNING)
def run(self):
self.tello.connect()
self.tello.set_speed(self.speed)
# In case streaming is on. This happens when we quit this program without the escape key.
self.tello.streamoff()
self.tello.streamon()
stream = cv2.VideoCapture(URL)
if stream.isOpened() == False:
print(f'[!] error opening {URL}')
self.screen.fill([0, 0, 0])
should_stop = False
counter = 0
detections = []
while not should_stop:
for event in pygame.event.get():
if event.type == pygame.USEREVENT + 1:
self.update()
elif event.type == pygame.QUIT:
should_stop = True
elif event.type == pygame.KEYDOWN:
if event.key == pygame.K_ESCAPE:
should_stop = True
else:
self.keydown(event.key)
elif event.type == pygame.KEYUP:
self.keyup(event.key)
ret, frame = stream.read()
if ret == True:
if(self.yolo_detection):
# run detection every tenth frame
counter += 1
if(counter >= 10):
counter = 0
detections = self.yolov7.detect(frame)
if len(detections) != 0:
print(f'\n{URL}:\n', json.dumps(detections, indent=4))
frame = draw(frame, detections)
# text overlays
text_battery = "Battery: {}%".format(self.tello.get_battery())
#text_pad_detection = f"Pad Detection: {self.pad_detection} {self.pad_direction}"
text_yolo_detection = f"Yolo Detection: {self.yolo_detection}"
cv2.putText(frame, text_battery, (5, 720 - 5), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
#cv2.putText(frame, text_pad_detection, (5, 25), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
cv2.putText(frame, text_yolo_detection, (5, 25), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
frame = np.rot90(frame)
frame = np.flipud(frame)
frame = pygame.surfarray.make_surface(frame)
self.screen.blit(frame, (0, 0))
pygame.display.update()
time.sleep(1 / FPS)
# Call it always before finishing. To deallocate resources.
self.tello.end()
self.yolov7.unload()
def keydown(self, key):
""" Update velocities based on key pressed
Arguments:
key: pygame key
"""
if key == pygame.K_UP: # set forward velocity
self.for_back_velocity = S
elif key == pygame.K_DOWN: # set backward velocity
self.for_back_velocity = -S
elif key == pygame.K_LEFT: # set left velocity
self.left_right_velocity = -S
elif key == pygame.K_RIGHT: # set right velocity
self.left_right_velocity = S
elif key == pygame.K_w: # set up velocity
self.up_down_velocity = S
elif key == pygame.K_s: # set down velocity
self.up_down_velocity = -S
elif key == pygame.K_a: # set yaw counter clockwise velocity
self.yaw_velocity = -S
elif key == pygame.K_d: # set yaw clockwise velocity
self.yaw_velocity = S
def keyup(self, key):
""" Update velocities based on key released
Arguments:
key: pygame key
"""
if key == pygame.K_UP or key == pygame.K_DOWN: # set zero forward/backward velocity
self.for_back_velocity = 0
elif key == pygame.K_LEFT or key == pygame.K_RIGHT: # set zero left/right velocity
self.left_right_velocity = 0
elif key == pygame.K_w or key == pygame.K_s: # set zero up/down velocity
self.up_down_velocity = 0
elif key == pygame.K_a or key == pygame.K_d: # set zero yaw velocity
self.yaw_velocity = 0
elif key == pygame.K_t: # takeoff
self.tello.takeoff()
self.send_rc_control = True
elif key == pygame.K_l: # land
not self.tello.land()
self.send_rc_control = False
elif key == pygame.K_p: # turn mission pad detection on/off
self.pad_detection = not self.pad_detection
if self.pad_detection:
self.tello.enable_mission_pads()
self.tello.set_mission_pad_detection_direction(self.pad_direction)
else:
self.tello.disable_mission_pads()
elif key == pygame.K_0: # mission pad detection direction downward
self.pad_direction = 0
elif key == pygame.K_1: # forward
self.pad_direction = 1
elif key == pygame.K_2: # both
self.pad_direction = 2
elif key == pygame.K_y: # turn yolo detection on/off
self.yolo_detection = not self.yolo_detection
def detect(self, frame):
detections = self.yolov7.detect(frame)
#detected_frame = draw(frame, detections)
return detections
def update(self):
""" Update routine. Send velocities to Tello.
"""
if self.send_rc_control:
self.tello.send_rc_control(self.left_right_velocity, self.for_back_velocity,
self.up_down_velocity, self.yaw_velocity)
if self.pad_detection:
pad = self.tello.get_mission_pad_id()
if(pad != -1 & pad != -2):
print(f"ID: {pad}")
def main():
frontend = FrontEnd()
frontend.run()
if __name__ == '__main__':
main()