diff --git a/README.md b/README.md index 48f4fc8..ac78618 100644 --- a/README.md +++ b/README.md @@ -8,16 +8,39 @@ Logo

-##
Installation
+##
Install
-```console +```bash npm install --save yolov5js ``` -##
Deployment
+##
Convert
+ +```bash +# clone YOLOv5 repository +git clone https://github.com/ultralytics/yolov5.git +cd yolov5 + +# create python virtual environment [recommended] +virtualenv venv +source venv/bin/activate + +# install dependencies +pip install -r requirements.txt +pip install tensorflowjs + +# convert model to tensorflow.js format +python export.py --weights yolov5s.pt --include tfjs +``` + +##
Zoo
+ +Use and share pretrained YOLOv5 tensorflow.js models with [yolov5.js-zoo](https://github.com/SkalskiP/yolov5js-zoo). + +##
Deploy
-Fetch from models library +Fetch from models zoo ```javascript import {load, YOLO_V5_N_COCO_MODEL_CONFIG} from 'yolov5js' @@ -33,7 +56,7 @@ const model = await load(YOLO_V5_N_COCO_MODEL_CONFIG) ```javascript import {load, ModelConfig} from 'yolov5js' -const config = { source: 'https://github.com/raw/SkalskiP/yolov5js/master/models/yolov5n/model.json' } +const config = { source: 'https://github.com/raw/SkalskiP/yolov5js-zoo/master/models/coco/yolov5n/model.json' } const model = await load(config) ``` @@ -54,6 +77,11 @@ const model = await load(config)
+##
Kudos
+ +Kudos to [ultralytics](https://ultralytics.com/) team as well as all other open-source contributors for building [YOLOv5](https://github.com/ultralytics/yolov5) project, and making it all possible. + + ##
License
Project is freely distributable under the terms of the [MIT license](LICENSE). diff --git a/docs/api.md b/docs/api.md deleted file mode 100644 index e69de29..0000000 diff --git a/docs/index.md b/docs/index.md index 9c561af..ff9c2f6 100644 --- a/docs/index.md +++ b/docs/index.md @@ -4,8 +4,18 @@ logo

-##
Installation
+##
Install
-```console +```bash npm install --save yolov5js ``` + +##
Deploy & Predict
+ +```javascript +import {load, YOLO_V5_N_COCO_MODEL_CONFIG} from 'yolov5js' + +const model = await load(YOLO_V5_N_COCO_MODEL_CONFIG) + +const detections = await model.detect(image) +``` diff --git a/docs/weights.md b/docs/weights.md index e95b867..2e9ecf8 100644 --- a/docs/weights.md +++ b/docs/weights.md @@ -1,9 +1,18 @@ -```console +## Convert + +```bash +# clone YOLOv5 repository git clone https://github.com/ultralytics/yolov5.git cd yolov5 + +# create python virtual environment [recommended] virtualenv venv source venv/bin/activate + +# install dependencies pip install -r requirements.txt pip install tensorflowjs + +# convert model to tensorflow.js format python export.py --weights yolov5s.pt --include tfjs ``` diff --git a/mkdocs.yml b/mkdocs.yml index 413914b..5168da1 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -26,8 +26,7 @@ extra: # Page tree nav: - Overview: index.md - - Docs: api.md - - Custom weights: weights.md + - Convert: weights.md # Configuration theme: diff --git a/package-lock.json b/package-lock.json index e62c415..92df390 100644 --- a/package-lock.json +++ b/package-lock.json @@ -1,5 +1,5 @@ { - "name": "yolov5-js", + "name": "yolov5js", "version": "0.2.0", "lockfileVersion": 1, "requires": true, diff --git a/package.json b/package.json index 46bf87a..3c330c0 100644 --- a/package.json +++ b/package.json @@ -1,5 +1,5 @@ { - "version": "0.2.0", + "version": "1.0.0", "license": "MIT", "main": "dist/index.js", "typings": "dist/index.d.ts", diff --git a/src/index.ts b/src/index.ts index 234767e..fbc5566 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,5 +1,6 @@ export { DetectedObject, + ModelConfig, YOLOv5, load, YOLO_V5_N_COCO_MODEL_CONFIG, diff --git a/src/yolov5.ts b/src/yolov5.ts index 9f0a216..171eccb 100644 --- a/src/yolov5.ts +++ b/src/yolov5.ts @@ -13,29 +13,23 @@ const COCO_NAMES = ['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', const INFERENCE_RESOLUTION: [number, number] = [640, 640]; -export declare type ObjectDetectionBaseModel = 'yolov5n' | 'yolov5s' | 'yolov5m'; - export interface ModelConfig { - base: ObjectDetectionBaseModel; - modelUrl: string; - classNames: string[]; + source: string | File[]; + classNames?: string[]; } export const YOLO_V5_N_COCO_MODEL_CONFIG: ModelConfig = { - base: 'yolov5n', - modelUrl: 'https://github.com/raw/SkalskiP/ILearnMachineLearning.js/master/models/yolov5n/model.json', + source: 'https://github.com/raw/SkalskiP/yolov5js-zoo/master/models/coco/yolov5n/model.json', classNames: COCO_NAMES } export const YOLO_V5_S_COCO_MODEL_CONFIG: ModelConfig = { - base: 'yolov5s', - modelUrl: 'https://github.com/raw/SkalskiP/ILearnMachineLearning.js/master/models/yolov5s/model.json', + source: 'https://github.com/raw/SkalskiP/yolov5js-zoo/master/models/coco/yolov5s/model.json', classNames: COCO_NAMES } export const YOLO_V5_M_COCO_MODEL_CONFIG: ModelConfig = { - base: 'yolov5m', - modelUrl: 'https://github.com/raw/SkalskiP/ILearnMachineLearning.js/master/models/yolov5m/model.json', + source: 'https://github.com/raw/SkalskiP/yolov5js-zoo/master/models/coco/yolov5m/model.json', classNames: COCO_NAMES } @@ -45,15 +39,16 @@ export interface DetectedObject { width: number; height: number; score: number; - class: string; + classId: number; + class?: string; } export class YOLOv5 { public model: GraphModel; public inferenceResolution: [number, number]; - public classNames: string[]; + public classNames?: string[]; - constructor(model: GraphModel, inferenceResolution: [number, number], classNames: string[]) { + constructor(model: GraphModel, inferenceResolution: [number, number], classNames?: string[]) { this.model = model; this.inferenceResolution = inferenceResolution; this.classNames = classNames; @@ -77,7 +72,7 @@ export class YOLOv5 { scores: Float32Array, classes: Float32Array, inputResolution: [number, number], - classNames: string[], + classNames?: string[], minScore?: number ): DetectedObject[] { const scoreThreshold: number = minScore !== undefined ? minScore : 0; @@ -98,14 +93,16 @@ export class YOLOv5 { const maxY = bbox[3] * inputHeight; const width = maxX - minX; const height = maxY - minY; - const className = classNames[classes[i]]; + const classId = classes[i]; + const className = classNames !== undefined ? classNames[classes[i]] : undefined; detections.push({ x: minX, y: minY, width: width, height: height, + score: score, + classId: classId, class: className, - score: score }); } return detections; @@ -128,7 +125,13 @@ export class YOLOv5 { } export async function load(config: ModelConfig, inputResolution: [number, number] = INFERENCE_RESOLUTION): Promise { - return tf.loadGraphModel(config.modelUrl).then((model: GraphModel) => { - return new YOLOv5(model, inputResolution, config.classNames); - }); + if (typeof config.source === 'string') { + return tf.loadGraphModel(config.source).then((model: GraphModel) => { + return new YOLOv5(model, inputResolution, config.classNames); + }); + } else { + return tf.loadGraphModel(tf.io.browserFiles(config.source)).then((model: GraphModel) => { + return new YOLOv5(model, inputResolution, config.classNames); + }); + } }