⚡️ Load your YOLO v5 or v8 model in browser

Run object detection models trained with YOLOv5 YOLOv8 in browser using tensorflow.js

Demo

check out a demo of Aquarium Dataset object detection

Install

Yarn

yarn add yolo-tfjs

Or NPM

npm install yolo-tfjs

Usage Example

import YOLOTf from "yolo-tfjs";

const CLASSES = ["fish", "jellyfish"]
const COLORS = ["#00C2FF", "#FF9D97"]
const imageRef = useRef<HTMLImageElement>(null)

// load model files
const yoloTf = await YOLOTf.loadYoloModel(`model_path/model.json`, CLASSES, {
    yoloVersion: 'v8', onProgress(fraction: number){
        console.log('loading model...')
    }})

// return dection results with detected boxes
const results = await yoloTf.predict(imageRef.current)

// draw boxes in the canvas element
yoloTf.renderBox(canvasRef.current, {
    ...results, ratio: [results["xRatio"],results["yRatio"]]
}, COLORS)

API Docs

loadYoloModel(model, classes, config): YOLOTf

Args

Param | Type | Description -- | -- | -- model | string | path to model.json file classes | string[] | classes of the trained model config | Object | see below model configuration

Config | Type | Default | Description -- | -- | -- | -- | [options.scoreThreshold] | Number | 0.5 | | | [options.iouThreshold] | Number | 0.45 | | | [options.maxOutputSize] | Number | 500 | | | [options.onProgress] | Callback | (fraction: number) => void | | | [options.yoloVersion] | YoloVersion | _ | selected version v5 or v8 |

YOLOTf

PredictionData: {boxes, classes, scores, xRatio, yRatio}

predict(image, preprocessImage): PredictionData

Param | Type | Description -- | -- | -- image | HTMLImageElement | preprocessImage | (image: HTMLImageElement) => PreprocessResult | this optional param to use custom image preprocessing

renderBox(canvas, predictionData, colors): {boxes, classes, scores, xRatio, yRatio}