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
This commit is contained in:
parent
c8ad606e7c
commit
a8bee6baf3
|
@ -51,9 +51,16 @@ export async function createTaskRunner<T extends TaskRunner>(
|
||||||
canvas: HTMLCanvasElement|OffscreenCanvas|null|undefined,
|
canvas: HTMLCanvasElement|OffscreenCanvas|null|undefined,
|
||||||
fileset: WasmFileset, options: TaskRunnerOptions): Promise<T> {
|
fileset: WasmFileset, options: TaskRunnerOptions): Promise<T> {
|
||||||
const fileLocator: FileLocator = {
|
const fileLocator: FileLocator = {
|
||||||
locateFile() {
|
locateFile(file): string {
|
||||||
// The only file loaded with this mechanism is the Wasm binary
|
const wasm = fileset.wasmBinaryPath.toString();
|
||||||
return fileset.wasmBinaryPath.toString();
|
if (wasm.includes(file)) {
|
||||||
|
return wasm;
|
||||||
|
}
|
||||||
|
const asset = fileset.assetBinaryPath?.toString();
|
||||||
|
if (asset?.includes(file)) {
|
||||||
|
return asset;
|
||||||
|
}
|
||||||
|
return file;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|
|
@ -20,11 +20,13 @@ import {InferenceCalculatorOptions} from '../../../calculators/tensor/inference_
|
||||||
import {BaseOptions as BaseOptionsProto} from '../../../tasks/cc/core/proto/base_options_pb';
|
import {BaseOptions as BaseOptionsProto} from '../../../tasks/cc/core/proto/base_options_pb';
|
||||||
import {TaskRunner} from '../../../tasks/web/core/task_runner';
|
import {TaskRunner} from '../../../tasks/web/core/task_runner';
|
||||||
import {createSpyWasmModule, SpyWasmModule} from '../../../tasks/web/core/task_runner_test_utils';
|
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';
|
import {ErrorListener} from '../../../web/graph_runner/graph_runner';
|
||||||
// Placeholder for internal dependency on trusted resource URL builder
|
// 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 {TaskRunnerOptions} from './task_runner_options';
|
||||||
|
import {WasmFileset} from './wasm_fileset';
|
||||||
|
|
||||||
type Writeable<T> = {
|
type Writeable<T> = {
|
||||||
-readonly[P in keyof T]: T[P]
|
-readonly[P in keyof T]: T[P]
|
||||||
|
@ -147,6 +149,9 @@ describe('TaskRunner', () => {
|
||||||
let fetchSpy: jasmine.Spy;
|
let fetchSpy: jasmine.Spy;
|
||||||
let taskRunner: TaskRunnerFake;
|
let taskRunner: TaskRunnerFake;
|
||||||
let fetchStatus: number;
|
let fetchStatus: number;
|
||||||
|
let locator: graphRunner.FileLocator|undefined;
|
||||||
|
|
||||||
|
let oldCreate = graphRunner.createMediaPipeLib;
|
||||||
|
|
||||||
beforeEach(() => {
|
beforeEach(() => {
|
||||||
fetchStatus = 200;
|
fetchStatus = 200;
|
||||||
|
@ -159,9 +164,70 @@ describe('TaskRunner', () => {
|
||||||
});
|
});
|
||||||
global.fetch = fetchSpy;
|
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();
|
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', () => {
|
it('handles errors during graph update', () => {
|
||||||
taskRunner.enqueueError('Test error');
|
taskRunner.enqueueError('Test error');
|
||||||
|
|
||||||
|
|
2
mediapipe/tasks/web/core/wasm_fileset.d.ts
vendored
2
mediapipe/tasks/web/core/wasm_fileset.d.ts
vendored
|
@ -24,4 +24,6 @@ export declare interface WasmFileset {
|
||||||
wasmBinaryPath: string;
|
wasmBinaryPath: string;
|
||||||
/** The optional path to the asset loader script. */
|
/** The optional path to the asset loader script. */
|
||||||
assetLoaderPath?: string;
|
assetLoaderPath?: string;
|
||||||
|
/** The optional path to the assets binary. */
|
||||||
|
assetBinaryPath?: string;
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue
Block a user