pose-update (#2)

- add pose model (92dc53a37b584fa311fcb5fad589ef012ddf2ca3)


Co-authored-by: Radamés Ajna <radames@users.noreply.huggingface.co>

build/m.d.ts

@@ -17,6 +17,23 @@ export class Model {
 */
   run(image: Uint8Array, conf_threshold: number, iou_threshold: number): string;
 }
+/**
+*/
+export class ModelPose {
+  free(): void;
+/**
+* @param {Uint8Array} data
+* @param {string} model_size
+*/
+  constructor(data: Uint8Array, model_size: string);
+/**
+* @param {Uint8Array} image
+* @param {number} conf_threshold
+* @param {number} iou_threshold
+* @returns {string}
+*/
+  run(image: Uint8Array, conf_threshold: number, iou_threshold: number): string;
+}
 
 export type InitInput = RequestInfo | URL | Response | BufferSource | WebAssembly.Module;
 
@@ -25,12 +42,14 @@ export interface InitOutput {
   readonly __wbg_model_free: (a: number) => void;
   readonly model_new: (a: number, b: number, c: number, d: number, e: number) => void;
   readonly model_run: (a: number, b: number, c: number, d: number, e: number, f: number) => void;
+  readonly __wbg_modelpose_free: (a: number) => void;
+  readonly modelpose_new: (a: number, b: number, c: number, d: number, e: number) => void;
+  readonly modelpose_run: (a: number, b: number, c: number, d: number, e: number, f: number) => void;
   readonly main: (a: number, b: number) => number;
   readonly __wbindgen_add_to_stack_pointer: (a: number) => number;
   readonly __wbindgen_malloc: (a: number, b: number) => number;
   readonly __wbindgen_realloc: (a: number, b: number, c: number, d: number) => number;
   readonly __wbindgen_free: (a: number, b: number, c: number) => void;
-  readonly __wbindgen_exn_store: (a: number) => void;
   readonly __wbindgen_start: () => void;
 }
 

build/m.js

@@ -33,20 +33,6 @@ function addHeapObject(obj) {
     return idx;
 }
 
-function getObject(idx) { return heap[idx]; }
-
-function dropObject(idx) {
-    if (idx < 132) return;
-    heap[idx] = heap_next;
-    heap_next = idx;
-}
-
-function takeObject(idx) {
-    const ret = getObject(idx);
-    dropObject(idx);
-    return ret;
-}
-
 let WASM_VECTOR_LEN = 0;
 
 function passArray8ToWasm0(arg, malloc) {
@@ -118,12 +104,18 @@ function getInt32Memory0() {
     return cachedInt32Memory0;
 }
 
-function handleError(f, args) {
-    try {
-        return f.apply(this, args);
-    } catch (e) {
-        wasm.__wbindgen_exn_store(addHeapObject(e));
-    }
+function getObject(idx) { return heap[idx]; }
+
+function dropObject(idx) {
+    if (idx < 132) return;
+    heap[idx] = heap_next;
+    heap_next = idx;
+}
+
+function takeObject(idx) {
+    const ret = getObject(idx);
+    dropObject(idx);
+    return ret;
 }
 /**
 */
@@ -204,6 +196,85 @@ export class Model {
         }
     }
 }
+/**
+*/
+export class ModelPose {
+
+    static __wrap(ptr) {
+        ptr = ptr >>> 0;
+        const obj = Object.create(ModelPose.prototype);
+        obj.__wbg_ptr = ptr;
+
+        return obj;
+    }
+
+    __destroy_into_raw() {
+        const ptr = this.__wbg_ptr;
+        this.__wbg_ptr = 0;
+
+        return ptr;
+    }
+
+    free() {
+        const ptr = this.__destroy_into_raw();
+        wasm.__wbg_modelpose_free(ptr);
+    }
+    /**
+    * @param {Uint8Array} data
+    * @param {string} model_size
+    */
+    constructor(data, model_size) {
+        try {
+            const retptr = wasm.__wbindgen_add_to_stack_pointer(-16);
+            const ptr0 = passArray8ToWasm0(data, wasm.__wbindgen_malloc);
+            const len0 = WASM_VECTOR_LEN;
+            const ptr1 = passStringToWasm0(model_size, wasm.__wbindgen_malloc, wasm.__wbindgen_realloc);
+            const len1 = WASM_VECTOR_LEN;
+            wasm.modelpose_new(retptr, ptr0, len0, ptr1, len1);
+            var r0 = getInt32Memory0()[retptr / 4 + 0];
+            var r1 = getInt32Memory0()[retptr / 4 + 1];
+            var r2 = getInt32Memory0()[retptr / 4 + 2];
+            if (r2) {
+                throw takeObject(r1);
+            }
+            return ModelPose.__wrap(r0);
+        } finally {
+            wasm.__wbindgen_add_to_stack_pointer(16);
+        }
+    }
+    /**
+    * @param {Uint8Array} image
+    * @param {number} conf_threshold
+    * @param {number} iou_threshold
+    * @returns {string}
+    */
+    run(image, conf_threshold, iou_threshold) {
+        let deferred3_0;
+        let deferred3_1;
+        try {
+            const retptr = wasm.__wbindgen_add_to_stack_pointer(-16);
+            const ptr0 = passArray8ToWasm0(image, wasm.__wbindgen_malloc);
+            const len0 = WASM_VECTOR_LEN;
+            wasm.modelpose_run(retptr, this.__wbg_ptr, ptr0, len0, conf_threshold, iou_threshold);
+            var r0 = getInt32Memory0()[retptr / 4 + 0];
+            var r1 = getInt32Memory0()[retptr / 4 + 1];
+            var r2 = getInt32Memory0()[retptr / 4 + 2];
+            var r3 = getInt32Memory0()[retptr / 4 + 3];
+            var ptr2 = r0;
+            var len2 = r1;
+            if (r3) {
+                ptr2 = 0; len2 = 0;
+                throw takeObject(r2);
+            }
+            deferred3_0 = ptr2;
+            deferred3_1 = len2;
+            return getStringFromWasm0(ptr2, len2);
+        } finally {
+            wasm.__wbindgen_add_to_stack_pointer(16);
+            wasm.__wbindgen_free(deferred3_0, deferred3_1, 1);
+        }
+    }
+}
 
 async function __wbg_load(module, imports) {
     if (typeof Response === 'function' && module instanceof Response) {
@@ -243,125 +314,12 @@ function __wbg_get_imports() {
         const ret = new Error(getStringFromWasm0(arg0, arg1));
         return addHeapObject(ret);
     };
-    imports.wbg.__wbg_log_20252ed396e16790 = function(arg0, arg1) {
+    imports.wbg.__wbg_log_598ccd735a33342c = function(arg0, arg1) {
         console.log(getStringFromWasm0(arg0, arg1));
     };
-    imports.wbg.__wbindgen_object_drop_ref = function(arg0) {
-        takeObject(arg0);
-    };
-    imports.wbg.__wbindgen_string_new = function(arg0, arg1) {
-        const ret = getStringFromWasm0(arg0, arg1);
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbindgen_object_clone_ref = function(arg0) {
-        const ret = getObject(arg0);
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbg_crypto_c48a774b022d20ac = function(arg0) {
-        const ret = getObject(arg0).crypto;
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbindgen_is_object = function(arg0) {
-        const val = getObject(arg0);
-        const ret = typeof(val) === 'object' && val !== null;
-        return ret;
-    };
-    imports.wbg.__wbg_process_298734cf255a885d = function(arg0) {
-        const ret = getObject(arg0).process;
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbg_versions_e2e78e134e3e5d01 = function(arg0) {
-        const ret = getObject(arg0).versions;
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbg_node_1cd7a5d853dbea79 = function(arg0) {
-        const ret = getObject(arg0).node;
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbindgen_is_string = function(arg0) {
-        const ret = typeof(getObject(arg0)) === 'string';
-        return ret;
-    };
-    imports.wbg.__wbg_msCrypto_bcb970640f50a1e8 = function(arg0) {
-        const ret = getObject(arg0).msCrypto;
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbg_require_8f08ceecec0f4fee = function() { return handleError(function () {
-        const ret = module.require;
-        return addHeapObject(ret);
-    }, arguments) };
-    imports.wbg.__wbindgen_is_function = function(arg0) {
-        const ret = typeof(getObject(arg0)) === 'function';
-        return ret;
-    };
-    imports.wbg.__wbg_getRandomValues_37fa2ca9e4e07fab = function() { return handleError(function (arg0, arg1) {
-        getObject(arg0).getRandomValues(getObject(arg1));
-    }, arguments) };
-    imports.wbg.__wbg_randomFillSync_dc1e9a60c158336d = function() { return handleError(function (arg0, arg1) {
-        getObject(arg0).randomFillSync(takeObject(arg1));
-    }, arguments) };
-    imports.wbg.__wbg_newnoargs_581967eacc0e2604 = function(arg0, arg1) {
-        const ret = new Function(getStringFromWasm0(arg0, arg1));
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbg_call_cb65541d95d71282 = function() { return handleError(function (arg0, arg1) {
-        const ret = getObject(arg0).call(getObject(arg1));
-        return addHeapObject(ret);
-    }, arguments) };
-    imports.wbg.__wbg_self_1ff1d729e9aae938 = function() { return handleError(function () {
-        const ret = self.self;
-        return addHeapObject(ret);
-    }, arguments) };
-    imports.wbg.__wbg_window_5f4faef6c12b79ec = function() { return handleError(function () {
-        const ret = window.window;
-        return addHeapObject(ret);
-    }, arguments) };
-    imports.wbg.__wbg_globalThis_1d39714405582d3c = function() { return handleError(function () {
-        const ret = globalThis.globalThis;
-        return addHeapObject(ret);
-    }, arguments) };
-    imports.wbg.__wbg_global_651f05c6a0944d1c = function() { return handleError(function () {
-        const ret = global.global;
-        return addHeapObject(ret);
-    }, arguments) };
-    imports.wbg.__wbindgen_is_undefined = function(arg0) {
-        const ret = getObject(arg0) === undefined;
-        return ret;
-    };
-    imports.wbg.__wbg_call_01734de55d61e11d = function() { return handleError(function (arg0, arg1, arg2) {
-        const ret = getObject(arg0).call(getObject(arg1), getObject(arg2));
-        return addHeapObject(ret);
-    }, arguments) };
-    imports.wbg.__wbg_buffer_085ec1f694018c4f = function(arg0) {
-        const ret = getObject(arg0).buffer;
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbg_newwithbyteoffsetandlength_6da8e527659b86aa = function(arg0, arg1, arg2) {
-        const ret = new Uint8Array(getObject(arg0), arg1 >>> 0, arg2 >>> 0);
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbg_new_8125e318e6245eed = function(arg0) {
-        const ret = new Uint8Array(getObject(arg0));
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbg_set_5cf90238115182c3 = function(arg0, arg1, arg2) {
-        getObject(arg0).set(getObject(arg1), arg2 >>> 0);
-    };
-    imports.wbg.__wbg_newwithlength_e5d69174d6984cd7 = function(arg0) {
-        const ret = new Uint8Array(arg0 >>> 0);
-        return addHeapObject(ret);
-    };
-    imports.wbg.__wbg_subarray_13db269f57aa838d = function(arg0, arg1, arg2) {
-        const ret = getObject(arg0).subarray(arg1 >>> 0, arg2 >>> 0);
-        return addHeapObject(ret);
-    };
     imports.wbg.__wbindgen_throw = function(arg0, arg1) {
         throw new Error(getStringFromWasm0(arg0, arg1));
     };
-    imports.wbg.__wbindgen_memory = function() {
-        const ret = wasm.memory;
-        return addHeapObject(ret);
-    };
 
     return imports;
 }

build/m_bg.wasm

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:755ae31980152f83b923e4f4b5e7f2e501c264539d0271f0455fdb6cf38b16c0
-size 1624018
+oid sha256:4ab25e3b74019ed2f5635491492ec52403ce009a36813e7829a4ed802ef2f637
+size 1577889

build/m_bg.wasm.d.ts

@@ -4,10 +4,12 @@ export const memory: WebAssembly.Memory;
 export function __wbg_model_free(a: number): void;
 export function model_new(a: number, b: number, c: number, d: number, e: number): void;
 export function model_run(a: number, b: number, c: number, d: number, e: number, f: number): void;
+export function __wbg_modelpose_free(a: number): void;
+export function modelpose_new(a: number, b: number, c: number, d: number, e: number): void;
+export function modelpose_run(a: number, b: number, c: number, d: number, e: number, f: number): void;
 export function main(a: number, b: number): number;
 export function __wbindgen_add_to_stack_pointer(a: number): number;
 export function __wbindgen_malloc(a: number, b: number): number;
 export function __wbindgen_realloc(a: number, b: number, c: number, d: number): number;
 export function __wbindgen_free(a: number, b: number, c: number): void;
-export function __wbindgen_exn_store(a: number): void;
 export function __wbindgen_start(): void;

index.html

@@ -54,8 +54,50 @@
           model_size: "x",
           url: "yolov8x.safetensors",
         },
+        yolov8n_pose: {
+          model_size: "n",
+          url: "yolov8n-pose.safetensors",
+        },
+        yolov8s_pose: {
+          model_size: "s",
+          url: "yolov8s-pose.safetensors",
+        },
+        yolov8m_pose: {
+          model_size: "m",
+          url: "yolov8m-pose.safetensors",
+        },
+        yolov8l_pose: {
+          model_size: "l",
+          url: "yolov8l-pose.safetensors",
+        },
+        yolov8x_pose: {
+          model_size: "x",
+          url: "yolov8x-pose.safetensors",
+        },
       };
 
+      const COCO_PERSON_SKELETON = [
+        [4, 0], // head
+        [3, 0],
+        [16, 14], // left lower leg
+        [14, 12], // left upper leg
+        [6, 12], // left torso
+        [6, 5], // top torso
+        [6, 8], // upper arm
+        [8, 10], // lower arm
+        [1, 2], // head
+        [1, 3], // right head
+        [2, 4], // left head
+        [3, 5], // right neck
+        [4, 6], // left neck
+        [5, 7], // right upper arm
+        [7, 9], // right lower arm
+        [5, 11], // right torso
+        [11, 12], // bottom torso
+        [11, 13], // right upper leg
+        [13, 15], // right lower leg
+      ];
+
       // init web worker
       const yoloWorker = new Worker("./yoloWorker.js", { type: "module" });
 
@@ -202,17 +244,28 @@
         ctx.fillStyle = "#0dff9a";
         const fontSize = 14 * scale;
         ctx.font = `${fontSize}px sans-serif`;
-        for (const [label, bbox] of output) {
-          const [x, y, w, h] = [
-            bbox.xmin,
-            bbox.ymin,
-            bbox.xmax - bbox.xmin,
-            bbox.ymax - bbox.ymin,
-          ];
-
-          const confidence = bbox.confidence;
+        for (const detection of output) {
+          // check keypoint for pose model data
+          let xmin, xmax, ymin, ymax, label, confidence, keypoints;
+          if ("keypoints" in detection) {
+            xmin = detection.xmin;
+            xmax = detection.xmax;
+            ymin = detection.ymin;
+            ymax = detection.ymax;
+            confidence = detection.confidence;
+            keypoints = detection.keypoints;
+          } else {
+            const [_label, bbox] = detection;
+            label = _label;
+            xmin = bbox.xmin;
+            xmax = bbox.xmax;
+            ymin = bbox.ymin;
+            ymax = bbox.ymax;
+            confidence = bbox.confidence;
+          }
+          const [x, y, w, h] = [xmin, ymin, xmax - xmin, ymax - ymin];
 
-          const text = `${label} ${confidence.toFixed(2)}`;
+          const text = `${label ? label + " " : ""}${confidence.toFixed(2)}`;
           const width = ctx.measureText(text).width;
           ctx.fillStyle = "#3c8566";
           ctx.fillRect(x - 2, y - fontSize, width + 4, fontSize);
@@ -220,6 +273,28 @@
 
           ctx.strokeRect(x, y, w, h);
           ctx.fillText(text, x, y - 2);
+          if (keypoints) {
+            ctx.save();
+            ctx.fillStyle = "magenta";
+            ctx.strokeStyle = "yellow";
+
+            for (const keypoint of keypoints) {
+              const { x, y } = keypoint;
+              ctx.beginPath();
+              ctx.arc(x, y, 3, 0, 2 * Math.PI);
+              ctx.fill();
+            }
+            ctx.beginPath();
+            for (const [xid, yid] of COCO_PERSON_SKELETON) {
+              //draw line between skeleton keypoitns
+              if (keypoints[xid] && keypoints[yid]) {
+                ctx.moveTo(keypoints[xid].x, keypoints[xid].y);
+                ctx.lineTo(keypoints[yid].x, keypoints[yid].y);
+              }
+            }
+            ctx.stroke();
+            ctx.restore();
+          }
         }
       });
 
@@ -229,12 +304,12 @@
           button.disabled = true;
           button.classList.add("bg-blue-700");
           button.classList.remove("bg-blue-950");
-          button.textContent = "Detecting...";
+          button.textContent = "Predicting...";
         } else if (statusMessage === "complete") {
           button.disabled = false;
           button.classList.add("bg-blue-950");
           button.classList.remove("bg-blue-700");
-          button.textContent = "Detect Objects";
+          button.textContent = "Predict";
           document.querySelector("#share-btn").hidden = false;
         }
       }
@@ -250,27 +325,31 @@
     </script>
   </head>
   <body class="container max-w-4xl mx-auto p-4">
-    <main class="grid grid-cols-1 gap-8">
+    <main class="grid grid-cols-1 gap-8 relative">
+      <span class="absolute text-5xl -ml-[1em]"> 🕯️ </span>
       <div>
         <h1 class="text-5xl font-bold">Candle YOLOv8</h1>
         <h2 class="text-2xl font-bold">Rust/WASM Demo</h2>
         <p class="max-w-lg">
-          Running an object detection model in the browser using rust/wasm with
-          an image. This demo uses the
+          This demo showcases object detection and pose estimation models in
+          your browser using Rust/WASM. It utilizes
           <a
             href="https://huggingface.co/lmz/candle-yolo-v8"
             target="_blank"
             class="underline hover:text-blue-500 hover:no-underline"
           >
-            Candle YOLOv8
+            safetensor's YOLOv8 models
           </a>
-          models to detect objects in images and WASM runtime built with
+          and a WASM runtime built with
           <a
             href="https://github.com/huggingface/candle/"
             target="_blank"
             class="underline hover:text-blue-500 hover:no-underline"
-            >Candle
-          </a>
+            >Candle </a
+          >.
+        </p>
+        <p>
+          To run pose estimation, select a yolo pose model from the dropdown
         </p>
       </div>
 
@@ -285,6 +364,12 @@
           <option value="yolov8m">yolov8m (51.9 MB)</option>
           <option value="yolov8l">yolov8l (87.5 MB)</option>
           <option value="yolov8x">yolov8x (137 MB)</option>
+          <!-- Pose models -->
+          <option value="yolov8n_pose">yolov8n_pose (6.65 MB)</option>
+          <option value="yolov8s_pose">yolov8s_pose (23.3 MB)</option>
+          <option value="yolov8m_pose">yolov8m_pose (53 MB)</option>
+          <option value="yolov8l_pose">yolov8l_pose (89.1 MB)</option>
+          <option value="yolov8x_pose">yolov8x_pose (139 MB)</option>
         </select>
       </div>
       <!-- drag and drop area -->
@@ -358,6 +443,10 @@
             src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/candle/examples/bike.jpeg"
             class="cursor-pointer w-24 h-24 object-cover"
           />
+          <img
+            src="https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/candle/examples/000000000077.jpg"
+            class="cursor-pointer w-24 h-24 object-cover"
+          />
         </div>
       </div>
       <div>
@@ -406,7 +495,7 @@
           disabled
           class="bg-blue-950 hover:bg-blue-700 text-white font-normal py-2 px-4 rounded disabled:opacity-75 disabled:hover:bg-blue-950"
         >
-          Detect Objects
+          Predict
         </button>
       </div>
     </main>

yoloWorker.js

@@ -1,5 +1,5 @@
 //load the candle yolo wasm module
-import init, { Model } from "./build/m.js";
+import init, { Model, ModelPose } from "./build/m.js";
 
 class Yolo {
   static instance = {};
@@ -14,7 +14,12 @@ class Yolo {
       const modelRes = await fetch(modelURL);
       const yoloArrayBuffer = await modelRes.arrayBuffer();
       const weightsArrayU8 = new Uint8Array(yoloArrayBuffer);
-      this.instance[modelID] = new Model(weightsArrayU8, modelSize);
+      if (/pose/.test(modelID)) {
+        // if pose model, use ModelPose
+        this.instance[modelID] = new ModelPose(weightsArrayU8, modelSize);
+      } else {
+        this.instance[modelID] = new Model(weightsArrayU8, modelSize);
+      }
     } else {
       self.postMessage({ status: "model already loaded" });
     }