【Tflite模型】yolov5 tflite模型推理

该博客介绍了一种使用TFLite_runtime库和预训练的TFLite模型进行实时目标检测的方法。通过加载模型,处理输入图像,执行推理并解析输出,实现了对图像中目标的检测和框选,对检测到的物体类别和置信度进行了阈值筛选,最后在图像上绘制了检测框并显示结果。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

#1.导入python库
#使用tflite_runtime来替换tensorflow,减少每次检测加载tf的时间
import cv2
import numpy as np
import tflite_runtime.interpreter as tflite
#import tensorflow as tf
#2.加载TFLite模型,具体模型文件根据模型的网络结构为主,解码不唯一,这个其中一种方法
interpreter = tflite.Interpreter(model_path=“detect.tflite”)
interpreter.allocate_tensors()
#3.获取模型输入、输出的数据的信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

#4.输入检测图像,处理图像矩阵(重点)
frame = cv2.imread("img/1.jpg")
#原图尺寸
imH, imW, _ = np.shape(frame)
#模型图尺寸
input_shape = input_details[0]['shape']
# RGB转BGR
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# 设置输入尺寸
frame_resized = cv2.resize(frame_rgb, (input_shape[1], input_shape[2]))
# 添加一维度
img = np.expand_dims(frame_resized, axis=0)
# float32浮点数,看模型量化类型unit8,不需要转换
#input_data = np.array(np.random.random_sample(input_shape), dtype=np.float32)
interpreter.set_tensor(input_details[0]['index'], img)

#5.创建检测网络层
interpreter.invoke()
#6.检测框/类别/置信度的检测
boxes = interpreter.get_tensor(output_details[0]['index'])[0]  
classes = interpreter.get_tensor(output_details[1]['index'])[0]  
scores = interpreter.get_tensor(output_details[2]['index'])[0]
#7.绘画检测框
# 对于概率大于 50%的进行显示
for i in range(len(scores)):
    if ((scores[i] > 0.35) and (scores[i] <= 1.0)):
        # 获取边框坐标
        ymin = int(max(1, (boxes[i][0] * imH)))
        xmin = int(max(1, (boxes[i][1] * imW)))
        ymax = int(min(imH, (boxes[i][2] * imH)))
        xmax = int(min(imW, (boxes[i][3] * imW)))
        # 画框
        cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (10, 255, 0), 2)
        # 获取检测标签
        object_name = str(int(classes[i]))  
        label = '%s 0.%d' % (object_name, int(scores[i] * 100))  
        print("label", label)
        # 显示标记
        frame = paint_chinese_opencv(frame, label, (xmin, ymin), (255, 0, 0))
cv2.imshow('object detect', img)
cv2.waitKey(0)
要将YOLOv5训练出来的PyTorch模型转换为TensorFlow Lite模型,可以按照以下步骤进行: 1. 安装依赖库 在转换之前,需要安装TensorFlow和PyTorch两个库。可以使用以下命令安装: ``` pip install tensorflow==2.5.0 pip install torch==1.9.0 ``` 2. 下载YOLOv5代码 从GitHub上下载YOLOv5的代码,并进入yolov5目录: ``` git clone https://github.com/ultralytics/yolov5.git cd yolov5 ``` 3. 下载预训练权重 从YOLOv5的官方网站上下载相应的预训练权重文件,例如yolov5s.pt。将权重文件保存到yolov5目录下。 4. 运行转换脚本 在yolov5目录下运行以下命令,将PyTorch模型转换为TensorFlow Lite模型: ``` python models/tf.py --weights yolov5s.pt --cfg models/yolov5s.yaml --img-size 640 ``` 其中,`--weights`选项指定PyTorch模型的权重文件,`--cfg`选项指定YOLOv5模型的配置文件,`--img-size`选项指定输入图像的大小(必须与训练时指定的大小相同)。 转换完成后,会在yolov5目录下生成一个`.tflite`文件,即为转换后的TensorFlow Lite模型5. 推理测试 可以使用TensorFlow Lite的Python API加载模型,并进行推理测试: ``` import tensorflow as tf import numpy as np interpreter = tf.lite.Interpreter(model_path="yolov5s.tflite") interpreter.allocate_tensors() input_details = interpreter.get_input_details() output_details = interpreter.get_output_details() input_shape = input_details[0]['shape'] input_data = np.array(np.random.random_sample(input_shape), dtype=np.float32) interpreter.set_tensor(input_details[0]['index'], input_data) interpreter.invoke() output_data = interpreter.get_tensor(output_details[0]['index']) ``` 以上代码中,首先加载`.tflite`文件,并分配内存空间,然后获取输入和输出的详细信息。接着,随机生成一个输入数据,并将其传递给模型进行推理。最后,获取模型的输出结果。
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值