毕设要做一个可以识别工业品缺陷的 Web 前后端项目,前端打算采用 vue 框架,后端打算使用 golang 的 gin 框架进行开发。 但是现在其中有个问题是如何调用 pytorch 框架训练的 YoloV5 模型。同时我希望可以高并发调用这个模型,就好像我可能有多个工厂同时部署了这个缺陷识别系统,每个工厂有多条流水线,都可以调用这个 Python 模型。

综合网络搜索的集中方案, 以下是目前已知的方案

golang 运行 Python 脚本调用模型

相当于在 golang 里面运行 python 对应 detect 的脚本然后捕获其输入了

os/exec 运行

类似于使用 cmd := exec.Command("python3", "predict.py", string(jsonData)) 这样的代码执行 python 程序 总感觉这样不太好,但事实上实验室原先的项目 java 配合 python 识别道路病害也是使用这样输入命令的方法

Go-Python 绑定库

在 Go 中嵌入 Python 解释器并运行 Python 代码,好像挺复杂的

TensortFlow

通过sdk加载模型方式

换个框架再训练一个 yolo 模型出来,然后 TF 官方有多种语言的 sdk 的,可以加载这个模型。通过使用 cgo 调用 但似乎存在一些问题,参考 note/ml_deploy.md at master · liyue201/note · GitHub

这种方法的优点是模型的训练和部署分离,模型的训练交给算法工程师去做,上线和部署交给开发工程师。因为是函数调用,没有网络损失,这种方式性能最好。众所周知,Go调C/C++都是通过cgo实现的。所以Go服务依赖TensortFlow编译出来的动态库文件。

而本人之前在部署随机森林时,发现官方编译出来的动态库里面有些模块是没有编译进去的,所以需要自己编译。而编译过程中需要下载一些依赖的第三方库代码,又由于众所周知的原因,有些可能会失败。所以,如果不是对性能要求太高,不建议用这种方式。

通过Tensortflow Serving方式

TensorFlow Serving 可以通过 gRPC 和 RESTful API 提供服务,使得外部应用可以通过网络调用模型进行推理。

基本工作流程

  • 训练并导出模型:在 TensorFlow 中训练模型并将其保存到磁盘。
  • 配置 TensorFlow Serving:TensorFlow Serving 读取保存的模型,并为其提供服务。
  • 客户端调用服务:客户端通过 gRPC 或 RESTful API 调用服务,并得到模型的预测结果。

有挺多教程的 Golang玩转TensorFlow深度学习模型 - Golang - Sunaloe 使用gRPC调用TensorFlow Serving后数据解析_grpc tensor-CSDN博客 Tensorflow Mnist手写数字识别Go实战

ONNX

ONNX(Open Neural Network Exchange)是一种开源的深度学习模型文件格式,它使不同的深度学习框架(如 PyTorch、TensorFlow、MXNet、Caffe2 等)可以互相兼容。ONNX 提供了一种通用的模型交换格式,使模型在不同平台和硬件加速器上更容易部署和优化。

可以把 Pytorxh 模型转换成 ONNX 格式文件,<— torch.onnx.export ONNX Runtime 可以调用

调用方式

有一个**go-onnxruntime**的包可以调用 onnxruntime package - github.com/ivansuteja96/go-onnxruntime - Go Packages 文档说明

存疑,golang 类似对这种模型的支持还是太少,而且没太多教程,现在就怕出问🙃

Python web 开发

实在不行学一学 Django 或者是 Flask 搞一个 python 开发好了

疑问

很早之前也用百度的文字识别 api 他们这种类似给了个 api 然后远程调用是如何实现的,有点想搞成这种类型。


最终方案

我跟朋友讨论了一下这个问题,脑子瓦特了,其实就类似 Tensortflow Serving 再部署一个网络服务一样。

我现在是两个一个前端一个后端 go web 服务器,那我可以后端再加一个服务器,专门用来运行 Python 的模型的服务器就可以了。收到 go 提交的 http 或者是 grpc 请求后运行模型对传输的图片进行查看,在返回对应格式的坐标。而且这个真的从实际出发考虑的话,如果有多个请求 detect 的请求也可以升级 AI 服务器的并发能力之类的。而且我可以先试试用 pytorch 进行部署。

实现

Python 部分

使用了 FastAPI 快速搭建了了一个可以调用 YOlO 模型的服务器 API。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import json

  

from fastapi import FastAPI

import cv2

import ml

import numpy as np

import requests

  

from fastapi import FastAPI, File, UploadFile

from fastapi.responses import JSONResponse

  

app = FastAPI()

  
  

# uvicorn main:app --reload

@app.get("")

  

@app.post("/steel/defectimg")

async def root(file: UploadFile):

    print('777')

    # return {"filename": file.filename}

    img = await file.read()

    # return {"file_size": len(file)}

    # 用opencv读取file字节流图片

  

    img = cv2.imdecode(np.frombuffer(img, np.uint8), cv2.IMREAD_COLOR)

    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    # cv2.imshow('img', img)

    # cv2.waitKey(0)

    # cv2.destroyAllWindows()

  

    # detect image in detecr function,and return the json data

    data = ml.detect(img)

    print(data)

    # data = {"class": "无缺陷", "conf": 0.0, "position": [0, 0, 0, 0],err_result:0}

    # add a field called err_result to the 'data'

    # if len(data)==0:

    #     data.append(err_result=0)

    # else:

    #     data.append(err_result=1)

    # print("json:")

    #convert data to json format

    # data = json.dumps(data)

    # print(data)

    #return response

    # return JSONResponse(content={"data":data,"status":0}, status_code=224)

    return data

Golang 部分

相当于 golang 造了一个 http 请求将图片数据发送过去

  1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
package util

import (
    "bytes"
    "encoding/json"
    "fmt"
    "image"
    "image/color"
    _ "image/jpeg"
    "io"
    "log"
    "mime/multipart"
    "net/http"
    "os"

    "github.com/gin-gonic/gin"
)

type DefectData struct {
    Class    string  `json:"class"`
    Conf     float64 `json:"conf"`
    Position []int   `json:"position"`
}

// 调用 YOLOv5 进行检测的函数
func detecttoyolov5(imageFile *os.File, filename string) (result *http.Response, err error) {

    // 创建一个新的 multipart 表单
    body := new(bytes.Buffer)
    writer := multipart.NewWriter(body)

    // 将 imageFile 添加到表单中
    part, err := writer.CreateFormFile("file", filename)
    if err != nil {
        fmt.Println(err.Error())
        fmt.Println("创建表单失败!")
        return
    }

    // 将 imageFile 复制到表单中
    _, err = io.Copy(part, imageFile)
    if err != nil {
        fmt.Println(err.Error())
        fmt.Println("复制文件到表单失败!")
        return
    }

    // 关闭 multipart writer
    err = writer.Close()
    if err != nil {
        fmt.Println(err.Error())
        fmt.Println("关闭表单失败!")
        return
    }

    // 创建一个包含表单数据的 POST 请求
    req, err := http.NewRequest("POST", "http://127.0.0.1:8000/steel/defectimg", body)
    if err != nil {
        fmt.Println(err.Error())
        fmt.Println("创建POST请求失败!")
    }

    req.Header.Set("Content-Type", writer.FormDataContentType())
    req.Header.Set("Accept-Encoding", "gzip, deflate")
    req.Header.Set("Accept", "*/*")

    // 发送请求
    client := &http.Client{}
    resp, err := client.Do(req)
    if err != nil {
        fmt.Println(err.Error())
        fmt.Println("发送POST请求失败!")
    }
    return resp, err
}

// Defectimg 是一个函数,用于从前端接收图片并发送到缺陷检测模型
func Defectimg(c *gin.Context) {

    // 从上下文中接收图片
    // 它返回一个 *multipart.FileHeader 和一个错误
    file, err := c.FormFile("file")

    if err != nil {
        fmt.Println(err.Error())
        fmt.Println("从前端获取图片失败!")
        return
    }

    // 使用 file.Header 获取文件信息
    fmt.Printf("上传的文件: %s (大小: %d 字节, 内容类型: %s)\n",
        file.Filename, file.Size, file.Header.Get("Content-Type"))
    filename := file.Filename

    // 保存图片到本地
    c.SaveUploadedFile(file, filename)

    // 上传图片到腾讯云 COS
    img_upload, err := os.Open(filename)
    if err != nil {
        fmt.Println(err.Error())
        fmt.Println("打开文件失败!")
        return
    }

    originalMode := "original"

    ori_url, err := UploadImageToTencentCloudCos(img_upload, originalMode)
    fmt.Println("原始图片URL: ", ori_url)
    if err != nil {
        fmt.Println(err.Error())
    }
    img_upload.Close()

    // 打开图片文件
    imageFile, err := os.Open(filename)
    if err != nil {
        fmt.Println(err.Error())
        fmt.Println("打开文件失败!")
        return
    }

    // 使用函数检测缺陷
    print(filename)
    resp, err := detecttoyolov5(imageFile, filename)
    if err != nil {
        fmt.Println(err.Error())
        fmt.Println("调用模型失败!")
        return
    }

    imageFile.Close()

    // 读取响应体
    respBody, err := io.ReadAll(resp.Body)
    if err != nil {
        fmt.Println(err.Error())
        fmt.Println("读取响应失败!")
    }

    // 解析 JSON 字符串为 Go 的 DefectData 结构体切片
    var defectData []DefectData
    if err := json.Unmarshal(respBody, &defectData); err != nil {
        c.String(http.StatusBadRequest, err.Error())
        fmt.Println("解析JSON失败!")
    }
    fmt.Println("缺陷数据: ", defectData)

    // 格式化打印数据
    for _, defect := range defectData {
        fmt.Printf("类别: %s, 置信度: %f, 位置: %v\n", defect.Class, defect.Conf, defect.Position)
    }

    // 使用 drawSquares 函数在图片上绘制方框
    // 解码图片
    imageDraw, err := os.Open(filename)
    if err != nil {
        fmt.Println(err.Error())
        fmt.Println("打开文件失败!")
        return
    }
    img, _, err := image.Decode(imageDraw)
    if err != nil {
        println("解码图片失败")
        log.Println(err)
    }

    // 使用 drawSquare 函数
    result_url := DrawSquare(img, defectData, color.RGBA{255, 0, 0, 255})
    imageDraw.Close()
    fmt.Println("结果图片URL: ", result_url)

    c.JSON(http.StatusOK, gin.H{
        "code":       200,
        "msg":        "success",
        "data":       defectData,
        "ori_url":    ori_url,
        "result_url": result_url,
    })
}

在操作里,图片先上传图床,在发送对应的链接,到时候返回回来直接是 json 数据格式(使用 http 请求调用)就可以了

另外的话需要注意 yolo 识别的一个坐标方位,这里还有点不一样,要具体看 yolo 模型里如何设置的 bounding box

YOLO模型识别的bounding box


2024.5.26 更新

现阶段针对这个系统的通信调用改变了,采用 grpc 通信,以 Protobuf 协议请求 Python 端的 AI 模型,然后使用在使用 json 格式返回。这样子 gin 后端向 AI 模型服务器请求调用检测的时候设置检测 token,(有点明白云平台是如何使用一个 token 就可以让很多类型语言的后端请求访问服务器的了)