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

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

 思考一下:点击浏览,选择了一张原始图片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页面代码:

yolo检测

上传图片

python 使用flask简单部署yolov5  第1张

python 使用flask简单部署yolov5  第2张

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

 

精彩链接

评论可见,请评论后查看内容,谢谢!!!评论后请刷新页面。