將 YOLO 帶到你的瀏覽器：如何快速部署YOLO模型

大貓咪收錄於 Javascript

2024-05-29 2024-05-29 約 2221 字預計閱讀 10 分鐘次閱讀 CC BY-NC 4.0

根據官方教學，你訓練好了 YOLO 模型

但是你要怎麼把模型部屬到網頁呢？

本篇就以 YOLOv8 來教你如何部屬 YOLO 模型

廣告 AD

YOLOv8

Github - ultralytics

本篇所使用的模型為 YOLOv8，負責找出驗證碼圖片中的數字，並辨別該數字。

模型都已經訓練過了，本篇著重在如何將模型部屬到網頁上進行推理。

有關於 YOLOv8 的訓練方法和使用教學請至 Github - ultralytics。

Captcha

我們這次要解決的問題是辨識驗證碼，驗證碼大概長得像下面這樣，圖片中的數字就是我們要辨識的目標，也就是 8402。圖片的原始長寬分別為 60 和 160。

Deploy

現在就來匯出我們訓練好的模型，這邊我習慣加上 simplify 的 flag 來簡化模型，根據 document 的說法這可能會增加效能和相容性。我們先讀取我們要輸出的模型，然後輸出成 ONNX 的格式，我們這篇將以 onnxruntime 的後端來跑模型。輸出成 ONNX 格式後，可以用 netron 來看整個模型的架構。

Python

from ultralytics import YOLO

model = YOLO('path/to/best.pt')
success = model.export(format='onnx', simplify=True)

有了模型後，接下來部屬時我們會需要先前處理圖片再跑模型，最後處理我們模型輸出的資料。

Load Image

首先我們來讀取圖片，我們有一個 id 為 input-field 的 input 元素，可以讓使用者選擇圖片，再來我們透過 FileReader 來讀取我們的檔案，並將檔案讀取成 Array Buffer，最後我們的圖片就會存在 img_buf 中。

const input_field = document.getElementById('input-field');
input_field.addEventListener('change', (e) => {
    const reader = new FileReader();
    reader.onload = async function () {
        const img_buf = reader.result;
        // ... //
    };
    reader.readAsArrayBuffer(e.target.files[0]);
});

Jimp

Github - Jimp

Jimp (JavaScript Image Manipulation Program) 是一個純 JS 的圖片處理函式庫，且沒有依賴其他任何第三方的函示庫。我們這次就使用這個 Library 來處理我們的圖片。

要使用 Jimp 之前要 include 檔案，放在 html 的 head 裡面。

html

<script 
  src="https://cdn.jsdelivr.net/npm/jimp/browser/lib/jimp.min.js">
</script>

透過 Jimp 讀取 image buffer 後，我們要調整圖片大小，因為模型的輸入為 160 x 160，因此我們將最長邊調整成 160，並對圖片做 padding 到 160 x 160，全部都用黑色來補。

const img = await Jimp.read(img_buf).then(image => {
    // 調整圖片大小
    if (image.getWidth() > image.getHeight()) {
        image.resize(160, Jimp.AUTO);
    } else {
        image.resize(Jimp.AUTO, 160);
    }

    // 塞入黑邊
    let bg_image = new Jimp(160, 160, '#000000', (err, image) => {
        if (err) throw err
    });
    bg_image.blit(image, 0, 0);

    return bg_image;
});

ONNX Runtime

npm - onnxruntime-web

處理好輸入的圖片後，接下來就可以跑模型了，我們要使用 onnxruntime-web 來在瀏覽器上跑 ONNX 模型。

html

<script 
  src="https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js">
</script>

調整 channel 和將圖片轉成 tensor 格式，首先我們的圖片格式為 [W, H, 4]，分別為寬高和 RGB 和透明度的 alpha 通道，但我們的輸入格式為 [3, W, H]，也就是 RGB 和寬高，並且資料要是一維的 Float32Array，因此我們先取得圖片 RGB 三個通道的數值，然後將資料攤平，並將資料轉成 float32，最後就可以轉成我們的 tensor 餵給模型了。

// 取得圖片的 RGB 通道
var imageBufferData = img.bitmap.data;
const [redArray, greenArray, blueArray] = new Array(new Array(), new Array(), new Array());
for (let i = 0; i < imageBufferData.length; i += 4) {
    redArray.push(imageBufferData[i]);
    greenArray.push(imageBufferData[i + 1]);
    blueArray.push(imageBufferData[i + 2]);
    // 跳過管控透明度的 alpha 通道=
}

// 轉成 RGB、寬、高
const transposedData = redArray.concat(greenArray).concat(blueArray);

// 轉成 float32
let i, l = transposedData.length;
const float32Data = new Float32Array(l);
for (i = 0; i < l; i++) {
    float32Data[i] = transposedData[i] / 255.0;
}

// 轉成 tensor 物件，shape 為 1x3x160x160
const inputTensor = new ort.Tensor("float32", float32Data, [1, 3, 160, 160]);

跑模型前讀取模型：

const session = await ort.InferenceSession.create(model_path);

跑模型得到結果：

// 設定輸入的名稱和資料
const feeds = {};
feeds[session.inputNames[0]] = inputTensor;
// 跑模型
const outputData = await session.run(feeds);
// 模型輸出結果
const output = outputData[session.outputNames[0]];

NMS

模型的輸出依然不是我們想要的答案，我們要對輸出做 NMS (Non-Maximum Suppression)，因為輸出的結果會抓到很多的物件，但這些物件又重疊在一起，每個信賴分數又不一樣，到底要選擇哪些的物件，於是我們要跑 NMS 的算法，找出那些信心分數較高的物件出來。

模型的輸出大小為 1x14*525，14 為 10 種 Label (數字 0 ~ 9) 加上 bounding box 的 4 項數據 (X 中心、Y 中心、寬度、高度)，525 為偵測的物件的數量。

我們會先轉換 bounding box 的格式變成記錄左上角的點的 XY 和右下角的點的 XY，方便後續算 IOU，另外也會過濾 confidence score 低於閥值的物件，一般設定是用 0.25。每一物件的 confidence score 的計算方式為該物件中，所有 label 下的最大值，且擁有最大 confidence score 的 label 就是模型判定該物件為該 label。

接著我們會跑 NMS，步驟如下：

將目前剩餘的物件(沒有被過濾掉的)加入候選物件
找出所有候選物件中 confidence score 最大的物件 (L)
將此物件 L 加入選定集合中
將除了物件 L 之外的其他候選物件計算與物件 L 的 IOU
移除 IOU 數值大於閥值 (預設為 0.5) 的候選物件
如果還有候選物件則回到步驟 1，否則選定集合為最終輸出結果

const conf_threshold = 0.25;
const iou_threshold = 0.5;

const dims = output.dims;
const rows = dims[2];
const clsNum = dims[1] - 4;
const rawData = output.data;

// 轉換 bounding box 的格式 
// (centerX, centerY, width, height) 
// --> 
// (left-top X, left-top Y, right-bottom X, right-bottom Y)
let data = [];
for (let i = 0; i < rows; ++i) {
    let arr = new Float32Array(clsNum + 4);
    for (let j = 0; j < clsNum + 4; ++j) {
        arr[j] = rawData[(j * rows) + i];
    }
    const center_x = arr[0], center_y = arr[1];
    const half_x = (arr[2] / 2), half_y = (arr[3] / 2);
    arr[0] = center_x - half_x;
    arr[1] = center_y - half_y;
    arr[2] = center_x + half_x;
    arr[3] = center_y + half_y;
    data.push(arr);
}

// 過濾 confidence score 低於閥值的物件，
let conf_arr = new Float32Array(rows);
let label_arr = new Uint8Array(rows);
let index = new Set();
for (let i = 0; i < rows; i++) {
    let conf = -1, idx;
    for (let j = 0; j < clsNum; ++j) {
        if (conf < data[i][4 + j]) {
            conf = data[i][4 + j];
            idx = j;
        }
    }
    if (conf > conf_threshold) {
        index.add(i);
    }
    conf_arr[i] = conf;
    label_arr[i] = idx;
}

// NMS
let select_idx = []
while (index.size > 0) {
    // 步驟 1：找出擁有最大 confidence score 的物件
    let max_conf = -1;
    let max_conf_idx = 0;
    for (let idx of index) {
        let conf = conf_arr[idx];
        if (max_conf < conf) {
            max_conf = conf;
            max_conf_idx = idx;
        }
    }

    // 步驟 2：加入選定集合
    select_idx.push(max_conf_idx)

    // 步驟 3：計算 IOU
    const select_width = (data[max_conf_idx][2] - data[max_conf_idx][0]);
    const select_height = (data[max_conf_idx][3] - data[max_conf_idx][1]);
    const select_area = select_width * select_height;
    let new_index = new Set();
    for (let idx of index) {
        const width = (data[idx][2] - data[idx][0]);
        const height = (data[idx][3] - data[idx][1]);
        const area = width * height;

        // 重疊面積
        const overlap_left = Math.max(data[idx][0], data[max_conf_idx][0]);
        const overlap_right = Math.min(data[idx][2], data[max_conf_idx][2]);
        const overlap_top = Math.max(data[idx][1], data[max_conf_idx][1]);
        const overlap_bottom = Math.min(data[idx][3], data[max_conf_idx][3]);
        const overlap_width = Math.max(overlap_right - overlap_left, 0);
        const overlap_height = Math.max(overlap_bottom - overlap_top, 0);
        const overlap_area = overlap_width * overlap_height;

        // IOU
        const iou = (overlap_area) / (select_area + area - overlap_area);

        // 步驟 4：根據 IOU 過濾物件
        if (iou < iou_threshold) {
            new_index.add(idx);
        }
    }

    // 更新候選名單
    index = new_index;
}

// 最終結果
let confidence = new Float32Array(select_idx.length);
let label = new Uint8Array(select_idx.length);
let box = new Float32Array(select_idx.length * 4);
for (let i = 0; i < select_idx.length; ++i) {
    const idx = select_idx[i];
    confidence[i] = conf_arr[idx];
    label[i] = label_arr[idx];
    box[i * 4 + 0] = data[idx][0];
    box[i * 4 + 1] = data[idx][1];
    box[i * 4 + 2] = data[idx][2];
    box[i * 4 + 3] = data[idx][3];
}

Combine objects

做完 NMS 找到物件後，我們要來找出驗證碼的答案了！由於驗證碼我們都是由左至右閱讀的，因此我們只要將物件根據 X 座標的數值來排列，並組合他們的 label 即可。

// 根據 X 座標的數值排列
let index = [];
for (let i = 0; i < label.length; ++i) {
    index.push(i);
}
index.sort((a, b) => box[a * 4] - box[b * 4]);

// 最終結果
let result = '';
for (let i = 0; i < index.length; ++i) {
    result += String(label[index[i]]);
}