二使用 F l a s k 构建 A I 模型部署应用 二 使用Flask构建AI模型部署应用 二使用Flask构建AI模型部署应用
前置
设计思路
一 环境配置
pip install flask -i https://pypi.tuna.tsinghua.edu.cn/simple
二 部署使用
项目结构
前端
index.html
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>AI部署界面</title>
</head>
<body>
<!--分界线-->
<hr/>
<p>上传模型</p>
<form action="/uploadModel" enctype="multipart/form-data" method="post">
<table>
<tbody>
<tr>
<td>算法类型：</td>
<td><input type="text" name="algoCategory"></td>
</tr>
<tr>
<td>任务类型：</td>
<td><input type="text" name="task"></td>
</tr>
<tr>
<td>图像输入大小：</td>
<td><input type="text" name="imgsz"></td>
</tr>
<tr>
<td>设备：</td>
<td><input type="text" name="device"></td>
</tr>
<tr>
<td>模型文件：</td>
<td><input type="file" name="model"></td>
</tr>
<tr>
<td>提交：</td>
<td><input type="submit" name="提交"></td>
</tr>
</tbody>
</table>
</form>
<!--分界线-->
<hr/>
<p>已注册模型列表</p>
<table>
<tbody>
<tr>
<td>唯一ID</td>
<td>算法类型</td>
<td>任务类型</td>
<td>图像输入大小</td>
<td>设备</td>
<td>模型文件</td>
</tr>
{% for item in content %}
<tr>
<form action="/ModifyModel" enctype="multipart/form-data" method="post">
<td><input type="text" name="uuid" value={{item}}></td>
<td><input type="text" name="algoCategory" value={{content[item]["algoCategory"]}}></td>
<td><input type="text" name="task" value={{content[item]["task"]}}></td>
<td><input type="text" name="imgsz" value={{content[item]["imgsz"]}}></td>
<td><input type="text" name="device" value={{content[item]["device"]}}></td>
<td><input type="text" name="ModelPath" value={{content[item]["ModelPath"]}}></td>
<td><input type="submit" name="delete" value="删除" onclick="return confirm('确认删除吗？')"></td>
<td><input type="submit" name="init" value="初始化模型"></td>
{% if alg_status[item] %}
<td><input type="text" style="background:Green" name="initStatus" value="初始化成功"}}></td>
{% endif %}
<td><input type="file" name="testImg"></td>
<td><input type="submit" name="test" value="测试"></td>
</form>
</tr>
{% endfor %}
</tbody>
</table>
</body>
</html>
showRes.html
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<title>显示结果图像</title>
</head>
<body>
<img src="{{ img_data }}" alt="图像路径错误" />
<hr/>
</body>
</html>
后端
flask_app.py
from flask import Flask, request,url_for,render_template
import os
from Helper.JsonHelper import jsonHelper
import uuid
import cv2
import numpy as np
from PIL import Image
from MajorAlgo import MajorModel
# flask实例
app = Flask(__name__)
# 导入配置文件
app.config.from_pyfile('config.py')
# 算法实例字典
dict_alg = {}
# 算法初始化状态栏
alg_status = {}
# 当前文件夹绝对路径
basedir = os.path.abspath(os.path.dirname(__file__))
# 数据库路径
databaseJsonPath = r"./modelData.json"
'''
主界面
'''
@app.route('/')
def index():
# 数据库json路径
dataBaseJsonPath = databaseJsonPath
# 判断json是否存在，不存在则返回空
if os.path.exists(dataBaseJsonPath):
# 加载json数据库
jsonRes = jsonHelper.read(dataBaseJsonPath)
# 跳转至主页面
return render_template('index.html',content=jsonRes,alg_status=alg_status)
else:
# 无数据跳转至主界面
return render_template('index.html',content=[],alg_status=alg_status)
'''
上传模型
'''
@app.route("/uploadModel",methods=['GET','POST'])
def uploadModel():
# 算法类型
algoCategory = request.form.get("algoCategory")
# 任务类型
task = request.form.get("task")
# 图像大小
imgsz = request.form.get("imgsz")
# 设备
device = request.form.get("device")
# 模型
model = request.files.get("model")
# 判断保存模型路径是否存在，如果不存在，则创建
if not os.path.exists("./AI_Model_Zoo"):
os.makedirs("./AI_Model_Zoo")
# 生成唯一id
uniqueId = str(uuid.uuid1())
# 保存模型
model.save(os.path.join("AI_Model_Zoo",uniqueId+"_"+model.filename))
# 更新数据信息到数据库
jsonHelper.update(
# 唯一id(键)
uniqueId,
# 模型数据（值）
{
"algoCategory":algoCategory,
"task":task,
"imgsz": imgsz,
"device": device,
"ModelPath": os.path.join("AI_Model_Zoo",uniqueId+"_"+model.filename)
},
# 数据库json路径（保存json路径）
databaseJsonPath)
# 读取数据库json
jsonRes = jsonHelper.read(databaseJsonPath)
# 更新到主界面
return render_template('index.html',content=jsonRes,alg_status=alg_status)
'''
模型操作（删除、初始化、删除）
'''
@app.route("/ModifyModel",methods=['GET','POST'])
def ModifyModel():
if request.method =="POST":
# 一 删除模型信息
if request.form.get("delete"):
# 获取提交表单的uuid
uuid = request.form.get("uuid")
# 根据uuid删除json数据
jsonRes = jsonHelper.delete(uuid,databaseJsonPath)
# 返回更新后的数据到主界面
return render_template('index.html',content=jsonRes,alg_status = alg_status)
# 二 初始化模型
elif request.form.get("init"):
# 从表单获取模型权重路径
weights = request.form.get("ModelPath")
# 获取算法类型
algoCategory = request.form.get("algoCategory")
# 获取图像大小
imgsz= request.form.get("imgsz")
# 获取uuid
uuid = request.form.get("uuid")
# 获取设备
device = request.form.get("device")
# 获取任务类型
task = request.form.get("task")
# 生成参数字典
parms = {
"task":task,
"device":device,
"uuid": uuid,
"weights": weights,
"imgsz":imgsz,
"algoCategory": algoCategory}
# 根据算法参数生成创建算法实例
alg = MajorModel(**parms)
# 将算法实例放入全局字典中
dict_alg[uuid] = alg
# 读取数据库信息返回至主界面
jsonres = jsonHelper.read(databaseJsonPath)
# 更新模型状态
alg_status[uuid] = True;
# 跳转至主界面
return render_template('index.html', content=jsonres, alg_status = alg_status)
# 三 测试模型
elif request.form.get("test"):
# 从表单中获取uuid
uuid = request.form.get("uuid")
# 获取测试图像
testImg = request.files.get("testImg")
# 转Pillow格式
image = Image.open(testImg)
# 模型推理
imgBase64 = dict_alg[uuid].test(source=image)
# 跳转至显示图像页面
img_base64 = "data:image/jpeg;base64,"+ imgBase64
return render_template("showRes.html",img_data=img_base64)
'''
推理部署
'''
@app.route("/predict",methods=['GET','POST'])
def predict():
# 获取唯一ID，确定调用哪个模型
uuid = request.args.get("uuid")
# 获取接受数据的byte流
file_byte = request.get_data()
# byte流转np数组
img_buffer_numpy = np.frombuffer(file_byte, dtype=np.uint8)
# 转opencv 0 灰度图 1表示彩色图。
img = cv2.imdecode(img_buffer_numpy, 1)
# 模型预测
res = dict_alg[uuid].predict(source=[img])
return str(res)
if __name__ == '__main__':
app.run()
MajorAlgo.py
from ultralytics import YOLO
import torch
import cv2
from PIL import Image
import os
import numpy as np
import base64
from PIL import ImageDraw
# 当前工作目录
basedir = os.path.abspath(os.path.dirname(__file__))
'''
算法类
'''
class MajorModel:
def __init__(self, **kwargs):
if kwargs["algoCategory"] == "v8":
self.algoCategory = "v8"
self.imgsz =int(kwargs['imgsz'])
self.uuid = kwargs['uuid']
if kwargs['device'] == "cpu":
self.device = "cpu"
else:
self.device = int(kwargs['device'])
self.task = kwargs['task']
self.weights = kwargs['weights']
self.model = YOLO(self.weights)
elif kwargs["algoCategory"] == "v5":
self.algoCategory = "v5"
self.imgsz = int(kwargs['imgsz'])
self.uuid = kwargs['uuid']
if kwargs['device'] == "cpu":
self.device = torch.device("cpu")
else:
self.device = torch.device('cuda:'+kwargs['device'])
self.task = kwargs['task']
self.weights = kwargs['weights']
self.model = torch.hub.load('./yolov5', 'custom', path=self.weights,source='local')
'''
部署预测
'''
def predict(self,source):
# yolov5 图像检测
if self.algoCategory == "v5" and self.task == "det":
res = self.v5_det_predict(source)
return res
# yolov8 图像分类
if self.algoCategory == "v8" and self.task == "cls":
res = self.v8_cls_predict(source)
return res
# yolov8 图像检测
if self.algoCategory == "v8" and self.task == "det":
res = self.v8_det_predict(source)
return res
# yolov8 图像分割
if self.algoCategory == "v8" and self.task == "seg":
res = self.v8_seg_predict(source)
return res
'''
测试功能
'''
def test(self,source):
# yolov5 图像检测
if self.algoCategory == "v5" and self.task == "det":
res = self.v5_det_test(source=source)
return res
# yolov8 图像分类
if self.algoCategory == "v8" and self.task == "cls":
res = self.v8_cls_test(source=source)
return res
# yolov8 图像检测
if self.algoCategory == "v8" and self.task == "det":
res = self.v8_det_test(source=source)
return res
# yolov8 图像分割
if self.algoCategory == "v8" and self.task == "seg":
res= self.v8_seg_test(source=source)
return res
'''
yolov8分类预测函数
'''
def v8_cls_predict(self, source):
pass
'''
yolov8检测预测函数
'''
def v8_det_predict(self, source):
pass
'''
yolov8分割预测函数
'''
def v8_seg_predict(self, source):
pass
'''
yolov5检测预测函数
'''
def v5_det_predict(self,source):
pass
'''
yolov8 分类测试函数
'''
def v8_cls_test(self, source):
results = self.model.predict(source=source, save=False, imgsz=self.imgsz, device=self.device)
# View results
for r in results:
# 获取预测类别
name = r.names[r.probs.top1]
# 绘制工具
draw = ImageDraw.Draw(source)
# 绘制文本
draw.text((5, 5), name, align="left")
# 转opencv
img = cv2.cvtColor(np.asarray(source), cv2.COLOR_RGB2BGR)
# np转byte流
ret, buffer = cv2.imencode('.jpg', img)
# byte流转base64编码
img_base64 = base64.b64encode(buffer).decode('utf-8')
return img_base64
'''
yolov8 检测测试函数
'''
def v8_det_test(self, source):
# 模型推理
results = self.model(source)
# 获取预测图像（np格式）
annotated_frame = results[0].plot()
# np转byte流
ret, buffer = cv2.imencode('.jpg', annotated_frame)
# byte流转base64编码
img_base64 = base64.b64encode(buffer).decode('utf-8')
return img_base64
'''
yolov8 分割测试函数
'''
def v8_seg_test(self, source):
# 模型推理
results = self.model(source)
# 获取预测图像（np格式）
annotated_frame = results[0].plot()
# np转byte流
ret, buffer = cv2.imencode('.jpg', annotated_frame)
# byte流转base64编码
img_base64 = base64.b64encode(buffer).decode('utf-8')
return img_base64
'''
yolov5 分类测试函数
'''
def v5_cls_test(self, source):
pass
'''
yolov5 检测测试函数
'''
def v5_det_test(self, source):
# # 图像预处理
# img0 = source
# im = letterbox(img0, self.imgsz, stride=32, auto=True)[0] # padded resize
# im = im.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
# im = np.ascontiguousarray(im) # contiguous
# im = torch.from_numpy(im).to(self.model.device)
# im = im.half() if self.model.fp16 else im.float() # uint8 to fp16/32
# im /= 255 # 0 - 255 to 0.0 - 1.0
# if len(im.shape) == 3:
# im = im[None]
# # 推理
# pred = self.model(im, augment=False, visualize=False)
# # 极大值抑制
# pred = non_max_suppression(pred, conf_thres=0.25, iou_thres=0.45, max_det=1000)
# # 处理
# for i, det in enumerate(pred): # per image
# if len(det):
# # 预测框大小返回到拉伸原图大小
# det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], img0.shape).round()
# # 保存结果图
# for *xyxy, conf, cls in reversed(det):
# if True: # Write to file
# box = xyxy
# p1, p2 = (int(box[0]), int(box[1])), (int(box[2]), int(box[3]))
# cv2.rectangle(img0, p1, p2, (0, 255, 255), thickness=1, lineType=cv2.LINE_AA)
# imgshow = Image.fromarray(cv2.cvtColor(img0, cv2.COLOR_BGR2RGB))
# # imgshow.show()
# # 保存到固定路径下，为显示做准备
# imgshow.save(os.path.join(basedir, 'static/images', "testRes.jpg"))
return os.path.join('static/images', "testRes.jpg")
JsonHepler.py
import json
import os
class JsonHelper:
'''
更新
'''
def update(self,key,data,data_path):
# 如果json文件不存在，则创建
if not os.path.exists(data_path):
with open(data_path, 'w') as f:
json.dump({"test":""}, f)
# 更新数据
with open(data_path, encoding='utf-8-sig', errors='ignore') as file:
content = json.load(file)
if key in content:
content[key].update(data)
else:
content[key] = data
# 保存数据
with open(data_path, 'w') as f:
json.dump(content, f, indent=2)
'''
读取
'''
def read(self,data_path):
with open(data_path, encoding='utf-8-sig', errors='ignore') as file:
content = json.load(file)
return content
'''
删除
'''
def delete(self,key,data_path):
with open(data_path, encoding='utf-8-sig', errors='ignore') as file:
content = json.load(file)
del content[key]
# 保存数据
with open(data_path, 'w') as f:
json.dump(content, f, indent=2)
return content
jsonHelper = JsonHelper()
if __name__ == '__main__':
jsonHelper = JsonHelper()
# jsonHelper.write({"tetset":1231})
print("写入json")
jsonHelper.data_path = r"../assets/data.json"
print(jsonHelper.read())
test= jsonHelper.read()
print(test["data"])
