feature/005-api-refinement (#6)

* initial api update * update package version * paths refer to yolov5 zoo, initial documentation
SkalskiP · Sep 29, 2022 · 9cff747 · 9cff747
1 parent 45bb78f
commit 9cff747
Show file tree

Hide file tree

Showing 9 changed files with 82 additions and 32 deletions.
diff --git a/README.md b/README.md
@@ -8,16 +8,39 @@
     <img width="200" src="https://yolov5js-images.s3.eu-central-1.amazonaws.com/yolov5js-logo.png" alt="Logo">
 </p>
 
-## <div align="center">Installation</div>
+## <div align="center">Install</div>
 
-```console
+```bash
 npm install --save yolov5js
 ```
 
-## <div align="center">Deployment</div>
+## <div align="center">Convert</div>
+
+```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
+```
+
+## <div align="center">Zoo</div>
+
+Use and share pretrained YOLOv5 tensorflow.js models with [yolov5.js-zoo](https://github.com/SkalskiP/yolov5js-zoo).
+
+## <div align="center">Deploy</div>
 
 <details open>
-<summary>Fetch from models library</summary>
+<summary>Fetch from models zoo</summary>
 
 ```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://raw.githubusercontent.com/SkalskiP/yolov5js/master/models/yolov5n/model.json' }
+const config = { source: 'https://raw.githubusercontent.com/SkalskiP/yolov5js-zoo/master/models/coco/yolov5n/model.json' }
 const model = await load(config)
 ```
 
@@ -54,6 +77,11 @@ const model = await load(config)
 
 </details>
 
+## <div align="center">Kudos</div>
+
+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.
+
+
 ## <div align="center">License</div>
 
 Project is freely distributable under the terms of the [MIT license](LICENSE).
diff --git a/docs/api.md b/docs/api.md
diff --git a/docs/index.md b/docs/index.md
@@ -4,8 +4,18 @@
     <img width="120" src="https://yolov5js-images.s3.eu-central-1.amazonaws.com/yolov5js-logo-black.png" alt="logo">
 </p>
 
-## <div align="center">Installation</div>
+## <div align="center">Install</div>
 
-```console
+```bash
 npm install --save yolov5js
 ```
+
+## <div align="center">Deploy & Predict</div>
+
+```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
@@ -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
@@ -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
diff --git 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
@@ -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
@@ -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://raw.githubusercontent.com/SkalskiP/ILearnMachineLearning.js/master/models/yolov5n/model.json',
+    source: 'https://raw.githubusercontent.com/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://raw.githubusercontent.com/SkalskiP/ILearnMachineLearning.js/master/models/yolov5s/model.json',
+    source: 'https://raw.githubusercontent.com/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://raw.githubusercontent.com/SkalskiP/ILearnMachineLearning.js/master/models/yolov5m/model.json',
+    source: 'https://raw.githubusercontent.com/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<YOLOv5> {
-    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);
+        });
+    }
 }