在计算机视觉领域,YOLOv5是一种常用的目标检测模型,用于快速而准确地识别图像中的目标。本文将介绍如何将YOLOv5模型应用于小程序,并展示了使用Flask框架实现的示例代码。
首先,我们需要从本地加载自定义的YOLOv5模型。示例代码中,使用了PyTorch库加载模型:
model_path = str(ROOT)
weight_path = str(ROOT / "dataset" / "best.pt")
model = torch.hub.load(model_path, "custom", weight_path, source="local", force_reload=True)
接下来,我们需要创建一个Flask应用,并定义一个路由用于接收上传的图像,并对图像进行目标识别:
@app.route('/upload', methods=['POST'])
def recognize_image():
if 'file' not in request.files:
return '未选择文件'
# 读取上传的图像文件
file = request.files['file']
if file.filename == '':
return '未选择文件'
# 保存文件到本地
save_path = model_path+'/dataset/upload/'
file.save(save_path + file.filename)
imgs = [str(save_path) + file.filename]
img_path = str(save_path) + file.filename
results = model(imgs)
...
# 对预测结果进行处理
...
return jsonify({'code': 200, 'msg': msg})
在这个示例中,我们使用了Flask的request
模块来处理上传的图像文件。图像文件保存在服务器上并传递给YOLOv5模型进行目标识别。然后,我们可以根据预测结果进行进一步的处理和返回。
除了上传图像进行目标识别外,我们还可以通过另一个路由来进行测试,读取本地图片文件并进行目标识别:
@app.route('/test', methods=['GET'])
def test_detect_objects():
imgs = [str(model_path) + '/dataset/t_1.jpg']
results = model(imgs)
...
# 对预测结果进行处理
...
return '识别完成物体为:' + tagname + '的可信度:' + confidence
最后,我们需要在Flask应用的if __name__ == '__main__':
部分添加代码来启动Web应用:
if __name__ == '__main__':
app.run(host='0.0.0.0')
现在,你可以将以上代码填充到你的博客中,并根据实际情况进行适当的调整和说明。此示例提供了一个简单的Web应用程序,可以上传图像并使用YOLOv5模型进行目标识别。你可以根据实际需求进行扩展和优化。
下面是接口文件的带代码
import torch
import sys
import os
from pathlib import Path
from flask import Flask, request,jsonify
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
app = Flask(__name__)
# 从本地加载自定义的YOLOv5模型
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
model_path = str(ROOT) # yolov5根目录,需要转换为字符串类型
weight_path = str(ROOT / "dataset" / "best.pt") # 设置正确的权重文件路径
model = torch.hub.load(model_path, "custom", weight_path, source="local", force_reload=True)
# results.show() # 这两句用于看一下模型检测结果
@app.route('/upload', methods=['POST'])
def recognize_image():
# 检查是否存在图像文件
if 'file' not in request.files:
return '未选择文件'
file = request.files['file']
if file.filename == '':
return '未选择文件'
# 保存文件到本地
save_path = model_path+'/dataset/upload/'
file.save(save_path + file.filename)
imgs = [str(save_path) + file.filename] # 设置正确的图像文件路径
img_path = str(save_path) + file.filename
results = model(imgs)
results.print()
results.save(save_path + 'detections')
res = results.pandas().xyxy
confidence = 0
tagname = ''
img = Image.open(img_path)
draw = ImageDraw.Draw(img)
for i, boxes in enumerate(res):
print(f"第 {i + 1} 个结果:")
for _, row in boxes.iterrows():
if row['confidence'] > 0.5:
# 计算模型准确率
confidence = '{:.2%}'.format(row['confidence'])
tagname = row['name']
print(row['confidence'])
# 打印预测结果
print(f"预测标签: {row['name']}")
# 计算模型准确率
# accuracy = calculate_accuracy([(img, pred_index)]) # 使用辅助函数计算准确率
# print("模型准确率: {:.2%}".format(accuracy))
# 处理检测结果并返回
# TODO: 返回处理后的检测结果
if confidence == 0:
msg = '当前没有脉动'
else:
msg = '识别完成物体为:'+tagname+'的可信度:' + confidence
return jsonify({'code': 200, 'msg': msg})
def get_image_path(image_folder, image_file):
# 获取当前脚本文件的路径
script_dir = Path(__file__).resolve().parent
# 构建图片文件的完整路径
image_path = script_dir / image_folder / image_file
return image_path
@app.route('/test', methods=['GET'])
def test_detect_objects():
# 读取本地图片文件
imgs = [str(model_path) + '/dataset/t_1.jpg'] # 设置正确的图像文件路径
results = model(imgs)
results.print()
res = results.pandas().xyxy
confidence = 0;
for i, boxes in enumerate(res):
print(f"第 {i + 1} 个结果:")
for _, row in boxes.iterrows():
if row['confidence']>0.5:
confidence = '{:.2%}'.format(row['confidence'])
print(row['confidence'])
# 打印预测结果
print(f"预测标签: {row['name']}")
# 计算模型准确率
# accuracy = calculate_accuracy([(img, pred_index)]) # 使用辅助函数计算准确率
# print("模型准确率: {:.2%}".format(accuracy))
# 处理检测结果并返回
# TODO: 返回处理后的检测结果
return '识别完成物体为:的可信度:'+confidence
if __name__ == '__main__':
app.run(host='0.0.0.0')