From db1a89324e6ffc100bf7723fbfaf2673b4f36ecc Mon Sep 17 00:00:00 2001 From: Sebastian Schmidt Date: Thu, 19 Jan 2023 10:39:05 -0800 Subject: [PATCH] Add mediapipe::Image output to the graph runner PiperOrigin-RevId: 503204918 --- .../graph_runner/graph_runner_image_lib.ts | 87 ++++++++++++++++--- 1 file changed, 74 insertions(+), 13 deletions(-) diff --git a/mediapipe/web/graph_runner/graph_runner_image_lib.ts b/mediapipe/web/graph_runner/graph_runner_image_lib.ts index 7a4ea09e..9608ebcc 100644 --- a/mediapipe/web/graph_runner/graph_runner_image_lib.ts +++ b/mediapipe/web/graph_runner/graph_runner_image_lib.ts @@ -1,4 +1,6 @@ -import {ImageSource, GraphRunner} from './graph_runner'; +import {GraphRunner, ImageSource} from './graph_runner'; + + /** * We extend from a GraphRunner constructor. This ensures our mixin has @@ -8,6 +10,12 @@ import {ImageSource, GraphRunner} from './graph_runner'; // tslint:disable-next-line:no-any type LibConstructor = new (...args: any[]) => GraphRunner; +/** An image returned from a MediaPipe graph. */ +export interface WasmImage { + data: Uint8Array|Float32Array; + width: number; + height: number; +} /** * Declarations for Emscripten's WebAssembly Module behavior, so TS compiler * doesn't break our JS/C++ bridge. @@ -16,26 +24,33 @@ export declare interface WasmImageModule { _addBoundTextureAsImageToStream: (streamNamePtr: number, width: number, height: number, timestamp: number) => void; + _attachImageListener: (streamNamePtr: number) => void; + _attachImageVectorListener: (streamNamePtr: number) => void; } /** * An implementation of GraphRunner that supports binding GPU image data as - * `mediapipe::Image` instances. We implement as a proper TS mixin, to allow for - * effective multiple inheritance. Example usage: - * `const GraphRunnerImageLib = SupportImage(GraphRunner);` + * `mediapipe::Image` instances. We implement as a proper TS mixin, to allow + * for effective multiple inheritance. Example usage: `const GraphRunnerImageLib + * = SupportImage(GraphRunner);` */ // tslint:disable-next-line:enforce-name-casing export function SupportImage(Base: TBase) { return class extends Base { + get wasmImageModule(): WasmImageModule { + return this.wasmModule as unknown as WasmImageModule; + } + /** - * Takes the relevant information from the HTML video or image element, and - * passes it into the WebGL-based graph for processing on the given stream - * at the given timestamp as a MediaPipe image. Processing will not occur - * until a blocking call (like processVideoGl or finishProcessing) is made. + * Takes the relevant information from the HTML video or image element, + * and passes it into the WebGL-based graph for processing on the given + * stream at the given timestamp as a MediaPipe image. Processing will not + * occur until a blocking call (like processVideoGl or finishProcessing) + * is made. * @param imageSource Reference to the video frame we wish to add into our * graph. - * @param streamName The name of the MediaPipe graph stream to add the frame - * to. + * @param streamName The name of the MediaPipe graph stream to add the + * frame to. * @param timestamp The timestamp of the input frame, in ms. */ addGpuBufferAsImageToStream( @@ -43,9 +58,55 @@ export function SupportImage(Base: TBase) { this.wrapStringPtr(streamName, (streamNamePtr: number) => { const [width, height] = this.bindTextureToStream(imageSource, streamNamePtr); - (this.wasmModule as unknown as WasmImageModule) - ._addBoundTextureAsImageToStream( - streamNamePtr, width, height, timestamp); + this.wasmImageModule._addBoundTextureAsImageToStream( + streamNamePtr, width, height, timestamp); + }); + } + + /** + * Attaches a mediapipe:Image packet listener to the specified output + * stream. + * @param outputStreamName The name of the graph output stream to grab + * mediapipe::Image data from. + * @param callbackFcn The function that will be called back with the data, + * as it is received. Note that the data is only guaranteed to exist + * for the duration of the callback, and the callback will be called + * inline, so it should not perform overly complicated (or any async) + * behavior. + */ + attachImageListener( + outputStreamName: string, + callbackFcn: (data: WasmImage, timestamp: number) => void): void { + // Set up our TS listener to receive any packets for this stream. + this.setListener(outputStreamName, callbackFcn); + + // Tell our graph to listen for mediapipe::Image packets on this stream. + this.wrapStringPtr(outputStreamName, (outputStreamNamePtr: number) => { + this.wasmImageModule._attachImageListener(outputStreamNamePtr); + }); + } + + /** + * Attaches a mediapipe:Image[] packet listener to the specified + * output_stream. + * @param outputStreamName The name of the graph output stream to grab + * std::vector data from. + * @param callbackFcn The function that will be called back with the data, + * as it is received. Note that the data is only guaranteed to exist + * for the duration of the callback, and the callback will be called + * inline, so it should not perform overly complicated (or any async) + * behavior. + */ + attachImageVectorListener( + outputStreamName: string, + callbackFcn: (data: WasmImage[], timestamp: number) => void): void { + // Set up our TS listener to receive any packets for this stream. + this.setVectorListener(outputStreamName, callbackFcn); + + // Tell our graph to listen for std::vector packets on + // this stream. + this.wrapStringPtr(outputStreamName, (outputStreamNamePtr: number) => { + this.wasmImageModule._attachImageVectorListener(outputStreamNamePtr); }); } };