当前位置: article > 正文

使用flask简单部署yolov5_yolov5 flask部署

作者：小蓝xlanll | 2024-02-15 09:33:57

踩

yolov5 flask部署

本文目前只实现检测图片功能，主要是为了学习加载本地模型以及实现后端逻辑。对我小白来说难点在于保存模型推理后的图片，路径搞对很重要！！！！

首先思考打开页面时应该有两个按钮，一个用来选择本地图片，另一个是将选择的图片提交至后端从而给模型进行推理检测。在进行推理检测后，应该把推理的结果以及原始图片放在页面展示出来。如下前端简单页面。如想实现更加复杂好看的页面可以自己渲染。

思考一下：点击浏览，选择了一张原始图片f，我们应该将选择的这个图像拿到，然后保存到项目目录中，这样再在html页面中根据路径将原始图片展示出来，可以通过下面代码：


# 1、获取前端提交的文件
f = request.files['file']
 
# 2、选择保存图片的路径
basepath = "/static/images/"
path = os.path.dirname(__file__)  # /home/nscrl/lhy/cv/yolov5_7.0
suffix = f.filename.split(".")[1]  # 获取.jpg
 
 
o_img_path = basepath + uniqueNum + "_o" + "." + suffix  # 这是/static/images/xxx.jpg
# 这个可以传参给前端拿到相对路径然后展示图片
 
f.save(path + o_img_path)  # 这是绝对路径，因为.save()貌似只能保存绝对路径，但保存的地址也是项目下的/static/images/xxx.jpg

之后，拿到你点击的图像并保存至项目的路径后，我们就可以加载yolov5检测模型了，我们使用model = torch.hub.load('./', 'custom', 'yolov5s.pt', source='local')加载模型，'./'就是你的模型文件放在哪里的，比如我的yolov5s.pt文件就放在项目的目录下，其他的参数可以不用动，直接这样写就行，想详细理解可以去查资料有关torch.hub.load的用法。至此，我们就加载好了模型，然后我们就可以直接img = model(img_path), 然后img.show()就能显示推理好的图像了。处理图像的代码如下：


# 1、加载模型
model = torch.hub.load('./', 'custom', 'yolov5s.pt', source='local')
 
# 2、拿到原始图像的路径，也就是上文说的model(img_path)中的img_path
img_detect_path = '.' + o_img_path  # ./static/images/xxx.jpg
img_detect = model(img_detect_path)  # 此时你就可以直接img_detect.show()查看检测的图片了

接下来，我们希望将推理后的图片也要保存到项目中，然后再前端也传入路径，就能显示图片了，注意这里我们不能直接使用img_detect.save()保存，因为会直接执行yolov5中的common.py文件下的save文件，需要改动很多才能保存到我们想要的目录中。对于如何保存处理后的图片也放到staitc/images/xxx.jpg，代码如下：


img_detect.render()
p_img_path = basepath + uniqueNum + "_p" + "." + suffix  # /static/images/xxx.jpg
 
Image.fromarray(img_detect.ims[0]).save(path+p_img_path)  # 也是保存绝对路径, /home/nscrl/lhy/cv/yolov5_7.0/static/images/xxx.jpg

原始图片的路径和处理后的图片的路径都可以拿到了，就是o_img_path(/static/images/xxx_o.jpg)，和p_img_path(/static/images/xxx_p.jpg)，然后传到html中就可以显示了。

整体的代码如下：

后端代码my_detect.py:


import datetime
import os.path
import random
import torch
from flask import Flask, render_template, request
from PIL import Image
 
app = Flask(__name__)
 
 
@app.route('/flask/demo/', methods=["POST", "GET"])
def flask_demo():
    # 如果post请求，也就是前端form提交表单，执行下面代码后返回index.html,get请求就直接返回index.html
    if request.method == "POST":
        # 1、生成一个随机的文件名：唯一
        nowTime = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        randomNum = random.randint(0, 100)
        if randomNum <= 10:
            randomNum = str(0) + str(randomNum)
        uniqueNum = str(nowTime) + str(randomNum)
 
        # 2、获取图像
 
        f = request.files['file']  # 获取前端提交的文件
        basepath = "/static/images/"
        path = os.path.dirname(__file__)   # /home/nscrl/lhy/cv/yolov5_7.0
        suffix = f.filename.split(".")[1]  # 获取.jpg
 
        # 3、保存原图像
        o_img_path = basepath + uniqueNum + "_o" + "." + suffix
        f.save(path + o_img_path)  # f.save只能保存绝对路径
        # 处理图像
        model = torch.hub.load('./', 'custom', 'yolov5s.pt', source='local')
        img_detect_path = '.' + o_img_path   # './static/images/xxx.jpg'
        img_detect = model(img_detect_path)
        img_detect.render()
        p_img_path = basepath + uniqueNum + "_p" + "." + suffix
        Image.fromarray(img_detect.ims[0]).save(path+p_img_path)
 
        return render_template('index.html', o_img_path='.' + o_img_path, p_img_path='.' + p_img_path)
    return render_template('index.html')
 
 
 
 
 
 
 
if __name__ == '__main__':
    app.run()

前端html页面代码：


<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <title>yolo检测</title>
</head>
<body>
    <h1>上传图片</h1>
    <form method="post" enctype="multipart/form-data">
        <div>
            <input type="file" name="file" value="选择图片">
        </div>
        <input type="submit" value="提交">
    </form>
    <img src={{ o_img_path }} width="300px"/>
    <img src={{ p_img_path }} width="300px" />
</body>
</html>

运行my_detect.py文件结果展示：

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/小蓝xlanll/article/detail/83220