@stream-io/video-filters-web 0.0.1

package/src/webgl2/jointBilateralFilterStage.ts

@@ -0,0 +1,159 @@
+import {
+  compileShader,
+  createPipelineStageProgram,
+  glsl,
+} from '../helpers/webglHelper';
+import { SegmentationParams } from '../segmentation';
+
+export function buildJointBilateralFilterStage(
+  gl: WebGL2RenderingContext,
+  vertexShader: WebGLShader,
+  positionBuffer: WebGLBuffer,
+  texCoordBuffer: WebGLBuffer,
+  inputTexture: WebGLTexture,
+  outputTexture: WebGLTexture,
+  canvas: HTMLCanvasElement,
+  segmentationConfig: SegmentationParams,
+) {
+  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 { width: segmentationWidth, height: segmentationHeight } =
+    segmentationConfig;
+  const { width: outputWidth, height: outputHeight } = canvas;
+  const texelWidth = 1 / outputWidth;
+  const texelHeight = 1 / outputHeight;
+
+  const fragmentShader = compileShader(
+    gl,
+    gl.FRAGMENT_SHADER,
+    fragmentShaderSource,
+  );
+  const program = createPipelineStageProgram(
+    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);
+
+  // Ensures default values are configured to prevent infinite
+  // loop in fragment shader
+  updateSigmaSpace(0);
+  updateSigmaColor(0);
+
+  function render() {
+    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);
+  }
+
+  function updateSigmaSpace(sigmaSpace: number) {
+    sigmaSpace *= Math.max(
+      outputWidth / segmentationWidth,
+      outputHeight / segmentationHeight,
+    );
+
+    const kSparsityFactor = 0.66; // Higher is sparser.
+    const step = Math.max(1, Math.sqrt(sigmaSpace) * kSparsityFactor);
+    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);
+  }
+
+  function updateSigmaColor(sigmaColor: number) {
+    gl.useProgram(program);
+    gl.uniform1f(sigmaColorLocation, sigmaColor);
+  }
+
+  function cleanUp() {
+    gl.deleteFramebuffer(frameBuffer);
+    gl.deleteProgram(program);
+    gl.deleteShader(fragmentShader);
+  }
+
+  return { render, updateSigmaSpace, updateSigmaColor, cleanUp };
+}

package/src/webgl2/resizingStage.ts

@@ -0,0 +1,103 @@
+import { TFLite } from '../tflite';
+import {
+  compileShader,
+  createPipelineStageProgram,
+  createTexture,
+  glsl,
+  readPixelsAsync,
+} from '../helpers/webglHelper';
+import { SegmentationParams } from '../segmentation';
+
+export function buildResizingStage(
+  gl: WebGL2RenderingContext,
+  vertexShader: WebGLShader,
+  positionBuffer: WebGLBuffer,
+  texCoordBuffer: WebGLBuffer,
+  tflite: TFLite,
+  segmentationConfig: SegmentationParams,
+) {
+  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);
+    }
+  `;
+
+  // TFLite memory will be accessed as float32
+  const tfliteInputMemoryOffset = tflite._getInputMemoryOffset() / 4;
+
+  const { width: outputWidth, height: outputHeight } = segmentationConfig;
+  const outputPixelCount = outputWidth * outputHeight;
+
+  const fragmentShader = compileShader(
+    gl,
+    gl.FRAGMENT_SHADER,
+    fragmentShaderSource,
+  );
+  const program = createPipelineStageProgram(
+    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);
+
+  function render() {
+    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.
+    // The pixels will be available in the next frame render which results
+    // in offsets in the segmentation output but increases the frame rate.
+    readPixelsAsync(
+      gl,
+      0,
+      0,
+      outputWidth,
+      outputHeight,
+      gl.RGBA,
+      gl.UNSIGNED_BYTE,
+      outputPixels,
+    );
+
+    for (let i = 0; i < outputPixelCount; i++) {
+      const tfliteIndex = tfliteInputMemoryOffset + i * 3;
+      const outputIndex = i * 4;
+      tflite.HEAPF32[tfliteIndex] = outputPixels[outputIndex] / 255;
+      tflite.HEAPF32[tfliteIndex + 1] = outputPixels[outputIndex + 1] / 255;
+      tflite.HEAPF32[tfliteIndex + 2] = outputPixels[outputIndex + 2] / 255;
+    }
+  }
+
+  function cleanUp() {
+    gl.deleteFramebuffer(frameBuffer);
+    gl.deleteTexture(outputTexture);
+    gl.deleteProgram(program);
+    gl.deleteShader(fragmentShader);
+  }
+
+  return { render, cleanUp };
+}

package/src/webgl2/softmaxStage.ts

@@ -0,0 +1,103 @@
+import { TFLite } from '../tflite';
+import {
+  compileShader,
+  createPipelineStageProgram,
+  createTexture,
+  glsl,
+} from '../helpers/webglHelper';
+import { SegmentationParams } from '../segmentation';
+
+export function buildSoftmaxStage(
+  gl: WebGL2RenderingContext,
+  vertexShader: WebGLShader,
+  positionBuffer: WebGLBuffer,
+  texCoordBuffer: WebGLBuffer,
+  tflite: TFLite,
+  outputTexture: WebGLTexture,
+  segmentationConfig: SegmentationParams,
+) {
+  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));
+    }
+  `;
+
+  // TFLite memory will be accessed as float32
+  const tfliteOutputMemoryOffset = tflite._getOutputMemoryOffset() / 4;
+  const { width: segmentationWidth, height: segmentationHeight } =
+    segmentationConfig;
+
+  const fragmentShader = compileShader(
+    gl,
+    gl.FRAGMENT_SHADER,
+    fragmentShaderSource,
+  );
+  const program = createPipelineStageProgram(
+    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);
+
+  function render() {
+    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,
+      tflite.HEAPF32,
+      tfliteOutputMemoryOffset,
+    );
+    gl.bindFramebuffer(gl.FRAMEBUFFER, frameBuffer);
+    gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4);
+  }
+
+  function cleanUp() {
+    gl.deleteFramebuffer(frameBuffer);
+    gl.deleteTexture(inputTexture);
+    gl.deleteProgram(program);
+    gl.deleteShader(fragmentShader);
+  }
+
+  return { render, cleanUp };
+}

package/src/webgl2/webgl2Pipeline.ts

@@ -0,0 +1,200 @@
+import { TFLite } from '../tflite';
+import { compileShader, createTexture, glsl } from '../helpers/webglHelper';
+import {
+  BackgroundBlurStage,
+  buildBackgroundBlurStage,
+} from './backgroundBlurStage';
+import {
+  BackgroundImageStage,
+  buildBackgroundImageStage,
+} from './backgroundImageStage';
+import { buildJointBilateralFilterStage } from './jointBilateralFilterStage';
+import { buildResizingStage } from './resizingStage';
+import { buildSoftmaxStage } from './softmaxStage';
+import { BackgroundBlurLevel, BackgroundFilter } from '../createRenderer';
+import { SegmentationParams } from '../segmentation';
+
+export function buildWebGL2Pipeline(
+  videoSource: HTMLVideoElement,
+  backgroundImage: HTMLImageElement | undefined,
+  blurLevel: BackgroundBlurLevel | undefined,
+  backgroundFilter: BackgroundFilter,
+  canvas: HTMLCanvasElement,
+  tflite: TFLite,
+  segmentationConfig: SegmentationParams,
+) {
+  const gl = canvas.getContext('webgl2')!;
+  if (!gl) throw new Error('WebGL2 is not supported');
+
+  const { width: frameWidth, height: frameHeight } = videoSource;
+  const { width: segmentationWidth, height: segmentationHeight } =
+    segmentationConfig;
+
+  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 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,
+    tflite,
+    segmentationConfig,
+  );
+  const loadSegmentationStage = buildSoftmaxStage(
+    gl,
+    vertexShader,
+    positionBuffer,
+    texCoordBuffer,
+    tflite,
+    segmentationTexture,
+    segmentationConfig,
+  );
+
+  const jointBilateralFilterStage = buildJointBilateralFilterStage(
+    gl,
+    vertexShader,
+    positionBuffer,
+    texCoordBuffer,
+    segmentationTexture,
+    personMaskTexture,
+    canvas,
+    segmentationConfig,
+  );
+  const backgroundStage =
+    backgroundFilter === 'blur'
+      ? buildBackgroundBlurStage(
+          gl,
+          vertexShader,
+          positionBuffer,
+          texCoordBuffer,
+          personMaskTexture,
+          canvas,
+          blurLevel || 'high',
+        )
+      : buildBackgroundImageStage(
+          gl,
+          positionBuffer,
+          texCoordBuffer,
+          personMaskTexture,
+          backgroundImage,
+          canvas,
+        );
+
+  function render() {
+    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,
+      videoSource,
+    );
+
+    gl.bindVertexArray(vertexArray);
+
+    resizingStage.render();
+
+    tflite._runInference();
+
+    loadSegmentationStage.render();
+    jointBilateralFilterStage.render();
+    backgroundStage.render();
+  }
+
+  function updatePostProcessingConfig() {
+    jointBilateralFilterStage.updateSigmaSpace(1);
+    jointBilateralFilterStage.updateSigmaColor(0.1);
+
+    if (backgroundFilter === 'image') {
+      const backgroundImageStage = backgroundStage as BackgroundImageStage;
+      backgroundImageStage.updateCoverage([0.5, 0.75]);
+      backgroundImageStage.updateLightWrapping(0.3);
+      backgroundImageStage.updateBlendMode('screen');
+    } else if (backgroundFilter === 'blur') {
+      const backgroundBlurStage = backgroundStage as BackgroundBlurStage;
+      backgroundBlurStage.updateCoverage([0.5, 0.75]);
+    } else {
+      // TODO Handle no background in a separate pipeline path
+      const backgroundImageStage = backgroundStage as BackgroundImageStage;
+      backgroundImageStage.updateCoverage([0, 0.9999]);
+      backgroundImageStage.updateLightWrapping(0);
+    }
+  }
+
+  function cleanUp() {
+    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 };
+}