Nhận diện vật thể chạy thẳng trong trình duyệt

Tôi tự chụp và gán nhãn một bộ ảnh đồ vật quanh bàn làm việc, train YOLOv5s trên đó, rồi mang model chạy hoàn toàn trong trình duyệt. Repo chia làm ba phần — train, yolov5-litert, yolov5-onnxruntime-web — vì mục tiêu thật sự là so hai runtime web trên cùng một model.

Dataset tự gán nhãn#

Tổng cộng 521 ảnh chụp bằng điện thoại, gán nhãn tay theo format YOLO cho sáu class: ban phim, but long, cay keo, con chuot, doi dua, tai nghe. Một script nhỏ xáo trộn danh sách với random.seed(42) rồi chia 80/20 thành 416 ảnh train và 105 ảnh val.

Dữ liệu tự làm thì luôn có rác. Lúc chạy val.py, YOLOv5 báo tám ảnh JPEG hỏng phải sửa lại và loại hẳn một ảnh vì toạ độ nhãn vượt ra ngoài biên (1,0912 thay vì ≤ 1) — nên phần đánh giá thực tế chỉ còn 104 ảnh với 270 đối tượng.

Train và export ra hai format#

Train từ checkpoint yolov5s.pt trên Colab với một Tesla T4: ảnh 640, batch 8, tối đa 300 epoch và --patience 30 để tự dừng khi không còn cải thiện.

!python train.py --img 640 --batch 8 --epochs 300 \
--data ../data.yml --weights yolov5s.pt \
--cache images --patience 30
 
!python export.py --weights runs/train/exp/weights/best.pt --include tflite --device cpu
 
# tham số simplify để gộp các operator thừa, giúp file gọn nhẹ tối ưu cho onnx runtime web
!python export.py --weights runs/train/exp/weights/best.pt --include onnx --simplify

Trên tập val, model đạt precision 0,810, recall 0,827, mAP50 0,843 và mAP50-95 0,508. Class cay keo tốt nhất với mAP50 0,910, but long kém nhất với 0,724 — hợp lý, vì bút mảnh và dễ lẫn vào nền bàn.

LiteRT.js bắt tôi tự chuyển tensor giữa các thiết bị#

Bản LiteRT khởi tạo theo hai bước: nạp runtime WASM trước, rồi mới biên dịch model với accelerator webgpu. Khác biệt lớn nhất so với ONNX là tensor không tự nằm đúng chỗ — tôi phải đẩy nó sang WebGPU trước khi chạy và kéo về WASM sau khi chạy.

await loadLiteRt(WASM_PATH);
const model = await loadAndCompile(MODEL_PATH, { accelerator: "webgpu" });
 
// TFLite nhận layout NHWC, không phải NCHW như bản ONNX.
const tensor = await new Tensor(inputData, [1, IMG_SIZE, IMG_SIZE, 3]).moveTo("webgpu");
const output = model.run([tensor])[0];
 
// Kết quả nằm trên GPU; phải kéo về WASM thì JS mới đọc được.
const cpuOutput = await output.moveTo("wasm");
const rawBoxes = parseYoloOutput(cpuOutput.toTypedArray(), classes);

Toàn bộ preprocessing ở đây là JavaScript thuần: vẽ frame vào canvas 640×640, gọi getImageData, rồi lặp tay để bỏ kênh alpha và chia 255. Không phụ thuộc thư viện ngoài nào, nhưng cũng không có letterbox — ảnh bị kéo giãn thẳng về hình vuông.

ONNX Runtime Web nặng hơn nhưng preprocessing đúng hơn#

Bản ONNX chạy trên execution provider wasm với simd và threads bật, kèm một lượt warmup bằng tensor rỗng để lần inference thật đầu tiên không bị khựng. File .onnx sau --simplify nặng 27,2 MB, nên tôi tải bằng XMLHttpRequest có onprogress để hiện phần trăm thay vì để người dùng nhìn màn hình trắng. Craco cũng phải thêm copy-webpack-plugin để đưa các file .wasm của onnxruntime-web ra thư mục build.

Đổi lại, preprocessing tốt hơn hẳn: OpenCV.js dùng copyMakeBorder chèn viền cho ảnh thành hình vuông rồi blobFromImage mới resize về 640×640 và chia 255. Tỉ lệ ảnh gốc được giữ nguyên, và tỉ lệ pad đó được dùng lại để quy box về toạ độ ảnh thật.

NMS vẫn nằm ở JavaScript#

Cả hai bản đều export không kèm NMS, nên output là tensor thô (1, 25200, 11) và cả hai phía web đều phải tự duyệt 25.200 anchor rồi chạy NMS viết tay. Ngưỡng hai bên cũng lệch: LiteRT lọc ở 0,45 và NMS theo từng class, ONNX lọc ở 0,25 với ngưỡng class 0,2, NMS không phân biệt class và giới hạn topk 100. Ở chế độ camera, tôi chặn nhịp inference ở 250 ms để vòng requestAnimationFrame không xếp chồng lên nhau.

Kết quả#

• Model sáu class đạt mAP50 0,843 trên 104 ảnh val — đủ dùng cho demo, nhưng vẫn là con số của một dataset 521 ảnh
• LiteRT.js gọn hơn về phụ thuộc: chỉ @litertjs/core, đi thẳng WebGPU, không cần cấu hình bundler — đổi lại tôi phải tự quản lý vị trí tensor và tự viết mọi bước preprocessing
• ONNX Runtime Web ít rủi ro hơn vì chỉ dựa vào WASM và có OpenCV.js lo phần letterbox, nhưng kéo theo 27,2 MB model phải tải về cộng cấu hình webpack cho các file .wasm
• Tôi chưa đo thời gian inference trong trình duyệt ở cả hai bản; số duy nhất tôi có là 9,0 ms inference và 5,3 ms NMS mỗi ảnh trên T4, nên chưa thể kết luận bên nào nhanh hơn trên web