From 023793651db5ddc6585b01bef7a9b1662127f072 Mon Sep 17 00:00:00 2001
From: raiot <raiot.lee@hotmail.com>
Date: Sun, 23 Jun 2024 14:57:02 +0800
Subject: [PATCH] feat: first demo
---
.gitignore | 197 ++++++++++++++++++++++++++++++++++++++++++++++++++++
agent.py | 52 +++++++++++---
control.py | 132 +++++++++++++++++++++++++++++------
picture.png | Bin 941 -> 0 bytes
test.py | 14 ++++
utils.py | 67 ++++++++++++++++++
video.avi | Bin 7172 -> 0 bytes
vision.py | 4 +-
8 files changed, 435 insertions(+), 31 deletions(-)
create mode 100644 .gitignore
delete mode 100644 picture.png
create mode 100644 test.py
create mode 100644 utils.py
delete mode 100644 video.avi
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..1c430c9
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,197 @@
+# Created by https://www.toptal.com/developers/gitignore/api/python,visualstudiocode
+# Edit at https://www.toptal.com/developers/gitignore?templates=python,visualstudiocode
+
+### Python ###
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+# For a library or package, you might want to ignore these files since the code is
+# intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+# However, in case of collaboration, if having platform-specific dependencies or dependencies
+# having no cross-platform support, pipenv may install dependencies that don't work, or not
+# install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+# This is especially recommended for binary packages to ensure reproducibility, and is more
+# commonly ignored for libraries.
+# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+# in version control.
+# https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+# and can be added to the global gitignore or merged into this file. For a more nuclear
+# option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+### Python Patch ###
+# Poetry local configuration file - https://python-poetry.org/docs/configuration/#local-configuration
+poetry.toml
+
+# ruff
+.ruff_cache/
+
+# LSP config files
+pyrightconfig.json
+
+### VisualStudioCode ###
+.vscode/*
+!.vscode/settings.json
+!.vscode/tasks.json
+!.vscode/launch.json
+!.vscode/extensions.json
+!.vscode/*.code-snippets
+
+# Local History for Visual Studio Code
+.history/
+
+# Built Visual Studio Code Extensions
+*.vsix
+
+### VisualStudioCode Patch ###
+# Ignore all local history of files
+.history
+.ionide
+
+data/
+
+# End of https://www.toptal.com/developers/gitignore/api/python,visualstudiocode
\ No newline at end of file
diff --git a/agent.py b/agent.py
index 18eba90..e0de1b2 100644
--- a/agent.py
+++ b/agent.py
@@ -6,9 +6,9 @@ import sys
import time
import logging
import cv2
-
+import numpy as np
from control import ManualControl, AutoControl
-from control import Status
+from utils import Status, DataRecorder
@@ -20,6 +20,11 @@ class UAV:
self.tello.connect()
print('Current_battery: ' + str(self.tello.get_battery()))
self.tello.streamon()
+ self.tello.set_video_resolution(Tello.RESOLUTION_480P)
+ self.tello.set_video_fps(Tello.FPS_30)
+ self.tello.set_video_bitrate(Tello.BITRATE_1MBPS)
+ self.tello.set_video_direction(Tello.CAMERA_FORWARD)
+
self.frame_read = self.tello.get_frame_read()
self.left_right_velocity = 0
@@ -33,44 +38,60 @@ class UAV:
self.status = Status.DETECT
self.frame_queue = Queue(maxsize=1) # 用于存放视频帧, 1用于避免帧堆积
+ record_name = input("请输入数据备注（英文）:")
+ if record_name:
+ self.recorder = DataRecorder(exp_note=record_name, need_record=True)
+ else:
+ self.recorder = DataRecorder()
+
+ self.front_cam_shape = (648, 478) # x, y
def update(self):
""" Update routine. Send velocities to Tello.
向Tello发送各方向速度信息
"""
+ if self.yaw_velocity > 50:
+ self.yaw_velocity = 50
+ elif self.yaw_velocity < -50:
+ self.yaw_velocity = -50
+
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)
+ print(self.left_right_velocity, self.for_back_velocity,
+ self.up_down_velocity, self.yaw_velocity)
def run(self):
""" 主程序 """
# self.init_flight()
+ # self.tello.turn_motor_on()
state_thread = Thread(target=self.video_stream)
state_thread.start()
- # self.send_rc_control = True
- # self.manual_control.start() # 启动手动控制键盘监听
+ self.manual_control.start() # 启动手动控制键盘监听
while True:
# self.up_down_velocity(10)
- # self.send_rc_control = True
- self.auto_control()
+ self.auto_track_land()
self.update()
- self.video_display()
+ # self.video_display()
time.sleep(0.05)
def signal_handler(self, signal, frame):
print ('\nSignal Catched! 执行急停！')
self.tello.emergency()
+ self.tello.streamoff()
sys.exit(0)
def init_flight(self):
self.tello.turn_motor_on()
- time.sleep(10)
+ time.sleep(5)
logging.warn('自检完毕，可以起飞')
self.tello.takeoff()
+ self.send_rc_control = True # 开启发送控制信号
+ self.update()
def video_stream(self):
height, width, _ = self.frame_read.frame.shape
@@ -78,6 +99,11 @@ class UAV:
frame = self.frame_read.frame
self.frame_queue.put(frame)
+ states = self.tello.get_current_state()
+ states['timestamp'] = time.time()
+ self.recorder.state_record(states)
+ self.recorder.frame_record(frame, 'origin')
+
time.sleep(1 / 30)
def video_display(self):
@@ -88,15 +114,23 @@ class UAV:
if cv2.waitKey(1) & 0xFF == ord('q'):
self.keep_recording = False
cv2.destroyAllWindows()
+ if cv2.waitKey(1) & 0xFF == ord('s'):
+ cv2.imwrite(str(time.time()) + '.jpg', frame)
def auto_track_land(self):
+ # print(self.status)
- if self.status == Status.INIT or self.status == Status.TRACKING:
+ if self.status == Status.TRACK or self.status == Status.DETECT:
front_img = self.frame_queue.get()
+ front_img = cv2.cvtColor(front_img, cv2.COLOR_RGB2BGR)
if self.status == Status.DETECT:
self.auto_control.detect(front_img=front_img)
elif self.status == Status.TRACK:
self.auto_control.track(front_img=front_img)
+
+ elif self.status == Status.LANDING:
+ bottom_img = self.frame_queue.get()
+ self.auto_control.landing(bottom_img=bottom_img)
diff --git a/control.py b/control.py
index c5a8161..4e3e96a 100644
--- a/control.py
+++ b/control.py
@@ -1,20 +1,18 @@
from pynput import keyboard
import time
-
+import logging
from vision import Tracker
import cv2
-from enum import Enum
+from utils import Status
+import pupil_apriltags as apriltag
+from djitellopy import Tello
TICKS_PER_SEC = 30
-class Status(Enum):
- INIT = 0
- DETECT = 1
- TRACK = 2
- LANDING = 3
+
class PIDController:
def __init__(self, kp, ki, kd, current_time):
@@ -25,7 +23,7 @@ class PIDController:
self.integral = 0
self.last_time = current_time
- def control(self, error, current_time):
+ def control(self, error, current_time) -> int:
current_time = current_time
delta_time = current_time - self.last_time
if delta_time == 0:
@@ -38,7 +36,7 @@ class PIDController:
self.previous_error = error
self.last_time = current_time
- return output
+ return int(output)
class ManualControl:
@@ -63,10 +61,12 @@ class ManualControl:
self.uav.up_down_velocity = 50
elif key.char == 's':
self.uav.up_down_velocity = -50
- elif key.char == 'q':
+ elif key.char == 'z':
self.uav.yaw_velocity = -50
- elif key.char == 'e':
+ elif key.char == 'c':
self.uav.yaw_velocity = 50
+ elif key.char == 'l':
+ self.uav.tello.land()
# print(f'left_right_velocity: {self.uav.left_right_velocity}, for_back_velocity: {self.uav.for_back_velocity}, up_down_velocity: {self.uav.up_down_velocity}, yaw_velocity: {self.uav.yaw_velocity}')
except AttributeError:
pass
@@ -79,7 +79,7 @@ class ManualControl:
self.uav.left_right_velocity = 0
elif key.char == 'w' or key.char == 's':
self.uav.up_down_velocity = 0
- elif key.char == 'q' or key.char == 'e':
+ elif key.char == 'z' or key.char == 'c':
self.uav.yaw_velocity = 0
except AttributeError:
pass
@@ -89,27 +89,117 @@ class AutoControl:
def __init__(self, uav):
self.uav = uav
self.target = [0, 0, 0, 0]
- self.pid = PIDController(1, 0, 0, time.time())
+ self.tag_detector = apriltag.Detector(families='tag36h11')
+
+ self.yaw_controller = PIDController(1, 0, 0.1, time.time())
+ self.for_back_controler = PIDController(100, 0, 0.1, time.time())
+ self.left_right_controler = PIDController(50, 0, 0.1, time.time()) # 摄像头左右视野较小
+ self.loss_target_times = 0
+ self.is_get_tag = False
def detect(self, front_img):
- cv2.imshow("ROI select", front_img[:, :, 0:3])
- self.gROI = cv2.selectROI("ROI select", front_img[:, :, 0:3], False)
+ cv2.imshow("ROI select", front_img)
+
+ # 在 selectROI 前等待按键（可选）
+ key = cv2.waitKey(10) # 等待按键
+ if key == ord('q'): # 如果按下 'q' 键
+ print("Quit selected before ROI.")
+ self.uav.tello.land()
+ cv2.destroyAllWindows()
+ exit()
+
+ self.gROI = cv2.selectROI("ROI select", front_img, False)
if (not self.gROI):
print("空框选，退出")
- quit()
+ self.uav.tello.land()
self.gTracker = Tracker(tracker_type="KCF")
- self.gTracker.initWorking(front_img[:, :, 0:3], self.gROI)
- print("start tracking")
+ self.gTracker.initWorking(front_img, self.gROI)
+ logging.info("start tracking")
self.uav.status = Status.TRACK
cv2.destroyWindow("ROI select")
- def track(self, front_img):
+ def track(self, front_img):
_item = self.gTracker.track(front_img)
- if _item.message:
+ if _item.getMessage():
+ self.loss_target_times = 0
+ print("tracking")
self.target = _item.getMessage()['target']
+ frame = _item.getFrame()
+ cv2.imshow("track result", frame)
+ self.uav.recorder.frame_record(frame, 'KCF_tracking')
+
+ # store the frame to video
+
+ _key = cv2.waitKey(1) & 0xFF
+ if (_key == ord('q')) | (_key == 27):
+ self.uav.status = Status.LANDING
+
+ if (_key == ord('r')) :
+ self.uav.status = Status.DETECT
+ # print(_item.getFrame())
+ # print(self.target)
+ yaw_vel = self.yaw_controller.control(self.target[0] - self.uav.front_cam_shape[0]/2, time.time())
+ self.uav.yaw_velocity = int(yaw_vel)
+ self.uav.for_back_velocity = 20
else:
- return None
+ self.loss_target_times += 1
+ if self.loss_target_times == 6:
+ # 等待6帧（0.2秒），仍丢失则重新检测并重置计数器
+ self.loss_target_times = 0 # 在函数头部还有一个
+ self.uav.status = Status.DETECT
+ self.uav.yaw_velocity = 0
+ self.uav.for_back_velocity = 0
+
+ if self.target[1] > self.uav.front_cam_shape[1] * 0.9:
+ print(self.target[1], self.uav.front_cam_shape[0]/10)
+ logging.warn("ready to land")
+ self.uav.status = Status.LANDING
+ self.uav.tello.set_video_direction(Tello.CAMERA_DOWNWARD)
+
+ def landing(self, bottom_img):
+ tag_large = self.tag_detector.detect(bottom_img[:, :, 0], estimate_tag_pose=True, camera_params=([353.836278849487, 353.077028029136, 163.745870517989, 115.130883974855]), tag_size=0.169)
+ tag_small = self.tag_detector.detect(bottom_img[:, :, 0], estimate_tag_pose=True, camera_params=([353.836278849487, 353.077028029136, 163.745870517989, 115.130883974855]), tag_size=0.018)
+ tag_list = []
+ for tag in tag_large:
+ if tag.tag_id == 1:
+ tag_list.append(tag)
+ for tag in tag_small:
+ if tag.tag_id == 0:
+ tag_list.append(tag)
+
+ if tag_list:
+ self.is_get_tag = True
+ for_vel = -1 * self.for_back_controler.control(tag_list[0].pose_t[0], time.time())
+ self.uav.for_back_velocity = int(for_vel)
+
+ left_vel = -1 * self.left_right_controler.control(tag_list[0].pose_t[1], time.time())
+ self.uav.left_right_velocity = int(left_vel)
+ # print(self.uav.for_back_velocity, self.uav.left_right_velocity)
+ if tag_list[0].pose_t[0] < 0.2 and tag_list[0].pose_t[1] < 0.2:
+ self.uav.up_down_velocity = -20
+ # print(tag_list[0].pose_t) # 后正左正高度
+ if tag_list[0].pose_t[2] < 0.3:
+ self.uav.tello.land()
+ else:
+ self.uav.up_down_velocity = 0
+ self.uav.left_right_velocity = 0
+ if self.is_get_tag:
+ self.uav.for_back_velocity = 0
+ logging.warn("no tag detected")
+
+ cv2.imshow('aa', bottom_img[:, :, 0])
+ cv2.waitKey(1)
+
+
+ # if tags:
+ # self.uav.status = Status.LANDED
+ # self.uav.tello.land()
+ # logging.warn('landed')
+ # else:
+ # self.uav.up_down_velocity = -20
+
+
class FakeUAV:
diff --git a/picture.png b/picture.png
deleted file mode 100644
index 4404e7dabb5bd7ed4734d6a3b10f91ba1051ac16..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001
literal 941
zcmeAS@N?(olHy`uVBq!ia0y~yV4MKNIvh+uk)*3dQyCbTOFUg1Ln;{G9%SSM3NkJD
zRliHN@fO3J3<f_F26d?hVcr9sY!6(R3l=e2TpSf14U^FXL2+(imb7Hlwf#_R3e0f~
Mp00i_>zopr0RL(is{jB1
diff --git a/test.py b/test.py
new file mode 100644
index 0000000..266df51
--- /dev/null
+++ b/test.py
@@ -0,0 +1,14 @@
+import pupil_apriltags as apriltag
+import cv2
+
+tag_detector = apriltag.Detector(families='tag36h11')
+
+img = cv2.imread('Snipaste_2024-06-22_21-30-43.png')
+# cv2.imshow('aaa', img)
+# print(img.shape)
+tag = tag_detector.detect(img[:,:,0], estimate_tag_pose=True, camera_params=([353.836278849487, 353.077028029136, 163.745870517989, 115.130883974855]), tag_size=0.018)
+
+for t in tag:
+ if t.tag_id == 0:
+ print(t.pose_t)
+ print(t.tag_id)
\ No newline at end of file
diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000..f454f4e
--- /dev/null
+++ b/utils.py
@@ -0,0 +1,67 @@
+from enum import Enum
+from pathlib import Path
+from datetime import datetime
+import time
+import cv2
+import logging
+import csv
+
+
+class Status(Enum):
+ # INIT = 0
+ DETECT = 1
+ TRACK = 2
+ LANDING = 3
+
+class DataRecorder:
+ def __init__(self, exp_note='default', need_record=False):
+ # use path lib to manage file path
+ self.need_record = need_record
+ if need_record:
+ self.root_folder = Path('data') / str(datetime.now().strftime('%Y-%m-%d_%H-%M-%S') + '_' + exp_note)
+ self.root_folder.mkdir(parents=True, exist_ok=False)
+ 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()
+
+ else:
+ logging.warn("实验数据将不被记录！")
+
+ def frame_record(self, frame, cam_name:str):
+ if self.need_record:
+ cam_folder = self.video_folder / cam_name # 分离不同来源的视频
+ if not cam_folder.exists():
+ # 如果文件夹不存在则创建
+ print('create folder:', str(cam_folder))
+ cam_folder.mkdir(parents=True, exist_ok=False)
+
+ frame_name = cam_folder / str(str(time.time()) + '.jpg')
+ if cv2.imwrite(str(frame_name), frame):
+ pass
+ else:
+ print('write frame failed')
+
+ def state_record(self, states):
+ 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'])
+ writer.writerow(states)
+
+
+
+if __name__ == '__main__':
+ recorder = DataRecorder('aa')
+ # print(recorder.root_folder)
+ camera = cv2.VideoCapture(0)
+ while True:
+ ret, frame = camera.read()
+ if ret:
+ print(frame.shape)
+ cv2.imshow('frame', frame)
+
+ recorder.frame_record(frame, 'front')
+ cv2.waitKey(1)
+ time.sleep(0.1)
\ No newline at end of file
diff --git a/video.avi b/video.avi
deleted file mode 100644
index 0ada87377c03c0f5b2d91a24a9cb33c8e4eaee0d..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001
literal 7172
zcmWIYbaVS7&A{Lo=BeQ0865IPkbxm1r6?z{EHlG`fq}uaiGg9kZeAdZfkA+ifq@Ck
zVqjo|h%ih5GISsus0@SxnKea#fuXphC<lnaW|d{86i0-4x*#hAX_bQtfXt3$WnlOZ
z1TZxd7<Cwd`qMOk0w4^tkC8!w0c4g#9Z&-r@Cx<wmH^t#0u?tfNJ&OhGRhqdfzc2c
z4S~@R7!85Z5Eu=C(GVC7fzc2c4S~@R7!85Z5E$km02&`)0S-Lmr{v}UQO1v9?q7V~
z1C7&20Tp=qx%qnryM^!p**=M7X{Hu>=4N_^1_lhE(E7oQ&x}#I(GVC7f#Dhgpb}ye
zaF8`Ozbq3pk_sB31&;r202ZK(TRVjzL&1y+3`cr97{2nfs7?WnwK8r3mLbW&GQ`l>
z1YC$PFm7fLU<XQU6G$j*vfci1+x=s=>!ox5<c=~&!(lWLjAn+>k^opD)-ixeSZ8om
q!niHqgE43-j7)H!Av2}IPyuKb$V}jL7<fJmRC>ea!_t6kkX-<J2I}en
diff --git a/vision.py b/vision.py
index bec6fe8..d81b5af 100644
--- a/vision.py
+++ b/vision.py
@@ -26,7 +26,7 @@ class Tracker(object):
'''
# 获得opencv版本
(major_ver, minor_ver, subminor_ver) = (cv2.__version__).split('.')
- self.tracker_types = ['BOOSTING', 'MIL', 'KCF', 'TLD', 'MEDIANFLOW', 'GOTURN']
+ self.tracker_types = ['BOOSTING', 'MIL', 'KCF', 'TLD', 'MEDIANFLOW', 'GOTURN', 'CSRT']
self.tracker_type = tracker_type
self.isWorking = False
self.draw_coord = draw_coord
@@ -46,6 +46,8 @@ class Tracker(object):
self.tracker = cv2.TrackerMedianFlow_create()
if tracker_type == 'GOTURN':
self.tracker = cv2.TrackerGOTURN_create()
+ if tracker_type == 'CSRT':
+ self.tracker = cv2.TrackerCSRT_create()
def initWorking(self, frame, box):
'''