Internal change
PiperOrigin-RevId: 529495239
This commit is contained in:
@@ -106,7 +106,7 @@ class Dataset(classification_dataset.ClassificationDataset):
|
||||
...
|
||||
Each <file0>.xml annotation file should have the following format:
|
||||
<annotation>
|
||||
<filename>file0.jpg<filename>
|
||||
<filename>file0.jpg</filename>
|
||||
<object>
|
||||
<name>kangaroo</name>
|
||||
<bndbox>
|
||||
@@ -114,6 +114,7 @@ class Dataset(classification_dataset.ClassificationDataset):
|
||||
<ymin>89</ymin>
|
||||
<xmax>386</xmax>
|
||||
<ymax>262</ymax>
|
||||
</bndbox>
|
||||
</object>
|
||||
<object>...</object>
|
||||
</annotation>
|
||||
|
||||
@@ -27,8 +27,6 @@ class HParams(hp.BaseHParams):
|
||||
learning_rate: Learning rate to use for gradient descent training.
|
||||
batch_size: Batch size for training.
|
||||
epochs: Number of training iterations over the dataset.
|
||||
do_fine_tuning: If true, the base module is trained together with the
|
||||
classification layer on top.
|
||||
cosine_decay_epochs: The number of epochs for cosine decay learning rate.
|
||||
See
|
||||
https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/schedules/CosineDecay
|
||||
@@ -39,13 +37,13 @@ class HParams(hp.BaseHParams):
|
||||
"""
|
||||
|
||||
# Parameters from BaseHParams class.
|
||||
learning_rate: float = 0.003
|
||||
batch_size: int = 32
|
||||
epochs: int = 10
|
||||
learning_rate: float = 0.3
|
||||
batch_size: int = 8
|
||||
epochs: int = 30
|
||||
|
||||
# Parameters for cosine learning rate decay
|
||||
cosine_decay_epochs: Optional[int] = None
|
||||
cosine_decay_alpha: float = 0.0
|
||||
cosine_decay_alpha: float = 1.0
|
||||
|
||||
|
||||
@dataclasses.dataclass
|
||||
@@ -67,8 +65,8 @@ class QATHParams:
|
||||
for more information.
|
||||
"""
|
||||
|
||||
learning_rate: float = 0.03
|
||||
batch_size: int = 32
|
||||
epochs: int = 10
|
||||
decay_steps: int = 231
|
||||
learning_rate: float = 0.3
|
||||
batch_size: int = 8
|
||||
epochs: int = 15
|
||||
decay_steps: int = 8
|
||||
decay_rate: float = 0.96
|
||||
|
||||
@@ -361,9 +361,10 @@ class FaceStylizerGraph : public core::ModelTaskGraph {
|
||||
|
||||
auto& tensors_to_image =
|
||||
graph.AddNode("mediapipe.tasks.TensorsToImageCalculator");
|
||||
ConfigureTensorsToImageCalculator(
|
||||
image_to_tensor_options,
|
||||
&tensors_to_image.GetOptions<TensorsToImageCalculatorOptions>());
|
||||
auto& tensors_to_image_options =
|
||||
tensors_to_image.GetOptions<TensorsToImageCalculatorOptions>();
|
||||
tensors_to_image_options.mutable_input_tensor_float_range()->set_min(-1);
|
||||
tensors_to_image_options.mutable_input_tensor_float_range()->set_max(1);
|
||||
face_alignment_image >> tensors_to_image.In(kTensorsTag);
|
||||
face_alignment = tensors_to_image.Out(kImageTag).Cast<Image>();
|
||||
|
||||
|
||||
@@ -41,7 +41,11 @@ mediapipe_ts_library(
|
||||
|
||||
mediapipe_ts_library(
|
||||
name = "image",
|
||||
srcs = ["image.ts"],
|
||||
srcs = [
|
||||
"image.ts",
|
||||
"image_converter.ts",
|
||||
"image_shader_context.ts",
|
||||
],
|
||||
)
|
||||
|
||||
mediapipe_ts_library(
|
||||
|
||||
@@ -16,7 +16,8 @@
|
||||
|
||||
import 'jasmine';
|
||||
|
||||
import {MPImage, MPImageShaderContext, MPImageType} from './image';
|
||||
import {MPImage, MPImageType} from './image';
|
||||
import {MPImageShaderContext} from './image_shader_context';
|
||||
|
||||
const WIDTH = 2;
|
||||
const HEIGHT = 2;
|
||||
|
||||
@@ -14,6 +14,9 @@
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
import {DefaultColorConverter} from '../../../../tasks/web/vision/core/image_converter';
|
||||
import {assertNotNull, MPImageShaderContext} from '../../../../tasks/web/vision/core/image_shader_context';
|
||||
|
||||
/** The underlying type of the image. */
|
||||
export enum MPImageType {
|
||||
/** Represents the native `UInt8ClampedArray` type. */
|
||||
@@ -34,235 +37,6 @@ export enum MPImageType {
|
||||
export type MPImageContainer =
|
||||
Uint8ClampedArray|Float32Array|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;
|
||||
}
|
||||
|
||||
// TODO: Move internal-only types to different module.
|
||||
|
||||
/**
|
||||
* Utility class that encapsulates the buffers used by `MPImageShaderContext`.
|
||||
* For internal use only.
|
||||
*/
|
||||
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.
|
||||
*
|
||||
* For internal use only.
|
||||
*/
|
||||
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();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/** A four channel color with a red, green, blue and alpha values. */
|
||||
export type RGBAColor = [number, number, number, number];
|
||||
|
||||
@@ -329,70 +103,6 @@ export interface MPImageChannelConverter {
|
||||
*/
|
||||
uint8ToFloatConverter?: (value: number) => number;
|
||||
}
|
||||
/**
|
||||
* Color converter that falls back to a default implementation if the
|
||||
* user-provided converter does not specify a conversion.
|
||||
*/
|
||||
class DefaultColorConverter implements Required<MPImageChannelConverter> {
|
||||
private static readonly WARNINGS_LOGGED = new Set<string>();
|
||||
|
||||
constructor(private readonly customConverter: MPImageChannelConverter) {}
|
||||
|
||||
floatToRGBAConverter(v: number): RGBAColor {
|
||||
if (this.customConverter.floatToRGBAConverter) {
|
||||
return this.customConverter.floatToRGBAConverter(v);
|
||||
}
|
||||
this.logWarningOnce('floatToRGBAConverter');
|
||||
return [v * 255, v * 255, v * 255, 255];
|
||||
}
|
||||
|
||||
uint8ToRGBAConverter(v: number): RGBAColor {
|
||||
if (this.customConverter.uint8ToRGBAConverter) {
|
||||
return this.customConverter.uint8ToRGBAConverter(v);
|
||||
}
|
||||
this.logWarningOnce('uint8ToRGBAConverter');
|
||||
return [v, v, v, 255];
|
||||
}
|
||||
|
||||
rgbaToFloatConverter(r: number, g: number, b: number, a: number): number {
|
||||
if (this.customConverter.rgbaToFloatConverter) {
|
||||
return this.customConverter.rgbaToFloatConverter(r, g, b, a);
|
||||
}
|
||||
this.logWarningOnce('rgbaToFloatConverter');
|
||||
return (r / 3 + g / 3 + b / 3) / 255;
|
||||
}
|
||||
|
||||
rgbaToUint8Converter(r: number, g: number, b: number, a: number): number {
|
||||
if (this.customConverter.rgbaToUint8Converter) {
|
||||
return this.customConverter.rgbaToUint8Converter(r, g, b, a);
|
||||
}
|
||||
this.logWarningOnce('rgbaToUint8Converter');
|
||||
return r / 3 + g / 3 + b / 3;
|
||||
}
|
||||
|
||||
floatToUint8Converter(v: number): number {
|
||||
if (this.customConverter.floatToUint8Converter) {
|
||||
return this.customConverter.floatToUint8Converter(v);
|
||||
}
|
||||
this.logWarningOnce('floatToUint8Converter');
|
||||
return v * 255;
|
||||
}
|
||||
|
||||
uint8ToFloatConverter(v: number): number {
|
||||
if (this.customConverter.uint8ToFloatConverter) {
|
||||
return this.customConverter.uint8ToFloatConverter(v);
|
||||
}
|
||||
this.logWarningOnce('uint8ToFloatConverter');
|
||||
return v / 255;
|
||||
}
|
||||
|
||||
private logWarningOnce(methodName: string): void {
|
||||
if (!DefaultColorConverter.WARNINGS_LOGGED.has(methodName)) {
|
||||
console.log(`Using default ${methodName}`);
|
||||
DefaultColorConverter.WARNINGS_LOGGED.add(methodName);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* The wrapper class for MediaPipe Image objects.
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
/**
|
||||
* Copyright 2023 The MediaPipe Authors.
|
||||
*
|
||||
* 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 {MPImageChannelConverter, RGBAColor} from '../../../../tasks/web/vision/core/image';
|
||||
|
||||
/**
|
||||
* Color converter that falls back to a default implementation if the
|
||||
* user-provided converter does not specify a conversion.
|
||||
*/
|
||||
export class DefaultColorConverter implements
|
||||
Required<MPImageChannelConverter> {
|
||||
private static readonly WARNINGS_LOGGED = new Set<string>();
|
||||
|
||||
constructor(private readonly customConverter: MPImageChannelConverter) {}
|
||||
|
||||
floatToRGBAConverter(v: number): RGBAColor {
|
||||
if (this.customConverter.floatToRGBAConverter) {
|
||||
return this.customConverter.floatToRGBAConverter(v);
|
||||
}
|
||||
this.logWarningOnce('floatToRGBAConverter');
|
||||
return [v * 255, v * 255, v * 255, 255];
|
||||
}
|
||||
|
||||
uint8ToRGBAConverter(v: number): RGBAColor {
|
||||
if (this.customConverter.uint8ToRGBAConverter) {
|
||||
return this.customConverter.uint8ToRGBAConverter(v);
|
||||
}
|
||||
this.logWarningOnce('uint8ToRGBAConverter');
|
||||
return [v, v, v, 255];
|
||||
}
|
||||
|
||||
rgbaToFloatConverter(r: number, g: number, b: number, a: number): number {
|
||||
if (this.customConverter.rgbaToFloatConverter) {
|
||||
return this.customConverter.rgbaToFloatConverter(r, g, b, a);
|
||||
}
|
||||
this.logWarningOnce('rgbaToFloatConverter');
|
||||
return (r / 3 + g / 3 + b / 3) / 255;
|
||||
}
|
||||
|
||||
rgbaToUint8Converter(r: number, g: number, b: number, a: number): number {
|
||||
if (this.customConverter.rgbaToUint8Converter) {
|
||||
return this.customConverter.rgbaToUint8Converter(r, g, b, a);
|
||||
}
|
||||
this.logWarningOnce('rgbaToUint8Converter');
|
||||
return r / 3 + g / 3 + b / 3;
|
||||
}
|
||||
|
||||
floatToUint8Converter(v: number): number {
|
||||
if (this.customConverter.floatToUint8Converter) {
|
||||
return this.customConverter.floatToUint8Converter(v);
|
||||
}
|
||||
this.logWarningOnce('floatToUint8Converter');
|
||||
return v * 255;
|
||||
}
|
||||
|
||||
uint8ToFloatConverter(v: number): number {
|
||||
if (this.customConverter.uint8ToFloatConverter) {
|
||||
return this.customConverter.uint8ToFloatConverter(v);
|
||||
}
|
||||
this.logWarningOnce('uint8ToFloatConverter');
|
||||
return v / 255;
|
||||
}
|
||||
|
||||
private logWarningOnce(methodName: string): void {
|
||||
if (!DefaultColorConverter.WARNINGS_LOGGED.has(methodName)) {
|
||||
console.log(`Using default ${methodName}`);
|
||||
DefaultColorConverter.WARNINGS_LOGGED.add(methodName);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,243 @@
|
||||
/**
|
||||
* Copyright 2023 The MediaPipe Authors.
|
||||
*
|
||||
* 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.
|
||||
*/
|
||||
|
||||
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);
|
||||
}
|
||||
`;
|
||||
|
||||
/** Helper to assert that `value` is not null. */
|
||||
export 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`.
|
||||
* For internal use only.
|
||||
*/
|
||||
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.
|
||||
*
|
||||
* For internal use only.
|
||||
*/
|
||||
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();
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -17,8 +17,9 @@
|
||||
import {NormalizedRect} from '../../../../framework/formats/rect_pb';
|
||||
import {TaskRunner} from '../../../../tasks/web/core/task_runner';
|
||||
import {WasmFileset} from '../../../../tasks/web/core/wasm_fileset';
|
||||
import {MPImage, MPImageShaderContext} from '../../../../tasks/web/vision/core/image';
|
||||
import {MPImage} from '../../../../tasks/web/vision/core/image';
|
||||
import {ImageProcessingOptions} from '../../../../tasks/web/vision/core/image_processing_options';
|
||||
import {MPImageShaderContext} from '../../../../tasks/web/vision/core/image_shader_context';
|
||||
import {GraphRunner, ImageSource, WasmMediaPipeConstructor} from '../../../../web/graph_runner/graph_runner';
|
||||
import {SupportImage, WasmImage} from '../../../../web/graph_runner/graph_runner_image_lib';
|
||||
import {isWebKit} from '../../../../web/graph_runner/platform_utils';
|
||||
|
||||
Reference in New Issue
Block a user