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

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

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