image-studio/renderer/image-texture-programs.ts

import type { MaskVisualizationMode } from "./image-textures";

export type MaskVisualizationResources = {
  program: WebGLProgram;
  positionLocation: number;
  texCoordLocation: number;
  samplerLocation: WebGLUniformLocation;
  modeLocation: WebGLUniformLocation;
  colorLocation: WebGLUniformLocation;
};

export function createProgram(gl: WebGL2RenderingContext) {
  const vertexShader = compileShader(
    gl,
    gl.VERTEX_SHADER,
    `#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 fragmentShader = compileShader(
    gl,
    gl.FRAGMENT_SHADER,
    `#version 300 es
    precision mediump float;
    uniform sampler2D u_image;
    uniform float u_opacity;
    in vec2 v_texCoord;
    out vec4 outColor;
    void main() {
      outColor = texture(u_image, v_texCoord) * u_opacity;
    }`,
  );
  const program = gl.createProgram();
  if (!program) throw new Error("Failed to create image shader program");

  gl.attachShader(program, vertexShader);
  gl.attachShader(program, fragmentShader);
  gl.linkProgram(program);
  gl.deleteShader(vertexShader);
  gl.deleteShader(fragmentShader);

  if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
    const message = gl.getProgramInfoLog(program) ?? "Unknown program link error";
    gl.deleteProgram(program);
    throw new Error(message);
  }

  return program;
}

export function getMaskVisualizationResources(
  gl: WebGL2RenderingContext,
  getResources: () => MaskVisualizationResources | "failed" | undefined,
  setResources: (resources: MaskVisualizationResources | "failed") => void,
): MaskVisualizationResources | undefined {
  const currentResources = getResources();
  if (currentResources === "failed") return undefined;
  if (currentResources) return currentResources;

  try {
    const resources = createMaskVisualizationProgram(gl);
    setResources(resources);
    return resources;
  } catch {
    setResources("failed");
    return undefined;
  }
}

export function maskVisualizationModeValue(mode: MaskVisualizationMode) {
  switch (mode) {
    case "blackWhite":
      return 0;
    case "alpha":
      return 1;
    case "hiddenOverlay":
      return 2;
  }
}

function createMaskVisualizationProgram(gl: WebGL2RenderingContext): MaskVisualizationResources {
  const vertexShader = compileShader(
    gl,
    gl.VERTEX_SHADER,
    `#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 fragmentShader = compileShader(
    gl,
    gl.FRAGMENT_SHADER,
    `#version 300 es
    precision mediump float;
    uniform sampler2D u_mask;
    uniform int u_mode;
    uniform vec4 u_color;
    in vec2 v_texCoord;
    out vec4 outColor;
    void main() {
      vec4 maskColor = texture(u_mask, v_texCoord);
      float maskAlpha = maskColor.a * dot(maskColor.rgb, vec3(0.2126, 0.7152, 0.0722));
      if (u_mode == 0) {
        float value = step(0.5, maskAlpha);
        outColor = vec4(value, value, value, 1.0);
        return;
      }
      if (u_mode == 1) {
        outColor = vec4(maskAlpha, maskAlpha, maskAlpha, 1.0);
        return;
      }

      float alpha = (1.0 - maskAlpha) * u_color.a;
      outColor = vec4(u_color.rgb * alpha, alpha);
    }`,
  );
  const program = gl.createProgram();
  if (!program) throw new Error("Failed to create mask visualization shader program");

  gl.attachShader(program, vertexShader);
  gl.attachShader(program, fragmentShader);
  gl.linkProgram(program);
  gl.deleteShader(vertexShader);
  gl.deleteShader(fragmentShader);

  if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
    const message = gl.getProgramInfoLog(program) ?? "Unknown mask visualization program link error";
    gl.deleteProgram(program);
    throw new Error(message);
  }

  const samplerLocation = gl.getUniformLocation(program, "u_mask");
  const modeLocation = gl.getUniformLocation(program, "u_mode");
  const colorLocation = gl.getUniformLocation(program, "u_color");
  if (!samplerLocation || !modeLocation || !colorLocation) {
    gl.deleteProgram(program);
    throw new Error("Failed to resolve mask visualization shader uniforms");
  }

  return {
    program,
    positionLocation: gl.getAttribLocation(program, "a_position"),
    texCoordLocation: gl.getAttribLocation(program, "a_texCoord"),
    samplerLocation,
    modeLocation,
    colorLocation,
  };
}

export function createTintedProgram(gl: WebGL2RenderingContext) {
  const vertexShader = compileShader(
    gl,
    gl.VERTEX_SHADER,
    `#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 fragmentShader = compileShader(
    gl,
    gl.FRAGMENT_SHADER,
    `#version 300 es
    precision mediump float;
    uniform sampler2D u_image;
    uniform vec4 u_color;
    in vec2 v_texCoord;
    out vec4 outColor;
    void main() {
      float alpha = u_color.a * texture(u_image, v_texCoord).a;
      outColor = vec4(u_color.rgb * alpha, alpha);
    }`,
  );
  const program = gl.createProgram();
  if (!program) throw new Error("Failed to create tinted image shader program");

  gl.attachShader(program, vertexShader);
  gl.attachShader(program, fragmentShader);
  gl.linkProgram(program);
  gl.deleteShader(vertexShader);
  gl.deleteShader(fragmentShader);

  if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
    const message = gl.getProgramInfoLog(program) ?? "Unknown tinted program link error";
    gl.deleteProgram(program);
    throw new Error(message);
  }

  return program;
}

export function createMaskRevealPreviewProgram(gl: WebGL2RenderingContext) {
  const vertexShader = compileShader(
    gl,
    gl.VERTEX_SHADER,
    `#version 300 es
    in vec2 a_position;
    in vec2 a_texCoord;
    in vec2 a_maskTexCoord;
    out vec2 v_texCoord;
    out vec2 v_maskTexCoord;
    void main() {
      gl_Position = vec4(a_position, 0.0, 1.0);
      v_texCoord = a_texCoord;
      v_maskTexCoord = a_maskTexCoord;
    }`,
  );
  const fragmentShader = compileShader(
    gl,
    gl.FRAGMENT_SHADER,
    `#version 300 es
    precision mediump float;
    uniform sampler2D u_image;
    uniform sampler2D u_mask;
    uniform float u_opacity;
    in vec2 v_texCoord;
    in vec2 v_maskTexCoord;
    out vec4 outColor;
    void main() {
      vec4 color = texture(u_image, v_texCoord);
      vec4 maskColor = texture(u_mask, v_maskTexCoord);
      float maskAlpha = maskColor.a * dot(maskColor.rgb, vec3(0.2126, 0.7152, 0.0722));
      float hiddenMaskAlpha = 1.0 - maskAlpha;
      float previewAlpha = clamp(hiddenMaskAlpha * u_opacity, 0.0, 1.0);
      outColor = vec4(color.rgb * previewAlpha, color.a * previewAlpha);
    }`,
  );
  const program = gl.createProgram();
  if (!program) throw new Error("Failed to create mask reveal preview shader program");

  gl.attachShader(program, vertexShader);
  gl.attachShader(program, fragmentShader);
  gl.linkProgram(program);
  gl.deleteShader(vertexShader);
  gl.deleteShader(fragmentShader);

  if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
    const message = gl.getProgramInfoLog(program) ?? "Unknown mask reveal preview program link error";
    gl.deleteProgram(program);
    throw new Error(message);
  }

  return program;
}

export function createMaskedProgram(gl: WebGL2RenderingContext) {
  const vertexShader = compileShader(
    gl,
    gl.VERTEX_SHADER,
    `#version 300 es
    in vec2 a_position;
    in vec2 a_texCoord;
    in vec2 a_maskTexCoord;
    out vec2 v_texCoord;
    out vec2 v_maskTexCoord;
    void main() {
      gl_Position = vec4(a_position, 0.0, 1.0);
      v_texCoord = a_texCoord;
      v_maskTexCoord = a_maskTexCoord;
    }`,
  );
  const fragmentShader = compileShader(
    gl,
    gl.FRAGMENT_SHADER,
    `#version 300 es
    precision mediump float;
    uniform sampler2D u_image;
    uniform sampler2D u_mask;
    uniform float u_opacity;
    in vec2 v_texCoord;
    in vec2 v_maskTexCoord;
    out vec4 outColor;
    void main() {
      vec4 color = texture(u_image, v_texCoord);
      vec4 maskColor = texture(u_mask, v_maskTexCoord);
      float maskAlpha = maskColor.a * dot(maskColor.rgb, vec3(0.2126, 0.7152, 0.0722));
      float alpha = color.a * maskAlpha;
      outColor = vec4(color.rgb * maskAlpha, alpha) * u_opacity;
    }`,
  );
  const program = gl.createProgram();
  if (!program) throw new Error("Failed to create masked image shader program");

  gl.attachShader(program, vertexShader);
  gl.attachShader(program, fragmentShader);
  gl.linkProgram(program);
  gl.deleteShader(vertexShader);
  gl.deleteShader(fragmentShader);

  if (!gl.getProgramParameter(program, gl.LINK_STATUS)) {
    const message = gl.getProgramInfoLog(program) ?? "Unknown masked program link error";
    gl.deleteProgram(program);
    throw new Error(message);
  }

  return program;
}

function compileShader(gl: WebGL2RenderingContext, type: number, source: string) {
  const shader = gl.createShader(type);
  if (!shader) throw new Error("Failed to create shader");

  gl.shaderSource(shader, source);
  gl.compileShader(shader);
  if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {
    const message = gl.getShaderInfoLog(shader) ?? "Unknown shader compile error";
    gl.deleteShader(shader);
    throw new Error(message);
  }

  return shader;
}