【发疯毕设日志day7】大疆tello火影结印起飞——模型测试代码

1. 写了一个类 yolo_test

启动tello,加载初始化模型

设置一个置信度阈值(设置的高一些)

import cv2
import numpy as np
from djitellopy import Tello

import torch

class yolo_test:
    def __init__(self):

        # tello初始化
        self.tello = Tello()
        self.tello.connect()
        self.tello.streamon()
        # 加载 yolov5模型
        self.model = torch.hub.load('./yolov5', 'custom', path='./weights/best.pt',source='local')
        # 置信度阈值
        self.model.conf = 0.7

2. 读取视频流 转化成RGB

解析目标检测结果 

3.  

"""
测试训练号的YOLO模型
"""

import cv2
import numpy as np
from djitellopy import Tello

import torch

class yolo_test:
    def __init__(self):

        # tello初始化
        self.tello = Tello()
        self.tello.connect()
        self.tello.streamon()

     
        # 加载 yolov5模型
        self.model = torch.hub.load('./yolov5', 'custom', path='./weights/best.pt',source='local')
        # 置信度阈值
        self.model.conf = 0.7

    def demo(self):
       
        while True:

            # 读取视频帧
            frame = self.tello.get_frame_read().frame
            # 镜像
            frame = cv2.flip(frame, 1)

            # 计算视频高度和宽度
            frame_h, frame_w = frame.shape[:2]

            # 转为RGB
            img_cvt = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)

            # 目标检测推理
            results = self.model(img_cvt)
            results_arr = results.pandas().xyxy[0].to_numpy()

            # 解析目标检测结果
            for item in results_arr:

                # 标签ID
                ret_label_id = item[-2]
                # 标签名称
                ret_label_text = item[-1]
                # 置信度
                ret_conf = item[-3]
                # 检测框坐标
                l,t,r,b = item[:4].astype('int')
                # 画矩形
                cv2.rectangle(frame,(l,t),(r,b),(0,255,0),2)
                # 绘制文字
                cv2.putText(frame,'{} {}%'.format(ret_label_text,round(ret_conf*100,2)),(l,t-10),cv2.FONT_ITALIC,1.5,(0,255,0),2)


            cv2.imshow('demo', frame)
            if cv2.waitKey(10) & 0xFF == ord('q'):
                break
            
        cv2.destroyAllWindows()



if __name__ == '__main__':

    # 实例化
    test = yolo_test()
    test.demo()
    

  • 4
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值