diff --git a/README.md b/README.md
index 48f4fc8..ac78618 100644
--- a/README.md
+++ b/README.md
@@ -8,16 +8,39 @@
-## 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 @@
-## 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);
+ });
+ }
}