项目gitee仓库完整代码:https://gitee.com/wsy-yjys/yolov5_-tello_-demo
B站视频demo:https://www.bilibili.com/video/BV1ZvytYZEqM/
一、环境配置注意事项
- robomaster适配python版本:3.6.6-3.8.9之间
- YOLOv5适配python版本:3.8.0以上
因此,项目环境python版本最好在3.8.0-3.8.9
之间
二、代码使用注意事项
- 运行前,已连接Tello无人机的WIFI,并确保公用网络的防火墙已关闭!
- self.model = 替换为自己的模型路径!
三、视频demo核心代码如下
import threading
import time
import cv2
import torch
import robomaster
from robomaster import robot
from robomaster.battery import TelloBattery as Battery
from queue import Queue
"""注意事项:
1. 运行前,已连接Tello无人机的WIFI,并确保公用网络的防火墙已关闭!
2. self.model = 替换为自己的模型路径!
"""
# For “NotImplementedError: cannot instantiate 'PosixPath' on your system" bug
import pathlib
temp = pathlib.PosixPath
pathlib.PosixPath = pathlib.WindowsPath
robomaster.config.LOCAL_IP_STR = "192.168.10.2"
class Tello():
def __init__(self,):
"""Initializes the YOLOv5 Detections class with image info, predictions, filenames, timing and normalization."""
super().__init__()
# 模型初始化
self.model = torch.hub.load(r'./', 'custom', path=r"runs/train/exp4/weights/best.pt", source="local") # 替换为自己的模型路径)
# 无人机初始化
self.drone = robot.Drone()
self.drone.initialize()
# 电池初始化
self.battery = Battery(self.drone)
# 相机初始化
self.camera = self.drone.camera
self.camera.stop_video_stream()
self.camera.start_video_stream(display=False)
self.camera.set_fps("high")
self.camera.set_resolution("high")
self.camera.set_bitrate(6)
def infer(self):
# 使用YOLOv5模型进行推理tello回传的视频流
img = self.camera.read_cv2_image(strategy="newest")
fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 指定视频编码器
video_writer = cv2.VideoWriter('output.mp4', fourcc, 30.0, (img.shape[1], img.shape[0])) # 创建视频写入对象
while True:
img = self.camera.read_cv2_image(strategy="newest") #直接读取最新帧,采用="pipeline"会在第二次启动时卡死
results = self.model(img) # Inference
# results.print()
image = results.show_cv()
cv2.imshow('results', image)
# 写入当前帧到视频文件
video_writer.write(image)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
if not q.empty() and q.get()["close"]:
break
video_writer.release()
cv2.destroyAllWindows()
print("infer end")
def get_verson(self):
# 获取飞机版本信息
print("Drone sdk version: {0}".format(self.drone.get_sdk_version()))
# 获取飞机SN信息
print("drone sn: {0}".format(self.drone.get_sn()))
def get_battery(self):
# 获取飞机电量信息
print("Battery: {}%".format(self.battery.get_battery()))
def close(self):
# 结束:释放资源
self.camera.stop_video_stream()
self.drone.close()
if __name__ == "__main__":
q = Queue(100)
# 初始化
tello = Tello()
tello.get_battery()
# YOLO推理线程
recv_thread = threading.Thread(target=tello.infer)
recv_thread.daemon = True
recv_thread.start()
# 主线程:用于实现预定动作
time.sleep(6)
q.put({"close":True})
# 结束
time.sleep(1) # 延迟1s,确保两个线程都正常结束
tello.close()
注意:这只是核心代码,项目完整代码请见gitee仓库