二使用Flask构建AI模型部署应用

Jumy_S

已于 2023-10-08 21:58:04 修改

阅读量352

点赞数 1

文章标签： python

于 2023-09-24 15:07:38 首次发布

本文链接：https://blog.csdn.net/Jumy_S/article/details/133241046

版权

$二使用 Fl a s k 构建 A I 模型部署应用$

前置

1.html、js、css

2.flask后端基本常识及操作

3.json文件常识及基本操作

设计思路

在这里插入图片描述

一环境配置

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"])