From 046efddf8b5789962c26782dd6371e2fcb781b10 Mon Sep 17 00:00:00 2001 From: Sebastian Schmidt Date: Fri, 3 Feb 2023 12:14:48 -0800 Subject: [PATCH] Add Empty Packet support to GraphRunner PiperOrigin-RevId: 506963970 --- mediapipe/web/graph_runner/graph_runner.ts | 125 ++++++++++-------- .../graph_runner/graph_runner_image_lib.ts | 8 +- 2 files changed, 75 insertions(+), 58 deletions(-) diff --git a/mediapipe/web/graph_runner/graph_runner.ts b/mediapipe/web/graph_runner/graph_runner.ts index 644d74918..eaec5c499 100644 --- a/mediapipe/web/graph_runner/graph_runner.ts +++ b/mediapipe/web/graph_runner/graph_runner.ts @@ -15,6 +15,25 @@ export declare interface FileLocator { locateFile: (filename: string) => string; } +/** + * A listener that receives the contents of a non-empty MediaPipe packet and + * its timestamp. + */ +export type SimpleListener = (data: T, timestamp: number) => void; + +/** + * A listener that receives the current MediaPipe packet timestamp. This is + * invoked even for empty packet. + */ +export type EmptyPacketListener = (timestamp: number) => void; + +/** + * A listener that receives a single element of vector-returning output packet. + * Intended for internal usage. + */ +export type VectorListener = + (data: T, index: number, length: number, timestamp: number) => void; + /** * Declarations for Emscripten's WebAssembly Module behavior, so TS compiler * doesn't break our JS/C++ bridge. @@ -71,14 +90,12 @@ export declare interface WasmModule { (dataPtr: number, dataSize: number, protoNamePtr: number, streamNamePtr: number) => void; - // Wasm Module output listener entrypoints. Also built as part of + // Map of output streams to packet listeners. Also built as part of // gl_graph_runner_internal_multi_input. simpleListeners?: - {[outputStreamName: string]: (data: unknown, timestamp: number) => void}; - vectorListeners?: { - [outputStreamName: string]: ( - data: unknown, index: number, length: number, timestamp: number) => void - }; + Record|VectorListener>; + // Map of output streams to empty packet listeners. + emptyPacketListeners?: Record; _attachBoolListener: (streamNamePtr: number) => void; _attachBoolVectorListener: (streamNamePtr: number) => void; _attachDoubleListener: (streamNamePtr: number) => void; @@ -147,7 +164,7 @@ export type WasmMediaPipeConstructor = * Simple class to run an arbitrary image-in/image-out MediaPipe graph (i.e. * as created by wasm_mediapipe_demo BUILD macro), and either render results * into canvas, or else return the output WebGLTexture. Takes a WebAssembly - * Module (must be instantiated to self.Module). + * Module. */ export class GraphRunner { // TODO: These should be protected/private, but are left exposed for @@ -419,12 +436,10 @@ export class GraphRunner { * Ensures existence of the simple listeners table and registers the callback. * Intended for internal usage. */ - setListener( - outputStreamName: string, - callbackFcn: (data: T, timestamp: number) => void) { + setListener(outputStreamName: string, callbackFcn: SimpleListener) { this.wasmModule.simpleListeners = this.wasmModule.simpleListeners || {}; this.wasmModule.simpleListeners[outputStreamName] = - callbackFcn as (data: unknown, timestamp: number) => void; + callbackFcn as SimpleListener; } /** @@ -432,20 +447,19 @@ export class GraphRunner { * Intended for internal usage. */ setVectorListener( - outputStreamName: string, - callbackFcn: (data: T[], timestamp: number) => void) { + outputStreamName: string, callbackFcn: SimpleListener) { let buffer: T[] = []; - this.wasmModule.vectorListeners = this.wasmModule.vectorListeners || {}; - this.wasmModule.vectorListeners[outputStreamName] = + this.wasmModule.simpleListeners = this.wasmModule.simpleListeners || {}; + this.wasmModule.simpleListeners[outputStreamName] = (data: unknown, index: number, length: number, timestamp: number) => { // The Wasm listener gets invoked once for each element. Once we - // receive all elements, we invoke the registered callback with the - // full array. + // receive all elements, we invoke the registered callback with + // the full array. buffer[index] = data as T; if (index === length - 1) { - // Invoke the user callback directly, as the Wasm layer may clean up - // the underlying data elements once we leave the scope of the - // listener. + // Invoke the user callback directly, as the Wasm layer may + // clean up the underlying data elements once we leave the scope + // of the listener. callbackFcn(buffer, timestamp); buffer = []; } @@ -460,6 +474,25 @@ export class GraphRunner { this.wasmModule.errorListener = callbackFcn; } + /** + * Attaches a listener that will be invoked when the MediaPipe framework + * receives an empty packet on the provided output stream. This can be used + * to receive the latest output timestamp. + * + * Empty packet listeners are only active if there is a corresponding packet + * listener. + * + * @param outputStreamName The name of the graph output stream to receive + * empty packets from. + * @param callbackFcn The callback to receive the timestamp. + */ + attachEmptyPacketListener( + outputStreamName: string, callbackFcn: EmptyPacketListener) { + this.wasmModule.emptyPacketListeners = + this.wasmModule.emptyPacketListeners || {}; + this.wasmModule.emptyPacketListeners[outputStreamName] = callbackFcn; + } + /** * Takes the raw data from a JS audio capture array, and sends it to C++ to be * processed. @@ -744,8 +777,7 @@ export class GraphRunner { * should not perform overly complicated (or any async) behavior. */ attachBoolListener( - outputStreamName: string, - callbackFcn: (data: boolean, timestamp: number) => void): void { + outputStreamName: string, callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setListener(outputStreamName, callbackFcn); @@ -765,8 +797,7 @@ export class GraphRunner { * should not perform overly complicated (or any async) behavior. */ attachBoolVectorListener( - outputStreamName: string, - callbackFcn: (data: boolean[], timestamp: number) => void): void { + outputStreamName: string, callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setVectorListener(outputStreamName, callbackFcn); @@ -786,8 +817,7 @@ export class GraphRunner { * should not perform overly complicated (or any async) behavior. */ attachIntListener( - outputStreamName: string, - callbackFcn: (data: number, timestamp: number) => void): void { + outputStreamName: string, callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setListener(outputStreamName, callbackFcn); @@ -807,8 +837,7 @@ export class GraphRunner { * should not perform overly complicated (or any async) behavior. */ attachIntVectorListener( - outputStreamName: string, - callbackFcn: (data: number[], timestamp: number) => void): void { + outputStreamName: string, callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setVectorListener(outputStreamName, callbackFcn); @@ -828,8 +857,7 @@ export class GraphRunner { * should not perform overly complicated (or any async) behavior. */ attachDoubleListener( - outputStreamName: string, - callbackFcn: (data: number, timestamp: number) => void): void { + outputStreamName: string, callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setListener(outputStreamName, callbackFcn); @@ -849,8 +877,7 @@ export class GraphRunner { * should not perform overly complicated (or any async) behavior. */ attachDoubleVectorListener( - outputStreamName: string, - callbackFcn: (data: number[], timestamp: number) => void): void { + outputStreamName: string, callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setVectorListener(outputStreamName, callbackFcn); @@ -870,8 +897,7 @@ export class GraphRunner { * should not perform overly complicated (or any async) behavior. */ attachFloatListener( - outputStreamName: string, - callbackFcn: (data: number, timestamp: number) => void): void { + outputStreamName: string, callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setListener(outputStreamName, callbackFcn); @@ -891,8 +917,7 @@ export class GraphRunner { * should not perform overly complicated (or any async) behavior. */ attachFloatVectorListener( - outputStreamName: string, - callbackFcn: (data: number[], timestamp: number) => void): void { + outputStreamName: string, callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setVectorListener(outputStreamName, callbackFcn); @@ -912,8 +937,7 @@ export class GraphRunner { * should not perform overly complicated (or any async) behavior. */ attachStringListener( - outputStreamName: string, - callbackFcn: (data: string, timestamp: number) => void): void { + outputStreamName: string, callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setListener(outputStreamName, callbackFcn); @@ -933,8 +957,7 @@ export class GraphRunner { * should not perform overly complicated (or any async) behavior. */ attachStringVectorListener( - outputStreamName: string, - callbackFcn: (data: string[], timestamp: number) => void): void { + outputStreamName: string, callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setVectorListener(outputStreamName, callbackFcn); @@ -964,8 +987,7 @@ export class GraphRunner { * with it). */ attachProtoListener( - outputStreamName: string, - callbackFcn: (data: Uint8Array, timestamp: number) => void, + outputStreamName: string, callbackFcn: SimpleListener, makeDeepCopy?: boolean): void { // Set up our TS listener to receive any packets for this stream. this.setListener(outputStreamName, callbackFcn); @@ -999,8 +1021,7 @@ export class GraphRunner { * with it). */ attachProtoVectorListener( - outputStreamName: string, - callbackFcn: (data: Uint8Array[], timestamp: number) => void, + outputStreamName: string, callbackFcn: SimpleListener, makeDeepCopy?: boolean): void { // Set up our TS listener to receive any packets for this stream. this.setVectorListener(outputStreamName, callbackFcn); @@ -1034,8 +1055,7 @@ export class GraphRunner { * with it). */ attachAudioListener( - outputStreamName: string, - callbackFcn: (data: Float32Array, timestamp: number) => void, + outputStreamName: string, callbackFcn: SimpleListener, makeDeepCopy?: boolean): void { if (!this.wasmModule._attachAudioListener) { console.warn( @@ -1045,13 +1065,12 @@ export class GraphRunner { // Set up our TS listener to receive any packets for this stream, and // additionally reformat our Uint8Array into a Float32Array for the user. - this.setListener( - outputStreamName, (data: Uint8Array, timestamp: number) => { - // Should be very fast - const floatArray = - new Float32Array(data.buffer, data.byteOffset, data.length / 4); - callbackFcn(floatArray, timestamp); - }); + this.setListener(outputStreamName, (data, timestamp) => { + // Should be very fast + const floatArray = + new Float32Array(data.buffer, data.byteOffset, data.length / 4); + callbackFcn(floatArray, timestamp); + }); // Tell our graph to listen for string packets on this stream. this.wrapStringPtr(outputStreamName, (outputStreamNamePtr: number) => { diff --git a/mediapipe/web/graph_runner/graph_runner_image_lib.ts b/mediapipe/web/graph_runner/graph_runner_image_lib.ts index 9608ebcc7..72d5ad965 100644 --- a/mediapipe/web/graph_runner/graph_runner_image_lib.ts +++ b/mediapipe/web/graph_runner/graph_runner_image_lib.ts @@ -1,6 +1,4 @@ -import {GraphRunner, ImageSource} from './graph_runner'; - - +import {GraphRunner, ImageSource, SimpleListener} from './graph_runner'; /** * We extend from a GraphRunner constructor. This ensures our mixin has @@ -76,7 +74,7 @@ export function SupportImage(Base: TBase) { */ attachImageListener( outputStreamName: string, - callbackFcn: (data: WasmImage, timestamp: number) => void): void { + callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setListener(outputStreamName, callbackFcn); @@ -99,7 +97,7 @@ export function SupportImage(Base: TBase) { */ attachImageVectorListener( outputStreamName: string, - callbackFcn: (data: WasmImage[], timestamp: number) => void): void { + callbackFcn: SimpleListener): void { // Set up our TS listener to receive any packets for this stream. this.setVectorListener(outputStreamName, callbackFcn);