Merge branch 'main' of https://github.com/raiots/USVLander

agent.py

@@ -64,7 +64,7 @@ class UAV:
 
     def run(self):
         """ 主程序 """
-        # self.init_flight()
+        self.init_flight()
         # self.tello.turn_motor_on()
 
         state_thread = Thread(target=self.video_stream)
@@ -92,6 +92,7 @@ class UAV:
         self.tello.takeoff()
         self.send_rc_control = True  # 开启发送控制信号
         self.update()
+        self.tello.go_xyz_speed(0, 0, 50, 20)  # 初始升高高度
 
     def video_stream(self):
         # 独立线程，用于获取视频流并记录状态
@@ -135,6 +136,10 @@ class UAV:
         elif self.status == Status.LANDING:
             bottom_img = self.frame_queue.get()
             self.auto_control.landing(bottom_img=bottom_img)
+
+        elif self.status == Status.FALLING:
+            # 自由落体阶段
+            pass
         time.sleep(1 / 30)
 
         

control.py

@@ -125,7 +125,9 @@ class AutoControl:
             self.loss_target_times = 0
             logging.info('tracking')
             self.target = _item.getMessage()['target']
-            print(self.target)
+            # print(self.target)
+            target_dict  = {'timestamp': time.time(), 'x': self.target[0], 'y':self.target[1]}
+            self.uav.recorder.anything_record(target_dict, 'track_result', ['timestamp', 'x', 'y'])
             
             frame = _item.getFrame()
             cv2.imshow("track result", frame)
@@ -178,7 +180,7 @@ class AutoControl:
         if tag_list:
             # 如果检测到标签
             self.get_tag_times += 1
-            if self.get_tag_times == 5:
+            if self.get_tag_times >= 5:
                 # 检测到5次后才认为大部分tag已经进入视野
                 self.is_get_tag = True
 
@@ -193,6 +195,7 @@ class AutoControl:
                 # print(tag_list[0].pose_t)  # 后正左正高度
                 if tag_list[0].pose_t[2] < 0.3:
                     self.uav.tello.land()
+                    self.uav.status = Status.FALLING
         else:
             self.uav.up_down_velocity = 0
             self.uav.left_right_velocity = 0

utils.py

@@ -12,6 +12,7 @@ class Status(Enum):
     DETECT = 1
     TRACK = 2
     LANDING = 3
+    FALLING =  4
 
 class DataRecorder:
     def __init__(self, exp_note='default', need_record=False):
@@ -23,9 +24,9 @@ class DataRecorder:
             self.video_folder = self.root_folder / 'video'
             self.video_folder.mkdir(parents=True, exist_ok=False)
 
-            # with open(str(self.root_folder) + '/sensor_data.csv', mode='a', newline='') as file:
-            #     self.writer = csv.DictWriter(file, fieldnames=['timestamp', 'mid', 'x', 'y', 'z', 'mpry', 'pitch', 'roll', 'yaw', 'vgx', 'vgy', 'vgz', 'templ', 'temph', 'tof', 'h', 'bat', 'baro', 'time', 'agx', 'agy', 'agz'])
-            #     self.writer.writeheader()
+            with open(str(self.root_folder) + '/sensor_data.csv', mode='a', newline='') as file:
+                writer = csv.DictWriter(file, fieldnames=['timestamp', 'mid', 'x', 'y', 'z', 'mpry', 'pitch', 'roll', 'yaw', 'vgx', 'vgy', 'vgz', 'templ', 'temph', 'tof', 'h', 'bat', 'baro', 'time', 'agx', 'agy', 'agz', 'status', 'cmd_vel'])
+                writer.writeheader()
 
         else:
             logging.warn("实验数据将不被记录！")
@@ -48,14 +49,14 @@ class DataRecorder:
         if self.need_record:
             with open(str(self.root_folder) + '/sensor_data.csv', mode='a', newline='') as file:
                 writer = csv.DictWriter(file, fieldnames=['timestamp', 'mid', 'x', 'y', 'z', 'mpry', 'pitch', 'roll', 'yaw', 'vgx', 'vgy', 'vgz', 'templ', 'temph', 'tof', 'h', 'bat', 'baro', 'time', 'agx', 'agy', 'agz', 'status', 'cmd_vel'])
-                writer.writeheader()
+                # writer.writeheader()
                 writer.writerow(states)
 
     def anything_record(self, data, data_name:str, header:list):
         if self.need_record:
             with open(str(self.root_folder) + '/' + data_name + '.csv', mode='a', newline='') as file:
                 writer = csv.DictWriter(file, fieldnames=header)
-                writer.writeheader()
+                # writer.writeheader()
                 writer.writerow(data)