From a8bee6baf3cf3cdf8c91d68362887b7c9cda4a91 Mon Sep 17 00:00:00 2001
From: MediaPipe Team <mediapipe-team@google.com>
Date: Mon, 14 Aug 2023 13:57:59 -0700
Subject: [PATCH] Updates the runners to support wasm-style binary assets
 files, and allows their URLs to be explicitly specified as part of the
 WasmFileset.

PiperOrigin-RevId: 556903356
---
 mediapipe/tasks/web/core/task_runner.ts      | 13 +++-
 mediapipe/tasks/web/core/task_runner_test.ts | 68 +++++++++++++++++++-
 mediapipe/tasks/web/core/wasm_fileset.d.ts   |  2 +
 3 files changed, 79 insertions(+), 4 deletions(-)

diff --git a/mediapipe/tasks/web/core/task_runner.ts b/mediapipe/tasks/web/core/task_runner.ts
index dde98192..e2690cde 100644
--- a/mediapipe/tasks/web/core/task_runner.ts
+++ b/mediapipe/tasks/web/core/task_runner.ts
@@ -51,9 +51,16 @@ export async function createTaskRunner<T extends TaskRunner>(
     canvas: HTMLCanvasElement|OffscreenCanvas|null|undefined,
     fileset: WasmFileset, options: TaskRunnerOptions): Promise<T> {
   const fileLocator: FileLocator = {
-    locateFile() {
-      // The only file loaded with this mechanism is the Wasm binary
-      return fileset.wasmBinaryPath.toString();
+    locateFile(file): string {
+      const wasm = fileset.wasmBinaryPath.toString();
+      if (wasm.includes(file)) {
+        return wasm;
+      }
+      const asset = fileset.assetBinaryPath?.toString();
+      if (asset?.includes(file)) {
+        return asset;
+      }
+      return file;
     }
   };
 
diff --git a/mediapipe/tasks/web/core/task_runner_test.ts b/mediapipe/tasks/web/core/task_runner_test.ts
index a68ba224..dd9b874b 100644
--- a/mediapipe/tasks/web/core/task_runner_test.ts
+++ b/mediapipe/tasks/web/core/task_runner_test.ts
@@ -20,11 +20,13 @@ import {InferenceCalculatorOptions} from '../../../calculators/tensor/inference_
 import {BaseOptions as BaseOptionsProto} from '../../../tasks/cc/core/proto/base_options_pb';
 import {TaskRunner} from '../../../tasks/web/core/task_runner';
 import {createSpyWasmModule, SpyWasmModule} from '../../../tasks/web/core/task_runner_test_utils';
+import * as graphRunner from '../../../web/graph_runner/graph_runner';
 import {ErrorListener} from '../../../web/graph_runner/graph_runner';
 // Placeholder for internal dependency on trusted resource URL builder
 
-import {CachedGraphRunner} from './task_runner';
+import {CachedGraphRunner, createTaskRunner} from './task_runner';
 import {TaskRunnerOptions} from './task_runner_options';
+import {WasmFileset} from './wasm_fileset';
 
 type Writeable<T> = {
   -readonly[P in keyof T]: T[P]
@@ -147,6 +149,9 @@ describe('TaskRunner', () => {
   let fetchSpy: jasmine.Spy;
   let taskRunner: TaskRunnerFake;
   let fetchStatus: number;
+  let locator: graphRunner.FileLocator|undefined;
+
+  let oldCreate = graphRunner.createMediaPipeLib;
 
   beforeEach(() => {
     fetchStatus = 200;
@@ -159,9 +164,70 @@ describe('TaskRunner', () => {
     });
     global.fetch = fetchSpy;
 
+    // Monkeypatch an exported static method for testing!
+    oldCreate = graphRunner.createMediaPipeLib;
+    locator = undefined;
+    (graphRunner as {createMediaPipeLib: Function}).createMediaPipeLib =
+        jasmine.createSpy().and.callFake(
+            (type, wasmLoaderPath, assetLoaderPath, canvas, fileLocator) => {
+              locator = fileLocator;
+              // tslint:disable-next-line:no-any Monkeypatching for test mocks.
+              return Promise.resolve(taskRunner as any);
+            });
+
     taskRunner = TaskRunnerFake.createFake();
   });
 
+  afterEach(() => {
+    // Restore the monkeypatch.
+    (graphRunner as {createMediaPipeLib: Function}).createMediaPipeLib =
+        oldCreate;
+  });
+
+  it('constructs with useful file locators for asset.data files', () => {
+    const fileset: WasmFileset = {
+      wasmLoaderPath: `wasm.js`,
+      wasmBinaryPath: `a/b/c/wasm.wasm`,
+      assetLoaderPath: `asset.js`,
+      assetBinaryPath: `a/b/c/asset.data`,
+    };
+
+    const options = {
+      baseOptions: {
+        modelAssetPath: `modelAssetPath`,
+      }
+    };
+
+    const runner = createTaskRunner(TaskRunnerFake, null, fileset, options);
+    expect(runner).toBeDefined();
+    expect(locator).toBeDefined();
+    expect(locator?.locateFile('wasm.wasm')).toEqual('a/b/c/wasm.wasm');
+    expect(locator?.locateFile('asset.data')).toEqual('a/b/c/asset.data');
+    expect(locator?.locateFile('unknown')).toEqual('unknown');
+  });
+
+  it('constructs without useful file locators with no asset.data file', () => {
+    const fileset: WasmFileset = {
+      wasmLoaderPath: `wasm.js`,
+      wasmBinaryPath: `a/b/c/wasm.wasm`,
+      assetLoaderPath: `asset.js`,
+      // No path to the assets binary.
+    };
+
+    const options = {
+      baseOptions: {
+        modelAssetPath: `modelAssetPath`,
+      }
+    };
+
+    const runner = createTaskRunner(TaskRunnerFake, null, fileset, options);
+    expect(runner).toBeDefined();
+    expect(locator).toBeDefined();
+    expect(locator?.locateFile('wasm.wasm')).toEqual('a/b/c/wasm.wasm');
+    expect(locator?.locateFile('asset.data')).toEqual('asset.data');
+    expect(locator?.locateFile('unknown')).toEqual('unknown');
+  });
+
   it('handles errors during graph update', () => {
     taskRunner.enqueueError('Test error');
 
diff --git a/mediapipe/tasks/web/core/wasm_fileset.d.ts b/mediapipe/tasks/web/core/wasm_fileset.d.ts
index dda466ad..e4cfbe1d 100644
--- a/mediapipe/tasks/web/core/wasm_fileset.d.ts
+++ b/mediapipe/tasks/web/core/wasm_fileset.d.ts
@@ -24,4 +24,6 @@ export declare interface WasmFileset {
   wasmBinaryPath: string;
   /** The optional path to the asset loader script. */
   assetLoaderPath?: string;
+  /** The optional path to the assets binary. */
+  assetBinaryPath?: string;
 }