pixijs实现绿幕抠图和视频

参考实现 obs studio

github.com/obsproject/...

最终实现

webgl and webgpu代码

webgl

in vec2 aPosition;
out vec2 vTextureCoord;

uniform vec4 uInputSize;
uniform vec4 uOutputFrame;
uniform vec4 uOutputTexture;

vec4 filterVertexPosition( void )
{
    vec2 position = aPosition * uOutputFrame.zw + uOutputFrame.xy;
    
    position.x = position.x * (2.0 / uOutputTexture.x) - 1.0;
    position.y = position.y * (2.0*uOutputTexture.z / uOutputTexture.y) - uOutputTexture.z;

    return vec4(position, 0.0, 1.0);
}

vec2 filterTextureCoord( void )
{
    return aPosition * (uOutputFrame.zw * uInputSize.zw);
}

void main(void)
{
    gl_Position = filterVertexPosition();
    vTextureCoord = filterTextureCoord();
}

precision highp float;
in vec2 vTextureCoord;
out vec4 finalColor;

uniform sampler2D uTexture;
uniform float opacity;
uniform vec4 color;
uniform float contrast;
uniform float brightness;
uniform float gamma;
uniform vec2 chromaKey;
uniform vec2 pixelSize;
uniform float similarity;
uniform float smoothness;
uniform float spill;
uniform vec4 cb_v4;
uniform vec4 cr_v4;

float getNonlinearChannel(float u) {
    if (u <= 0.0031308) {
        return 12.92 * u;
    } else {
        return (1.055 * pow(u, 1.0 / 2.4)) - 0.055;
    }
}

vec3 getNonlinearColor(vec3 rgb) {
    return vec3(
        getNonlinearChannel(rgb.r),
        getNonlinearChannel(rgb.g),
        getNonlinearChannel(rgb.b)
    );
}

vec4 calcColor(vec4 rgba) {
    return vec4(
        pow(rgba.rgb, vec3(gamma)) * contrast + brightness,
        rgba.a
    );
}

float getChromaDist(vec3 rgb) {
    float cb = dot(rgb, cb_v4.xyz) + cb_v4.w;
    float cr = dot(rgb, cr_v4.xyz) + cr_v4.w;
    return distance(chromaKey, vec2(cb, cr));
}

vec3 sampleTexture(vec2 uv) {
    vec3 sampled = texture2D(uTexture, uv).rgb;
    return getNonlinearColor(sampled);
}

float getBoxFilteredChromaDist(vec3 rgb, vec2 texCoord) {
    vec2 hPixelSize = pixelSize / 2.0;
    vec2 point0 = vec2(pixelSize.x, hPixelSize.y);
    vec2 point1 = vec2(hPixelSize.x, -pixelSize.y);

    float distVal = getChromaDist(sampleTexture(texCoord - point0));
    distVal += getChromaDist(sampleTexture(texCoord + point0));
    distVal += getChromaDist(sampleTexture(texCoord - point1));
    distVal += getChromaDist(sampleTexture(texCoord + point1));
    distVal *= 2.0;
    distVal += getChromaDist(getNonlinearColor(rgb));

    return distVal / 9.0;
}

void main(void)
{
    vec4 rgba = texture2D(uTexture, vTextureCoord);

    // 修复向量分量赋值错误：使用临时变量计算新的rgb值
    vec3 newRgb = rgba.rgb;
    if (rgba.a > 0.0) {
        newRgb *= 1.0 / rgba.a;
    } else {
        newRgb *= 0.0;
    }
    // 整体更新rgba
    rgba = vec4(newRgb, rgba.a);

    float chromaDist = getBoxFilteredChromaDist(rgba.rgb, vTextureCoord);
    float baseMask = chromaDist - similarity;

    float fullMaskValue = baseMask / smoothness;
    float clampedFullMask = clamp(fullMaskValue, 0.0, 1.0);
    float fullMask = pow(clampedFullMask, 1.5);

    float spillValue = baseMask / spill;
    float clampedSpill = clamp(spillValue, 0.0, 1.0);
    float spillVal = pow(clampedSpill, 1.5);

    // 处理result的创建和修改
    vec4 result = rgba * color;
    result = vec4(result.rgb, result.a * opacity * fullMask);

    // 处理颜色溢出修正
    float desat = dot(result.rgb, vec3(0.2126, 0.7152, 0.0722));
    vec3 mixedRgb = mix(vec3(desat), result.rgb, spillVal);
    result = vec4(mixedRgb, result.a);

    // 处理最终颜色计算
    vec4 tempFinalColor = calcColor(result);
    vec3 finalRgb = tempFinalColor.rgb * tempFinalColor.a;
    finalColor = vec4(finalRgb, tempFinalColor.a);
}

webgpu

struct GlobalFilterUniforms {
  uInputSize:vec4<f32>,
  uInputPixel:vec4<f32>,
  uInputClamp:vec4<f32>,
  uOutputFrame:vec4<f32>,
  uGlobalFrame:vec4<f32>,
  uOutputTexture:vec4<f32>,
};

@group(0) @binding(0) var<uniform> gfu: GlobalFilterUniforms;

struct VSOutput {
    @builtin(position) position: vec4<f32>,
    @location(0) uv : vec2<f32>
  };

fn filterVertexPosition(aPosition:vec2<f32>) -> vec4<f32>
{
    var position = aPosition * gfu.uOutputFrame.zw + gfu.uOutputFrame.xy;

    position.x = position.x * (2.0 / gfu.uOutputTexture.x) - 1.0;
    position.y = position.y * (2.0*gfu.uOutputTexture.z / gfu.uOutputTexture.y) - gfu.uOutputTexture.z;

    return vec4(position, 0.0, 1.0);
}

fn filterTextureCoord( aPosition:vec2<f32> ) -> vec2<f32>
{
    return aPosition * (gfu.uOutputFrame.zw * gfu.uInputSize.zw);
}

fn globalTextureCoord( aPosition:vec2<f32> ) -> vec2<f32>
{
  return  (aPosition.xy / gfu.uGlobalFrame.zw) + (gfu.uGlobalFrame.xy / gfu.uGlobalFrame.zw);  
}

fn getSize() -> vec2<f32>
{
  return gfu.uGlobalFrame.zw;
}
  
@vertex
fn mainVertex(
  @location(0) aPosition : vec2<f32>, 
) -> VSOutput {
  return VSOutput(
   filterVertexPosition(aPosition),
   filterTextureCoord(aPosition)
  );
}

struct MainUniforms {
    opacity: f32,
    color: vec4<f32>,
    contrast: f32,
    brightness: f32,
    gamma: f32,
    chromaKey: vec2<f32>,
    pixelSize: vec2<f32>,
    similarity: f32,
    smoothness: f32,
    spill: f32,
    cb_v4: vec4<f32>,
    cr_v4: vec4<f32>,
};

@group(1) @binding(0) var<uniform> uniforms: MainUniforms;

@group(0) @binding(1) var uTexture: texture_2d<f32>;
@group(0) @binding(2) var uSampler: sampler;

fn getNonlinearChannel(u: f32) -> f32 {
    if u <= 0.0031308 {
        return 12.92 * u;
    } else {
        return (1.055 * pow(u, 1.0 / 2.4)) - 0.055;
    }
}

fn getNonlinearColor(rgb: vec3<f32>) -> vec3<f32> {
    return vec3<f32>(
        getNonlinearChannel(rgb.r),
        getNonlinearChannel(rgb.g),
        getNonlinearChannel(rgb.b)
    );
}

fn calcColor(rgba: vec4<f32>) -> vec4<f32> {
    return vec4<f32>(
        pow(rgba.rgb, vec3<f32>(uniforms.gamma)) * uniforms.contrast + uniforms.brightness,
        rgba.a
    );
}

fn getChromaDist(rgb: vec3<f32>) -> f32 {
    let cb = dot(rgb, uniforms.cb_v4.xyz) + uniforms.cb_v4.w;
    let cr = dot(rgb, uniforms.cr_v4.xyz) + uniforms.cr_v4.w;
    return distance(uniforms.chromaKey, vec2<f32>(cb, cr));
}

fn sampleTexture(uv: vec2<f32>) -> vec3<f32> {
    let sampled = textureSample(uTexture, uSampler, uv).rgb;
    return getNonlinearColor(sampled);
}

fn getBoxFilteredChromaDist(rgb: vec3<f32>, texCoord: vec2<f32>) -> f32 {
    let hPixelSize = uniforms.pixelSize / 2.0;
    let point0 = vec2<f32>(uniforms.pixelSize.x, hPixelSize.y);
    let point1 = vec2<f32>(hPixelSize.x, -uniforms.pixelSize.y);

    var distVal = getChromaDist(sampleTexture(texCoord - point0));
    distVal += getChromaDist(sampleTexture(texCoord + point0));
    distVal += getChromaDist(sampleTexture(texCoord - point1));
    distVal += getChromaDist(sampleTexture(texCoord + point1));
    distVal *= 2.0;
    distVal += getChromaDist(getNonlinearColor(rgb));

    return distVal / 9.0;
}

@fragment
fn mainFragment(
    @builtin(position) position: vec4<f32>,
    @location(0) uv: vec2<f32>,
) -> @location(0) vec4<f32> {
    var rgba = textureSample(uTexture, uSampler, uv);

    // 修复向量分量赋值错误：使用临时变量计算新的rgb值
    var newRgb = rgba.rgb;
    if rgba.a > 0.0 {
        newRgb *= 1.0 / rgba.a;
    } else {
        newRgb *= 0.0;
    }
    // 整体更新rgba
    rgba = vec4<f32>(newRgb, rgba.a);

    let chromaDist = getBoxFilteredChromaDist(rgba.rgb, uv);
    let baseMask = chromaDist - uniforms.similarity;

    let fullMaskValue = baseMask / uniforms.smoothness;
    let clampedFullMask = clamp(fullMaskValue, 0.0, 1.0);
    let fullMask = pow(clampedFullMask, 1.5);

    let spillValue = baseMask / uniforms.spill;
    let clampedSpill = clamp(spillValue, 0.0, 1.0);
    let spillVal = pow(clampedSpill, 1.5);

    // 处理result的创建和修改
    var result = rgba * uniforms.color;
    result = vec4<f32>(result.rgb, result.a * uniforms.opacity * fullMask);

    // 处理颜色溢出修正
    let desat = dot(result.rgb, vec3<f32>(0.2126, 0.7152, 0.0722));
    let mixedRgb = mix(vec3<f32>(desat), result.rgb, spillVal);
    result = vec4<f32>(mixedRgb, result.a);

    // 处理最终颜色计算
    let finalColor = calcColor(result);
    let finalRgb = finalColor.rgb * finalColor.a;
    return vec4<f32>(finalRgb, finalColor.a);
}

滤镜封装

import * as Pixi from "pixi.js";
import { wgslVertex, vertex } from "../default";
import wgsl from "./source.wgsl";
import frag from "./source.frag";

export interface ChromaKeyFilterV2Options {
  opacity: number;
  color: [number, number, number, number];
  contrast: number;
  brightness: number;
  gamma: number;
  chromaKey: [number, number];
  pixelSize: [number, number];
  similarity: number;
  smoothness: number;
  spill: number;
  cb_v4: [number, number, number, number];
  cr_v4: [number, number, number, number];
}

export class ChromaKeyFilterV2 extends Pixi.Filter {
  public static readonly DEFAULT_OPTIONS: ChromaKeyFilterV2Options = {
    opacity: 1,
    color: [1, 1, 1, 1],
    contrast: 1.0,
    brightness: 0.0,
    gamma: 1.0,
    chromaKey: [0.163389, 0.10301899999999997],
    pixelSize: [1 / 1080, 1 / 1920],
    similarity: 0.4,
    smoothness: 0.08,
    spill: 0.1,
    cb_v4: [-0.100644, -0.338572, 0.439216, 0.501961],
    cr_v4: [0.439216, -0.398942, -0.040274, 0.501961],
  };

  public uniforms: {
    opacity: number;
    color: Float32Array;
    contrast: number;
    brightness: number;
    gamma: number;
    chromaKey: Float32Array;
    pixelSize: Float32Array;
    similarity: number;
    smoothness: number;
    spill: number;
    cb_v4: Float32Array;
    cr_v4: Float32Array;
  };

  constructor(options?: Partial<ChromaKeyFilterV2Options>) {
    options = { ...ChromaKeyFilterV2.DEFAULT_OPTIONS, ...(options ?? {}) };

    const gpuProgram = Pixi.GpuProgram.from({
      vertex: {
        source: wgslVertex,
        entryPoint: "mainVertex",
      },
      fragment: {
        source: wgsl,
        entryPoint: "mainFragment",
      },
    });

    const glProgram = Pixi.GlProgram.from({
      vertex: vertex,
      fragment: frag,
      name: "chroma-key-filter-v2",
    });

    super({
      gpuProgram,
      glProgram,
      resources: {
        uniforms: {
          opacity: { value: options.opacity, type: "f32" },
          color: { value: options.color, type: "vec4<f32>" },
          contrast: { value: options.contrast, type: "f32" },
          brightness: { value: options.brightness, type: "f32" },
          gamma: { value: options.gamma, type: "f32" },
          chromaKey: { value: options.chromaKey, type: "vec2<f32>" },
          pixelSize: { value: options.pixelSize, type: "vec2<f32>" },
          similarity: { value: options.similarity, type: "f32" },
          smoothness: { value: options.smoothness, type: "f32" },
          spill: { value: options.spill, type: "f32" },
          cb_v4: { value: options.cb_v4, type: "vec4<f32>" },
          cr_v4: { value: options.cr_v4, type: "vec4<f32>" },
        },
      },
    });

    this.uniforms = this.resources.uniforms.uniforms;
  }

  update(options: Partial<ChromaKeyFilterV2Options>) {
    this.uniforms.opacity = options.opacity ?? this.uniforms.opacity;
    this.uniforms.color = new Float32Array(
      options.color ?? this.uniforms.color
    );
    this.uniforms.contrast = options.contrast ?? this.uniforms.contrast;
    this.uniforms.brightness = options.brightness ?? this.uniforms.brightness;
    this.uniforms.gamma = options.gamma ?? this.uniforms.gamma;
    this.uniforms.chromaKey = new Float32Array(
      options.chromaKey ?? this.uniforms.chromaKey
    );
    this.uniforms.pixelSize = new Float32Array(
      options.pixelSize ?? this.uniforms.pixelSize
    );
    this.uniforms.similarity = options.similarity ?? this.uniforms.similarity;
    this.uniforms.smoothness = options.smoothness ?? this.uniforms.smoothness;
    this.uniforms.spill = options.spill ?? this.uniforms.spill;
    this.uniforms.cb_v4 = new Float32Array(
      options.cb_v4 ?? this.uniforms.cb_v4
    );
    this.uniforms.cr_v4 = new Float32Array(
      options.cr_v4 ?? this.uniforms.cr_v4
    );
  }
}

核心参数计算公式

let rgba = [0.0, 1.0, 0.0, 1.0];
if (SettingStore.type === "red") {
  rgba = [1.0, 0.0, 0.0, 1.0];
} else if (SettingStore.type === "blue") {
  rgba = [0.0, 0.0, 1.0, 1.0];
} else if (SettingStore.type === "custom") {
  rgba = [
    SettingStore.color[0] / 255,
    SettingStore.color[1] / 255,
    SettingStore.color[2] / 255,
    SettingStore.color[3] / 255,
  ];
}
const cb =
  rgba[0] * -0.100644 +
  rgba[1] * -0.338572 +
  rgba[2] * 0.439216 +
  rgba[3] * 0.501961;

const cr =
  rgba[0] * 0.439216 +
  rgba[1] * -0.398942 +
  rgba[2] * -0.040274 +
  rgba[3] * 0.501961;

shader.current.update({
  opacity: SettingStore.opacity,
  contrast:
    SettingStore.contrast < 0.0
      ? 1.0 / (-SettingStore.contrast + 1.0)
      : SettingStore.contrast + 1.0,
  brightness: SettingStore.brightness * 0.5,
  gamma:
    SettingStore.gamma < 0.0
      ? -SettingStore.gamma + 1.0
      : 1.0 / (SettingStore.gamma + 1.0),
  pixelSize: [1 / width, 1 / height],
  similarity: SettingStore.similarity * 0.001,
  smoothness: SettingStore.smoothness * 0.001,
  spill: SettingStore.spill * 0.001,
  chromaKey: [cb, cr],
});

ps 自定义颜色还未实现。

ui库没带，懒得找，基本上计算公式都是一致的

github

github.com/tetap/obs-c...

github pages

tetap.github.io/obs-chroma-...