Create MPImage type for Web
PiperOrigin-RevId: 525873209
This commit is contained in:
Sebastian Schmidt
Copybara-Service
parent
e9bb849503
commit
9be748db00
|
|
@ -20,6 +20,7 @@ mediapipe_files(srcs = [
|
|||
VISION_LIBS = [
|
||||
"//mediapipe/tasks/web/core:fileset_resolver",
|
||||
"//mediapipe/tasks/web/vision/core:drawing_utils",
|
||||
"//mediapipe/tasks/web/vision/core:image",
|
||||
"//mediapipe/tasks/web/vision/face_detector",
|
||||
"//mediapipe/tasks/web/vision/face_landmarker",
|
||||
"//mediapipe/tasks/web/vision/face_stylizer",
|
||||
|
|
|
|||
|
|
@ -39,6 +39,23 @@ mediapipe_ts_library(
|
|||
],
|
||||
)
|
||||
|
||||
mediapipe_ts_library(
|
||||
name = "image",
|
||||
srcs = ["image.ts"],
|
||||
)
|
||||
|
||||
mediapipe_ts_library(
|
||||
name = "image_test_lib",
|
||||
testonly = True,
|
||||
srcs = ["image.test.ts"],
|
||||
deps = [":image"],
|
||||
)
|
||||
|
||||
jasmine_node_test(
|
||||
name = "image_test",
|
||||
deps = [":image_test_lib"],
|
||||
)
|
||||
|
||||
mediapipe_ts_library(
|
||||
name = "vision_task_runner",
|
||||
srcs = ["vision_task_runner.ts"],
|
||||
|
|
|
|||
287
mediapipe/tasks/web/vision/core/image.test.ts
Normal file
287
mediapipe/tasks/web/vision/core/image.test.ts
Normal file
|
|
@ -0,0 +1,287 @@
|
|||
/**
|
||||
* 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';
|
||||
|
||||
import {MPImage, MPImageShaderContext, MPImageStorageType} from './image';
|
||||
|
||||
const WIDTH = 2;
|
||||
const HEIGHT = 2;
|
||||
|
||||
const skip = typeof document === 'undefined';
|
||||
if (skip) {
|
||||
console.log('These tests must be run in a browser.');
|
||||
}
|
||||
|
||||
/** The image types supported by MPImage. */
|
||||
type ImageType = ImageData|ImageBitmap|WebGLTexture;
|
||||
|
||||
async function createTestData(
|
||||
gl: WebGL2RenderingContext, data: number[], width: number,
|
||||
height: number): Promise<[ImageData, ImageBitmap, WebGLTexture]> {
|
||||
const imageData = new ImageData(new Uint8ClampedArray(data), width, height);
|
||||
const imageBitmap = await createImageBitmap(imageData);
|
||||
const webGlTexture = gl.createTexture()!;
|
||||
|
||||
gl.bindTexture(gl.TEXTURE_2D, webGlTexture);
|
||||
gl.texImage2D(
|
||||
gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, imageBitmap);
|
||||
gl.bindTexture(gl.TEXTURE_2D, null);
|
||||
|
||||
return [imageData, imageBitmap, webGlTexture];
|
||||
}
|
||||
|
||||
(skip ? xdescribe : describe)('MPImage', () => {
|
||||
let canvas: OffscreenCanvas;
|
||||
let gl: WebGL2RenderingContext;
|
||||
let imageData: ImageData;
|
||||
let imageBitmap: ImageBitmap;
|
||||
let webGlTexture: WebGLTexture;
|
||||
|
||||
beforeEach(async () => {
|
||||
canvas = new OffscreenCanvas(WIDTH, HEIGHT);
|
||||
gl = canvas.getContext('webgl2') as WebGL2RenderingContext;
|
||||
|
||||
const images = await createTestData(
|
||||
gl, [1, 0, 0, 255, 2, 0, 0, 255, 3, 0, 0, 255, 4, 0, 0, 255], WIDTH,
|
||||
HEIGHT);
|
||||
imageData = images[0];
|
||||
imageBitmap = images[1];
|
||||
webGlTexture = images[2];
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
gl.deleteTexture(webGlTexture);
|
||||
imageBitmap.close();
|
||||
});
|
||||
|
||||
function readPixelsFromImageBitmap(imageBitmap: ImageBitmap): ImageData {
|
||||
const canvas = new OffscreenCanvas(imageBitmap.width, imageBitmap.height);
|
||||
const ctx = canvas.getContext('2d') as OffscreenCanvasRenderingContext2D;
|
||||
ctx.drawImage(imageBitmap, 0, 0);
|
||||
return ctx.getImageData(0, 0, imageBitmap.width, imageBitmap.height);
|
||||
}
|
||||
|
||||
function readPixelsFromWebGLTexture(texture: WebGLTexture): Uint8Array {
|
||||
const pixels = new Uint8Array(WIDTH * WIDTH * 4);
|
||||
|
||||
gl.bindTexture(gl.TEXTURE_2D, texture);
|
||||
|
||||
const framebuffer = gl.createFramebuffer()!;
|
||||
gl.bindFramebuffer(gl.FRAMEBUFFER, framebuffer);
|
||||
gl.framebufferTexture2D(
|
||||
gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, texture, 0);
|
||||
gl.readPixels(0, 0, WIDTH, HEIGHT, gl.RGBA, gl.UNSIGNED_BYTE, pixels);
|
||||
gl.bindFramebuffer(gl.FRAMEBUFFER, null);
|
||||
gl.deleteFramebuffer(framebuffer);
|
||||
|
||||
gl.bindTexture(gl.TEXTURE_2D, null);
|
||||
|
||||
return pixels;
|
||||
}
|
||||
|
||||
function assertEquality(image: MPImage, expected: ImageType): void {
|
||||
if (expected instanceof ImageData) {
|
||||
const result = image.getImage(MPImageStorageType.IMAGE_DATA);
|
||||
expect(result).toEqual(expected);
|
||||
} else if (expected instanceof ImageBitmap) {
|
||||
const result = image.getImage(MPImageStorageType.IMAGE_BITMAP);
|
||||
expect(readPixelsFromImageBitmap(result))
|
||||
.toEqual(readPixelsFromImageBitmap(expected));
|
||||
} else { // WebGLTexture
|
||||
const result = image.getImage(MPImageStorageType.WEBGL_TEXTURE);
|
||||
expect(readPixelsFromWebGLTexture(result))
|
||||
.toEqual(readPixelsFromWebGLTexture(expected));
|
||||
}
|
||||
}
|
||||
|
||||
function createImage(
|
||||
shaderContext: MPImageShaderContext, input: ImageType, width: number,
|
||||
height: number): MPImage {
|
||||
return new MPImage(
|
||||
input instanceof ImageData ? input : null,
|
||||
input instanceof ImageBitmap ? input : null,
|
||||
input instanceof WebGLTexture ? input : null,
|
||||
/* ownsImageBitmap= */ false, /* ownsWebGLTexture= */ false, canvas,
|
||||
shaderContext, width, height);
|
||||
}
|
||||
|
||||
function runConversionTest(
|
||||
input: ImageType, output: ImageType, width = WIDTH,
|
||||
height = HEIGHT): void {
|
||||
const shaderContext = new MPImageShaderContext();
|
||||
const image = createImage(shaderContext, input, width, height);
|
||||
assertEquality(image, output);
|
||||
image.close();
|
||||
shaderContext.close();
|
||||
}
|
||||
|
||||
function runCloneTest(input: ImageType): void {
|
||||
const shaderContext = new MPImageShaderContext();
|
||||
const image = createImage(shaderContext, input, WIDTH, HEIGHT);
|
||||
const clone = image.clone();
|
||||
assertEquality(clone, input);
|
||||
clone.close();
|
||||
shaderContext.close();
|
||||
}
|
||||
|
||||
it(`converts from ImageData to ImageData`, () => {
|
||||
runConversionTest(imageData, imageData);
|
||||
});
|
||||
|
||||
it(`converts from ImageData to ImageBitmap`, () => {
|
||||
runConversionTest(imageData, imageBitmap);
|
||||
});
|
||||
|
||||
it(`converts from ImageData to WebGLTexture`, () => {
|
||||
runConversionTest(imageData, webGlTexture);
|
||||
});
|
||||
|
||||
it(`converts from ImageBitmap to ImageData`, () => {
|
||||
runConversionTest(imageBitmap, imageData);
|
||||
});
|
||||
|
||||
it(`converts from ImageBitmap to ImageBitmap`, () => {
|
||||
runConversionTest(imageBitmap, imageBitmap);
|
||||
});
|
||||
|
||||
it(`converts from ImageBitmap to WebGLTexture`, () => {
|
||||
runConversionTest(imageBitmap, webGlTexture);
|
||||
});
|
||||
|
||||
it(`converts from WebGLTexture to ImageData`, () => {
|
||||
runConversionTest(webGlTexture, imageData);
|
||||
});
|
||||
|
||||
it(`converts from WebGLTexture to ImageBitmap`, () => {
|
||||
runConversionTest(webGlTexture, imageBitmap);
|
||||
});
|
||||
|
||||
it(`converts from WebGLTexture to WebGLTexture`, () => {
|
||||
runConversionTest(webGlTexture, webGlTexture);
|
||||
});
|
||||
|
||||
it(`clones ImageData`, () => {
|
||||
runCloneTest(imageData);
|
||||
});
|
||||
|
||||
it(`clones ImageBitmap`, () => {
|
||||
runCloneTest(imageBitmap);
|
||||
});
|
||||
|
||||
it(`clones WebGLTextures`, () => {
|
||||
runCloneTest(webGlTexture);
|
||||
});
|
||||
|
||||
it(`does not flip textures twice`, async () => {
|
||||
const [imageData, , webGlTexture] = await createTestData(
|
||||
gl, [1, 0, 0, 255, 2, 0, 0, 255, 3, 0, 0, 255, 4, 0, 0, 255], WIDTH,
|
||||
HEIGHT);
|
||||
|
||||
const shaderContext = new MPImageShaderContext();
|
||||
const image = new MPImage(
|
||||
/* imageData= */ null, /* imageBitmap= */ null, webGlTexture,
|
||||
/* ownsImageBitmap= */ false, /* ownsWebGLTexture= */ false, canvas,
|
||||
shaderContext, WIDTH, HEIGHT);
|
||||
|
||||
const result = image.clone().getImage(MPImageStorageType.IMAGE_DATA);
|
||||
expect(result).toEqual(imageData);
|
||||
|
||||
gl.deleteTexture(webGlTexture);
|
||||
shaderContext.close();
|
||||
});
|
||||
|
||||
it(`can clone and get image`, async () => {
|
||||
const [imageData, , webGlTexture] = await createTestData(
|
||||
gl, [1, 0, 0, 255, 2, 0, 0, 255, 3, 0, 0, 255, 4, 0, 0, 255], WIDTH,
|
||||
HEIGHT);
|
||||
|
||||
const shaderContext = new MPImageShaderContext();
|
||||
const image = new MPImage(
|
||||
/* imageData= */ null, /* imageBitmap= */ null, webGlTexture,
|
||||
/* ownsImageBitmap= */ false, /* ownsWebGLTexture= */ false, canvas,
|
||||
shaderContext, WIDTH, HEIGHT);
|
||||
|
||||
// Verify that we can mix the different shader modes by running them out of
|
||||
// order.
|
||||
let result = image.getImage(MPImageStorageType.IMAGE_DATA);
|
||||
expect(result).toEqual(imageData);
|
||||
|
||||
result = image.clone().getImage(MPImageStorageType.IMAGE_DATA);
|
||||
expect(result).toEqual(imageData);
|
||||
|
||||
result = image.getImage(MPImageStorageType.IMAGE_DATA);
|
||||
expect(result).toEqual(imageData);
|
||||
|
||||
gl.deleteTexture(webGlTexture);
|
||||
shaderContext.close();
|
||||
});
|
||||
|
||||
it('supports hasType()', async () => {
|
||||
const shaderContext = new MPImageShaderContext();
|
||||
const image = createImage(shaderContext, imageData, WIDTH, HEIGHT);
|
||||
|
||||
expect(image.hasType(MPImageStorageType.IMAGE_DATA)).toBe(true);
|
||||
expect(image.hasType(MPImageStorageType.WEBGL_TEXTURE)).toBe(false);
|
||||
expect(image.hasType(MPImageStorageType.IMAGE_BITMAP)).toBe(false);
|
||||
|
||||
image.getImage(MPImageStorageType.WEBGL_TEXTURE);
|
||||
|
||||
expect(image.hasType(MPImageStorageType.IMAGE_DATA)).toBe(true);
|
||||
expect(image.hasType(MPImageStorageType.WEBGL_TEXTURE)).toBe(true);
|
||||
expect(image.hasType(MPImageStorageType.IMAGE_BITMAP)).toBe(false);
|
||||
|
||||
await image.getImage(MPImageStorageType.IMAGE_BITMAP);
|
||||
|
||||
expect(image.hasType(MPImageStorageType.IMAGE_DATA)).toBe(true);
|
||||
expect(image.hasType(MPImageStorageType.WEBGL_TEXTURE)).toBe(true);
|
||||
expect(image.hasType(MPImageStorageType.IMAGE_BITMAP)).toBe(true);
|
||||
|
||||
image.close();
|
||||
shaderContext.close();
|
||||
});
|
||||
|
||||
it('supports image that is smaller than the canvas', async () => {
|
||||
const [imageData, imageBitmap, webGlTexture] = await createTestData(
|
||||
gl, [1, 0, 0, 255, 2, 0, 0, 255], /* width= */ 2, /* height= */ 1);
|
||||
|
||||
runConversionTest(imageData, webGlTexture, /* width= */ 2, /* height= */ 1);
|
||||
runConversionTest(
|
||||
webGlTexture, imageBitmap, /* width= */ 2, /* height= */ 1);
|
||||
runConversionTest(imageBitmap, imageData, /* width= */ 2, /* height= */ 1);
|
||||
|
||||
gl.deleteTexture(webGlTexture);
|
||||
imageBitmap.close();
|
||||
});
|
||||
|
||||
it('supports image that is larger than the canvas', async () => {
|
||||
const [imageData, imageBitmap, webGlTexture] = await createTestData(
|
||||
gl,
|
||||
[
|
||||
1, 0, 0, 255, 2, 0, 0, 255, 3, 0, 0, 255,
|
||||
4, 0, 0, 255, 5, 0, 0, 255, 6, 0, 0, 255
|
||||
],
|
||||
/* width= */ 2, /* height= */ 3);
|
||||
|
||||
runConversionTest(imageData, webGlTexture, /* width= */ 2, /* height= */ 3);
|
||||
runConversionTest(
|
||||
webGlTexture, imageBitmap, /* width= */ 2, /* height= */ 3);
|
||||
runConversionTest(imageBitmap, imageData, /* width= */ 2, /* height= */ 3);
|
||||
|
||||
gl.deleteTexture(webGlTexture);
|
||||
imageBitmap.close();
|
||||
});
|
||||
});
|
||||
595
mediapipe/tasks/web/vision/core/image.ts
Normal file
595
mediapipe/tasks/web/vision/core/image.ts
Normal file
|
|
@ -0,0 +1,595 @@
|
|||
/**
|
||||
* Copyright 2023 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.
|
||||
*/
|
||||
|
||||
/** The underlying type of the image. */
|
||||
export enum MPImageStorageType {
|
||||
/** Represents the native `ImageData` type. */
|
||||
IMAGE_DATA,
|
||||
/** Represents the native `ImageBitmap` type. */
|
||||
IMAGE_BITMAP,
|
||||
/** Represents the native `WebGLTexture` type. */
|
||||
WEBGL_TEXTURE
|
||||
}
|
||||
|
||||
type MPImageNativeContainer = ImageData|ImageBitmap|WebGLTexture;
|
||||
|
||||
const VERTEX_SHADER = `
|
||||
attribute vec2 aVertex;
|
||||
attribute vec2 aTex;
|
||||
varying vec2 vTex;
|
||||
void main(void) {
|
||||
gl_Position = vec4(aVertex, 0.0, 1.0);
|
||||
vTex = aTex;
|
||||
}`;
|
||||
|
||||
const FRAGMENT_SHADER = `
|
||||
precision mediump float;
|
||||
varying vec2 vTex;
|
||||
uniform sampler2D inputTexture;
|
||||
void main() {
|
||||
gl_FragColor = texture2D(inputTexture, vTex);
|
||||
}
|
||||
`;
|
||||
|
||||
function assertNotNull<T>(value: T|null, msg: string): T {
|
||||
if (value === null) {
|
||||
throw new Error(`Unable to obtain required WebGL resource: ${msg}`);
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
/**
|
||||
* Utility class that encapsulates the buffers used by `MPImageShaderContext`.
|
||||
*/
|
||||
class MPImageShaderBuffers {
|
||||
constructor(
|
||||
private readonly gl: WebGL2RenderingContext,
|
||||
private readonly vertexArrayObject: WebGLVertexArrayObject,
|
||||
private readonly vertexBuffer: WebGLBuffer,
|
||||
private readonly textureBuffer: WebGLBuffer) {}
|
||||
|
||||
bind() {
|
||||
this.gl.bindVertexArray(this.vertexArrayObject);
|
||||
}
|
||||
|
||||
unbind() {
|
||||
this.gl.bindVertexArray(null);
|
||||
}
|
||||
|
||||
close() {
|
||||
this.gl.deleteVertexArray(this.vertexArrayObject);
|
||||
this.gl.deleteBuffer(this.vertexBuffer);
|
||||
this.gl.deleteBuffer(this.textureBuffer);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* A class that encapsulates the shaders used by an MPImage. Can be re-used
|
||||
* across MPImages that use the same WebGL2Rendering context.
|
||||
*/
|
||||
export class MPImageShaderContext {
|
||||
private gl?: WebGL2RenderingContext;
|
||||
private framebuffer?: WebGLFramebuffer;
|
||||
private program?: WebGLProgram;
|
||||
private vertexShader?: WebGLShader;
|
||||
private fragmentShader?: WebGLShader;
|
||||
private aVertex?: GLint;
|
||||
private aTex?: GLint;
|
||||
|
||||
/**
|
||||
* The shader buffers used for passthrough renders that don't modify the
|
||||
* input texture.
|
||||
*/
|
||||
private shaderBuffersPassthrough?: MPImageShaderBuffers;
|
||||
|
||||
/**
|
||||
* The shader buffers used for passthrough renders that flip the input texture
|
||||
* vertically before conversion to a different type. This is used to flip the
|
||||
* texture to the expected orientation for drawing in the browser.
|
||||
*/
|
||||
private shaderBuffersFlipVertically?: MPImageShaderBuffers;
|
||||
|
||||
private compileShader(source: string, type: number): WebGLShader {
|
||||
const gl = this.gl!;
|
||||
const shader =
|
||||
assertNotNull(gl.createShader(type), 'Failed to create WebGL shader');
|
||||
gl.shaderSource(shader, source);
|
||||
gl.compileShader(shader);
|
||||
if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {
|
||||
const info = gl.getShaderInfoLog(shader);
|
||||
throw new Error(`Could not compile WebGL shader: ${info}`);
|
||||
}
|
||||
gl.attachShader(this.program!, shader);
|
||||
return shader;
|
||||
}
|
||||
|
||||
private setupShaders(): void {
|
||||
const gl = this.gl!;
|
||||
this.program =
|
||||
assertNotNull(gl.createProgram()!, 'Failed to create WebGL program');
|
||||
|
||||
this.vertexShader = this.compileShader(VERTEX_SHADER, gl.VERTEX_SHADER);
|
||||
this.fragmentShader =
|
||||
this.compileShader(FRAGMENT_SHADER, gl.FRAGMENT_SHADER);
|
||||
|
||||
gl.linkProgram(this.program);
|
||||
const linked = gl.getProgramParameter(this.program, gl.LINK_STATUS);
|
||||
if (!linked) {
|
||||
const info = gl.getProgramInfoLog(this.program);
|
||||
throw new Error(`Error during program linking: ${info}`);
|
||||
}
|
||||
|
||||
this.aVertex = gl.getAttribLocation(this.program, 'aVertex');
|
||||
this.aTex = gl.getAttribLocation(this.program, 'aTex');
|
||||
}
|
||||
|
||||
private createBuffers(flipVertically: boolean): MPImageShaderBuffers {
|
||||
const gl = this.gl!;
|
||||
const vertexArrayObject =
|
||||
assertNotNull(gl.createVertexArray(), 'Failed to create vertex array');
|
||||
gl.bindVertexArray(vertexArrayObject);
|
||||
|
||||
const vertexBuffer =
|
||||
assertNotNull(gl.createBuffer(), 'Failed to create buffer');
|
||||
gl.bindBuffer(gl.ARRAY_BUFFER, vertexBuffer);
|
||||
gl.enableVertexAttribArray(this.aVertex!);
|
||||
gl.vertexAttribPointer(this.aVertex!, 2, gl.FLOAT, false, 0, 0);
|
||||
gl.bufferData(
|
||||
gl.ARRAY_BUFFER, new Float32Array([-1, -1, -1, 1, 1, 1, 1, -1]),
|
||||
gl.STATIC_DRAW);
|
||||
|
||||
const textureBuffer =
|
||||
assertNotNull(gl.createBuffer(), 'Failed to create buffer');
|
||||
gl.bindBuffer(gl.ARRAY_BUFFER, textureBuffer);
|
||||
gl.enableVertexAttribArray(this.aTex!);
|
||||
gl.vertexAttribPointer(this.aTex!, 2, gl.FLOAT, false, 0, 0);
|
||||
|
||||
const bufferData =
|
||||
flipVertically ? [0, 1, 0, 0, 1, 0, 1, 1] : [0, 0, 0, 1, 1, 1, 1, 0];
|
||||
gl.bufferData(
|
||||
gl.ARRAY_BUFFER, new Float32Array(bufferData), gl.STATIC_DRAW);
|
||||
|
||||
gl.bindBuffer(gl.ARRAY_BUFFER, null);
|
||||
gl.bindVertexArray(null);
|
||||
|
||||
return new MPImageShaderBuffers(
|
||||
gl, vertexArrayObject, vertexBuffer, textureBuffer);
|
||||
}
|
||||
|
||||
private getShaderBuffers(flipVertically: boolean): MPImageShaderBuffers {
|
||||
if (flipVertically) {
|
||||
if (!this.shaderBuffersFlipVertically) {
|
||||
this.shaderBuffersFlipVertically =
|
||||
this.createBuffers(/* flipVertically= */ true);
|
||||
}
|
||||
return this.shaderBuffersFlipVertically;
|
||||
} else {
|
||||
if (!this.shaderBuffersPassthrough) {
|
||||
this.shaderBuffersPassthrough =
|
||||
this.createBuffers(/* flipVertically= */ false);
|
||||
}
|
||||
return this.shaderBuffersPassthrough;
|
||||
}
|
||||
}
|
||||
|
||||
private maybeInitGL(gl: WebGL2RenderingContext): void {
|
||||
if (!this.gl) {
|
||||
this.gl = gl;
|
||||
} else if (gl !== this.gl) {
|
||||
throw new Error('Cannot change GL context once initialized');
|
||||
}
|
||||
}
|
||||
|
||||
/** Runs the callback using the shader. */
|
||||
run<T>(
|
||||
gl: WebGL2RenderingContext, flipVertically: boolean,
|
||||
callback: () => T): T {
|
||||
this.maybeInitGL(gl);
|
||||
|
||||
if (!this.program) {
|
||||
this.setupShaders();
|
||||
}
|
||||
|
||||
const shaderBuffers = this.getShaderBuffers(flipVertically);
|
||||
gl.useProgram(this.program!);
|
||||
shaderBuffers.bind();
|
||||
const result = callback();
|
||||
shaderBuffers.unbind();
|
||||
|
||||
return result;
|
||||
}
|
||||
/**
|
||||
* Binds a framebuffer to the canvas. If the framebuffer does not yet exist,
|
||||
* creates it first. Binds the provided texture to the framebuffer.
|
||||
*/
|
||||
bindFramebuffer(gl: WebGL2RenderingContext, texture: WebGLTexture): void {
|
||||
this.maybeInitGL(gl);
|
||||
if (!this.framebuffer) {
|
||||
this.framebuffer =
|
||||
assertNotNull(gl.createFramebuffer(), 'Failed to create framebuffe.');
|
||||
}
|
||||
gl.bindFramebuffer(gl.FRAMEBUFFER, this.framebuffer);
|
||||
gl.framebufferTexture2D(
|
||||
gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, texture, 0);
|
||||
}
|
||||
|
||||
unbindFramebuffer(): void {
|
||||
this.gl?.bindFramebuffer(this.gl.FRAMEBUFFER, null);
|
||||
}
|
||||
|
||||
close() {
|
||||
if (this.program) {
|
||||
const gl = this.gl!;
|
||||
gl.deleteProgram(this.program);
|
||||
gl.deleteShader(this.vertexShader!);
|
||||
gl.deleteShader(this.fragmentShader!);
|
||||
}
|
||||
if (this.framebuffer) {
|
||||
this.gl!.deleteFramebuffer(this.framebuffer);
|
||||
}
|
||||
if (this.shaderBuffersPassthrough) {
|
||||
this.shaderBuffersPassthrough.close();
|
||||
}
|
||||
if (this.shaderBuffersFlipVertically) {
|
||||
this.shaderBuffersFlipVertically.close();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* The wrapper class for MediaPipe Image objects.
|
||||
*
|
||||
* Images are stored as `ImageData`, `ImageBitmap` or `WebGLTexture` objects.
|
||||
* You can convert the underlying type to any other type by passing the
|
||||
* desired type to `getImage()`. As type conversions can be expensive, it is
|
||||
* recommended to limit these conversions. You can verify what underlying
|
||||
* types are already available by invoking `hasType()`.
|
||||
*
|
||||
* Images that are returned from a MediaPipe Tasks are owned by by the
|
||||
* underlying C++ Task. If you need to extend the lifetime of these objects,
|
||||
* you can invoke the `clone()` method. To free up the resources obtained
|
||||
* during any clone or type conversion operation, it is important to invoke
|
||||
* `close()` on the `MPImage` instance.
|
||||
*
|
||||
* Converting to and from ImageBitmap requires that the MediaPipe task is
|
||||
* initialized with an `OffscreenCanvas`. As we require WebGL2 support, this
|
||||
* places some limitations on Browser support as outlined here:
|
||||
* https://developer.mozilla.org/en-US/docs/Web/API/OffscreenCanvas/getContext
|
||||
*/
|
||||
export class MPImage {
|
||||
private gl?: WebGL2RenderingContext;
|
||||
|
||||
/** @hideconstructor */
|
||||
constructor(
|
||||
private imageData: ImageData|null,
|
||||
private imageBitmap: ImageBitmap|null,
|
||||
private webGLTexture: WebGLTexture|null,
|
||||
private ownsImageBitmap: boolean,
|
||||
private ownsWebGLTexture: boolean,
|
||||
/** Returns the canvas element that the image is bound to. */
|
||||
readonly canvas: HTMLCanvasElement|OffscreenCanvas|undefined,
|
||||
private shaderContext: MPImageShaderContext|undefined,
|
||||
/** Returns the width of the image. */
|
||||
readonly width: number,
|
||||
/** Returns the height of the image. */
|
||||
readonly height: number,
|
||||
) {}
|
||||
|
||||
/**
|
||||
* Returns whether this `MPImage` stores the image in the desired format.
|
||||
* This method can be called to reduce expensive conversion before invoking
|
||||
* `getType()`.
|
||||
*/
|
||||
hasType(type: MPImageStorageType): boolean {
|
||||
if (type === MPImageStorageType.IMAGE_DATA) {
|
||||
return !!this.imageData;
|
||||
} else if (type === MPImageStorageType.IMAGE_BITMAP) {
|
||||
return !!this.imageBitmap;
|
||||
} else if (type === MPImageStorageType.WEBGL_TEXTURE) {
|
||||
return !!this.webGLTexture;
|
||||
} else {
|
||||
throw new Error(`Type is not supported: ${type}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Returns the underlying image as an `ImageData` object. Note that this
|
||||
* involves an expensive GPU to CPU transfer if the current image is only
|
||||
* available as an `ImageBitmap` or `WebGLTexture`.
|
||||
*
|
||||
* @return The current image as an ImageData object.
|
||||
*/
|
||||
getImage(type: MPImageStorageType.IMAGE_DATA): ImageData;
|
||||
/**
|
||||
* Returns the underlying image as an `ImageBitmap`. Note that
|
||||
* conversions to `ImageBitmap` are expensive, especially if the data
|
||||
* currently resides on CPU.
|
||||
*
|
||||
* Processing with `ImageBitmap`s requires that the MediaPipe Task was
|
||||
* initialized with an `OffscreenCanvas` with WebGL2 support. See
|
||||
* https://developer.mozilla.org/en-US/docs/Web/API/OffscreenCanvas/getContext
|
||||
* for a list of supported platforms.
|
||||
*
|
||||
* @return The current image as an ImageBitmap object.
|
||||
*/
|
||||
getImage(type: MPImageStorageType.IMAGE_BITMAP): ImageBitmap;
|
||||
/**
|
||||
* Returns the underlying image as a `WebGLTexture` object. Note that this
|
||||
* involves a CPU to GPU transfer if the current image is only available as
|
||||
* an `ImageData` object. The returned texture is bound to the current
|
||||
* canvas (see `.canvas`).
|
||||
*
|
||||
* @return The current image as a WebGLTexture.
|
||||
*/
|
||||
getImage(type: MPImageStorageType.WEBGL_TEXTURE): WebGLTexture;
|
||||
getImage(type?: MPImageStorageType): MPImageNativeContainer {
|
||||
if (type === MPImageStorageType.IMAGE_DATA) {
|
||||
return this.convertToImageData();
|
||||
} else if (type === MPImageStorageType.IMAGE_BITMAP) {
|
||||
return this.convertToImageBitmap();
|
||||
} else if (type === MPImageStorageType.WEBGL_TEXTURE) {
|
||||
return this.convertToWebGLTexture();
|
||||
} else {
|
||||
throw new Error(`Type is not supported: ${type}`);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Creates a copy of the resources stored in this `MPImage`. You can invoke
|
||||
* this method to extend the lifetime of an image returned by a MediaPipe
|
||||
* Task. Note that performance critical applications should aim to only use
|
||||
* the `MPImage` within the MediaPipe Task callback so that copies can be
|
||||
* avoided.
|
||||
*/
|
||||
clone(): MPImage {
|
||||
// TODO: We might only want to clone one backing datastructure
|
||||
// even if multiple are defined.
|
||||
let destinationImageData: ImageData|null = null;
|
||||
let destinationImageBitmap: ImageBitmap|null = null;
|
||||
let destinationWebGLTexture: WebGLTexture|null = null;
|
||||
|
||||
if (this.imageData) {
|
||||
destinationImageData =
|
||||
new ImageData(this.imageData.data, this.width, this.height);
|
||||
}
|
||||
|
||||
if (this.webGLTexture) {
|
||||
const gl = this.getGL();
|
||||
const shaderContext = this.getShaderContext();
|
||||
|
||||
// Create a new texture and use it to back a framebuffer
|
||||
gl.activeTexture(gl.TEXTURE1);
|
||||
destinationWebGLTexture =
|
||||
assertNotNull(gl.createTexture(), 'Failed to create texture');
|
||||
gl.bindTexture(gl.TEXTURE_2D, destinationWebGLTexture);
|
||||
|
||||
gl.texImage2D(
|
||||
gl.TEXTURE_2D, 0, gl.RGBA, this.width, this.height, 0, gl.RGBA,
|
||||
gl.UNSIGNED_BYTE, null);
|
||||
|
||||
shaderContext.bindFramebuffer(gl, destinationWebGLTexture);
|
||||
shaderContext.run(gl, /* flipVertically= */ false, () => {
|
||||
this.bindTexture(); // This activates gl.TEXTURE0
|
||||
gl.clearColor(0, 0, 0, 0);
|
||||
gl.clear(gl.COLOR_BUFFER_BIT);
|
||||
gl.drawArrays(gl.TRIANGLE_FAN, 0, 4);
|
||||
this.unbindTexture();
|
||||
});
|
||||
shaderContext.unbindFramebuffer();
|
||||
|
||||
this.unbindTexture();
|
||||
}
|
||||
|
||||
if (this.imageBitmap) {
|
||||
this.convertToWebGLTexture();
|
||||
this.bindTexture();
|
||||
destinationImageBitmap = this.copyTextureToBitmap();
|
||||
this.unbindTexture();
|
||||
}
|
||||
|
||||
return new MPImage(
|
||||
destinationImageData, destinationImageBitmap, destinationWebGLTexture,
|
||||
!!destinationImageBitmap, !!destinationWebGLTexture, this.canvas,
|
||||
this.shaderContext, this.width, this.height);
|
||||
}
|
||||
|
||||
|
||||
private getOffscreenCanvas(): OffscreenCanvas {
|
||||
if (!(this.canvas instanceof OffscreenCanvas)) {
|
||||
throw new Error(
|
||||
'Conversion to ImageBitmap requires that the MediaPipe Tasks is ' +
|
||||
'initialized with an OffscreenCanvas');
|
||||
}
|
||||
return this.canvas;
|
||||
}
|
||||
|
||||
private getGL(): WebGL2RenderingContext {
|
||||
if (!this.canvas) {
|
||||
throw new Error(
|
||||
'Conversion to different image formats require that a canvas ' +
|
||||
'is passed when iniitializing the image.');
|
||||
}
|
||||
if (!this.gl) {
|
||||
this.gl = assertNotNull(
|
||||
this.canvas.getContext('webgl2') as WebGL2RenderingContext | null,
|
||||
'You cannot use a canvas that is already bound to a different ' +
|
||||
'type of rendering context.');
|
||||
}
|
||||
return this.gl;
|
||||
}
|
||||
|
||||
private getShaderContext(): MPImageShaderContext {
|
||||
if (!this.shaderContext) {
|
||||
this.shaderContext = new MPImageShaderContext();
|
||||
}
|
||||
return this.shaderContext;
|
||||
}
|
||||
|
||||
private convertToImageBitmap(): ImageBitmap {
|
||||
if (!this.imageBitmap) {
|
||||
if (!this.webGLTexture) {
|
||||
this.webGLTexture = this.convertToWebGLTexture();
|
||||
}
|
||||
this.imageBitmap = this.convertWebGLTextureToImageBitmap();
|
||||
this.ownsImageBitmap = true;
|
||||
}
|
||||
|
||||
return this.imageBitmap;
|
||||
}
|
||||
|
||||
private convertToImageData(): ImageData {
|
||||
if (!this.imageData) {
|
||||
const gl = this.getGL();
|
||||
const shaderContext = this.getShaderContext();
|
||||
const pixels = new Uint8Array(this.width * this.height * 4);
|
||||
|
||||
// Create texture if needed
|
||||
this.convertToWebGLTexture();
|
||||
|
||||
// Create a framebuffer from the texture and read back pixels
|
||||
shaderContext.bindFramebuffer(gl, this.webGLTexture!);
|
||||
gl.readPixels(
|
||||
0, 0, this.width, this.height, gl.RGBA, gl.UNSIGNED_BYTE, pixels);
|
||||
shaderContext.unbindFramebuffer();
|
||||
|
||||
this.imageData = new ImageData(
|
||||
new Uint8ClampedArray(pixels.buffer), this.width, this.height);
|
||||
}
|
||||
|
||||
return this.imageData;
|
||||
}
|
||||
|
||||
private convertToWebGLTexture(): WebGLTexture {
|
||||
if (!this.webGLTexture) {
|
||||
const gl = this.getGL();
|
||||
this.bindTexture();
|
||||
const source = (this.imageBitmap || this.imageData)!;
|
||||
gl.texImage2D(
|
||||
gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, source);
|
||||
this.unbindTexture();
|
||||
}
|
||||
|
||||
return this.webGLTexture!;
|
||||
}
|
||||
|
||||
/**
|
||||
* Binds the backing texture to the canvas. If the texture does not yet
|
||||
* exist, creates it first.
|
||||
*/
|
||||
private bindTexture() {
|
||||
const gl = this.getGL();
|
||||
|
||||
gl.viewport(0, 0, this.width, this.height);
|
||||
|
||||
gl.activeTexture(gl.TEXTURE0);
|
||||
if (!this.webGLTexture) {
|
||||
this.webGLTexture =
|
||||
assertNotNull(gl.createTexture(), 'Failed to create texture');
|
||||
this.ownsWebGLTexture = true;
|
||||
}
|
||||
|
||||
gl.bindTexture(gl.TEXTURE_2D, this.webGLTexture);
|
||||
// TODO: Ideally, we would only set these once per texture and
|
||||
// not once every frame.
|
||||
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_S, gl.CLAMP_TO_EDGE);
|
||||
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_WRAP_T, gl.CLAMP_TO_EDGE);
|
||||
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.LINEAR);
|
||||
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.LINEAR);
|
||||
}
|
||||
|
||||
private unbindTexture(): void {
|
||||
this.gl!.bindTexture(this.gl!.TEXTURE_2D, null);
|
||||
}
|
||||
|
||||
/**
|
||||
* Invokes a shader to render the current texture and return it as an
|
||||
* ImageBitmap
|
||||
*/
|
||||
private copyTextureToBitmap(): ImageBitmap {
|
||||
const gl = this.getGL();
|
||||
const shaderContext = this.getShaderContext();
|
||||
|
||||
return shaderContext.run(gl, /* flipVertically= */ true, () => {
|
||||
return this.runWithResizedCanvas(() => {
|
||||
// Unbind any framebuffer that may be bound since
|
||||
// `transferToImageBitmap()` requires rendering into the display (null)
|
||||
// framebuffer.
|
||||
gl.bindFramebuffer(gl.FRAMEBUFFER, null);
|
||||
|
||||
gl.clearColor(0, 0, 0, 0);
|
||||
gl.clear(gl.COLOR_BUFFER_BIT);
|
||||
gl.drawArrays(gl.TRIANGLE_FAN, 0, 4);
|
||||
return this.getOffscreenCanvas().transferToImageBitmap();
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
private convertWebGLTextureToImageBitmap(): ImageBitmap {
|
||||
this.bindTexture();
|
||||
const result = this.copyTextureToBitmap();
|
||||
this.unbindTexture();
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Temporarily resizes the underlying canvas to match the dimensions of the
|
||||
* image. Runs the provided callback on the resized canvas.
|
||||
*
|
||||
* Note that while resizing is an expensive operation, it allows us to use
|
||||
* the synchronous `transferToImageBitmap()` API.
|
||||
*/
|
||||
private runWithResizedCanvas<T>(callback: () => T): T {
|
||||
const canvas = this.canvas!;
|
||||
|
||||
if (canvas.width === this.width && canvas.height === this.height) {
|
||||
return callback();
|
||||
}
|
||||
|
||||
const originalWidth = canvas.width;
|
||||
const originalHeight = canvas.height;
|
||||
canvas.width = this.width;
|
||||
canvas.height = this.height;
|
||||
|
||||
const result = callback();
|
||||
|
||||
canvas.width = originalWidth;
|
||||
canvas.height = originalHeight;
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
/**
|
||||
* Frees up any resources owned by this `MPImage` instance.
|
||||
*
|
||||
* Note that this method does not free images that are owned by the C++
|
||||
* Task, as these are freed automatically once you leave the MediaPipe
|
||||
* callback. Additionally, some shared state is freed only once you invoke the
|
||||
* Task's `close()` method.
|
||||
*/
|
||||
close(): void {
|
||||
if (this.ownsImageBitmap) {
|
||||
this.imageBitmap!.close();
|
||||
}
|
||||
|
||||
if (!this.gl) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (this.ownsWebGLTexture) {
|
||||
this.gl.deleteTexture(this.webGLTexture!);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -16,6 +16,7 @@
|
|||
|
||||
import {FilesetResolver as FilesetResolverImpl} from '../../../tasks/web/core/fileset_resolver';
|
||||
import {DrawingUtils as DrawingUtilsImpl} from '../../../tasks/web/vision/core/drawing_utils';
|
||||
import {MPImage as MPImageImpl} from '../../../tasks/web/vision/core/image';
|
||||
import {FaceDetector as FaceDetectorImpl} from '../../../tasks/web/vision/face_detector/face_detector';
|
||||
import {FaceLandmarker as FaceLandmarkerImpl, FaceLandmarksConnections as FaceLandmarksConnectionsImpl} from '../../../tasks/web/vision/face_landmarker/face_landmarker';
|
||||
import {FaceStylizer as FaceStylizerImpl} from '../../../tasks/web/vision/face_stylizer/face_stylizer';
|
||||
|
|
@ -31,6 +32,7 @@ import {ObjectDetector as ObjectDetectorImpl} from '../../../tasks/web/vision/ob
|
|||
// as exports.
|
||||
const DrawingUtils = DrawingUtilsImpl;
|
||||
const FilesetResolver = FilesetResolverImpl;
|
||||
const MPImage = MPImageImpl;
|
||||
const FaceDetector = FaceDetectorImpl;
|
||||
const FaceLandmarker = FaceLandmarkerImpl;
|
||||
const FaceLandmarksConnections = FaceLandmarksConnectionsImpl;
|
||||
|
|
@ -46,6 +48,7 @@ const ObjectDetector = ObjectDetectorImpl;
|
|||
export {
|
||||
DrawingUtils,
|
||||
FilesetResolver,
|
||||
MPImage,
|
||||
FaceDetector,
|
||||
FaceLandmarker,
|
||||
FaceLandmarksConnections,
|
||||
|
|
|
|||
|
|
@ -16,6 +16,7 @@
|
|||
|
||||
export * from '../../../tasks/web/core/fileset_resolver';
|
||||
export * from '../../../tasks/web/vision/core/drawing_utils';
|
||||
export * from '../../../tasks/web/vision/core/image';
|
||||
export * from '../../../tasks/web/vision/face_detector/face_detector';
|
||||
export * from '../../../tasks/web/vision/face_landmarker/face_landmarker';
|
||||
export * from '../../../tasks/web/vision/face_stylizer/face_stylizer';
|
||||
|
|
|
|||
Loading…
Reference in New Issue
Block a user