From 874cc9dea36e1aef9c6f0cfd3c48d3cdc1dc6d3a Mon Sep 17 00:00:00 2001 From: Sebastian Schmidt Date: Fri, 28 Apr 2023 17:11:03 -0700 Subject: [PATCH] Update PoseLandmarker to return MPImage PiperOrigin-RevId: 528022223 --- mediapipe/tasks/web/vision/core/render_utils.ts | 14 -------------- mediapipe/tasks/web/vision/pose_landmarker/BUILD | 2 ++ .../web/vision/pose_landmarker/pose_landmarker.ts | 2 +- .../pose_landmarker/pose_landmarker_result.d.ts | 3 ++- .../vision/pose_landmarker/pose_landmarker_test.ts | 11 +++++------ 5 files changed, 10 insertions(+), 22 deletions(-) diff --git a/mediapipe/tasks/web/vision/core/render_utils.ts b/mediapipe/tasks/web/vision/core/render_utils.ts index 066494f57..05f2a4df1 100644 --- a/mediapipe/tasks/web/vision/core/render_utils.ts +++ b/mediapipe/tasks/web/vision/core/render_utils.ts @@ -35,20 +35,6 @@ const COLOR_MAP: Array<[number, number, number, number]> = [ [255, 255, 255, CM_ALPHA] // class 11 is white; could do black instead? ]; -/** Helper function to draw a confidence mask */ -export function drawConfidenceMask( - ctx: CanvasRenderingContext2D, image: Float32Array, width: number, - height: number): void { - const uint8ClampedArray = new Uint8ClampedArray(width * height * 4); - for (let i = 0; i < image.length; i++) { - uint8ClampedArray[4 * i] = 128; - uint8ClampedArray[4 * i + 1] = 0; - uint8ClampedArray[4 * i + 2] = 0; - uint8ClampedArray[4 * i + 3] = image[i] * 255; - } - ctx.putImageData(new ImageData(uint8ClampedArray, width, height), 0, 0); -} - /** The color converter we use in our demos. */ export const RENDER_UTIL_CONVERTER: MPImageChannelConverter = { floatToRGBAConverter: v => [128, 0, 0, v * 255], diff --git a/mediapipe/tasks/web/vision/pose_landmarker/BUILD b/mediapipe/tasks/web/vision/pose_landmarker/BUILD index 932fce8bc..8d128ac1a 100644 --- a/mediapipe/tasks/web/vision/pose_landmarker/BUILD +++ b/mediapipe/tasks/web/vision/pose_landmarker/BUILD @@ -45,6 +45,7 @@ mediapipe_ts_declaration( "//mediapipe/tasks/web/components/containers:category", "//mediapipe/tasks/web/components/containers:landmark", "//mediapipe/tasks/web/core", + "//mediapipe/tasks/web/vision/core:image", "//mediapipe/tasks/web/vision/core:vision_task_options", ], ) @@ -62,6 +63,7 @@ mediapipe_ts_library( "//mediapipe/tasks/web/components/processors:landmark_result", "//mediapipe/tasks/web/core", "//mediapipe/tasks/web/core:task_runner_test_utils", + "//mediapipe/tasks/web/vision/core:image", "//mediapipe/tasks/web/vision/core:vision_task_runner", ], ) diff --git a/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker.ts b/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker.ts index 0b119780c..44effa879 100644 --- a/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker.ts +++ b/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker.ts @@ -417,7 +417,7 @@ export class PoseLandmarker extends VisionTaskRunner { this.graphRunner.attachImageVectorListener( SEGMENTATION_MASK_STREAM, (masks, timestamp) => { this.result.segmentationMasks = - masks.map(m => m.data) as Float32Array[] | WebGLBuffer[]; + masks.map(wasmImage => this.convertToMPImage(wasmImage)); this.setLatestOutputTimestamp(timestamp); }); this.graphRunner.attachEmptyPacketListener( diff --git a/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker_result.d.ts b/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker_result.d.ts index 4ddd085df..66d0498a6 100644 --- a/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker_result.d.ts +++ b/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker_result.d.ts @@ -16,6 +16,7 @@ import {Category} from '../../../../tasks/web/components/containers/category'; import {Landmark, NormalizedLandmark} from '../../../../tasks/web/components/containers/landmark'; +import {MPImage} from '../../../../tasks/web/vision/core/image'; export {Category, Landmark, NormalizedLandmark}; @@ -34,5 +35,5 @@ export declare interface PoseLandmarkerResult { auxilaryLandmarks: NormalizedLandmark[]; /** Segmentation mask for the detected pose. */ - segmentationMasks?: Float32Array[]|WebGLTexture[]; + segmentationMasks?: MPImage[]; } diff --git a/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker_test.ts b/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker_test.ts index 1a5afba50..2d76f656f 100644 --- a/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker_test.ts +++ b/mediapipe/tasks/web/vision/pose_landmarker/pose_landmarker_test.ts @@ -18,6 +18,7 @@ import 'jasmine'; import {CalculatorGraphConfig} from '../../../../framework/calculator_pb'; import {createLandmarks, createWorldLandmarks} from '../../../../tasks/web/components/processors/landmark_result_test_lib'; import {addJasmineCustomFloatEqualityTester, createSpyWasmModule, MediapipeTasksFake, SpyWasmModule, verifyGraph} from '../../../../tasks/web/core/task_runner_test_utils'; +import {MPImage} from '../../../../tasks/web/vision/core/image'; import {VisionGraphRunner} from '../../../../tasks/web/vision/core/vision_task_runner'; import {PoseLandmarker} from './pose_landmarker'; @@ -221,12 +222,10 @@ describe('PoseLandmarker', () => { .toHaveBeenCalledTimes(1); expect(poseLandmarker.fakeWasmModule._waitUntilIdle).toHaveBeenCalled(); - expect(result).toEqual({ - 'landmarks': [{'x': 0, 'y': 0, 'z': 0}], - 'worldLandmarks': [{'x': 0, 'y': 0, 'z': 0}], - 'auxilaryLandmarks': [{'x': 0, 'y': 0, 'z': 0}], - 'segmentationMasks': [new Float32Array([0, 1, 2, 3])], - }); + expect(result.landmarks).toEqual([{'x': 0, 'y': 0, 'z': 0}]); + expect(result.worldLandmarks).toEqual([{'x': 0, 'y': 0, 'z': 0}]); + expect(result.auxilaryLandmarks).toEqual([{'x': 0, 'y': 0, 'z': 0}]); + expect(result.segmentationMasks![0]).toBeInstanceOf(MPImage); done(); }); });