From 52c1d44561129698be2e83136f84026d5bb959c1 Mon Sep 17 00:00:00 2001 From: MediaPipe Team Date: Wed, 20 Dec 2023 13:14:10 -0800 Subject: [PATCH] TypeScript: adding VideoFrame typings support to video input PiperOrigin-RevId: 592640146 --- .../web/vision/core/vision_task_runner.ts | 20 +--------- mediapipe/web/graph_runner/graph_runner.ts | 37 ++++++++++++------- .../web/graph_runner/graph_runner_api.d.ts | 4 +- 3 files changed, 27 insertions(+), 34 deletions(-) diff --git a/mediapipe/tasks/web/vision/core/vision_task_runner.ts b/mediapipe/tasks/web/vision/core/vision_task_runner.ts index 292a37eec..4353d51e4 100644 --- a/mediapipe/tasks/web/vision/core/vision_task_runner.ts +++ b/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; diff --git a/mediapipe/web/graph_runner/graph_runner.ts b/mediapipe/web/graph_runner/graph_runner.ts index 06b615502..510e9908a 100644 --- a/mediapipe/web/graph_runner/graph_runner.ts +++ b/mediapipe/web/graph_runner/graph_runner.ts @@ -52,6 +52,26 @@ declare global { */ declare function importScripts(...urls: Array): 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; + 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; - 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; - } + const [width, height] = getImageSourceSize(imageSource); 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( diff --git a/mediapipe/web/graph_runner/graph_runner_api.d.ts b/mediapipe/web/graph_runner/graph_runner_api.d.ts index f6318416f..13ad0456f 100644 --- a/mediapipe/web/graph_runner/graph_runner_api.d.ts +++ b/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 = - HTMLCanvasElement|HTMLVideoElement|HTMLImageElement|ImageData|ImageBitmap; +export type ImageSource = HTMLCanvasElement|HTMLVideoElement|HTMLImageElement| + ImageData|ImageBitmap|VideoFrame; /** * Simple interface for a class to run an arbitrary MediaPipe graph on web, and