diff --git a/.eslintignore b/.eslintignore index 57d0f64c..3f673d82 100644 --- a/.eslintignore +++ b/.eslintignore @@ -1 +1,3 @@ -/common/component/component.js \ No newline at end of file +/common/component/component.js +/selfie_segmentation/tflite-support/ +/selfie_segmentation/lib/ \ No newline at end of file diff --git a/code/index.html b/code/index.html index fc845356..67a631f3 100644 --- a/code/index.html +++ b/code/index.html @@ -61,7 +61,7 @@ + + + + + + + + + \ No newline at end of file diff --git a/selfie_segmentation/lib/README.md b/selfie_segmentation/lib/README.md new file mode 100644 index 00000000..ae2dc647 --- /dev/null +++ b/selfie_segmentation/lib/README.md @@ -0,0 +1 @@ +Resources under `helpers` and `webgl2` foloders are used for postprocessing and referred from https://github.com/Volcomix/virtual-background/tree/main/src/pipelines/ under the license of Apache License 2.0. \ No newline at end of file diff --git a/selfie_segmentation/lib/helpers/webglHelper.js b/selfie_segmentation/lib/helpers/webglHelper.js new file mode 100644 index 00000000..76da05cc --- /dev/null +++ b/selfie_segmentation/lib/helpers/webglHelper.js @@ -0,0 +1,111 @@ +export const glsl = String.raw; + +/** + * Create stage program for WebGL pipeline + * @param {WebGL2RenderingContext} gl, gl object + * @param {WebGLShader} vertexShader, vertex shader + * @param {WebGLShader} fragmentShader, fragment shader + * @param {WebGLBuffer} positionBuffer, buffer of position + * @param {WebGLBuffer} texCoordBuffer, buffer of tex coord + */ +export function createPiplelineStageProgram( + gl, vertexShader, fragmentShader, positionBuffer, texCoordBuffer) { + const program = createProgram(gl, vertexShader, fragmentShader); + + const positionAttributeLocation = gl.getAttribLocation(program, 'a_position'); + gl.enableVertexAttribArray(positionAttributeLocation); + gl.bindBuffer(gl.ARRAY_BUFFER, positionBuffer); + gl.vertexAttribPointer(positionAttributeLocation, 2, gl.FLOAT, false, 0, 0); + + const texCoordAttributeLocation = gl.getAttribLocation(program, 'a_texCoord'); + gl.enableVertexAttribArray(texCoordAttributeLocation); + gl.bindBuffer(gl.ARRAY_BUFFER, texCoordBuffer); + gl.vertexAttribPointer(texCoordAttributeLocation, 2, gl.FLOAT, false, 0, 0); + + return program; +} + +export function createProgram(gl, vertexShader, fragmentShader) { + const program = gl.createProgram(); + + gl.attachShader(program, vertexShader); + gl.attachShader(program, fragmentShader); + gl.linkProgram(program); + if (!gl.getProgramParameter(program, gl.LINK_STATUS)) { + throw new Error( + `Could not link WebGL program: ${gl.getProgramInfoLog(program)}`); + } + return program; +} + +export function compileShader(gl, shaderType, shaderSource) { + const shader = gl.createShader(shaderType); + gl.shaderSource(shader, shaderSource); + gl.compileShader(shader); + if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) { + throw new Error(`Could not compile shader: ${gl.getShaderInfoLog(shader)}`); + } + return shader; +} + +export function createTexture( + gl, internalformat, width, height, minFilter, magFilter) { + const texture = gl.createTexture() + minFilter = minFilter === undefined ? gl.NEAREST : minFilter; + magFilter = magFilter === undefined ? gl.NEAREST : magFilter; + + gl.bindTexture(gl.TEXTURE_2D, texture); + 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, minFilter); + gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, magFilter); + gl.texStorage2D(gl.TEXTURE_2D, 1, internalformat, width, height); + + return texture; +} + +export async function readPixelsAsync( + gl, x, y, width, height, format, type, dest) { + const buf = gl.createBuffer(); + gl.bindBuffer(gl.PIXEL_PACK_BUFFER, buf); + gl.bufferData(gl.PIXEL_PACK_BUFFER, dest.byteLength, gl.STREAM_READ); + gl.readPixels(x, y, width, height, format, type, 0); + gl.bindBuffer(gl.PIXEL_PACK_BUFFER, null); + + await getBufferSubDataAsync(gl, gl.PIXEL_PACK_BUFFER, buf, 0, dest); + + gl.deleteBuffer(buf); + return dest; +} + +async function getBufferSubDataAsync( + gl, target, buffer, srcByteOffset, dstBuffer, dstOffset = undefined, length = undefined) { + const sync = gl.fenceSync(gl.SYNC_GPU_COMMANDS_COMPLETE, 0); + gl.flush(); + const res = await clientWaitAsync(gl, sync); + gl.deleteSync(sync); + + if (res !== gl.WAIT_FAILED) { + gl.bindBuffer(target, buffer); + gl.getBufferSubData(target, srcByteOffset, dstBuffer, dstOffset, length); + gl.bindBuffer(target, null); + } +} + +function clientWaitAsync(gl, sync) { + return new Promise((resolve) => { + function test() { + const res = gl.clientWaitSync(sync, 0, 0); + if (res === gl.WAIT_FAILED) { + resolve(res); + return; + } + if (res === gl.TIMEOUT_EXPIRED) { + requestAnimationFrame(test); + return; + } + resolve(res); + } + requestAnimationFrame(test); + }); +} diff --git a/selfie_segmentation/lib/webgl2/backgroundBlurStage.js b/selfie_segmentation/lib/webgl2/backgroundBlurStage.js new file mode 100644 index 00000000..52371a76 --- /dev/null +++ b/selfie_segmentation/lib/webgl2/backgroundBlurStage.js @@ -0,0 +1,269 @@ +import { + compileShader, + createPiplelineStageProgram, + createTexture, + glsl, +} from '../helpers/webglHelper.js'; + +export function buildBackgroundBlurStage( + gl, vertexShader, positionBuffer, texCoordBuffer, personMaskTexture, canvas) { + const blurPass = buildBlurPass( + gl, + vertexShader, + positionBuffer, + texCoordBuffer, + personMaskTexture, + canvas, + ); + const blendPass = buildBlendPass(gl, positionBuffer, texCoordBuffer, canvas); + + const render = function() { + blurPass.render(); + blendPass.render(); + }; + + const updateCoverage = function(coverage) { + blendPass.updateCoverage(coverage); + }; + + const cleanUp = function() { + blendPass.cleanUp(); + blurPass.cleanUp(); + }; + + return { + render, + updateCoverage, + cleanUp, + }; +} + +function buildBlurPass( + gl, vertexShader, positionBuffer, texCoordBuffer, personMaskTexture, canvas) { + const fragmentShaderSource = glsl`#version 300 es + + precision highp float; + + uniform sampler2D u_inputFrame; + uniform sampler2D u_personMask; + uniform vec2 u_texelSize; + + in vec2 v_texCoord; + + out vec4 outColor; + + const float offset[5] = float[](0.0, 1.0, 2.0, 3.0, 4.0); + const float weight[5] = float[](0.2270270270, 0.1945945946, 0.1216216216, + 0.0540540541, 0.0162162162); + + void main() { + vec4 centerColor = texture(u_inputFrame, v_texCoord); + float personMask = texture(u_personMask, v_texCoord).a; + + vec4 frameColor = centerColor * weight[0] * (1.0 - personMask); + + for (int i = 1; i < 5; i++) { + vec2 offset = vec2(offset[i]) * u_texelSize; + + vec2 texCoord = v_texCoord + offset; + frameColor += texture(u_inputFrame, texCoord) * weight[i] * + (1.0 - texture(u_personMask, texCoord).a); + + texCoord = v_texCoord - offset; + frameColor += texture(u_inputFrame, texCoord) * weight[i] * + (1.0 - texture(u_personMask, texCoord).a); + } + outColor = vec4(frameColor.rgb + (1.0 - frameColor.a) * centerColor.rgb, 1.0); + } + `; + + const scale = 0.5; + const outputWidth = canvas.width * scale; + const outputHeight = canvas.height * scale; + const texelWidth = 1 / outputWidth; + const texelHeight = 1 / outputHeight; + + const fragmentShader = compileShader( + gl, + gl.FRAGMENT_SHADER, + fragmentShaderSource, + ); + const program = createPiplelineStageProgram( + gl, + vertexShader, + fragmentShader, + positionBuffer, + texCoordBuffer, + ); + const inputFrameLocation = gl.getUniformLocation(program, 'u_inputFrame'); + const personMaskLocation = gl.getUniformLocation(program, 'u_personMask'); + const texelSizeLocation = gl.getUniformLocation(program, 'u_texelSize'); + const texture1 = createTexture( + gl, + gl.RGBA8, + outputWidth, + outputHeight, + gl.NEAREST, + gl.LINEAR, + ); + const texture2 = createTexture( + gl, + gl.RGBA8, + outputWidth, + outputHeight, + gl.NEAREST, + gl.LINEAR, + ); + + const frameBuffer1 = gl.createFramebuffer(); + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer1); + gl.framebufferTexture2D( + gl.FRAMEBUFFER, + gl.COLOR_ATTACHMENT0, + gl.TEXTURE_2D, + texture1, + 0, + ); + + const frameBuffer2 = gl.createFramebuffer(); + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer2); + gl.framebufferTexture2D( + gl.FRAMEBUFFER, + gl.COLOR_ATTACHMENT0, + gl.TEXTURE_2D, + texture2, + 0, + ); + + gl.useProgram(program); + gl.uniform1i(personMaskLocation, 1); + + const render = function() { + gl.viewport(0, 0, outputWidth, outputHeight) + gl.useProgram(program) + gl.uniform1i(inputFrameLocation, 0) + gl.activeTexture(gl.TEXTURE1) + gl.bindTexture(gl.TEXTURE_2D, personMaskTexture) + + for (let i = 0; i < 3; i++) { + gl.uniform2f(texelSizeLocation, 0, texelHeight) + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer1) + gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4) + + gl.activeTexture(gl.TEXTURE2) + gl.bindTexture(gl.TEXTURE_2D, texture1) + gl.uniform1i(inputFrameLocation, 2) + + gl.uniform2f(texelSizeLocation, texelWidth, 0) + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer2) + gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4) + + gl.bindTexture(gl.TEXTURE_2D, texture2) + } + } + + const cleanUp = function() { + gl.deleteFramebuffer(frameBuffer2) + gl.deleteFramebuffer(frameBuffer1) + gl.deleteTexture(texture2) + gl.deleteTexture(texture1) + gl.deleteProgram(program) + gl.deleteShader(fragmentShader) + } + + return { + render, + cleanUp, + }; +} + +function buildBlendPass(gl, positionBuffer, texCoordBuffer, canvas) { + const vertexShaderSource = glsl`#version 300 es + + in vec2 a_position; + in vec2 a_texCoord; + + out vec2 v_texCoord; + + void main() { + // Flipping Y is required when rendering to canvas + gl_Position = vec4(a_position * vec2(1.0, -1.0), 0.0, 1.0); + v_texCoord = a_texCoord; + } + `; + + const fragmentShaderSource = glsl`#version 300 es + + precision highp float; + + uniform sampler2D u_inputFrame; + uniform sampler2D u_personMask; + uniform sampler2D u_blurredInputFrame; + uniform vec2 u_coverage; + + in vec2 v_texCoord; + + out vec4 outColor; + + void main() { + vec3 color = texture(u_inputFrame, v_texCoord).rgb; + vec3 blurredColor = texture(u_blurredInputFrame, v_texCoord).rgb; + float personMask = texture(u_personMask, v_texCoord).a; + personMask = smoothstep(u_coverage.x, u_coverage.y, personMask); + outColor = vec4(mix(blurredColor, color, personMask), 1.0); + } + `; + + const { width: outputWidth, height: outputHeight } = canvas; + + const vertexShader = compileShader(gl, gl.VERTEX_SHADER, vertexShaderSource); + const fragmentShader = compileShader( + gl, + gl.FRAGMENT_SHADER, + fragmentShaderSource, + ); + const program = createPiplelineStageProgram( + gl, + vertexShader, + fragmentShader, + positionBuffer, + texCoordBuffer, + ); + const inputFrameLocation = gl.getUniformLocation(program, 'u_inputFrame'); + const personMaskLocation = gl.getUniformLocation(program, 'u_personMask'); + const blurredInputFrame = gl.getUniformLocation( + program, + 'u_blurredInputFrame', + ); + const coverageLocation = gl.getUniformLocation(program, 'u_coverage'); + + gl.useProgram(program); + gl.uniform1i(inputFrameLocation, 0); + gl.uniform1i(personMaskLocation, 1); + gl.uniform1i(blurredInputFrame, 2); + gl.uniform2f(coverageLocation, 0, 1); + + const render = function() { + gl.viewport(0, 0, outputWidth, outputHeight); + gl.useProgram(program); + gl.bindFramebuffer(gl.FRAMEBUFFER, null); + gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4); + } + + const updateCoverage = function(coverage) { + gl.useProgram(program); + gl.uniform2f(coverageLocation, coverage[0], coverage[1]); + } + + const cleanUp = function() { + gl.deleteProgram(program); + gl.deleteShader(fragmentShader); + gl.deleteShader(vertexShader); + } + + return { + render, + updateCoverage, + cleanUp, + }; +} diff --git a/selfie_segmentation/lib/webgl2/backgroundImageStage.js b/selfie_segmentation/lib/webgl2/backgroundImageStage.js new file mode 100644 index 00000000..a59440b9 --- /dev/null +++ b/selfie_segmentation/lib/webgl2/backgroundImageStage.js @@ -0,0 +1,206 @@ +import { + compileShader, + createPiplelineStageProgram, + createTexture, + glsl, +} from '../helpers/webglHelper.js'; + +export function buildBackgroundImageStage( + gl, positionBuffer, texCoordBuffer, personMaskTexture, backgroundImage, canvas) { + const vertexShaderSource = glsl`#version 300 es + + uniform vec2 u_backgroundScale; + uniform vec2 u_backgroundOffset; + + in vec2 a_position; + in vec2 a_texCoord; + + out vec2 v_texCoord; + out vec2 v_backgroundCoord; + + void main() { + // Flipping Y is required when rendering to canvas + gl_Position = vec4(a_position * vec2(1.0, -1.0), 0.0, 1.0); + v_texCoord = a_texCoord; + v_backgroundCoord = a_texCoord * u_backgroundScale + u_backgroundOffset; + } + `; + + const fragmentShaderSource = glsl`#version 300 es + + precision highp float; + + uniform sampler2D u_inputFrame; + uniform sampler2D u_personMask; + uniform sampler2D u_background; + uniform vec2 u_coverage; + uniform float u_lightWrapping; + uniform float u_blendMode; + + in vec2 v_texCoord; + in vec2 v_backgroundCoord; + + out vec4 outColor; + + vec3 screen(vec3 a, vec3 b) { + return 1.0 - (1.0 - a) * (1.0 - b); + } + + vec3 linearDodge(vec3 a, vec3 b) { + return a + b; + } + + void main() { + vec3 frameColor = texture(u_inputFrame, v_texCoord).rgb; + vec3 backgroundColor = texture(u_background, v_backgroundCoord).rgb; + float personMask = texture(u_personMask, v_texCoord).a; + float lightWrapMask = 1.0 - max(0.0, personMask - u_coverage.y) / (1.0 - u_coverage.y); + vec3 lightWrap = u_lightWrapping * lightWrapMask * backgroundColor; + frameColor = u_blendMode * linearDodge(frameColor, lightWrap) + + (1.0 - u_blendMode) * screen(frameColor, lightWrap); + personMask = smoothstep(u_coverage.x, u_coverage.y, personMask); + outColor = vec4(frameColor * personMask + backgroundColor * (1.0 - personMask), 1.0); + } + `; + + const { width: outputWidth, height: outputHeight } = canvas; + const outputRatio = outputWidth / outputHeight; + + const vertexShader = compileShader(gl, gl.VERTEX_SHADER, vertexShaderSource); + const fragmentShader = compileShader( + gl, + gl.FRAGMENT_SHADER, + fragmentShaderSource, + ); + const program = createPiplelineStageProgram( + gl, + vertexShader, + fragmentShader, + positionBuffer, + texCoordBuffer, + ); + const backgroundScaleLocation = gl.getUniformLocation( + program, + 'u_backgroundScale', + ); + const backgroundOffsetLocation = gl.getUniformLocation( + program, + 'u_backgroundOffset', + ); + const inputFrameLocation = gl.getUniformLocation(program, 'u_inputFrame'); + const personMaskLocation = gl.getUniformLocation(program, 'u_personMask'); + const backgroundLocation = gl.getUniformLocation(program, 'u_background'); + const coverageLocation = gl.getUniformLocation(program, 'u_coverage'); + const lightWrappingLocation = gl.getUniformLocation( + program, + 'u_lightWrapping', + ); + const blendModeLocation = gl.getUniformLocation(program, 'u_blendMode'); + + gl.useProgram(program); + gl.uniform2f(backgroundScaleLocation, 1, 1); + gl.uniform2f(backgroundOffsetLocation, 0, 0); + gl.uniform1i(inputFrameLocation, 0); + gl.uniform1i(personMaskLocation, 1); + gl.uniform2f(coverageLocation, 0, 1); + gl.uniform1f(lightWrappingLocation, 0); + gl.uniform1f(blendModeLocation, 0); + + let backgroundTexture = null; + // TODO Find a better to handle background being loaded + if (backgroundImage === null || backgroundImage === void 0 ? void 0 : backgroundImage.complete) { + updateBackgroundImage(backgroundImage); + } else if (backgroundImage) { + backgroundImage.onload = () => { + updateBackgroundImage(backgroundImage); + }; + } + + const render = function() { + gl.viewport(0, 0, outputWidth, outputHeight); + gl.useProgram(program); + gl.activeTexture(gl.TEXTURE1); + gl.bindTexture(gl.TEXTURE_2D, personMaskTexture); + if (backgroundTexture !== null) { + gl.activeTexture(gl.TEXTURE2); + gl.bindTexture(gl.TEXTURE_2D, backgroundTexture); + // TODO Handle correctly the background not loaded yet + gl.uniform1i(backgroundLocation, 2); + } + gl.bindFramebuffer(gl.FRAMEBUFFER, null); + gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4); + }; + + function updateBackgroundImage(backgroundImage) { + backgroundTexture = createTexture( + gl, + gl.RGBA8, + backgroundImage.naturalWidth, + backgroundImage.naturalHeight, + gl.LINEAR, + gl.LINEAR, + ); + gl.texSubImage2D( + gl.TEXTURE_2D, + 0, + 0, + 0, + backgroundImage.naturalWidth, + backgroundImage.naturalHeight, + gl.RGBA, + gl.UNSIGNED_BYTE, + backgroundImage, + ); + + let xOffset = 0; + let yOffset = 0; + let backgroundWidth = backgroundImage.naturalWidth; + let backgroundHeight = backgroundImage.naturalHeight; + const backgroundRatio = backgroundWidth / backgroundHeight; + if (backgroundRatio < outputRatio) { + backgroundHeight = backgroundWidth / outputRatio; + yOffset = (backgroundImage.naturalHeight - backgroundHeight) / 2; + } else { + backgroundWidth = backgroundHeight * outputRatio; + xOffset = (backgroundImage.naturalWidth - backgroundWidth) / 2; + } + + const xScale = backgroundWidth / backgroundImage.naturalWidth; + const yScale = backgroundHeight / backgroundImage.naturalHeight; + xOffset /= backgroundImage.naturalWidth; + yOffset /= backgroundImage.naturalHeight; + + gl.uniform2f(backgroundScaleLocation, xScale, yScale); + gl.uniform2f(backgroundOffsetLocation, xOffset, yOffset); + } + + const updateCoverage = function(coverage) { + gl.useProgram(program); + gl.uniform2f(coverageLocation, coverage[0], coverage[1]); + }; + + const updateLightWrapping = function(lightWrapping) { + gl.useProgram(program); + gl.uniform1f(lightWrappingLocation, lightWrapping); + }; + + const updateBlendMode = function(blendMode) { + gl.useProgram(program); + gl.uniform1f(blendModeLocation, blendMode === 'screen' ? 0 : 1); + }; + + const cleanUp = function() { + gl.deleteTexture(backgroundTexture); + gl.deleteProgram(program); + gl.deleteShader(fragmentShader); + gl.deleteShader(vertexShader); + }; + + return { + render, + updateCoverage, + updateLightWrapping, + updateBlendMode, + cleanUp, + }; +} diff --git a/selfie_segmentation/lib/webgl2/jointBilateralFilterStage.js b/selfie_segmentation/lib/webgl2/jointBilateralFilterStage.js new file mode 100644 index 00000000..6da40af6 --- /dev/null +++ b/selfie_segmentation/lib/webgl2/jointBilateralFilterStage.js @@ -0,0 +1,150 @@ +import { + compileShader, + createPiplelineStageProgram, + glsl, +} from '../helpers/webglHelper.js'; + +export function buildJointBilateralFilterStage( + gl, vertexShader, positionBuffer, texCoordBuffer, inputTexture, inputResolution, outputTexture, canvas) { + const fragmentShaderSource = glsl`#version 300 es + + precision highp float; + + uniform sampler2D u_inputFrame; + uniform sampler2D u_segmentationMask; + uniform vec2 u_texelSize; + uniform float u_step; + uniform float u_radius; + uniform float u_offset; + uniform float u_sigmaTexel; + uniform float u_sigmaColor; + + in vec2 v_texCoord; + + out vec4 outColor; + + float gaussian(float x, float sigma) { + float coeff = -0.5 / (sigma * sigma * 4.0 + 1.0e-6); + return exp((x * x) * coeff); + } + + void main() { + vec2 centerCoord = v_texCoord; + vec3 centerColor = texture(u_inputFrame, centerCoord).rgb; + float newVal = 0.0; + + float spaceWeight = 0.0; + float colorWeight = 0.0; + float totalWeight = 0.0; + + // Subsample kernel space. + for (float i = -u_radius + u_offset; i <= u_radius; i += u_step) { + for (float j = -u_radius + u_offset; j <= u_radius; j += u_step) { + vec2 shift = vec2(j, i) * u_texelSize; + vec2 coord = vec2(centerCoord + shift); + vec3 frameColor = texture(u_inputFrame, coord).rgb; + float outVal = texture(u_segmentationMask, coord).a; + + spaceWeight = gaussian(distance(centerCoord, coord), u_sigmaTexel); + colorWeight = gaussian(distance(centerColor, frameColor), u_sigmaColor); + totalWeight += spaceWeight * colorWeight; + + newVal += spaceWeight * colorWeight * outVal; + } + } + newVal /= totalWeight; + + outColor = vec4(vec3(0.0), newVal); + } + `; + + const [segmentationWidth, segmentationHeight] = inputResolution; + const { width: outputWidth, height: outputHeight } = canvas; + const texelWidth = 1 / outputWidth; + const texelHeight = 1 / outputHeight; + + const fragmentShader = compileShader( + gl, + gl.FRAGMENT_SHADER, + fragmentShaderSource, + ); + const program = createPiplelineStageProgram( + gl, + vertexShader, + fragmentShader, + positionBuffer, + texCoordBuffer, + ); + const inputFrameLocation = gl.getUniformLocation(program, 'u_inputFrame'); + const segmentationMaskLocation = gl.getUniformLocation( + program, + 'u_segmentationMask', + ); + const texelSizeLocation = gl.getUniformLocation(program, 'u_texelSize'); + const stepLocation = gl.getUniformLocation(program, 'u_step'); + const radiusLocation = gl.getUniformLocation(program, 'u_radius'); + const offsetLocation = gl.getUniformLocation(program, 'u_offset'); + const sigmaTexelLocation = gl.getUniformLocation(program, 'u_sigmaTexel'); + const sigmaColorLocation = gl.getUniformLocation(program, 'u_sigmaColor'); + + const frameBuffer = gl.createFramebuffer(); + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer); + gl.framebufferTexture2D( + gl.FRAMEBUFFER, + gl.COLOR_ATTACHMENT0, + gl.TEXTURE_2D, + outputTexture, + 0, + ); + + gl.useProgram(program); + gl.uniform1i(inputFrameLocation, 0); + gl.uniform1i(segmentationMaskLocation, 1); + gl.uniform2f(texelSizeLocation, texelWidth, texelHeight); + + const render = function() { + gl.viewport(0, 0, outputWidth, outputHeight); + gl.useProgram(program); + gl.activeTexture(gl.TEXTURE1); + gl.bindTexture(gl.TEXTURE_2D, inputTexture); + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer); + gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4); + }; + + const updateSigmaSpace = function(sigmaSpace) { + sigmaSpace *= Math.max( + outputWidth / segmentationWidth, + outputHeight / segmentationHeight, + ); + + const kSparsityFactor = 0.66; // Higher is more sparse. + const sparsity = Math.max(1, Math.sqrt(sigmaSpace) * kSparsityFactor); + const step = sparsity; + const radius = sigmaSpace; + const offset = step > 1 ? step * 0.5 : 0; + const sigmaTexel = Math.max(texelWidth, texelHeight) * sigmaSpace; + + gl.useProgram(program); + gl.uniform1f(stepLocation, step); + gl.uniform1f(radiusLocation, radius); + gl.uniform1f(offsetLocation, offset); + gl.uniform1f(sigmaTexelLocation, sigmaTexel); + }; + + const updateSigmaColor = function(sigmaColor) { + gl.useProgram(program); + gl.uniform1f(sigmaColorLocation, sigmaColor); + }; + + const cleanUp = function() { + gl.deleteFramebuffer(frameBuffer); + gl.deleteProgram(program); + gl.deleteShader(fragmentShader); + }; + // Ensures default values are configured to prevent infinite + // loop in fragment shader + updateSigmaSpace(0); + updateSigmaColor(0); + + return { render, updateSigmaSpace, updateSigmaColor, cleanUp } +} diff --git a/selfie_segmentation/lib/webgl2/loadSegmentationStage.js b/selfie_segmentation/lib/webgl2/loadSegmentationStage.js new file mode 100644 index 00000000..3d0e79c4 --- /dev/null +++ b/selfie_segmentation/lib/webgl2/loadSegmentationStage.js @@ -0,0 +1,89 @@ +import { + compileShader, + createPiplelineStageProgram, + createTexture, + glsl, +} from '../helpers/webglHelper.js'; + +export function buildLoadSegmentationStage( + gl, vertexShader, positionBuffer, texCoordBuffer, inputResolution, outputBuffer, outputTexture) { + const fragmentShaderSource = glsl`#version 300 es + + precision highp float; + + uniform sampler2D u_inputSegmentation; + + in vec2 v_texCoord; + + out vec4 outColor; + + void main() { + float segmentation = texture(u_inputSegmentation, v_texCoord).r; + outColor = vec4(vec3(0.0), segmentation); + } + `; + + const [segmentationWidth, segmentationHeight] = inputResolution; + + const fragmentShader = compileShader( + gl, + gl.FRAGMENT_SHADER, + fragmentShaderSource, + ); + const program = createPiplelineStageProgram( + gl, + vertexShader, + fragmentShader, + positionBuffer, + texCoordBuffer, + ); + const inputLocation = gl.getUniformLocation(program, 'u_inputSegmentation'); + const inputTexture = createTexture( + gl, + gl.R32F, + segmentationWidth, + segmentationHeight, + ); + + const frameBuffer = gl.createFramebuffer(); + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer); + gl.framebufferTexture2D( + gl.FRAMEBUFFER, + gl.COLOR_ATTACHMENT0, + gl.TEXTURE_2D, + outputTexture, + 0, + ); + + gl.useProgram(program); + gl.uniform1i(inputLocation, 1); + + const render = function() { + gl.viewport(0, 0, segmentationWidth, segmentationHeight); + gl.useProgram(program); + gl.activeTexture(gl.TEXTURE1); + gl.bindTexture(gl.TEXTURE_2D, inputTexture); + gl.texSubImage2D( + gl.TEXTURE_2D, + 0, + 0, + 0, + segmentationWidth, + segmentationHeight, + gl.RED, + gl.FLOAT, + outputBuffer, + ); + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer); + gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4); + }; + + const cleanUp = function() { + gl.deleteFramebuffer(frameBuffer); + gl.deleteTexture(inputTexture); + gl.deleteProgram(program); + gl.deleteShader(fragmentShader); + }; + + return {render, cleanUp}; +} diff --git a/selfie_segmentation/lib/webgl2/resizingStage.js b/selfie_segmentation/lib/webgl2/resizingStage.js new file mode 100644 index 00000000..798c494f --- /dev/null +++ b/selfie_segmentation/lib/webgl2/resizingStage.js @@ -0,0 +1,85 @@ +import { + compileShader, + createPiplelineStageProgram, + createTexture, + glsl, + readPixelsAsync, +} from '../helpers/webglHelper.js'; + +export function buildResizingStage( + gl, vertexShader, positionBuffer, texCoordBuffer, inputResolution) { + const fragmentShaderSource = glsl`#version 300 es + + precision highp float; + + uniform sampler2D u_inputFrame; + + in vec2 v_texCoord; + + out vec4 outColor; + + void main() { + outColor = texture(u_inputFrame, v_texCoord); + } + `; + + const [outputWidth, outputHeight] = inputResolution; + const outputPixelCount = outputWidth * outputHeight; + + const fragmentShader = compileShader( + gl, + gl.FRAGMENT_SHADER, + fragmentShaderSource, + ); + const program = createPiplelineStageProgram( + gl, + vertexShader, + fragmentShader, + positionBuffer, + texCoordBuffer, + ); + const inputFrameLocation = gl.getUniformLocation(program, 'u_inputFrame'); + const outputTexture = createTexture(gl, gl.RGBA8, outputWidth, outputHeight); + + const frameBuffer = gl.createFramebuffer(); + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer); + gl.framebufferTexture2D( + gl.FRAMEBUFFER, + gl.COLOR_ATTACHMENT0, + gl.TEXTURE_2D, + outputTexture, + 0, + ); + const outputPixels = new Uint8Array(outputPixelCount * 4); + + gl.useProgram(program); + gl.uniform1i(inputFrameLocation, 0); + + const render = function() { + gl.viewport(0, 0, outputWidth, outputHeight); + gl.useProgram(program); + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer); + gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4); + + // Downloads pixels asynchronously from GPU while rendering the current frame + readPixelsAsync( + gl, + 0, + 0, + outputWidth, + outputHeight, + gl.RGBA, + gl.UNSIGNED_BYTE, + outputPixels, + ); + }; + + const cleanUp = function() { + gl.deleteFramebuffer(frameBuffer); + gl.deleteTexture(outputTexture); + gl.deleteProgram(program); + gl.deleteShader(fragmentShader); + }; + + return {render, cleanUp}; +} diff --git a/selfie_segmentation/lib/webgl2/softmaxStage.js b/selfie_segmentation/lib/webgl2/softmaxStage.js new file mode 100644 index 00000000..e6919fd7 --- /dev/null +++ b/selfie_segmentation/lib/webgl2/softmaxStage.js @@ -0,0 +1,92 @@ +import { + compileShader, + createPiplelineStageProgram, + createTexture, + glsl, +} from '../helpers/webglHelper.js'; + +export function buildSoftmaxStage( + gl, vertexShader, positionBuffer, texCoordBuffer, inputResolution, outputBuffer, outputTexture) { + const fragmentShaderSource = glsl`#version 300 es + + precision highp float; + + uniform sampler2D u_inputSegmentation; + + in vec2 v_texCoord; + + out vec4 outColor; + + void main() { + vec2 segmentation = texture(u_inputSegmentation, v_texCoord).rg; + float shift = max(segmentation.r, segmentation.g); + float backgroundExp = exp(segmentation.r - shift); + float personExp = exp(segmentation.g - shift); + outColor = vec4(vec3(0.0), personExp / (backgroundExp + personExp)); + } + `; + + const [segmentationWidth, segmentationHeight] = inputResolution; + + const fragmentShader = compileShader( + gl, + gl.FRAGMENT_SHADER, + fragmentShaderSource, + ); + const program = createPiplelineStageProgram( + gl, + vertexShader, + fragmentShader, + positionBuffer, + texCoordBuffer, + ); + const inputLocation = gl.getUniformLocation(program, 'u_inputSegmentation'); + const inputTexture = createTexture( + gl, + gl.RG32F, + segmentationWidth, + segmentationHeight, + ); + + const frameBuffer = gl.createFramebuffer(); + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer); + gl.framebufferTexture2D( + gl.FRAMEBUFFER, + gl.COLOR_ATTACHMENT0, + gl.TEXTURE_2D, + outputTexture, + 0, + ); + + gl.useProgram(program); + gl.uniform1i(inputLocation, 1); + + const render = function() { + gl.viewport(0, 0, segmentationWidth, segmentationHeight); + gl.useProgram(program); + gl.activeTexture(gl.TEXTURE1); + gl.bindTexture(gl.TEXTURE_2D, inputTexture); + gl.texSubImage2D( + gl.TEXTURE_2D, + 0, + 0, + 0, + segmentationWidth, + segmentationHeight, + gl.RG, + gl.FLOAT, + outputBuffer, + ); + gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer); + gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4); + }; + + const cleanUp = function() { + gl.deleteFramebuffer(frameBuffer); + gl.deleteTexture(inputTexture); + gl.deleteProgram(program); + gl.deleteShader(fragmentShader); + }; + + return {render, cleanUp}; +} diff --git a/selfie_segmentation/lib/webgl2/webgl2Pipeline.js b/selfie_segmentation/lib/webgl2/webgl2Pipeline.js new file mode 100644 index 00000000..e7b0827c --- /dev/null +++ b/selfie_segmentation/lib/webgl2/webgl2Pipeline.js @@ -0,0 +1,184 @@ +import {compileShader, createTexture, glsl} from '../helpers/webglHelper.js'; +import {buildBackgroundBlurStage} from './backgroundBlurStage.js'; +import {buildBackgroundImageStage} from './backgroundImageStage.js'; +import {buildJointBilateralFilterStage} from './jointBilateralFilterStage.js'; +import {buildLoadSegmentationStage} from './loadSegmentationStage.js'; +import {buildResizingStage} from './resizingStage.js'; + +export function buildWebGL2Pipeline( + sourcePlayback, backgroundImage, backgroundType, inputResolution, canvas, outputBuffer) { + const vertexShaderSource = glsl`#version 300 es + + in vec2 a_position; + in vec2 a_texCoord; + + out vec2 v_texCoord; + + void main() { + gl_Position = vec4(a_position, 0.0, 1.0); + v_texCoord = a_texCoord; + } + `; + + const {width: frameWidth, height: frameHeight} = sourcePlayback; + const [segmentationWidth, segmentationHeight] = inputResolution; + + const gl = canvas.getContext('webgl2'); + + const vertexShader = compileShader(gl, gl.VERTEX_SHADER, vertexShaderSource); + + const vertexArray = gl.createVertexArray(); + gl.bindVertexArray(vertexArray); + + const positionBuffer = gl.createBuffer(); + gl.bindBuffer(gl.ARRAY_BUFFER, positionBuffer); + gl.bufferData( + gl.ARRAY_BUFFER, + new Float32Array([-1.0, -1.0, 1.0, -1.0, -1.0, 1.0, 1.0, 1.0]), + gl.STATIC_DRAW, + ); + + const texCoordBuffer = gl.createBuffer(); + gl.bindBuffer(gl.ARRAY_BUFFER, texCoordBuffer); + gl.bufferData( + gl.ARRAY_BUFFER, + new Float32Array([0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 1.0, 1.0]), + gl.STATIC_DRAW, + ); + + // We don't use texStorage2D here because texImage2D seems faster + // to upload video texture than texSubImage2D even though the latter + // is supposed to be the recommended way: + // https://developer.mozilla.org/en-US/docs/Web/API/WebGL_API/WebGL_best_practices#use_texstorage_to_create_textures + const inputFrameTexture = gl.createTexture(); + gl.bindTexture(gl.TEXTURE_2D, inputFrameTexture); + 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.NEAREST); + gl.texParameteri(gl.TEXTURE_2D, gl.TEXTURE_MAG_FILTER, gl.NEAREST); + + // TODO Rename segmentation and person mask to be more specific + const segmentationTexture = createTexture( + gl, + gl.RGBA8, + segmentationWidth, + segmentationHeight, + ); + const personMaskTexture = createTexture( + gl, + gl.RGBA8, + frameWidth, + frameHeight, + ); + + const resizingStage = buildResizingStage( + gl, + vertexShader, + positionBuffer, + texCoordBuffer, + inputResolution, + ); + const loadSegmentationStage = buildLoadSegmentationStage( + gl, + vertexShader, + positionBuffer, + texCoordBuffer, + inputResolution, + outputBuffer, + segmentationTexture, + ); + const jointBilateralFilterStage = buildJointBilateralFilterStage( + gl, + vertexShader, + positionBuffer, + texCoordBuffer, + segmentationTexture, + inputResolution, + personMaskTexture, + canvas, + ); + const backgroundStage = + backgroundType === 'blur' + ? buildBackgroundBlurStage( + gl, + vertexShader, + positionBuffer, + texCoordBuffer, + personMaskTexture, + canvas, + ) + : buildBackgroundImageStage( + gl, + positionBuffer, + texCoordBuffer, + personMaskTexture, + backgroundImage, + canvas, + ); + + const render = async function() { + gl.clearColor(0, 0, 0, 0); + gl.clear(gl.COLOR_BUFFER_BIT); + + gl.activeTexture(gl.TEXTURE0); + gl.bindTexture(gl.TEXTURE_2D, inputFrameTexture); + + // texImage2D seems faster than texSubImage2D to upload + // video texture + gl.texImage2D( + gl.TEXTURE_2D, + 0, + gl.RGBA, + gl.RGBA, + gl.UNSIGNED_BYTE, + sourcePlayback, + ); + + gl.bindVertexArray(vertexArray); + + resizingStage.render(); + + loadSegmentationStage.render(); + jointBilateralFilterStage.render(); + backgroundStage.render(); + }; + + const updatePostProcessingConfig = function(postProcessingConfig) { + jointBilateralFilterStage.updateSigmaSpace( + postProcessingConfig.jointBilateralFilter.sigmaSpace); + jointBilateralFilterStage.updateSigmaColor( + postProcessingConfig.jointBilateralFilter.sigmaColor); + + if (backgroundType === 'image') { + // const backgroundImageStage = backgroundStage as BackgroundImageStage + backgroundStage.updateCoverage(postProcessingConfig.coverage); + backgroundStage.updateLightWrapping(postProcessingConfig.lightWrapping); + backgroundStage.updateBlendMode(postProcessingConfig.blendMode); + } else if (backgroundType === 'blur') { + // const backgroundBlurStage = backgroundStage as BackgroundBlurStage + backgroundStage.updateCoverage(postProcessingConfig.coverage); + } else { + // TODO Handle no background in a separate pipeline path + // const backgroundImageStage = backgroundStage as BackgroundImageStage + backgroundStage.updateCoverage([0, 0.9999]); + backgroundStage.updateLightWrapping(0); + } + }; + + const cleanUp = function() { + backgroundStage.cleanUp(); + jointBilateralFilterStage.cleanUp(); + loadSegmentationStage.cleanUp(); + resizingStage.cleanUp(); + + gl.deleteTexture(personMaskTexture); + gl.deleteTexture(segmentationTexture); + gl.deleteTexture(inputFrameTexture); + gl.deleteBuffer(texCoordBuffer); + gl.deleteBuffer(positionBuffer); + gl.deleteVertexArray(vertexArray); + gl.deleteShader(vertexShader); + }; + + return {render, updatePostProcessingConfig, cleanUp}; +} diff --git a/selfie_segmentation/main.js b/selfie_segmentation/main.js new file mode 100644 index 00000000..bca1ca57 --- /dev/null +++ b/selfie_segmentation/main.js @@ -0,0 +1,246 @@ +'use strict'; + +import * as utils from '../common/utils.js'; +import {buildWebGL2Pipeline} from './lib/webgl2/webgl2Pipeline.js'; +import * as ui from '../common/ui.js'; +const worker = new Worker('./builtin_delegate_worker.js'); + +const imgElement = document.getElementById('feedElement'); +imgElement.src = './images/test.jpg'; +const camElement = document.getElementById('feedMediaElement'); +const outputCanvas = document.getElementById('outputCanvas'); +let modelName = ''; +let rafReq; +let isFirstTimeLoad = true; +let inputType = 'image'; +let stream = null; +let loadTime = 0; +let computeTime = 0; +let outputBuffer; +let modelChanged = false; +let backgroundImageSource = document.getElementById('00-img'); +let backgroundType = 'img'; // 'none', 'blur', 'image' +const inputOptions = { + mean: [127.5, 127.5, 127.5], + std: [127.5, 127.5, 127.5], + scaledFlag: false, + inputLayout: 'nhwc', +}; +let enableWebnnDelegate = false; +const disabledSelectors = ['#tabs > li', '.btn']; + +$(document).ready(() => { + $('.icdisplay').hide(); +}); + +$('#modelBtns .btn').on('change', async (e) => { + modelChanged = true; + modelName = $(e.target).attr('id'); + if (modelName.includes('landscape')) { + inputOptions.inputDimensions = [1, 144, 256, 3]; + inputOptions.inputResolution = [256, 144]; + } else { + inputOptions.inputDimensions = [1, 256, 256, 3]; + inputOptions.inputResolution = [256, 256]; + } + if (inputType === 'camera') utils.stopCameraStream(rafReq, stream); + await main(); +}); + +$('#webnnDelegate').on('change', async (e) => { + modelChanged = true; + console.log($(e.target)); + enableWebnnDelegate = $(e.target)[0].checked; + if (inputType === 'camera') utils.stopCameraStream(rafReq, stream); + await main(); +}); + +// Click trigger to do inference with element +$('#img').click(async () => { + if (inputType === 'camera') utils.stopCameraStream(rafReq, stream); + inputType = 'image'; + $('#pickimage').show(); + $('.shoulddisplay').hide(); + await main(); +}); + +$('#imageFile').change((e) => { + const files = e.target.files; + if (files.length > 0) { + $('#feedElement').removeAttr('height'); + $('#feedElement').removeAttr('width'); + imgElement.src = URL.createObjectURL(files[0]); + } +}); + +$('#feedElement').on('load', async () => { + await main(); +}); + +// Click trigger to do inference with