YOLOv5+Flask+Vue实现图片检测

发布于:2022-12-30 ⋅ 阅读:(447) ⋅ 点赞:(0)

检测效果图

在这里插入图片描述

实现过程

1.前端获取

<el-card>
  <el-upload
    ref="uploadRef"
    class="upload-demo"
    action="#"
    :auto-upload="false"
  >
    <template #trigger>
      <el-button type="primary">选择文件</el-button>
    </template>
  </el-upload>
  <div class="main-card">
    <div class="sourceImage">
      <img v-if="imgUrl" :src="imgUrl" alt="" />
    </div>
    <div class="sourceImage">
      <!-- <template #file="{file}">
          <img :src="file.url" alt="">
        </template> -->
      <img v-if="resultUrl" :src="resultUrl" alt="" />
    </div>
  </div>
  <el-button type="primary" @click="getData()">检测</el-button>
  <el-table :data="tableData" stripe style="width: 100%">
    <el-table-column prop="name" label="Name" width="180" />
    <el-table-column prop="age" label="Age" width="180" />
  </el-table>

  <el-upload
    action="#"
    list-type="picture-card"
    :auto-upload="false"
    :on-change="handleSuccess"
  >
    <el-icon><Plus /></el-icon>
    <template #file="{ file }">
      <div>
        <img class="el-upload-list__item-thumbnail" :src="file.url" alt="" />
        <span class="el-upload-list__item-actions">
          <span
            class="el-upload-list__item-preview"
            @click="handlePictureCardPreview(file)"
          >
            <el-icon><zoom-in /></el-icon>
          </span>
          <span
            v-if="!disabled"
            class="el-upload-list__item-delete"
            @click="handleDownload(file)"
          >
            <el-icon><Download /></el-icon>
          </span>
          <span
            v-if="!disabled"
            class="el-upload-list__item-delete"
            @click="handleRemove(file)"
          >
            <el-icon><Delete /></el-icon>
          </span>
        </span>
      </div>
    </template>
  </el-upload>

  <el-dialog v-model="dialogVisible">
    <img w-full :src="dialogImageUrl" alt="Preview Image" />
  </el-dialog>
</el-card>
</template>
<script lang="ts" setup>
import { ref, onMounted, toRaw, renderList } from 'vue'
import type { UploadInstance } from 'element-plus'
import axios from 'axios'
const tableData = ref([])
const imgUrl = ref('')
const resultUrl = ref('')

const uploadRef = ref<UploadInstance>()

onMounted(() => {
//   getData()
})

import { Delete, Download, Plus, ZoomIn } from '@element-plus/icons-vue'
import type { UploadFile } from 'element-plus'
let image_base64 = ref('')
const handleSuccess = (file, filelist) => {
//   file = toRaw(file)
imgUrl.value = file.url
//   console.log(imgUrl.value)
console.log(file)
var reader = new FileReader()
reader.readAsDataURL(file.raw)
reader.onload = () => {
  //   转换成base64的结果
  image_base64.value = reader.result
  // console.log(reader.result)
}
reader.onerror = function (error) {
  console.log('Error: ', error)
}
}

const beforeUpload = (res, file: UploadFile) => {
console.log(file)
}
const dialogImageUrl = ref('')
const dialogVisible = ref(false)
const disabled = ref(false)

const handlePictureCardPreview = (file: UploadFile) => {
//   console.log(file.url)
dialogImageUrl.value = file.url!
dialogVisible.value = true
}

const handleRemove = (file: UploadFile) => {
console.log(file)
}

const handleDownload = (file: UploadFile) => {
console.log(file)
}

const getData = async () => {
console.log(1)
const message = await axios
  .post('/api/students', {
    img: image_base64.value
  })
  .then((res) => {
    //   tableData.value = res.data
    // res.data是base64编码。再将这个转成图片。
    resultUrl.value = res.data
    console.log(res)
  })
  .catch((err) => {
    console.log(err)
  })
}
</script>

<style lang="less" scoped>
.main-card {
display: flex;
justify-content: space-around;
.sourceImage img {
  height: 300px;
  width: 300px;
  background-color: palegoldenrod;
}
}
</style>

使用vue+element-plus来搭建前端页面,使用el-upload组件上传图片,其次将图片转为base64使用axios进行请求。

2.后端处理

使用Flask作为后端,简单方便,接收到前端传过来的base64,再转为图片,输入到YOLOv5模型当中进行识别,然后会得到识别完成的base64图像,再返回给前端。
使用yolov5s.pt模型。直接使用torch.hub.load()进行模型加载。从本地加载可以避免由于github下载导致的问题。图像传递的过程为web->flask。传递方式都为base64。

# 加载模型,输入图片进行detect,返回图片的base64编码。
import base64

import torch
import os
from PIL import Image
import cv2
from io import BytesIO, StringIO


def getbase64(img):
    # os.environ["CUDA_VISIBLE_DEVICES"]
    model = torch.hub.load('./', 'custom', path='./yolov5s.pt', source='local')
    # model.cuda()
    # print(model)
    img1 = Image.open(BytesIO(img))
    # img2 = cv2.imread('./data/images/zidane.jpg')[:,:,::-1]
    imgs = [img1]

    results = model(imgs, size=640)
    print('识别成功')
    results.imgs
    results.render()
    base64_image = ''
    for img in results.imgs:
        buffered = BytesIO()
        img_base64 = Image.fromarray(img)
        img_base64.save(buffered, format='JPEG')
        base64_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
        # with open('./result.txt', 'a+') as f:
        #     f.write(base64_image)
        #     f.close()
    print(results)
    results.save()
    base64_image = 'data:image/jpeg;base64,%s' % base64_image
    return base64_image
    # print(results.xyxy[0])
    # print(results.pandas().xyxy[0])

最后将编码返回前端,前端:src="检测完成的base64编码",就能将图像显示出来。


# -*- coding: utf-8 -*-
# @Time    : 2022/9/3
# @Author  : rickHan
# @Software: PyCharm
# @File    : app.py
import base64
import os
import json
from PIL import Image
import cv2
from io import BytesIO, StringIO
from flask import Flask, request
from yolov5 import getbase64

app = Flask(__name__)

@app.route('/')
def hello_world():
    return 'hello World'

@app.route('/students', methods=['POST'])
def get_students():
    # getbase64(img)
    if request.method == 'POST':
        base64_data = request.json['img']
        print('已经收到了')
        head,context = base64_data.split(',')
        img_data = base64.b64decode(context)
        image =getbase64(img_data)
        print('success')
        return image
        # image.show()
    # img = ''
    # detect_image = getbase64(img)
    # return detect_image
    return 'false'
if __name__ == '__main__':
    app.run()

识别速度上还是可以的,几秒钟就能识别出来了。
github:https://github.com/itouchzh/yolov5_flask_vue