赞
踩
在计算机视觉领域,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')
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。