image-studio/platform/browser/maskRaster.ts

import type { Rect } from "@core/geometry";
import { assertCanvasRasterSize, assertProcessingRasterSize } from "./rasterLimits";

export type MaskFill = "white" | "black" | "clear";

export type MaskRasterOperation =
  | { type: "invert" }
  | { type: "fill"; fill: MaskFill }
  | { type: "feather"; radius: number }
  | { type: "expand"; radius: number }
  | { type: "contract"; radius: number }
  | { type: "blur"; radius: number }
  | { type: "despeckle"; strength: number };

export type MaskAnalysis = {
  width: number;
  height: number;
  revealedPixels: number;
  hiddenPixels: number;
  coverage: number;
  hiddenCoverage: number;
  bounds?: Rect;
  hiddenBounds?: Rect;
  thumbnail: string;
};

export type NormalizedMaskOptions = {
  polarity: "hidden" | "revealed";
  expand?: number;
  feather?: number;
  blur?: number;
  despeckle?: number;
};

export async function createSolidMaskSource(width: number, height: number, fill: MaskFill): Promise<string> {
  assertProcessingRasterSize(width, height, "Mask");
  const canvas = createCanvas(width, height);
  const context = require2dContext(canvas);
  context.clearRect(0, 0, canvas.width, canvas.height);

  if (fill === "white") {
    context.fillStyle = "#ffffff";
    context.fillRect(0, 0, canvas.width, canvas.height);
  } else if (fill === "black") {
    context.fillStyle = "#000000";
    context.fillRect(0, 0, canvas.width, canvas.height);
  }

  return canvas.toDataURL("image/png");
}

export async function applyMaskRasterOperation(source: string, width: number, height: number, operation: MaskRasterOperation): Promise<string> {
  assertProcessingRasterSize(width, height, "Mask");
  if (operation.type === "fill") return createSolidMaskSource(width, height, operation.fill);

  const mask = await loadMaskValues(source, width, height);
  const next = applyMaskValueOperation(mask.values, mask.width, mask.height, operation);
  return maskValuesToDataUrl(next, mask.width, mask.height);
}

export async function analyzeMaskSource(source: string, width: number, height: number): Promise<MaskAnalysis> {
  assertProcessingRasterSize(width, height, "Mask");
  const mask = await loadMaskValues(source, width, height);
  const revealed = analyzeValues(mask.values, mask.width, mask.height, false);
  const hiddenValues = invertMaskValues(mask.values);
  const hidden = analyzeValues(hiddenValues, mask.width, mask.height, false);
  return {
    width: mask.width,
    height: mask.height,
    revealedPixels: revealed.pixels,
    hiddenPixels: hidden.pixels,
    coverage: revealed.coverage,
    hiddenCoverage: hidden.coverage,
    bounds: revealed.bounds,
    hiddenBounds: hidden.bounds,
    thumbnail: maskValuesToDataUrl(mask.values, mask.width, mask.height, 72, 48),
  };
}

export async function createNormalizedMaskSource(source: string, width: number, height: number, options: NormalizedMaskOptions): Promise<{ source: string; values: Uint8ClampedArray; bounds?: Rect }> {
  assertProcessingRasterSize(width, height, "Mask");
  const mask = await loadMaskValues(source, width, height);
  let values = options.polarity === "hidden" ? invertMaskValues(mask.values) : new Uint8ClampedArray(mask.values);

  const despeckle = Math.round(clampNumber(options.despeckle ?? 0, 0, 64));
  const expand = Math.round(clampNumber(options.expand ?? 0, -256, 256));
  const feather = Math.round(clampNumber(options.feather ?? 0, 0, 256));
  const blur = Math.round(clampNumber(options.blur ?? 0, 0, 256));

  if (despeckle > 0) values = despeckleMaskValues(values, mask.width, mask.height, despeckle);
  if (expand > 0) values = dilateMaskValues(values, mask.width, mask.height, expand);
  if (expand < 0) values = erodeMaskValues(values, mask.width, mask.height, -expand);
  if (feather > 0) values = blurMaskValues(values, mask.width, mask.height, feather);
  if (blur > 0) values = blurMaskValues(values, mask.width, mask.height, blur);

  const analysis = analyzeValues(values, mask.width, mask.height, false);
  return { source: maskValuesToDataUrl(values, mask.width, mask.height), values, bounds: analysis.bounds };
}

export async function loadImageCanvas(source: string, width?: number, height?: number): Promise<HTMLCanvasElement> {
  const image = await loadImage(source);
  assertCanvasRasterSize(width ?? image.naturalWidth, height ?? image.naturalHeight);
  const canvas = createCanvas(width ?? image.naturalWidth, height ?? image.naturalHeight);
  const context = require2dContext(canvas);
  context.clearRect(0, 0, canvas.width, canvas.height);
  context.drawImage(image, 0, 0, canvas.width, canvas.height);
  return canvas;
}

export async function imageSourceToPngDataUrl(source: string): Promise<string> {
  if (source.startsWith("data:image/png;base64,")) return source;
  const canvas = await loadImageCanvas(source);
  return canvas.toDataURL("image/png");
}

export function cropCanvas(sourceCanvas: HTMLCanvasElement, crop: Rect, outputWidth = crop.w, outputHeight = crop.h): string {
  assertProcessingRasterSize(outputWidth, outputHeight, "Crop");
  const canvas = createCanvas(outputWidth, outputHeight);
  const context = require2dContext(canvas);
  context.clearRect(0, 0, outputWidth, outputHeight);
  context.drawImage(sourceCanvas, crop.x, crop.y, crop.w, crop.h, 0, 0, crop.w, crop.h);
  return canvas.toDataURL("image/png");
}

export function cropMaskValuesToDataUrl(values: Uint8ClampedArray, width: number, height: number, crop: Rect, outputWidth = crop.w, outputHeight = crop.h): string {
  assertProcessingRasterSize(outputWidth, outputHeight, "Mask crop");
  const canvas = createCanvas(outputWidth, outputHeight);
  const context = require2dContext(canvas);
  const imageData = context.createImageData(outputWidth, outputHeight);
  imageData.data.set(cropMaskValuesToRgba(values, width, height, crop, outputWidth, outputHeight));
  context.putImageData(imageData, 0, 0);
  return canvas.toDataURL("image/png");
}

export function cropMaskValuesToRgba(values: Uint8ClampedArray, width: number, height: number, crop: Rect, outputWidth = crop.w, outputHeight = crop.h): Uint8ClampedArray {
  const safeOutputWidth = Math.max(1, Math.round(outputWidth));
  const safeOutputHeight = Math.max(1, Math.round(outputHeight));
  const data = new Uint8ClampedArray(safeOutputWidth * safeOutputHeight * 4);
  for (let pixel = 0; pixel < safeOutputWidth * safeOutputHeight; pixel += 1) data[pixel * 4 + 3] = 255;

  for (let y = 0; y < Math.min(crop.h, safeOutputHeight); y += 1) {
    for (let x = 0; x < Math.min(crop.w, safeOutputWidth); x += 1) {
      const sourceX = crop.x + x;
      const sourceY = crop.y + y;
      if (sourceX < 0 || sourceY < 0 || sourceX >= width || sourceY >= height) continue;
      const value = values[sourceY * width + sourceX] ?? 0;
      const index = (y * safeOutputWidth + x) * 4;
      data[index] = value;
      data[index + 1] = value;
      data[index + 2] = value;
    }
  }

  return data;
}

export function expandRectWithinBounds(rect: Rect, padding: number, bounds: { w: number; h: number }, multiple = 1, minSize = 1): Rect {
  const padded = Math.max(0, Math.round(padding));
  let x1 = Math.max(0, Math.floor(rect.x) - padded);
  let y1 = Math.max(0, Math.floor(rect.y) - padded);
  let x2 = Math.min(bounds.w, Math.ceil(rect.x + rect.w) + padded);
  let y2 = Math.min(bounds.h, Math.ceil(rect.y + rect.h) + padded);

  const safeMinSize = Math.max(1, Math.round(minSize));
  const targetWidth = Math.min(bounds.w, roundUp(Math.max(safeMinSize, x2 - x1), multiple));
  const targetHeight = Math.min(bounds.h, roundUp(Math.max(safeMinSize, y2 - y1), multiple));

  const extraWidth = targetWidth - (x2 - x1);
  const extraHeight = targetHeight - (y2 - y1);
  x1 = Math.max(0, x1 - Math.floor(extraWidth / 2));
  y1 = Math.max(0, y1 - Math.floor(extraHeight / 2));
  x2 = Math.min(bounds.w, x1 + targetWidth);
  y2 = Math.min(bounds.h, y1 + targetHeight);
  x1 = Math.max(0, x2 - targetWidth);
  y1 = Math.max(0, y2 - targetHeight);

  return { x: x1, y: y1, w: Math.max(1, x2 - x1), h: Math.max(1, y2 - y1) };
}

export function maskValueFromRgba(data: Uint8ClampedArray, index: number): number {
  const red = data[index] ?? 0;
  const green = data[index + 1] ?? 0;
  const blue = data[index + 2] ?? 0;
  const alpha = data[index + 3] ?? 0;
  const luminance = 0.2126 * red + 0.7152 * green + 0.0722 * blue;
  return Math.round((alpha * luminance) / 255);
}

export function applyMaskValueOperation(values: Uint8ClampedArray, width: number, height: number, operation: Exclude<MaskRasterOperation, { type: "fill" }>): Uint8ClampedArray {
  switch (operation.type) {
    case "invert":
      return invertMaskValues(values);
    case "feather":
    case "blur":
      return blurMaskValues(values, width, height, Math.round(clampNumber(operation.radius, 0, 256)));
    case "expand":
      return dilateMaskValues(values, width, height, Math.round(clampNumber(operation.radius, 0, 256)));
    case "contract":
      return erodeMaskValues(values, width, height, Math.round(clampNumber(operation.radius, 0, 256)));
    case "despeckle":
      return despeckleMaskValues(values, width, height, Math.round(clampNumber(operation.strength, 0, 64)));
  }
}

export function invertMaskValues(values: Uint8ClampedArray): Uint8ClampedArray {
  const next = new Uint8ClampedArray(values.length);
  for (let index = 0; index < values.length; index += 1) next[index] = 255 - (values[index] ?? 0);
  return next;
}

export function erodeMaskValues(values: Uint8ClampedArray, width: number, height: number, radius: number): Uint8ClampedArray {
  const safeRadius = Math.round(clampNumber(radius, 0, 256));
  if (safeRadius <= 0) return new Uint8ClampedArray(values);
  return separableExtrema(values, width, height, safeRadius, Math.min, 255);
}

export function dilateMaskValues(values: Uint8ClampedArray, width: number, height: number, radius: number): Uint8ClampedArray {
  const safeRadius = Math.round(clampNumber(radius, 0, 256));
  if (safeRadius <= 0) return new Uint8ClampedArray(values);
  return separableExtrema(values, width, height, safeRadius, Math.max, 0);
}

export function blurMaskValues(values: Uint8ClampedArray, width: number, height: number, radius: number): Uint8ClampedArray {
  const safeRadius = Math.round(clampNumber(radius, 0, 256));
  if (safeRadius <= 0) return new Uint8ClampedArray(values);
  const horizontal = new Float64Array(values.length);
  const next = new Uint8ClampedArray(values.length);
  const span = safeRadius * 2 + 1;
  for (let y = 0; y < height; y += 1) {
    let total = 0;
    for (let ox = -safeRadius; ox <= safeRadius; ox += 1) total += values[y * width + clampInt(ox, 0, width - 1)] ?? 0;
    for (let x = 0; x < width; x += 1) {
      horizontal[y * width + x] = total / span;
      total += (values[y * width + clampInt(x + safeRadius + 1, 0, width - 1)] ?? 0) - (values[y * width + clampInt(x - safeRadius, 0, width - 1)] ?? 0);
    }
  }
  for (let x = 0; x < width; x += 1) {
    let total = 0;
    for (let oy = -safeRadius; oy <= safeRadius; oy += 1) total += horizontal[clampInt(oy, 0, height - 1) * width + x] ?? 0;
    for (let y = 0; y < height; y += 1) {
      next[y * width + x] = Math.round(total / span);
      total += (horizontal[clampInt(y + safeRadius + 1, 0, height - 1) * width + x] ?? 0) - (horizontal[clampInt(y - safeRadius, 0, height - 1) * width + x] ?? 0);
    }
  }
  return next;
}

function separableExtrema(values: Uint8ClampedArray, width: number, height: number, radius: number, combine: (a: number, b: number) => number, initial: number) {
  const horizontal = new Uint8ClampedArray(values.length);
  const next = new Uint8ClampedArray(values.length);
  for (let y = 0; y < height; y += 1) for (let x = 0; x < width; x += 1) {
    let value = initial;
    for (let offset = -radius; offset <= radius; offset += 1) value = combine(value, values[y * width + clampInt(x + offset, 0, width - 1)] ?? 0);
    horizontal[y * width + x] = value;
  }
  for (let y = 0; y < height; y += 1) for (let x = 0; x < width; x += 1) {
    let value = initial;
    for (let offset = -radius; offset <= radius; offset += 1) value = combine(value, horizontal[clampInt(y + offset, 0, height - 1) * width + x] ?? 0);
    next[y * width + x] = value;
  }
  return next;
}

export function despeckleMaskValues(values: Uint8ClampedArray, width: number, height: number, strength: number): Uint8ClampedArray {
  const safeStrength = Math.round(clampNumber(strength, 0, 64));
  if (safeStrength <= 0) return new Uint8ClampedArray(values);

  const radius = Math.max(1, Math.ceil(safeStrength / 6));
  const threshold = Math.max(1, Math.round(safeStrength / 2));
  const next = new Uint8ClampedArray(values);

  for (let y = 0; y < height; y += 1) {
    for (let x = 0; x < width; x += 1) {
      const index = y * width + x;
      const visible = (values[index] ?? 0) > 127;
      let same = 0;
      for (let oy = -radius; oy <= radius; oy += 1) {
        for (let ox = -radius; ox <= radius; ox += 1) {
          if (ox === 0 && oy === 0) continue;
          const sample = values[clampInt(y + oy, 0, height - 1) * width + clampInt(x + ox, 0, width - 1)] ?? 0;
          if ((sample > 127) === visible) same += 1;
        }
      }
      if (same <= threshold) next[index] = visible ? 0 : 255;
    }
  }

  return next;
}

async function loadMaskValues(source: string, width: number, height: number): Promise<{ width: number; height: number; values: Uint8ClampedArray }> {
  const canvas = await loadImageCanvas(source, width, height);
  const context = require2dContext(canvas);
  const data = context.getImageData(0, 0, canvas.width, canvas.height);
  const values = new Uint8ClampedArray(canvas.width * canvas.height);
  for (let pixel = 0; pixel < values.length; pixel += 1) values[pixel] = maskValueFromRgba(data.data, pixel * 4);
  return { width: canvas.width, height: canvas.height, values };
}

function maskValuesToDataUrl(values: Uint8ClampedArray, width: number, height: number, outputWidth = width, outputHeight = height): string {
  const canvas = createCanvas(outputWidth, outputHeight);
  const context = require2dContext(canvas);
  const imageData = context.createImageData(outputWidth, outputHeight);

  for (let y = 0; y < outputHeight; y += 1) {
    for (let x = 0; x < outputWidth; x += 1) {
      const sourceX = Math.floor((x / outputWidth) * width);
      const sourceY = Math.floor((y / outputHeight) * height);
      const value = values[clampInt(sourceY, 0, height - 1) * width + clampInt(sourceX, 0, width - 1)] ?? 0;
      const index = (y * outputWidth + x) * 4;
      imageData.data[index] = 255;
      imageData.data[index + 1] = 255;
      imageData.data[index + 2] = 255;
      imageData.data[index + 3] = value;
    }
  }

  context.putImageData(imageData, 0, 0);
  return canvas.toDataURL("image/png");
}

function analyzeValues(values: Uint8ClampedArray, width: number, height: number, includeSoftPixels: boolean): { pixels: number; coverage: number; bounds?: Rect } {
  let pixels = 0;
  let minX = Number.POSITIVE_INFINITY;
  let minY = Number.POSITIVE_INFINITY;
  let maxX = Number.NEGATIVE_INFINITY;
  let maxY = Number.NEGATIVE_INFINITY;

  for (let y = 0; y < height; y += 1) {
    for (let x = 0; x < width; x += 1) {
      const value = values[y * width + x] ?? 0;
      const active = includeSoftPixels ? value > 0 : value > 127;
      if (!active) continue;
      pixels += 1;
      minX = Math.min(minX, x);
      minY = Math.min(minY, y);
      maxX = Math.max(maxX, x + 1);
      maxY = Math.max(maxY, y + 1);
    }
  }

  return {
    pixels,
    coverage: pixels / Math.max(1, width * height),
    bounds: pixels > 0 ? { x: minX, y: minY, w: maxX - minX, h: maxY - minY } : undefined,
  };
}

function createCanvas(width: number, height: number): HTMLCanvasElement {
  const canvas = document.createElement("canvas");
  canvas.width = Math.max(1, Math.round(width));
  canvas.height = Math.max(1, Math.round(height));
  return canvas;
}

function require2dContext(canvas: HTMLCanvasElement): CanvasRenderingContext2D {
  const context = canvas.getContext("2d");
  if (!context) throw new Error("Unable to create mask canvas");
  return context;
}

function loadImage(source: string): Promise<HTMLImageElement> {
  return new Promise((resolve, reject) => {
    const image = new Image();
    image.onload = () => resolve(image);
    image.onerror = () => reject(new Error("Failed to load image"));
    image.src = source;
  });
}

function roundUp(value: number, multiple: number): number {
  const safeMultiple = Math.max(1, Math.round(multiple));
  return Math.ceil(value / safeMultiple) * safeMultiple;
}

function clampNumber(value: number, min: number, max: number): number {
  if (!Number.isFinite(value)) return min;
  return Math.max(min, Math.min(max, value));
}

function clampInt(value: number, min: number, max: number): number {
  return Math.round(Math.max(min, Math.min(max, value)));
}