TypeScript: adding VideoFrame typings support to video input

PiperOrigin-RevId: 592640146

mediapipe/tasks/web/vision/core/vision_task_runner.ts

@@ -21,7 +21,7 @@ import {MPImage} from '../../../../tasks/web/vision/core/image';
 import {ImageProcessingOptions} from '../../../../tasks/web/vision/core/image_processing_options';
 import {MPImageShaderContext} from '../../../../tasks/web/vision/core/image_shader_context';
 import {MPMask} from '../../../../tasks/web/vision/core/mask';
-import {GraphRunner, ImageSource, WasmMediaPipeConstructor} from '../../../../web/graph_runner/graph_runner';
+import {getImageSourceSize, GraphRunner, ImageSource, WasmMediaPipeConstructor} from '../../../../web/graph_runner/graph_runner';
 import {SupportImage, WasmImage} from '../../../../web/graph_runner/graph_runner_image_lib';
 import {isWebKit} from '../../../../web/graph_runner/platform_utils';
 import {SupportModelResourcesGraphService} from '../../../../web/graph_runner/register_model_resources_graph_service';
@@ -134,22 +134,6 @@ export abstract class VisionTaskRunner extends TaskRunner {
     this.process(imageFrame, imageProcessingOptions, timestamp);
   }
 
-  private getImageSourceSize(imageSource: ImageSource): [number, number] {
-    if ((imageSource as HTMLVideoElement).videoWidth !== undefined) {
-      return [
-        (imageSource as HTMLVideoElement).videoWidth,
-        (imageSource as HTMLVideoElement).videoHeight
-      ];
-    } else if ((imageSource as HTMLImageElement).naturalWidth !== undefined) {
-      return [
-        (imageSource as HTMLImageElement).naturalWidth,
-        (imageSource as HTMLImageElement).naturalHeight
-      ];
-    } else {
-      return [imageSource.width, imageSource.height];
-    }
-  }
-
   private convertToNormalizedRect(
       imageSource: ImageSource,
       imageProcessingOptions?: ImageProcessingOptions): NormalizedRect {
@@ -199,7 +183,7 @@ export abstract class VisionTaskRunner extends TaskRunner {
       //   uses this for cropping,
       // - then finally rotates this back.
       if (imageProcessingOptions?.rotationDegrees % 180 !== 0) {
-        const [imageWidth, imageHeight] = this.getImageSourceSize(imageSource);
+        const [imageWidth, imageHeight] = getImageSourceSize(imageSource);
         // tslint:disable:no-unnecessary-type-assertion
         const width = normalizedRect.getHeight()! * imageHeight / imageWidth;
         const height = normalizedRect.getWidth()! * imageWidth / imageHeight;

mediapipe/web/graph_runner/graph_runner.ts

@@ -52,6 +52,26 @@ declare global {
  */
 declare function importScripts(...urls: Array<string|URL>): void;
 
+/**
+ * Detects image source size.
+ */
+export function getImageSourceSize(imageSource: ImageSource): [number, number] {
+  if ((imageSource as HTMLVideoElement).videoWidth !== undefined) {
+    const videoElement = imageSource as HTMLVideoElement;
+    return [videoElement.videoWidth, videoElement.videoHeight];
+  } else if ((imageSource as HTMLImageElement).naturalWidth !== undefined) {
+    // TODO: Ensure this works with SVG images
+    const imageElement = imageSource as HTMLImageElement;
+    return [imageElement.naturalWidth, imageElement.naturalHeight];
+  } else if ((imageSource as VideoFrame).displayWidth !== undefined) {
+    const videoFrame = imageSource as VideoFrame;
+    return [videoFrame.displayWidth, videoFrame.displayHeight];
+  } else {
+    const notVideoFrame = imageSource as Exclude<ImageSource, VideoFrame>;
+    return [notVideoFrame.width, notVideoFrame.height];
+  }
+}
+
 /**
  * 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
@@ -64,7 +84,7 @@ export class GraphRunner implements GraphRunnerApi {
   //   should be somewhat fixed when we create our .d.ts files.
   readonly wasmModule: WasmModule;
   readonly hasMultiStreamSupport: boolean;
-  autoResizeCanvas: boolean = true;
+  autoResizeCanvas = true;
   audioPtr: number|null;
   audioSize: number;
 
@@ -196,18 +216,7 @@ export class GraphRunner implements GraphRunnerApi {
       gl.pixelStorei(gl.UNPACK_FLIP_Y_WEBGL, false);
     }
 
-    let width, height;
+    const [width, height] = getImageSourceSize(imageSource);
-    if ((imageSource as HTMLVideoElement).videoWidth) {
-      width = (imageSource as HTMLVideoElement).videoWidth;
-      height = (imageSource as HTMLVideoElement).videoHeight;
-    } else if ((imageSource as HTMLImageElement).naturalWidth) {
-      // TODO: Ensure this works with SVG images
-      width = (imageSource as HTMLImageElement).naturalWidth;
-      height = (imageSource as HTMLImageElement).naturalHeight;
-    } else {
-      width = imageSource.width;
-      height = imageSource.height;
-    }
 
     if (this.autoResizeCanvas &&
         (width !== this.wasmModule.canvas.width ||
@@ -295,7 +304,7 @@ export class GraphRunner implements GraphRunnerApi {
    * format).
    *
    * Consumers must deserialize the binary representation themselves as this
-   * avoids addding a direct dependency on the Protobuf JSPB target in the graph
+   * avoids adding a direct dependency on the Protobuf JSPB target in the graph
    * library.
    */
   getCalculatorGraphConfig(

mediapipe/web/graph_runner/graph_runner_api.d.ts

@@ -26,8 +26,8 @@ export {
 /**
  * Valid types of image sources which we can run our GraphRunner over.
  */
-export type ImageSource =
+export type ImageSource = HTMLCanvasElement|HTMLVideoElement|HTMLImageElement|
-    HTMLCanvasElement|HTMLVideoElement|HTMLImageElement|ImageData|ImageBitmap;
+    ImageData|ImageBitmap|VideoFrame;
 
 /**
  * Simple interface for a class to run an arbitrary MediaPipe graph on web, and