Add the FaceStylizer Web API

PiperOrigin-RevId: 518409812
This commit is contained in:
Sebastian Schmidt 2023-03-21 16:13:58 -07:00 committed by Copybara-Service
parent b71e1d14d3
commit bbd21e9a6d
10 changed files with 546 additions and 1 deletions

View File

@ -19,6 +19,7 @@ mediapipe_files(srcs = [
VISION_LIBS = [
"//mediapipe/tasks/web/core:fileset_resolver",
"//mediapipe/tasks/web/vision/face_stylizer",
"//mediapipe/tasks/web/vision/gesture_recognizer",
"//mediapipe/tasks/web/vision/hand_landmarker",
"//mediapipe/tasks/web/vision/image_classifier",

View File

@ -2,6 +2,21 @@
This package contains the vision tasks for MediaPipe.
## Face Stylizer
The MediaPipe Face Stylizer lets you perform face stylization on images.
```
const vision = await FilesetResolver.forVisionTasks(
"https://cdn.jsdelivr.net/npm/@mediapipe/tasks-vision@latest/wasm"
);
const faceStylizer = await FaceStylizer.createFromModelPath(vision,
"model.tflite"
);
const image = document.getElementById("image") as HTMLImageElement;
const stylizedImage = faceStylizer.stylize(image);
```
## Gesture Recognition
The MediaPipe Gesture Recognizer task lets you recognize hand gestures in real

View File

@ -35,6 +35,17 @@ export type SegmentationMask = Uint8ClampedArray|Float32Array|WebGLTexture;
export type SegmentationMaskCallback =
(masks: SegmentationMask[], width: number, height: number) => void;
/**
* A callback that receives an `ImageData` object from a Vision task. The
* lifetime of the underlying data is limited to the duration of the callback.
* If asynchronous processing is needed, all data needs to be copied before the
* callback returns.
*
* The `WebGLTexture` output type is reserved for future usage.
*/
export type ImageCallback =
(image: ImageData|WebGLTexture, width: number, height: number) => void;
/** A Region-Of-Interest (ROI) to represent a region within an image. */
export declare interface RegionOfInterest {
/** The ROI in keypoint format. */

View File

@ -18,7 +18,7 @@ import {NormalizedRect} from '../../../../framework/formats/rect_pb';
import {TaskRunner} from '../../../../tasks/web/core/task_runner';
import {ImageProcessingOptions} from '../../../../tasks/web/vision/core/image_processing_options';
import {GraphRunner, ImageSource} from '../../../../web/graph_runner/graph_runner';
import {SupportImage} from '../../../../web/graph_runner/graph_runner_image_lib';
import {SupportImage, WasmImage} from '../../../../web/graph_runner/graph_runner_image_lib';
import {SupportModelResourcesGraphService} from '../../../../web/graph_runner/register_model_resources_graph_service';
import {VisionTaskOptions} from './vision_task_options';
@ -148,6 +148,31 @@ export abstract class VisionTaskRunner extends TaskRunner {
imageSource, this.imageStreamName, timestamp ?? performance.now());
this.finishProcessing();
}
/** Converts the RGB or RGBA Uint8Array of a WasmImage to ImageData. */
protected convertToImageData(wasmImage: WasmImage): ImageData {
const {data, width, height} = wasmImage;
if (!(data instanceof Uint8ClampedArray)) {
throw new Error(
'Only Uint8ClampedArray-based images can be converted to ImageData');
}
if (data.length === width * height * 4) {
return new ImageData(data, width, height);
} else if (data.length === width * height * 3) {
const rgba = new Uint8ClampedArray(width * height * 4);
for (let i = 0; i < width * height; ++i) {
rgba[4 * i] = data[3 * i];
rgba[4 * i + 1] = data[3 * i + 1];
rgba[4 * i + 2] = data[3 * i + 2];
rgba[4 * i + 3] = 255;
}
return new ImageData(rgba, width, height);
} else {
throw new Error(
`Unsupported channel count: ${data.length / width / height}`);
}
}
}

View File

@ -0,0 +1,57 @@
# This contains the MediaPipe Face Stylizer Task.
load("//mediapipe/framework/port:build_config.bzl", "mediapipe_ts_declaration", "mediapipe_ts_library")
load("@npm//@bazel/jasmine:index.bzl", "jasmine_node_test")
package(default_visibility = ["//mediapipe/tasks:internal"])
licenses(["notice"])
mediapipe_ts_library(
name = "face_stylizer",
srcs = ["face_stylizer.ts"],
deps = [
":face_stylizer_types",
"//mediapipe/framework:calculator_jspb_proto",
"//mediapipe/framework:calculator_options_jspb_proto",
"//mediapipe/tasks/cc/core/proto:base_options_jspb_proto",
"//mediapipe/tasks/cc/vision/face_stylizer/proto:face_stylizer_graph_options_jspb_proto",
"//mediapipe/tasks/web/core",
"//mediapipe/tasks/web/vision/core:image_processing_options",
"//mediapipe/tasks/web/vision/core:types",
"//mediapipe/tasks/web/vision/core:vision_task_runner",
"//mediapipe/web/graph_runner:graph_runner_ts",
],
)
mediapipe_ts_declaration(
name = "face_stylizer_types",
srcs = ["face_stylizer_options.d.ts"],
deps = [
"//mediapipe/tasks/web/core",
"//mediapipe/tasks/web/core:classifier_options",
"//mediapipe/tasks/web/vision/core:vision_task_options",
],
)
mediapipe_ts_library(
name = "face_stylizer_test_lib",
testonly = True,
srcs = [
"face_stylizer_test.ts",
],
deps = [
":face_stylizer",
":face_stylizer_types",
"//mediapipe/framework:calculator_jspb_proto",
"//mediapipe/tasks/web/core",
"//mediapipe/tasks/web/core:task_runner_test_utils",
"//mediapipe/web/graph_runner:graph_runner_image_lib_ts",
],
)
jasmine_node_test(
name = "face_stylizer_test",
tags = ["nomsan"],
deps = [":face_stylizer_test_lib"],
)

View File

@ -0,0 +1,298 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import {CalculatorGraphConfig} from '../../../../framework/calculator_pb';
import {CalculatorOptions} from '../../../../framework/calculator_options_pb';
import {BaseOptions as BaseOptionsProto} from '../../../../tasks/cc/core/proto/base_options_pb';
import {FaceStylizerGraphOptions as FaceStylizerGraphOptionsProto} from '../../../../tasks/cc/vision/face_stylizer/proto/face_stylizer_graph_options_pb';
import {WasmFileset} from '../../../../tasks/web/core/wasm_fileset';
import {ImageProcessingOptions} from '../../../../tasks/web/vision/core/image_processing_options';
import {ImageCallback} from '../../../../tasks/web/vision/core/types';
import {VisionGraphRunner, VisionTaskRunner} from '../../../../tasks/web/vision/core/vision_task_runner';
import {ImageSource, WasmModule} from '../../../../web/graph_runner/graph_runner';
// Placeholder for internal dependency on trusted resource url
import {FaceStylizerOptions} from './face_stylizer_options';
export * from './face_stylizer_options';
export {ImageSource}; // Used in the public API
const IMAGE_STREAM = 'image_in';
const NORM_RECT_STREAM = 'norm_rect';
const STYLIZED_IMAGE_STREAM = 'stylized_image';
const FACE_STYLIZER_GRAPH =
'mediapipe.tasks.vision.face_stylizer.FaceStylizerGraph';
// The OSS JS API does not support the builder pattern.
// tslint:disable:jspb-use-builder-pattern
export {ImageCallback};
/** Performs face stylization on images. */
export class FaceStylizer extends VisionTaskRunner {
private userCallback: ImageCallback = () => {};
private readonly options: FaceStylizerGraphOptionsProto;
/**
* Initializes the Wasm runtime and creates a new Face Stylizer from the
* provided options.
* @param wasmFileset A configuration object that provides the location of
* the Wasm binary and its loader.
* @param faceStylizerOptions The options for the Face Stylizer. Note
* that either a path to the model asset or a model buffer needs to be
* provided (via `baseOptions`).
*/
static createFromOptions(
wasmFileset: WasmFileset,
faceStylizerOptions: FaceStylizerOptions): Promise<FaceStylizer> {
return VisionTaskRunner.createInstance(
FaceStylizer, /* initializeCanvas= */ true, wasmFileset,
faceStylizerOptions);
}
/**
* Initializes the Wasm runtime and creates a new Face Stylizer based on
* the provided model asset buffer.
* @param wasmFileset A configuration object that provides the location of
* the Wasm binary and its loader.
* @param modelAssetBuffer A binary representation of the model.
*/
static createFromModelBuffer(
wasmFileset: WasmFileset,
modelAssetBuffer: Uint8Array): Promise<FaceStylizer> {
return VisionTaskRunner.createInstance(
FaceStylizer, /* initializeCanvas= */ true, wasmFileset,
{baseOptions: {modelAssetBuffer}});
}
/**
* Initializes the Wasm runtime and creates a new Face Stylizer based on
* the path to the model asset.
* @param wasmFileset A configuration object that provides the location of
* the Wasm binary and its loader.
* @param modelAssetPath The path to the model asset.
*/
static createFromModelPath(
wasmFileset: WasmFileset,
modelAssetPath: string): Promise<FaceStylizer> {
return VisionTaskRunner.createInstance(
FaceStylizer, /* initializeCanvas= */ true, wasmFileset,
{baseOptions: {modelAssetPath}});
}
/** @hideconstructor */
constructor(
wasmModule: WasmModule,
glCanvas?: HTMLCanvasElement|OffscreenCanvas|null) {
super(
new VisionGraphRunner(wasmModule, glCanvas), IMAGE_STREAM,
NORM_RECT_STREAM, /* roiAllowed= */ true);
this.options = new FaceStylizerGraphOptionsProto();
this.options.setBaseOptions(new BaseOptionsProto());
}
protected override get baseOptions(): BaseOptionsProto {
return this.options.getBaseOptions()!;
}
protected override set baseOptions(proto: BaseOptionsProto) {
this.options.setBaseOptions(proto);
}
/**
* Sets new options for the Face Stylizer.
*
* Calling `setOptions()` with a subset of options only affects those
* options. You can reset an option back to its default value by
* explicitly setting it to `undefined`.
*
* @param options The options for the Face Stylizer.
*/
override setOptions(options: FaceStylizerOptions): Promise<void> {
return super.applyOptions(options);
}
/**
* Performs face stylization on the provided single image. The method returns
* synchronously once the callback returns. Only use this method when the
* FaceStylizer is created with the image running mode.
*
* The input image can be of any size. To ensure that the output image has
* reasonable quailty, the stylized output image size is determined by the
* model output size.
*
* @param image An image to process.
* @param callback The callback that is invoked with the stylized image. The
* lifetime of the returned data is only guaranteed for the duration of the
* callback.
*/
stylize(image: ImageSource, callback: ImageCallback): void;
/**
* Performs face stylization on the provided single image. The method returns
* synchronously once the callback returns. Only use this method when the
* FaceStylizer is created with the image running mode.
*
* The 'imageProcessingOptions' parameter can be used to specify one or all
* of:
* - the rotation to apply to the image before performing stylization, by
* setting its 'rotationDegrees' property.
* - the region-of-interest on which to perform stylization, by setting its
* 'regionOfInterest' property. If not specified, the full image is used.
* If both are specified, the crop around the region-of-interest is extracted
* first, then the specified rotation is applied to the crop.
*
* The input image can be of any size. To ensure that the output image has
* reasonable quailty, the stylized output image size is the smaller of the
* model output size and the size of the 'regionOfInterest' specified in
* 'imageProcessingOptions'.
*
* @param image An image to process.
* @param imageProcessingOptions the `ImageProcessingOptions` specifying how
* to process the input image before running inference.
* @param callback The callback that is invoked with the stylized image. The
* lifetime of the returned data is only guaranteed for the duration of the
* callback.
*/
stylize(
image: ImageSource, imageProcessingOptions: ImageProcessingOptions,
callback: ImageCallback): void;
stylize(
image: ImageSource,
imageProcessingOptionsOrCallback: ImageProcessingOptions|ImageCallback,
callback?: ImageCallback): void {
const imageProcessingOptions =
typeof imageProcessingOptionsOrCallback !== 'function' ?
imageProcessingOptionsOrCallback :
{};
this.userCallback = typeof imageProcessingOptionsOrCallback === 'function' ?
imageProcessingOptionsOrCallback :
callback!;
this.processImageData(image, imageProcessingOptions ?? {});
this.userCallback = () => {};
}
/**
* Performs face stylization on the provided video frame. Only use this method
* when the FaceStylizer is created with the video running mode.
*
* The input frame can be of any size. It's required to provide the video
* frame's timestamp (in milliseconds). The input timestamps must be
* monotonically increasing.
*
* To ensure that the output image has reasonable quality, the stylized
* output image size is determined by the model output size.
*
* @param videoFrame A video frame to process.
* @param timestamp The timestamp of the current frame, in ms.
* @param callback The callback that is invoked with the stylized image. The
* lifetime of the returned data is only guaranteed for the duration of
* the callback.
*/
stylizeForVideo(
videoFrame: ImageSource, timestamp: number,
callback: ImageCallback): void;
/**
* Performs face stylization on the provided video frame. Only use this
* method when the FaceStylizer is created with the video running mode.
*
* The 'imageProcessingOptions' parameter can be used to specify one or all
* of:
* - the rotation to apply to the image before performing stylization, by
* setting its 'rotationDegrees' property.
* - the region-of-interest on which to perform stylization, by setting its
* 'regionOfInterest' property. If not specified, the full image is used.
* If both are specified, the crop around the region-of-interest is
* extracted first, then the specified rotation is applied to the crop.
*
* The input frame can be of any size. It's required to provide the video
* frame's timestamp (in milliseconds). The input timestamps must be
* monotonically increasing.
*
* To ensure that the output image has reasonable quailty, the stylized
* output image size is the smaller of the model output size and the size of
* the 'regionOfInterest' specified in 'imageProcessingOptions'.
*
* @param videoFrame A video frame to process.
* @param imageProcessingOptions the `ImageProcessingOptions` specifying how
* to process the input image before running inference.
* @param timestamp The timestamp of the current frame, in ms.
* @param callback The callback that is invoked with the stylized image. The
* lifetime of the returned data is only guaranteed for the duration of
* the callback.
*/
stylizeForVideo(
videoFrame: ImageSource, imageProcessingOptions: ImageProcessingOptions,
timestamp: number, callback: ImageCallback): void;
stylizeForVideo(
videoFrame: ImageSource,
timestampOrImageProcessingOptions: number|ImageProcessingOptions,
timestampOrCallback: number|ImageCallback,
callback?: ImageCallback): void {
const imageProcessingOptions =
typeof timestampOrImageProcessingOptions !== 'number' ?
timestampOrImageProcessingOptions :
{};
const timestamp = typeof timestampOrImageProcessingOptions === 'number' ?
timestampOrImageProcessingOptions :
timestampOrCallback as number;
this.userCallback = typeof timestampOrCallback === 'function' ?
timestampOrCallback :
callback!;
this.processVideoData(videoFrame, imageProcessingOptions, timestamp);
this.userCallback = () => {};
}
/** Updates the MediaPipe graph configuration. */
protected override refreshGraph(): void {
const graphConfig = new CalculatorGraphConfig();
graphConfig.addInputStream(IMAGE_STREAM);
graphConfig.addInputStream(NORM_RECT_STREAM);
graphConfig.addOutputStream(STYLIZED_IMAGE_STREAM);
const calculatorOptions = new CalculatorOptions();
calculatorOptions.setExtension(
FaceStylizerGraphOptionsProto.ext, this.options);
const segmenterNode = new CalculatorGraphConfig.Node();
segmenterNode.setCalculator(FACE_STYLIZER_GRAPH);
segmenterNode.addInputStream('IMAGE:' + IMAGE_STREAM);
segmenterNode.addInputStream('NORM_RECT:' + NORM_RECT_STREAM);
segmenterNode.addOutputStream('STYLIZED_IMAGE:' + STYLIZED_IMAGE_STREAM);
segmenterNode.setOptions(calculatorOptions);
graphConfig.addNode(segmenterNode);
this.graphRunner.attachImageListener(
STYLIZED_IMAGE_STREAM, (image, timestamp) => {
const imageData = this.convertToImageData(image);
this.userCallback(imageData, image.width, image.height);
this.setLatestOutputTimestamp(timestamp);
});
this.graphRunner.attachEmptyPacketListener(
STYLIZED_IMAGE_STREAM, timestamp => {
this.setLatestOutputTimestamp(timestamp);
});
const binaryGraph = graphConfig.serializeBinary();
this.setGraph(new Uint8Array(binaryGraph), /* isBinary= */ true);
}
}

View File

@ -0,0 +1,20 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import {VisionTaskOptions} from '../../../../tasks/web/vision/core/vision_task_options';
/** Options to configure the MediaPipe Face Stylizer Task */
export interface FaceStylizerOptions extends VisionTaskOptions {}

View File

@ -0,0 +1,114 @@
/**
* Copyright 2022 The MediaPipe Authors. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import 'jasmine';
// Placeholder for internal dependency on encodeByteArray
import {CalculatorGraphConfig} from '../../../../framework/calculator_pb';
import {addJasmineCustomFloatEqualityTester, createSpyWasmModule, MediapipeTasksFake, SpyWasmModule, verifyGraph, verifyListenersRegistered} from '../../../../tasks/web/core/task_runner_test_utils';
import {WasmImage} from '../../../../web/graph_runner/graph_runner_image_lib';
import {FaceStylizer} from './face_stylizer';
class FaceStylizerFake extends FaceStylizer implements MediapipeTasksFake {
calculatorName = 'mediapipe.tasks.vision.face_stylizer.FaceStylizerGraph';
attachListenerSpies: jasmine.Spy[] = [];
graph: CalculatorGraphConfig|undefined;
fakeWasmModule: SpyWasmModule;
imageListener: ((images: WasmImage, timestamp: number) => void)|undefined;
constructor() {
super(createSpyWasmModule(), /* glCanvas= */ null);
this.fakeWasmModule =
this.graphRunner.wasmModule as unknown as SpyWasmModule;
this.attachListenerSpies[0] =
spyOn(this.graphRunner, 'attachImageListener')
.and.callFake((stream, listener) => {
expect(stream).toEqual('stylized_image');
this.imageListener = listener;
});
spyOn(this.graphRunner, 'setGraph').and.callFake(binaryGraph => {
this.graph = CalculatorGraphConfig.deserializeBinary(binaryGraph);
});
spyOn(this.graphRunner, 'addGpuBufferAsImageToStream');
}
}
describe('FaceStylizer', () => {
let faceStylizer: FaceStylizerFake;
beforeEach(async () => {
addJasmineCustomFloatEqualityTester();
faceStylizer = new FaceStylizerFake();
await faceStylizer.setOptions(
{baseOptions: {modelAssetBuffer: new Uint8Array([])}});
});
it('initializes graph', async () => {
verifyGraph(faceStylizer);
verifyListenersRegistered(faceStylizer);
});
it('can use custom models', async () => {
const newModel = new Uint8Array([0, 1, 2, 3, 4]);
const newModelBase64 = Buffer.from(newModel).toString('base64');
await faceStylizer.setOptions({
baseOptions: {
modelAssetBuffer: newModel,
}
});
verifyGraph(
faceStylizer,
/* expectedCalculatorOptions= */ undefined,
/* expectedBaseOptions= */
[
'modelAsset', {
fileContent: newModelBase64,
fileName: undefined,
fileDescriptorMeta: undefined,
filePointerMeta: undefined
}
]);
});
it('invokes callback', (done) => {
if (typeof ImageData === 'undefined') {
console.log('ImageData tests are not supported on Node');
done();
return;
}
// Pass the test data to our listener
faceStylizer.fakeWasmModule._waitUntilIdle.and.callFake(() => {
verifyListenersRegistered(faceStylizer);
faceStylizer.imageListener!
({data: new Uint8ClampedArray([1, 1, 1, 1]), width: 1, height: 1},
/* timestamp= */ 1337);
});
// Invoke the face stylizeer
faceStylizer.stylize({} as HTMLImageElement, (image, width, height) => {
expect(faceStylizer.fakeWasmModule._waitUntilIdle).toHaveBeenCalled();
expect(image).toBeInstanceOf(ImageData);
expect(width).toEqual(1);
expect(height).toEqual(1);
done();
});
});
});

View File

@ -15,6 +15,7 @@
*/
import {FilesetResolver as FilesetResolverImpl} from '../../../tasks/web/core/fileset_resolver';
import {FaceStylizer as FaceStylizerImpl} from '../../../tasks/web/vision/face_stylizer/face_stylizer';
import {GestureRecognizer as GestureRecognizerImpl} from '../../../tasks/web/vision/gesture_recognizer/gesture_recognizer';
import {HandLandmarker as HandLandmarkerImpl} from '../../../tasks/web/vision/hand_landmarker/hand_landmarker';
import {ImageClassifier as ImageClassifierImpl} from '../../../tasks/web/vision/image_classifier/image_classifier';
@ -26,6 +27,7 @@ import {ObjectDetector as ObjectDetectorImpl} from '../../../tasks/web/vision/ob
// Declare the variables locally so that Rollup in OSS includes them explicitly
// as exports.
const FilesetResolver = FilesetResolverImpl;
const FaceStylizer = FaceStylizerImpl;
const GestureRecognizer = GestureRecognizerImpl;
const HandLandmarker = HandLandmarkerImpl;
const ImageClassifier = ImageClassifierImpl;
@ -36,6 +38,7 @@ const ObjectDetector = ObjectDetectorImpl;
export {
FilesetResolver,
FaceStylizer,
GestureRecognizer,
HandLandmarker,
ImageClassifier,

View File

@ -15,6 +15,7 @@
*/
export * from '../../../tasks/web/core/fileset_resolver';
export * from '../../../tasks/web/vision/face_stylizer/face_stylizer';
export * from '../../../tasks/web/vision/gesture_recognizer/gesture_recognizer';
export * from '../../../tasks/web/vision/hand_landmarker/hand_landmarker';
export * from '../../../tasks/web/vision/image_classifier/image_classifier';