YOLOV5 使用API方式用在FLASK網頁上

yolo yolov5 python html flask

neil410912021 2023-05-21 00:33:02 ‧ 1043 瀏覽

分享至

之前還沒有很瞭接怎麼把YOLO偵測的東西用在網頁上，想問一下現在這樣的架構是否是對的呢?

app.py

from flask import Flask, jsonify, request, render_template, send_file
import torch
import numpy as np
import cv2
import io

app = Flask(__name__)


model = torch.hub.load('ultralytics/yolov5', 'yolov5s')


@app.route('/api/yolo', methods=['POST'])
def yolo_api():
    # 獲取圖像資料
    image_file = request.files['image']
    image = cv2.imdecode(np.frombuffer(image_file.read(), np.uint8), cv2.IMREAD_COLOR)
    image = cv2.resize(image, (800, 480))

    # 執行模型
    results = model(image)

    # 獲取檢測結果的圖像
    output_image = np.squeeze(results.render())

    # 將圖像轉換為位元組流
    _, image_stream = cv2.imencode('.jpg', output_image)

    # 建立記憶體檔案物件
    output_stream = io.BytesIO(image_stream.tobytes())

    # 將圖像檔案傳送給模板進行顯示
    return send_file(output_stream, mimetype='image/jpeg')

@app.route('/')
def index():
    return render_template('index.html')
    

if __name__ == '__main__':
    app.run(debug=True)

index.html

<!DOCTYPE html>
<html>
<head>
    <title>YOLO API</title>
    <style>
        body {
            font-family: Arial, sans-serif;
        }
        h1 {
            color: #333;
        }
    </style>
</head>
<body>
    <h1>來偵測囉</h1>
    <form action="/api/yolo" method="post" enctype="multipart/form-data">
        <input type="file" name="image">
        <button type="submit">Submit</button>
    </form>

</body>
</html>