cassdlcm/mediapipe
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update MPImage to use containers

Browse Source 
PiperOrigin-RevId: 527596164
This commit is contained in:
Sebastian Schmidt 2023-04-27 09:15:24 -07:00 committed by Copybara-Service
parent 7c70c62465
commit bc3434108e
2 changed files with 109 additions and 90 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
mediapipe/tasks/web/vision/core/image.test.ts
View File

@ -112,9 +112,7 @@ async function createTestData(
shaderContext: MPImageShaderContext, input: ImageType, width: number, shaderContext: MPImageShaderContext, input: ImageType, width: number,
height: number): MPImage { height: number): MPImage {
return new MPImage( return new MPImage(
input instanceof ImageData ? input : null, [input],
input instanceof ImageBitmap ? input : null,
input instanceof WebGLTexture ? input : null,
/* ownsImageBitmap= */ false, /* ownsWebGLTexture= */ false, canvas, /* ownsImageBitmap= */ false, /* ownsWebGLTexture= */ false, canvas,
shaderContext, width, height); shaderContext, width, height);
} }
@ -193,7 +191,7 @@ async function createTestData(
const shaderContext = new MPImageShaderContext(); const shaderContext = new MPImageShaderContext();
const image = new MPImage( const image = new MPImage(
/* imageData= */ null, /* imageBitmap= */ null, webGlTexture, [webGlTexture],
/* ownsImageBitmap= */ false, /* ownsWebGLTexture= */ false, canvas, /* ownsImageBitmap= */ false, /* ownsWebGLTexture= */ false, canvas,
shaderContext, WIDTH, HEIGHT); shaderContext, WIDTH, HEIGHT);
@ -211,7 +209,7 @@ async function createTestData(
const shaderContext = new MPImageShaderContext(); const shaderContext = new MPImageShaderContext();
const image = new MPImage( const image = new MPImage(
/* imageData= */ null, /* imageBitmap= */ null, webGlTexture, [webGlTexture],
/* ownsImageBitmap= */ false, /* ownsWebGLTexture= */ false, canvas, /* ownsImageBitmap= */ false, /* ownsWebGLTexture= */ false, canvas,
shaderContext, WIDTH, HEIGHT); shaderContext, WIDTH, HEIGHT);

147
mediapipe/tasks/web/vision/core/image.ts
View File

@ -24,7 +24,9 @@ export enum MPImageStorageType {
WEBGL_TEXTURE WEBGL_TEXTURE
} }
type MPImageNativeContainer = ImageData|ImageBitmap|WebGLTexture; /** The supported image formats. For internal usage. */
export type MPImageNativeContainer =
Uint8ClampedArray|Float32Array|ImageData|ImageBitmap|WebGLTexture;
const VERTEX_SHADER = ` const VERTEX_SHADER = `
attribute vec2 aVertex; attribute vec2 aVertex;
@ -274,9 +276,7 @@ export class MPImage {
/** @hideconstructor */ /** @hideconstructor */
constructor( constructor(
private imageData: ImageData|null, private readonly containers: MPImageNativeContainer[],
private imageBitmap: ImageBitmap|null,
private webGLTexture: WebGLTexture|null,
private ownsImageBitmap: boolean, private ownsImageBitmap: boolean,
private ownsWebGLTexture: boolean, private ownsWebGLTexture: boolean,
/** Returns the canvas element that the image is bound to. */ /** Returns the canvas element that the image is bound to. */
@ -294,15 +294,7 @@ export class MPImage {
* `getType()`. * `getType()`.
*/ */
hasType(type: MPImageStorageType): boolean { hasType(type: MPImageStorageType): boolean {
if (type === MPImageStorageType.IMAGE_DATA) { return !!this.getContainer(type);
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}`);
}
} }
/** /**
@ -336,13 +328,36 @@ export class MPImage {
*/ */
getImage(type: MPImageStorageType.WEBGL_TEXTURE): WebGLTexture; getImage(type: MPImageStorageType.WEBGL_TEXTURE): WebGLTexture;
getImage(type?: MPImageStorageType): MPImageNativeContainer { getImage(type?: MPImageStorageType): MPImageNativeContainer {
if (type === MPImageStorageType.IMAGE_DATA) { switch (type) {
case MPImageStorageType.IMAGE_DATA:
return this.convertToImageData(); return this.convertToImageData();
} else if (type === MPImageStorageType.IMAGE_BITMAP) { case MPImageStorageType.IMAGE_BITMAP:
return this.convertToImageBitmap(); return this.convertToImageBitmap();
} else if (type === MPImageStorageType.WEBGL_TEXTURE) { case MPImageStorageType.WEBGL_TEXTURE:
return this.convertToWebGLTexture(); return this.convertToWebGLTexture();
} else { default:
throw new Error(`Type is not supported: ${type}`);
}
}
private getContainer(type: MPImageStorageType.IMAGE_DATA): ImageData
|undefined;
private getContainer(type: MPImageStorageType.IMAGE_BITMAP): ImageBitmap
|undefined;
private getContainer(type: MPImageStorageType.WEBGL_TEXTURE): WebGLTexture
|undefined;
private getContainer(type: MPImageStorageType): MPImageNativeContainer
|undefined;
private getContainer(type: MPImageStorageType): MPImageNativeContainer
|undefined {
switch (type) {
case MPImageStorageType.IMAGE_DATA:
return this.containers.find(img => img instanceof ImageData);
case MPImageStorageType.IMAGE_BITMAP:
return this.containers.find(img => img instanceof ImageBitmap);
case MPImageStorageType.WEBGL_TEXTURE:
return this.containers.find(img => img instanceof WebGLTexture);
default:
throw new Error(`Type is not supported: ${type}`); throw new Error(`Type is not supported: ${type}`);
} }
} }
@ -355,32 +370,31 @@ export class MPImage {
* avoided. * avoided.
*/ */
clone(): MPImage { clone(): MPImage {
const destinationContainers: MPImageNativeContainer[] = [];
// TODO: We might only want to clone one backing datastructure // TODO: We might only want to clone one backing datastructure
// even if multiple are defined. // even if multiple are defined;
let destinationImageData: ImageData|null = null; for (const container of this.containers) {
let destinationImageBitmap: ImageBitmap|null = null; let destinationContainer: MPImageNativeContainer;
let destinationWebGLTexture: WebGLTexture|null = null;
if (this.imageData) { if (container instanceof ImageData) {
destinationImageData = destinationContainer =
new ImageData(this.imageData.data, this.width, this.height); new ImageData(container.data, this.width, this.height);
} } else if (container instanceof WebGLTexture) {
if (this.webGLTexture) {
const gl = this.getGL(); const gl = this.getGL();
const shaderContext = this.getShaderContext(); const shaderContext = this.getShaderContext();
// Create a new texture and use it to back a framebuffer // Create a new texture and use it to back a framebuffer
gl.activeTexture(gl.TEXTURE1); gl.activeTexture(gl.TEXTURE1);
destinationWebGLTexture = destinationContainer =
assertNotNull(gl.createTexture(), 'Failed to create texture'); assertNotNull(gl.createTexture(), 'Failed to create texture');
gl.bindTexture(gl.TEXTURE_2D, destinationWebGLTexture); gl.bindTexture(gl.TEXTURE_2D, destinationContainer);
gl.texImage2D( gl.texImage2D(
gl.TEXTURE_2D, 0, gl.RGBA, this.width, this.height, 0, gl.RGBA, gl.TEXTURE_2D, 0, gl.RGBA, this.width, this.height, 0, gl.RGBA,
gl.UNSIGNED_BYTE, null); gl.UNSIGNED_BYTE, null);
shaderContext.bindFramebuffer(gl, destinationWebGLTexture); shaderContext.bindFramebuffer(gl, destinationContainer);
shaderContext.run(gl, /* flipVertically= */ false, () => { shaderContext.run(gl, /* flipVertically= */ false, () => {
this.bindTexture(); // This activates gl.TEXTURE0 this.bindTexture(); // This activates gl.TEXTURE0
gl.clearColor(0, 0, 0, 0); gl.clearColor(0, 0, 0, 0);
@ -391,18 +405,21 @@ export class MPImage {
shaderContext.unbindFramebuffer(); shaderContext.unbindFramebuffer();
this.unbindTexture(); this.unbindTexture();
} } else if (container instanceof ImageBitmap) {
if (this.imageBitmap) {
this.convertToWebGLTexture(); this.convertToWebGLTexture();
this.bindTexture(); this.bindTexture();
destinationImageBitmap = this.copyTextureToBitmap(); destinationContainer = this.copyTextureToBitmap();
this.unbindTexture(); this.unbindTexture();
} else {
throw new Error(`Type is not supported: ${container}`);
}
destinationContainers.push(destinationContainer);
} }
return new MPImage( return new MPImage(
destinationImageData, destinationImageBitmap, destinationWebGLTexture, destinationContainers, this.hasType(MPImageStorageType.IMAGE_BITMAP),
!!destinationImageBitmap, !!destinationWebGLTexture, this.canvas, this.hasType(MPImageStorageType.WEBGL_TEXTURE), this.canvas,
this.shaderContext, this.width, this.height); this.shaderContext, this.width, this.height);
} }
@ -439,75 +456,82 @@ export class MPImage {
} }
private convertToImageBitmap(): ImageBitmap { private convertToImageBitmap(): ImageBitmap {
if (!this.imageBitmap) { let imageBitmap = this.getContainer(MPImageStorageType.IMAGE_BITMAP);
if (!this.webGLTexture) { if (!imageBitmap) {
this.webGLTexture = this.convertToWebGLTexture(); this.convertToWebGLTexture();
} imageBitmap = this.convertWebGLTextureToImageBitmap();
this.imageBitmap = this.convertWebGLTextureToImageBitmap(); this.containers.push(imageBitmap);
this.ownsImageBitmap = true; this.ownsImageBitmap = true;
} }
return imageBitmap;
return this.imageBitmap;
} }
private convertToImageData(): ImageData { private convertToImageData(): ImageData {
if (!this.imageData) { let imageData = this.getContainer(MPImageStorageType.IMAGE_DATA);
if (!imageData) {
const gl = this.getGL(); const gl = this.getGL();
const shaderContext = this.getShaderContext(); const shaderContext = this.getShaderContext();
const pixels = new Uint8Array(this.width * this.height * 4); const pixels = new Uint8Array(this.width * this.height * 4);
// Create texture if needed // Create texture if needed
this.convertToWebGLTexture(); const webGLTexture = this.convertToWebGLTexture();
// Create a framebuffer from the texture and read back pixels // Create a framebuffer from the texture and read back pixels
shaderContext.bindFramebuffer(gl, this.webGLTexture!); shaderContext.bindFramebuffer(gl, webGLTexture);
gl.readPixels( gl.readPixels(
0, 0, this.width, this.height, gl.RGBA, gl.UNSIGNED_BYTE, pixels); 0, 0, this.width, this.height, gl.RGBA, gl.UNSIGNED_BYTE, pixels);
shaderContext.unbindFramebuffer(); shaderContext.unbindFramebuffer();
this.imageData = new ImageData( imageData = new ImageData(
new Uint8ClampedArray(pixels.buffer), this.width, this.height); new Uint8ClampedArray(pixels.buffer), this.width, this.height);
this.containers.push(imageData);
} }
return this.imageData; return imageData;
} }
private convertToWebGLTexture(): WebGLTexture { private convertToWebGLTexture(): WebGLTexture {
if (!this.webGLTexture) { let webGLTexture = this.getContainer(MPImageStorageType.WEBGL_TEXTURE);
if (!webGLTexture) {
const gl = this.getGL(); const gl = this.getGL();
this.bindTexture(); webGLTexture = this.bindTexture();
const source = (this.imageBitmap || this.imageData)!; const source = this.getContainer(MPImageStorageType.IMAGE_BITMAP) ||
this.convertToImageData();
gl.texImage2D( gl.texImage2D(
gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, source); gl.TEXTURE_2D, 0, gl.RGBA, gl.RGBA, gl.UNSIGNED_BYTE, source);
this.unbindTexture(); this.unbindTexture();
} }
return this.webGLTexture!; return webGLTexture;
} }
/** /**
* Binds the backing texture to the canvas. If the texture does not yet * Binds the backing texture to the canvas. If the texture does not yet
* exist, creates it first. * exist, creates it first.
*/ */
private bindTexture() { private bindTexture(): WebGLTexture {
const gl = this.getGL(); const gl = this.getGL();
gl.viewport(0, 0, this.width, this.height); gl.viewport(0, 0, this.width, this.height);
gl.activeTexture(gl.TEXTURE0); gl.activeTexture(gl.TEXTURE0);
if (!this.webGLTexture) {
this.webGLTexture = let webGLTexture = this.getContainer(MPImageStorageType.WEBGL_TEXTURE);
if (!webGLTexture) {
webGLTexture =
assertNotNull(gl.createTexture(), 'Failed to create texture'); assertNotNull(gl.createTexture(), 'Failed to create texture');
this.containers.push(webGLTexture);
this.ownsWebGLTexture = true; this.ownsWebGLTexture = true;
} }
gl.bindTexture(gl.TEXTURE_2D, this.webGLTexture); gl.bindTexture(gl.TEXTURE_2D, webGLTexture);
// TODO: Ideally, we would only set these once per texture and // TODO: Ideally, we would only set these once per texture and
// not once every frame. // 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_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_WRAP_T, gl.CLAMP_TO_EDGE);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.LINEAR); gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MIN_FILTER, gl.LINEAR);
gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.LINEAR); gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.LINEAR);
return webGLTexture;
} }
private unbindTexture(): void { private unbindTexture(): void {
@ -581,15 +605,12 @@ export class MPImage {
*/ */
close(): void { close(): void {
if (this.ownsImageBitmap) { if (this.ownsImageBitmap) {
this.imageBitmap!.close(); this.getContainer(MPImageStorageType.IMAGE_BITMAP)!.close();
}
if (!this.gl) {
return;
} }
if (this.ownsWebGLTexture) { if (this.ownsWebGLTexture) {
this.gl.deleteTexture(this.webGLTexture!); const gl = this.getGL();
gl.deleteTexture(this.getContainer(MPImageStorageType.WEBGL_TEXTURE)!);
} }
} }
} }