Add mediapipe::Image output to the graph runner

PiperOrigin-RevId: 503204918

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
+ * `mediapipe::Image` instances. We implement as a proper TS mixin, to allow
- * effective multiple inheritance. Example usage:
+ * for effective multiple inheritance. Example usage: `const GraphRunnerImageLib
- * `const GraphRunnerImageLib = SupportImage(GraphRunner);`
+ * = 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
+     * Takes the relevant information from the HTML video or image element,
-     * passes it into the WebGL-based graph for processing on the given stream
+     * and passes it into the WebGL-based graph for processing on the given
-     * at the given timestamp as a MediaPipe image. Processing will not occur
+     * stream at the given timestamp as a MediaPipe image. Processing will not
-     * until a blocking call (like processVideoGl or finishProcessing) is made.
+     * 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
+     * @param streamName The name of the MediaPipe graph stream to add the
-     *     to.
+     *     frame to.
      * @param timestamp The timestamp of the input frame, in ms.
      */
     addGpuBufferAsImageToStream(
@@ -43,9 +58,55 @@ 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)
+        this.wasmImageModule._addBoundTextureAsImageToStream(
-            ._addBoundTextureAsImageToStream(
+            streamNamePtr, width, height, timestamp);
-                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);
       });
     }
   };