rad.vg
this repo has no description
1from djitellopy import Tello
2import cv2
3import pygame
4import numpy as np
5import time
6
7MOTION = 50
8FPS = 120
9
10class TelloController():
11 def __init__(self):
12 pygame.init()
13 pygame.display.set_caption("Tello controller")
14 self.screen = pygame.display.set_mode([960, 720])
15 self.has_controller = pygame.joystick.get_count() > 0
16 if self.has_controller:
17 self.joystick = pygame.joystick.Joystick(0)
18 self.tello = Tello()
19
20 self.for_back_velocity = 0
21 self.left_right_velocity = 0
22 self.up_down_velocity = 0
23 self.yaw_velocity = 0
24
25 self.send_rc_control = False
26
27 pygame.time.set_timer(pygame.USEREVENT + 1, 1000 // FPS)
28
29 def run(self):
30 self.tello.connect()
31
32 # Reset stream in case it was left on by accident
33 self.tello.streamoff()
34 self.tello.streamon()
35
36 frame_read = self.tello.get_frame_read()
37
38 should_stop = False
39 while not should_stop:
40 for event in pygame.event.get():
41 if event.type == pygame.USEREVENT + 1:
42 self.update()
43 elif event.type == pygame.QUIT:
44 should_stop = True
45 elif event.type == pygame.KEYDOWN:
46 if event.key == pygame.K_ESCAPE:
47 should_stop = True
48 else:
49 self.keydown(event.key)
50 elif event.type == pygame.KEYUP:
51 self.keyup(event.key)
52 elif event.type == pygame.JOYAXISMOTION and self.has_controller:
53 self.stickmotion(event)
54 elif event.type == pygame.JOYBUTTONDOWN and self.has_controller:
55 self.buttonup(event.button)
56
57 if frame_read.stopped:
58 break
59
60 self.screen.fill([0, 0, 0])
61
62 frame = frame_read.frame
63
64 # Display battery status
65 text = "{}%".format(self.tello.get_battery())
66 cv2.putText(frame, text, (5, 720 - 5),
67 cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
68 frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
69 frame = np.rot90(frame)
70 frame = np.flipud(frame)
71
72 frame = pygame.surfarray.make_surface(frame)
73 self.screen.blit(frame, (0, 0))
74 pygame.display.update()
75
76 time.sleep(1 / FPS)
77
78 self.tello.end()
79
80 def keydown(self, key):
81 if key == pygame.K_UP: # Forward
82 self.for_back_velocity = MOTION
83 elif key == pygame.K_DOWN: # Backward
84 self.for_back_velocity = -MOTION
85 elif key == pygame.K_LEFT: # Left
86 self.left_right_velocity = -MOTION
87 elif key == pygame.K_RIGHT: # Right
88 self.left_right_velocity = MOTION
89 elif key == pygame.K_w: # Up
90 self.up_down_velocity = MOTION
91 elif key == pygame.K_s: # Down
92 self.up_down_velocity = -MOTION
93 elif key == pygame.K_a: # Rotate counter clockwise
94 self.yaw_velocity = -MOTION
95 elif key == pygame.K_d: # Rotate clockwise
96 self.yaw_velocity = MOTION
97
98 def keyup(self, key):
99 if key == pygame.K_UP or key == pygame.K_DOWN: # Reset forward/backward
100 self.for_back_velocity = 0
101 elif key == pygame.K_LEFT or key == pygame.K_RIGHT: # Reset left/right
102 self.left_right_velocity = 0
103 elif key == pygame.K_w or key == pygame.K_s: # Reset up/down
104 self.up_down_velocity = 0
105 elif key == pygame.K_a or key == pygame.K_d: # Reset yaw
106 self.yaw_velocity = 0
107 elif key == pygame.K_t: # Take off
108 self.tello.takeoff()
109 self.send_rc_control = True
110 elif key == pygame.K_l: # Land
111 not self.tello.land()
112 self.send_rc_control = False
113 elif key == pygame.K_h:
114 self.tello.flip_forward()
115 elif key == pygame.K_j:
116 self.tello.flip_back()
117
118 def stickmotion(self, event):
119 if event.axis == 0: # left stick left/right
120 self.yaw_velocity = round(self.joystick.get_axis(event.axis), 3) * SPEED
121 elif event.axis == 3: # right stick left/right
122 self.left_right_velocity = round(self.joystick.get_axis(event.axis), 3) * SPEED
123 elif event.axis == 2: # right stick up/down
124 self.for_back_velocity = -round(self.joystick.get_axis(event.axis), 3) * SPEED
125 elif event.axis == 2: # left trigger
126 self.up_down_velocity = -min(round(self.joystick.get_axis(event.axis), 3) + 1, 1) * SPEED
127 elif event.axis == 5: # right trigger
128 self.up_down_velocity = min(round(self.joystick.get_axis(event.axis), 3) + 1, 1) * SPEED
129
130 def buttonup(self, button):
131 if button == 0:
132 self.tello.land()
133 self.send_rc_control = False
134 elif button == 3:
135 self.tello.takeoff()
136 self.send_rc_control = True
137
138 def update(self):
139 if self.send_rc_control:
140 self.tello.send_rc_control(
141 # Whilst this is technically less accurate, I'm not sure how well
142 # these drones handle repeating numbers like 0.333333333 and I doubt
143 # that level of precision is necessary at all
144 round(self.left_right_velocity, 3),
145 round(self.for_back_velocity, 3),
146 round(self.up_down_velocity, 3),
147 round(self.yaw_velocity, 3)
148 )
149
150if __name__ == '__main__':
151 TelloController().run()