Use WebWorker for MediaPipe Holistic - Fixes #24

angular.json

@@ -40,6 +40,11 @@
                 "input": "./node_modules/flag-icons/flags/1x1/",
                 "output": "./assets/flags/1x1/"
               },
+              {
+                "glob": "**/*",
+                "input": "./node_modules/@mediapipe/holistic/",
+                "output": "./"
+              },
               {
                 "glob": "*.ttf",
                 "input": "./node_modules/@sutton-signwriting/font-ttf/font",

src/app/components/video/video.component.ts

@@ -1,7 +1,7 @@
 import {AfterViewInit, Component, ElementRef, HostBinding, Input, ViewChild} from '@angular/core';
 import {Store} from '@ngxs/store';
 import {combineLatest, firstValueFrom} from 'rxjs';
-import {VideoSettings, VideoStateModel} from '../../core/modules/ngxs/store/video/video.state';
+import {VideoStateModel} from '../../core/modules/ngxs/store/video/video.state';
 import Stats from 'stats.js';
 import {distinctUntilChanged, filter, map, takeUntil, tap} from 'rxjs/operators';
 import {BaseComponent} from '../base/base.component';
@@ -112,14 +112,15 @@ export class VideoComponent extends BaseComponent implements AfterViewInit {
       .pipe(
         map(state => state.videoSettings),
         filter(Boolean),
-        tap(({width, height}) => {
+        tap(({width, height, aspectRatio}) => {
+          this.aspectRatio = 'aspect-' + aspectRatio;
+
           this.canvasEl.nativeElement.width = width;
           this.canvasEl.nativeElement.height = height;
 
           // It is required to wait for next frame, as grid element might still be resizing
           requestAnimationFrame(this.scaleCanvas.bind(this));
         }),
-        tap((settings: VideoSettings) => (this.aspectRatio = 'aspect-' + settings.aspectRatio)),
         takeUntil(this.ngUnsubscribe)
       )
       .subscribe();

src/app/modules/animation/animation.service.ts

@@ -4,6 +4,7 @@ import {LayersModel} from '@tensorflow/tfjs-layers';
 import {Injectable} from '@angular/core';
 import {TensorflowService} from '../../core/services/tfjs/tfjs.service';
 import {MediapipeHolisticService} from '../../core/services/holistic.service';
+import {POSE_LANDMARKS} from '@mediapipe/holistic';
 
 const ANIMATION_KEYS = [
   'mixamorigHead.quaternion',
@@ -80,8 +81,7 @@ export class AnimationService {
   }
 
   normalizePose(pose: Pose): Tensor {
-    const bodyLandmarks =
-      pose.poseLandmarks || new Array(Object.keys(this.holistic.POSE_LANDMARKS).length).fill(EMPTY_LANDMARK);
+    const bodyLandmarks = pose.poseLandmarks || new Array(Object.keys(POSE_LANDMARKS).length).fill(EMPTY_LANDMARK);
     const leftHandLandmarks = pose.leftHandLandmarks || new Array(21).fill(EMPTY_LANDMARK);
     const rightHandLandmarks = pose.rightHandLandmarks || new Array(21).fill(EMPTY_LANDMARK);
     const landmarks = bodyLandmarks.concat(leftHandLandmarks, rightHandLandmarks);
@@ -90,8 +90,8 @@ export class AnimationService {
       .tensor(landmarks.map(l => [l.x, l.y, l.z]))
       .mul(this.tf.tensor([pose.image.width, pose.image.height, pose.image.width]));
 
-    const p1 = tensor.slice(this.holistic.POSE_LANDMARKS.LEFT_SHOULDER, 1);
-    const p2 = tensor.slice(this.holistic.POSE_LANDMARKS.RIGHT_SHOULDER, 1);
+    const p1 = tensor.slice(POSE_LANDMARKS.LEFT_SHOULDER, 1);
+    const p2 = tensor.slice(POSE_LANDMARKS.RIGHT_SHOULDER, 1);
 
     const d = this.tf.sqrt(this.tf.pow(p2.sub(p1), 2).sum());
     let normTensor = this.tf.sub(tensor, p1.add(p2).div(2)).div(d);

src/app/modules/detector/detector.service.ts

@@ -2,8 +2,8 @@ import {Tensor} from '@tensorflow/tfjs';
 import {EMPTY_LANDMARK, Pose, PoseLandmark} from '../pose/pose.state';
 import {LayersModel} from '@tensorflow/tfjs-layers';
 import {Injectable} from '@angular/core';
+import {POSE_LANDMARKS} from '@mediapipe/holistic';
 import {TensorflowService} from '../../core/services/tfjs/tfjs.service';
-import {MediapipeHolisticService} from '../../core/services/holistic.service';
 
 const WINDOW_SIZE = 20;
 
@@ -19,16 +19,11 @@ export class DetectorService {
 
   sequentialModel: LayersModel;
 
-  constructor(private tf: TensorflowService, private holistic: MediapipeHolisticService) {}
+  constructor(private tf: TensorflowService) {}
 
   async loadModel() {
-    return Promise.all([
-      this.holistic.load(),
-      this.tf
-        .load()
-        .then(() => this.tf.loadLayersModel('assets/models/sign-detector/model.json'))
-        .then(model => (this.sequentialModel = model as unknown as LayersModel)),
-    ]);
+    await this.tf.load();
+    this.sequentialModel = await this.tf.loadLayersModel('assets/models/sign-detector/model.json');
   }
 
   distance(p1: PoseLandmark, p2: PoseLandmark): number {
@@ -38,16 +33,15 @@ export class DetectorService {
   }
 
   normalizePose(pose: Pose): PoseLandmark[] {
-    const bodyLandmarks =
-      pose.poseLandmarks || new Array(Object.keys(this.holistic.POSE_LANDMARKS).length).fill(EMPTY_LANDMARK);
+    const bodyLandmarks = pose.poseLandmarks || new Array(Object.keys(POSE_LANDMARKS).length).fill(EMPTY_LANDMARK);
     const leftHandLandmarks = pose.leftHandLandmarks || new Array(21).fill(EMPTY_LANDMARK);
     const rightHandLandmarks = pose.leftHandLandmarks || new Array(21).fill(EMPTY_LANDMARK);
     const landmarks = bodyLandmarks
       .concat(leftHandLandmarks, rightHandLandmarks)
       .map(l => (this.isValidLandmark(l) ? l : EMPTY_LANDMARK));
 
-    const p1 = landmarks[this.holistic.POSE_LANDMARKS.LEFT_SHOULDER];
-    const p2 = landmarks[this.holistic.POSE_LANDMARKS.RIGHT_SHOULDER];
+    const p1 = landmarks[POSE_LANDMARKS.LEFT_SHOULDER];
+    const p2 = landmarks[POSE_LANDMARKS.RIGHT_SHOULDER];
 
     if (p1.x > 0 && p2.x > 0) {
       this.shoulderWidth[this.shoulderWidthIndex % WINDOW_SIZE] = this.distance(p1, p2);
@@ -69,36 +63,29 @@ export class DetectorService {
 
     // TODO remove, this is to be compliant with openpose
     const neck = {
-      x:
-        (newPose[this.holistic.POSE_LANDMARKS.LEFT_SHOULDER].x +
-          newPose[this.holistic.POSE_LANDMARKS.RIGHT_SHOULDER].x) /
-        2,
-      y:
-        (newPose[this.holistic.POSE_LANDMARKS.LEFT_SHOULDER].y +
-          newPose[this.holistic.POSE_LANDMARKS.RIGHT_SHOULDER].y) /
-        2,
+      x: (newPose[POSE_LANDMARKS.LEFT_SHOULDER].x + newPose[POSE_LANDMARKS.RIGHT_SHOULDER].x) / 2,
+      y: (newPose[POSE_LANDMARKS.LEFT_SHOULDER].y + newPose[POSE_LANDMARKS.RIGHT_SHOULDER].y) / 2,
     };
-
     return [
-      newPose[this.holistic.POSE_LANDMARKS.NOSE],
+      newPose[POSE_LANDMARKS.NOSE],
       neck,
-      newPose[this.holistic.POSE_LANDMARKS.RIGHT_SHOULDER],
-      newPose[this.holistic.POSE_LANDMARKS.RIGHT_ELBOW],
-      newPose[this.holistic.POSE_LANDMARKS.RIGHT_WRIST],
-      newPose[this.holistic.POSE_LANDMARKS.LEFT_SHOULDER],
-      newPose[this.holistic.POSE_LANDMARKS.LEFT_ELBOW],
-      newPose[this.holistic.POSE_LANDMARKS.LEFT_WRIST],
+      newPose[POSE_LANDMARKS.RIGHT_SHOULDER],
+      newPose[POSE_LANDMARKS.RIGHT_ELBOW],
+      newPose[POSE_LANDMARKS.RIGHT_WRIST],
+      newPose[POSE_LANDMARKS.LEFT_SHOULDER],
+      newPose[POSE_LANDMARKS.LEFT_ELBOW],
+      newPose[POSE_LANDMARKS.LEFT_WRIST],
       EMPTY_LANDMARK,
       EMPTY_LANDMARK,
       EMPTY_LANDMARK,
       EMPTY_LANDMARK,
       EMPTY_LANDMARK,
       EMPTY_LANDMARK,
       EMPTY_LANDMARK,
-      newPose[this.holistic.POSE_LANDMARKS.RIGHT_EYE],
-      newPose[this.holistic.POSE_LANDMARKS.LEFT_EYE],
-      newPose[this.holistic.POSE_LANDMARKS.RIGHT_EAR],
-      newPose[this.holistic.POSE_LANDMARKS.LEFT_EAR],
+      newPose[POSE_LANDMARKS.RIGHT_EYE],
+      newPose[POSE_LANDMARKS.LEFT_EYE],
+      newPose[POSE_LANDMARKS.RIGHT_EAR],
+      newPose[POSE_LANDMARKS.LEFT_EAR],
       EMPTY_LANDMARK,
       EMPTY_LANDMARK,
       EMPTY_LANDMARK,

src/app/modules/pose/pose.actions.ts

@@ -1,5 +1,9 @@
 import {Pose} from './pose.state';
 
+export class LoadPoseModel {
+  static readonly type = '[Pose] Load Pose Model';
+}
+
 export class PoseVideoFrame {
   static readonly type = '[Pose] Pose Video Frame';
 

src/app/modules/pose/pose.service.ts

@@ -1,55 +1,72 @@
 import {Injectable} from '@angular/core';
+import {
+  FACEMESH_FACE_OVAL,
+  FACEMESH_LEFT_EYE,
+  FACEMESH_LEFT_EYEBROW,
+  FACEMESH_LIPS,
+  FACEMESH_RIGHT_EYE,
+  FACEMESH_RIGHT_EYEBROW,
+  FACEMESH_TESSELATION,
+  HAND_CONNECTIONS,
+  POSE_CONNECTIONS,
+  POSE_LANDMARKS,
+} from '@mediapipe/holistic';
 import * as drawing from '@mediapipe/drawing_utils/drawing_utils.js';
 import {Pose, PoseLandmark} from './pose.state';
 import {GoogleAnalyticsService} from '../../core/modules/google-analytics/google-analytics.service';
-import {MediapipeHolisticService} from '../../core/services/holistic.service';
+import * as comlink from 'comlink';
+import {transferableImage} from '../../core/helpers/image/transferable';
 
 const IGNORED_BODY_LANDMARKS = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 15, 16, 17, 18, 19, 20, 21, 22];
 
 @Injectable({
   providedIn: 'root',
 })
 export class PoseService {
-  model?: any;
   isFirstFrame = true;
-  onResultsCallbacks = [];
 
-  constructor(private ga: GoogleAnalyticsService, private holistic: MediapipeHolisticService) {}
+  worker: comlink.Remote<{
+    loadModel: () => Promise<void>;
+    pose: (imageBitmap: ImageBitmap | ImageData) => Promise<Pose>;
+  }>;
 
-  onResults(onResultsCallback) {
-    this.onResultsCallbacks.push(onResultsCallback);
-  }
+  constructor(private ga: GoogleAnalyticsService) {}
 
   async load(): Promise<void> {
-    if (this.model) {
+    if (this.worker) {
       return;
     }
 
-    await this.holistic.load();
-
-    await this.ga.trace('pose', 'load', () => {
-      this.model = new this.holistic.Holistic({locateFile: file => `assets/models/holistic/${file}`});
-
-      this.model.setOptions({
-        upperBodyOnly: false,
-        modelComplexity: 1,
-      });
-
-      this.model.onResults(results => {
-        for (const callback of this.onResultsCallbacks) {
-          callback(results);
-        }
-      });
+    await this.ga.trace('pose', 'load', async () => {
+      this.worker = comlink.wrap(new Worker(new URL('./pose.worker', import.meta.url), {type: 'module'}));
+      await this.worker.loadModel();
     });
   }
 
-  async predict(video: HTMLVideoElement | HTMLImageElement): Promise<void> {
-    await this.load();
+  async predict(video: HTMLVideoElement | HTMLImageElement): Promise<Pose> {
+    const width = (video as HTMLVideoElement).videoWidth ?? video.width;
+    if (!this.worker || width === 0) {
+      return null;
+    }
 
     const frameType = this.isFirstFrame ? 'first-frame' : 'frame';
-    await this.ga.trace('pose', frameType, () => {
+    const image = await transferableImage(video);
+
+    return this.ga.trace('pose', frameType, async () => {
       this.isFirstFrame = false;
-      return this.model.send({image: video});
+      const result: Pose = await this.worker.pose(image);
+      if (!result) {
+        return null;
+      }
+
+      // TODO not sure if this is needed
+      // const newImage = document.createElement('canvas');
+      // newImage.width = image.width;
+      // newImage.height = image.height;
+      // const ctx = newImage.getContext('2d');
+      // ctx.drawImage(image as any, 0, 0);
+      // result.image = newImage;
+      return result;
     });
   }
 
@@ -59,7 +76,7 @@ export class PoseService {
       delete filteredLandmarks[l];
     }
 
-    drawing.drawConnectors(ctx, filteredLandmarks, this.holistic.POSE_CONNECTIONS, {color: '#00FF00'});
+    drawing.drawConnectors(ctx, filteredLandmarks, POSE_CONNECTIONS, {color: '#00FF00'});
     drawing.drawLandmarks(ctx, filteredLandmarks, {color: '#00FF00', fillColor: '#FF0000'});
   }
 
@@ -70,7 +87,7 @@ export class PoseService {
     dotColor: string,
     dotFillColor: string
   ): void {
-    drawing.drawConnectors(ctx, landmarks, this.holistic.HAND_CONNECTIONS, {color: lineColor});
+    drawing.drawConnectors(ctx, landmarks, HAND_CONNECTIONS, {color: lineColor});
     drawing.drawLandmarks(ctx, landmarks, {
       color: dotColor,
       fillColor: dotFillColor,
@@ -82,13 +99,13 @@ export class PoseService {
   }
 
   drawFace(landmarks: PoseLandmark[], ctx: CanvasRenderingContext2D): void {
-    drawing.drawConnectors(ctx, landmarks, this.holistic.FACEMESH_TESSELATION, {color: '#C0C0C070', lineWidth: 1});
-    drawing.drawConnectors(ctx, landmarks, this.holistic.FACEMESH_RIGHT_EYE, {color: '#FF3030'});
-    drawing.drawConnectors(ctx, landmarks, this.holistic.FACEMESH_RIGHT_EYEBROW, {color: '#FF3030'});
-    drawing.drawConnectors(ctx, landmarks, this.holistic.FACEMESH_LEFT_EYE, {color: '#30FF30'});
-    drawing.drawConnectors(ctx, landmarks, this.holistic.FACEMESH_LEFT_EYEBROW, {color: '#30FF30'});
-    drawing.drawConnectors(ctx, landmarks, this.holistic.FACEMESH_FACE_OVAL, {color: '#E0E0E0'});
-    drawing.drawConnectors(ctx, landmarks, this.holistic.FACEMESH_LIPS, {color: '#E0E0E0'});
+    drawing.drawConnectors(ctx, landmarks, FACEMESH_TESSELATION, {color: '#C0C0C070', lineWidth: 1});
+    drawing.drawConnectors(ctx, landmarks, FACEMESH_RIGHT_EYE, {color: '#FF3030'});
+    drawing.drawConnectors(ctx, landmarks, FACEMESH_RIGHT_EYEBROW, {color: '#FF3030'});
+    drawing.drawConnectors(ctx, landmarks, FACEMESH_LEFT_EYE, {color: '#30FF30'});
+    drawing.drawConnectors(ctx, landmarks, FACEMESH_LEFT_EYEBROW, {color: '#30FF30'});
+    drawing.drawConnectors(ctx, landmarks, FACEMESH_FACE_OVAL, {color: '#E0E0E0'});
+    drawing.drawConnectors(ctx, landmarks, FACEMESH_LIPS, {color: '#E0E0E0'});
   }
 
   drawConnect(connectors: PoseLandmark[][], ctx: CanvasRenderingContext2D): void {
@@ -112,15 +129,12 @@ export class PoseService {
 
     if (pose.rightHandLandmarks) {
       ctx.strokeStyle = '#00FF00';
-      this.drawConnect(
-        [[pose.poseLandmarks[this.holistic.POSE_LANDMARKS.RIGHT_ELBOW], pose.rightHandLandmarks[0]]],
-        ctx
-      );
+      this.drawConnect([[pose.poseLandmarks[POSE_LANDMARKS.RIGHT_ELBOW], pose.rightHandLandmarks[0]]], ctx);
     }
 
     if (pose.leftHandLandmarks) {
       ctx.strokeStyle = '#FF0000';
-      this.drawConnect([[pose.poseLandmarks[this.holistic.POSE_LANDMARKS.LEFT_ELBOW], pose.leftHandLandmarks[0]]], ctx);
+      this.drawConnect([[pose.poseLandmarks[POSE_LANDMARKS.LEFT_ELBOW], pose.leftHandLandmarks[0]]], ctx);
     }
   }