cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add mediapipe::Image output to the graph runner

Browse Source 
PiperOrigin-RevId: 503204918
This commit is contained in:
Sebastian Schmidt 2023-01-19 10:39:05 -08:00 committed by Copybara-Service
parent a02097ea08
commit db1a89324e
1 changed files with 74 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

85
mediapipe/web/graph_runner/graph_runner_image_lib.ts
View File

@ -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<TBase extends LibConstructor>(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,10 +58,56 @@ export function SupportImage<TBase extends LibConstructor>(Base: TBase) {
this.wrapStringPtr(streamName, (streamNamePtr: number) => {
const [width, height] =
this.bindTextureToStream(imageSource, streamNamePtr);
(this.wasmModule as unknown as WasmImageModule)
._addBoundTextureAsImageToStream(
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<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.
*/
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<mediapipe::Image> packets on
// this stream.
this.wrapStringPtr(outputStreamName, (outputStreamNamePtr: number) => {
this.wasmImageModule._attachImageVectorListener(outputStreamNamePtr);
});
}
};
}