image-studio/view/bottom-controls/ChromaKeyControls.tsx

import { useEffect } from "react";
import { DropHalf } from "@phosphor-icons/react";
import { commandIds } from "@commands/ids";
import type { ImageDocument } from "@core/document";
import type { Layer } from "@core/layer";
import { getLayerMask } from "@core/layer-mask-utils";
import { resolveTransformTargetBounds } from "@editor/transform-targets";
import type { AppStore } from "@editor/store";
import type { ChromaKeySettings } from "@editor/tools";
import type { SelectionState } from "@editor/state";
import { blurMaskValues, despeckleMaskValues, dilateMaskValues, erodeMaskValues } from "../mask/maskRaster";
import { BottomControlColorPicker } from "./ColorPicker";
import { BottomControlDivider } from "./Divider";
import { BottomControlSlider } from "./Slider";
import { bottomControlFieldClass, bottomControlIconSlotClass, bottomControlLabelClass, bottomControlMenuClass } from "./styles";

export type ChromaKeyControlsProps = {
  document: ImageDocument;
  selection: SelectionState;
  settings: ChromaKeySettings;
  dispatch: AppStore["dispatch"];
};

export function ChromaKeyControls({ document, selection, settings, dispatch }: ChromaKeyControlsProps) {
  const target = resolveChromaKeyTarget(document, selection);

  useEffect(() => {
    return () => {
      dispatch(commandIds.toolSetBrushStrokePreview, undefined);
    };
  }, [dispatch]);

  useEffect(() => {
    let cancelled = false;

    if (!target) {
      dispatch(commandIds.toolSetBrushStrokePreview, undefined);
      return;
    }

    void chromaKeySource(target.asset.source, target.asset.intrinsicSize.w, target.asset.intrinsicSize.h, settings).then((source) => {
      if (cancelled) return;
      dispatch(commandIds.toolSetBrushStrokePreview, { layerId: target.layer.id, assetId: target.asset.id, source });
    });

    return () => {
      cancelled = true;
    };
  }, [dispatch, settings.choke, settings.color, settings.despeckle, settings.feather, settings.softness, settings.spill, settings.tolerance, target?.asset.id, target?.asset.source, target?.asset.intrinsicSize.w, target?.asset.intrinsicSize.h, target?.layer.id]);

  return (
    <div className={bottomControlMenuClass()}>
      <span className={bottomControlIconSlotClass()} title="Chroma key">
        <DropHalf size={24} weight="regular" />
      </span>
      <BottomControlDivider />
      <label className={bottomControlFieldClass()}>
        <span className={bottomControlLabelClass()}>Key</span>
        <BottomControlColorPicker
          value={settings.color}
          aria-label="Chroma key color"
          onValueChange={(color) => dispatch(commandIds.toolSetChromaKeySettings, { color })}
        />
      </label>
      <BottomControlDivider />
      <label className={bottomControlFieldClass()}>
        <span className={bottomControlLabelClass()}>Tol</span>
        <BottomControlSlider
          min={0}
          max={255}
          value={settings.tolerance}
          className="w-36"
          aria-label="Chroma key tolerance"
          onValueChange={(tolerance) => dispatch(commandIds.toolSetChromaKeySettings, { tolerance })}
        />
        <span className="w-10 text-right text-base text-white">{Math.round(settings.tolerance)}</span>
      </label>
      <BottomControlDivider />
      <label className={bottomControlFieldClass()}>
        <span className={bottomControlLabelClass()}>Soft</span>
        <BottomControlSlider
          min={0}
          max={255}
          value={settings.softness}
          className="w-36"
          aria-label="Chroma key edge softness"
          onValueChange={(softness) => dispatch(commandIds.toolSetChromaKeySettings, { softness })}
        />
        <span className="w-10 text-right text-base text-white">{Math.round(settings.softness)}</span>
      </label>
      <BottomControlDivider />
      <label className={bottomControlFieldClass()}>
        <span className={bottomControlLabelClass()}>Feather</span>
        <BottomControlSlider
          min={0}
          max={20}
          value={settings.feather}
          className="w-32"
          aria-label="Chroma key mask feather"
          onValueChange={(feather) => dispatch(commandIds.toolSetChromaKeySettings, { feather })}
        />
        <span className="w-8 text-right text-base text-white">{Math.round(settings.feather)}</span>
      </label>
      <BottomControlDivider />
      <label className={bottomControlFieldClass()}>
        <span className={bottomControlLabelClass()}>Choke</span>
        <BottomControlSlider
          min={-20}
          max={20}
          value={settings.choke}
          className="w-32"
          aria-label="Chroma key mask choke or expand"
          onValueChange={(choke) => dispatch(commandIds.toolSetChromaKeySettings, { choke })}
        />
        <span className="w-8 text-right text-base text-white">{Math.round(settings.choke)}</span>
      </label>
      <BottomControlDivider />
      <label className={bottomControlFieldClass()}>
        <span className={bottomControlLabelClass()}>Clean</span>
        <BottomControlSlider
          min={0}
          max={20}
          value={settings.despeckle}
          className="w-32"
          aria-label="Chroma key mask despeckle"
          onValueChange={(despeckle) => dispatch(commandIds.toolSetChromaKeySettings, { despeckle })}
        />
        <span className="w-8 text-right text-base text-white">{Math.round(settings.despeckle)}</span>
      </label>
      <BottomControlDivider />
      <label className={bottomControlFieldClass()}>
        <span className={bottomControlLabelClass()}>Spill</span>
        <BottomControlSlider
          min={0}
          max={100}
          value={settings.spill}
          className="w-36"
          aria-label="Chroma key spill suppression"
          onValueChange={(spill) => dispatch(commandIds.toolSetChromaKeySettings, { spill })}
        />
        <span className="w-10 text-right text-base text-white">{Math.round(settings.spill)}</span>
      </label>
      <BottomControlDivider />
      <button
        type="button"
        className="rounded-full bg-white px-5 py-2 text-base font-medium text-black transition hover:bg-white/90 disabled:pointer-events-none disabled:opacity-35"
        disabled={!target}
        title={target ? "Create or update a layer mask from the chroma key" : "Select one image or raster layer to create a chroma key mask"}
        onClick={() => target && void applyChromaKeyMask(target, settings, dispatch)}
      >
        {target?.maskAsset ? "Update Mask" : "Create Mask"}
      </button>
    </div>
  );
}

function resolveChromaKeyTarget(document: ImageDocument, selection: SelectionState) {
  const layerId = selection.layerIds[0];
  if (selection.layerIds.length !== 1 || !layerId) return undefined;
  const layer = findLayer(document.artboards.find((artboard) => artboard.id === selection.artboardId)?.layers ?? [], layerId);
  if (!layer || layer.type === "group") return undefined;
  const asset = document.assets.find((candidate) => candidate.id === layer.assetId);
  const bounds = resolveTransformTargetBounds(document, { type: "layer", id: layer.id });
  const layerMask = getLayerMask(layer);
  const maskLayer = layerMask?.enabled ? findLayer(document.artboards.flatMap((artboard) => artboard.layers), layerMask.maskLayerId) : undefined;
  const maskAsset = maskLayer && maskLayer.type !== "group" ? document.assets.find((candidate) => candidate.id === maskLayer.assetId) : undefined;
  return asset && bounds ? { layer, asset, bounds, maskLayer, maskAsset } : undefined;
}

function findLayer(layers: readonly Layer[], layerId: string): Layer | undefined {
  for (const layer of layers) {
    if (layer.id === layerId) return layer;
    if (layer.type === "group") {
      const found = findLayer(layer.children, layerId);
      if (found) return found;
    }
  }
  return undefined;
}

async function applyChromaKeyMask(target: NonNullable<ReturnType<typeof resolveChromaKeyTarget>>, settings: ChromaKeySettings, dispatch: AppStore["dispatch"]) {
  const source = await chromaKeyMaskSource(target.asset.source, target.asset.intrinsicSize.w, target.asset.intrinsicSize.h, settings);
  dispatch(commandIds.toolSetBrushStrokePreview, undefined);

  if (target.maskAsset && target.maskLayer && target.maskLayer.type !== "group") {
    dispatch(commandIds.documentApplyLayerMaskOperation, { maskLayerId: target.maskLayer.id, source, mimeType: "image/png", operation: { type: "chromaKey" } });
    return;
  }

  const assetId = crypto.randomUUID();
  const maskLayerId = crypto.randomUUID();
  const width = Math.max(1, Math.round(target.asset.intrinsicSize.w));
  const height = Math.max(1, Math.round(target.asset.intrinsicSize.h));
  dispatch(commandIds.documentAddLayerMask, {
    layerId: target.layer.id,
    asset: {
      id: assetId,
      name: `${target.layer.name} Chroma Mask`,
      mimeType: "image/png",
      source,
      intrinsicSize: { w: width, h: height },
    },
    maskLayer: {
      id: maskLayerId,
      type: "raster",
      name: `${target.layer.name} Chroma Mask`,
      visible: true,
      locked: false,
      opacity: 1,
      assetId,
      transform: {
        position: { x: target.bounds.x, y: target.bounds.y },
        scale: { x: target.bounds.w / width, y: target.bounds.h / height },
        rotation: target.layer.transform.rotation,
      },
    },
  });
  dispatch(commandIds.toolExitMaskEdit, undefined);
  dispatch(commandIds.toolSetActive, { tool: "chromaKey" });
}

async function chromaKeySource(source: string, width: number, height: number, settings: ChromaKeySettings) {
  const canvas = document.createElement("canvas");
  canvas.width = Math.max(1, Math.round(width));
  canvas.height = Math.max(1, Math.round(height));
  const context = canvas.getContext("2d");
  if (!context) return source;
  const image = await loadImage(source);
  context.drawImage(image, 0, 0, canvas.width, canvas.height);
  const data = context.getImageData(0, 0, canvas.width, canvas.height);
  const alpha = chromaKeyAlpha(data, canvas.width, canvas.height, settings);
  for (let pixel = 0; pixel < alpha.length; pixel++) data.data[pixel * 4 + 3] = alpha[pixel] ?? 255;
  context.putImageData(data, 0, 0);
  return canvas.toDataURL("image/png");
}

async function chromaKeyMaskSource(source: string, width: number, height: number, settings: ChromaKeySettings) {
  const canvas = document.createElement("canvas");
  canvas.width = Math.max(1, Math.round(width));
  canvas.height = Math.max(1, Math.round(height));
  const context = canvas.getContext("2d");
  if (!context) return source;
  const image = await loadImage(source);
  context.drawImage(image, 0, 0, canvas.width, canvas.height);
  const data = context.getImageData(0, 0, canvas.width, canvas.height);
  const alpha = chromaKeyAlpha(data, canvas.width, canvas.height, settings);

  for (let pixel = 0; pixel < alpha.length; pixel++) {
    const index = pixel * 4;
    data.data[index] = 255;
    data.data[index + 1] = 255;
    data.data[index + 2] = 255;
    data.data[index + 3] = alpha[pixel] ?? 255;
  }

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

function hexToRgb(color: string) {
  const hex = color.replace("#", "");
  return { r: Number.parseInt(hex.slice(0, 2), 16), g: Number.parseInt(hex.slice(2, 4), 16), b: Number.parseInt(hex.slice(4, 6), 16) };
}

function chromaKeyAlpha(data: ImageData, width: number, height: number, settings: ChromaKeySettings) {
  const key = hexToRgb(settings.color);
  const alpha = new Uint8ClampedArray(width * height);
  for (let pixel = 0; pixel < alpha.length; pixel++) {
    const index = pixel * 4;
    const red = data.data[index] ?? 0;
    const green = data.data[index + 1] ?? 0;
    const blue = data.data[index + 2] ?? 0;
    const sourceAlpha = data.data[index + 3] ?? 255;
    alpha[pixel] = Math.round(sourceAlpha * chromaKeyKeepFactor(red, green, blue, key, settings));
  }
  return postProcessAlpha(alpha, width, height, settings);
}

function chromaKeyKeepFactor(red: number, green: number, blue: number, key: { r: number; g: number; b: number }, settings: ChromaKeySettings) {
  const tolerance = Math.max(0, Math.min(255, settings.tolerance));
  const softness = Math.max(0, Math.min(255, settings.softness));
  const spill = Math.max(0, Math.min(100, settings.spill)) / 100;
  const distance = Math.hypot(red - key.r, green - key.g, blue - key.b);
  const edgeKeep = distance <= tolerance ? 0 : softness > 0 && distance < tolerance + softness ? (distance - tolerance) / softness : 1;
  if (spill <= 0) return edgeKeep;

  const dominant = key.g >= key.r && key.g >= key.b ? green : key.r >= key.b ? red : blue;
  const neutral = key.g >= key.r && key.g >= key.b ? Math.max(red, blue) : key.r >= key.b ? Math.max(green, blue) : Math.max(red, green);
  const spillAmount = Math.max(0, dominant - neutral) / 255;
  return Math.max(0, Math.min(edgeKeep, 1 - spillAmount * spill));
}

function postProcessAlpha(alpha: Uint8ClampedArray, width: number, height: number, settings: ChromaKeySettings) {
  let next = alpha;
  const despeckle = Math.round(Math.max(0, Math.min(20, settings.despeckle)));
  const choke = Math.round(Math.max(-20, Math.min(20, settings.choke)));
  const feather = Math.round(Math.max(0, Math.min(20, settings.feather)));

  if (despeckle > 0) next = despeckleMaskValues(next, width, height, despeckle);
  if (choke > 0) next = erodeMaskValues(next, width, height, choke);
  if (choke < 0) next = dilateMaskValues(next, width, height, -choke);
  if (feather > 0) next = blurMaskValues(next, width, height, feather);
  return next;
}

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