DroneGPT/tello_wrapper.py at main · ubermensch19/DroneGPT · GitHub

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

import sys
import threading
import time


from djitellopy import Tello
import math
import numpy as np
import io
import cv2
import random
from PIL import Image


BOX_TRESHOLD = 0.25
TEXT_TRESHOLD = 0.25


class TelloWrapper:
    def __init__(self):
        self.client = Tello()
        self.client.connect()


        # Start video stream
        self.client.streamon()  # Enable video transmission
        t = threading.Thread(target=self.get_stream)
        t.setDaemon(True)
        t.start()

        self.head_img = None # Tello front camera


    def get_stream(self):
        while True:
            self.head_img = self.client.get_frame_read().frame
            time.sleep(0.01)

    def takeoff(self):
        self.client.takeoff()

    def land(self):
        self.client.land()


    def turn_left(self,degree=10):
        """
        Turn left by degree degrees
        :return:
        """
        degree = 10*degree # tello base is 0.1 degrees
        self.client.rotate_counter_clockwise(degree)


    def turn_right(self,degree=10):
        """
        Turn right by degree degrees
        :return:
        """
        degree = 10*degree # tello base is 0.1 degrees
        self.client.rotate_clockwise(degree)


    def forward(self, distance):
        """
        Move forward, too small movements won't work
        distance: distance in meters
        :return:
        """
        distance = distance*100 # tello base is 1cm
        self.client.move_forward(distance) # Move forward


    def back(self, distance):
        """
        Move backward, too small movements won't work
        distance: distance in meters
        :return:
        """
        distance = distance*100 # tello base is 1cm
        self.client.move_back(distance) # Move backward


    def up(self, distance):
        """
        Move up, too small movements won't work
        distance: distance in meters
        :return:
        """
        distance = distance*100 # tello base is 1cm
        self.client.move_up(distance) # Move up


    def down(self, distance):
        """
        Move down, too small movements won't work
        distance: distance in meters
        :return:
        """
        distance = distance*100 # tello base is 1cm
        self.client.move_down(distance) # Move down


    def get_image(self):
        """
        Get front camera rendered image
        :return:
        """
        return self.head_img

    def get_drone_state(self):
        """
        Get drone state
        :return:{'pitch': int, 'roll': int, 'yaw': int}
        """
        return self.client.query_attitude()


    def ob_objects(self,obj_name_list):
        """
        Note: need to execute get_image first,
        Run object detection model on image img, get target list [<object_name, distance, angle_in_degrees>,...]
        :return: object name list, object info list, bbox image
        """
        pass


    def ob_objects_llm(self,obj_name_list):
        """
        Note: need to execute get_image first, provide observation results for llm
        Run object detection model on image img, get target list [<object_name, distance, angle_in_degrees>,...] for llm reasoning
        :return: [<object_name, distance, angle_in_degrees>,...] e.g. [(door, 0.53, 22), (chair, 4.84, -21)]
        """
        # Get recognition results
        ob_list, final_obj_list, annotated_frame = self.ob_objects(obj_name_list)

        final_result = []

        for obj_info in final_obj_list:
            item = (obj_info[0], obj_info[1], obj_info[3]) #obj_name, camera_distance, angel_degree
            final_result.append(item)

        return final_result