image-studio/operations/masks/magic-wand.ts

import { commandIds } from "@commands/ids";
import type { ImageDocument } from "@core/document";
import type { Vec2D } from "@core/geometry";
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 { EditorState } from "@editor/state";
import { createWandMask } from "@platform/browser/magicWandRaster";

export async function applyMagicWandAt(store: AppStore, point: Vec2D, modeOverride?: EditorState["tools"]["magicWand"]["mode"]) {
  const state = store.getState();
  if (state.editor.tools.activeTool !== "magicWand") return false;
  const target = resolveMaskSelectionTarget(state.document, state.editor);
  if (!target) return true;
  const assetPoint = documentPointToAssetPoint(point, target.layer, target.asset.intrinsicSize);
  const x = Math.floor(assetPoint.x);
  const y = Math.floor(assetPoint.y);
  if (x < 0 || y < 0 || x >= target.asset.intrinsicSize.w || y >= target.asset.intrinsicSize.h) return true;
  const source = await createWandMask(target.asset.source, target.maskAsset?.source, Math.round(target.asset.intrinsicSize.w), Math.round(target.asset.intrinsicSize.h), x, y, { ...state.editor.tools.magicWand, mode: modeOverride ?? state.editor.tools.magicWand.mode, target: target.inpaintRegion ? "inpaint" : "visibility" });
  if (target.inpaintRegion && target.maskAsset) {
    store.dispatch(commandIds.documentApplyInpaintRegionMaskOperation, { regionId: target.inpaintRegion.id, source, mimeType: "image/png", operation: { type: "magicWand" } });
    return true;
  }
  if (target.maskAsset && target.maskLayer && (target.maskLayer.type === "image" || target.maskLayer.type === "raster")) {
    store.dispatch(commandIds.documentApplyLayerMaskOperation, { maskLayerId: target.maskLayer.id, source, mimeType: "image/png", operation: { type: "magicWand" } });
    return true;
  }
  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));
  store.dispatch(commandIds.documentAddLayerMask, {
    layerId: target.layer.id,
    asset: { id: assetId, name: `${target.layer.name} Wand Mask`, mimeType: "image/png", source, intrinsicSize: { w: width, h: height } },
    maskLayer: { id: maskLayerId, type: "raster", name: `${target.layer.name} Wand 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 } },
  });
  store.dispatch(commandIds.toolExitMaskEdit, undefined);
  store.dispatch(commandIds.toolSetActive, { tool: "magicWand" });
  return true;
}

export function resolveMaskSelectionTarget(document: ImageDocument, editor: EditorState) {
  const layerId = editor.selection.layerIds[0];
  if (!layerId || editor.selection.layerIds.length !== 1) return undefined;
  const layer = findLayer(document.artboards.flatMap((artboard) => artboard.layers), layerId);
  if (!layer || (layer.type !== "image" && layer.type !== "raster")) 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 inpaintRegion = editor.maskEdit?.kind === "inpaintRegion" ? document.inpaintRegions.find((candidate) => candidate.id === editor.maskEdit?.inpaintRegionId && candidate.targetLayerId === layer.id) : undefined;
  const maskLayer = layerMask?.enabled ? findLayer(document.artboards.flatMap((artboard) => artboard.layers), layerMask.maskLayerId) : undefined;
  const maskAsset = inpaintRegion ? document.assets.find((candidate) => candidate.id === inpaintRegion.maskAssetId) : maskLayer && (maskLayer.type === "image" || maskLayer.type === "raster") ? document.assets.find((candidate) => candidate.id === maskLayer.assetId) : undefined;
  return asset && bounds ? { layer, asset, bounds, maskLayer, maskAsset, inpaintRegion } : undefined;
}

export function documentPointToAssetPoint(point: Vec2D, layer: Extract<Layer, { type: "image" | "raster" }>, intrinsicSize: { w: number; h: number }): Vec2D {
  const source = layer.sourceRect ?? { x: 0, y: 0, ...intrinsicSize };
  const destination = { x: layer.transform.position.x + source.x * layer.transform.scale.x, y: layer.transform.position.y + source.y * layer.transform.scale.y, w: source.w * layer.transform.scale.x, h: source.h * layer.transform.scale.y };
  const center = { x: destination.x + destination.w / 2, y: destination.y + destination.h / 2 };
  const dx = point.x - center.x;
  const dy = point.y - center.y;
  const cos = Math.cos(-layer.transform.rotation);
  const sin = Math.sin(-layer.transform.rotation);
  const x = center.x + dx * cos - dy * sin;
  const y = center.y + dx * sin + dy * cos;
  return { x: source.x + (x - destination.x) / Math.max(0.0001, layer.transform.scale.x), y: source.y + (y - destination.y) / Math.max(0.0001, layer.transform.scale.y) };
}

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;
}