261 lines
11 KiB
TypeScript
261 lines
11 KiB
TypeScript
import type { Asset } from "@core/asset";
|
|
import type { ImageDocument } from "@core/document";
|
|
import type { Rect } from "@core/geometry";
|
|
import type { Layer } from "@core/layer";
|
|
import { getLayerMask } from "@core/layer-mask-utils";
|
|
import type { SelectionState } from "@editor/state";
|
|
import type { GenerateSettings } from "@editor/tools";
|
|
import { createDocumentReadIndex, resolveIndexedLayerBounds } from "@editor/document-indexes";
|
|
import { createNormalizedMaskSource, cropCanvas, cropMaskValuesToDataUrl, expandRectWithinBounds, loadImageCanvas } from "@platform/browser/maskRaster";
|
|
|
|
export type InpaintBundle = {
|
|
inputImage: string;
|
|
maskImage: string;
|
|
width: number;
|
|
height: number;
|
|
targetLayerId: string;
|
|
maskLayerId: string;
|
|
sourceAssetId: string;
|
|
maskAssetId: string;
|
|
crop: {
|
|
assetBounds: Rect;
|
|
documentBounds: Rect;
|
|
padding: number;
|
|
maskedAreaOnly: boolean;
|
|
};
|
|
mask: {
|
|
polarity: GenerateSettings["inpaint"]["maskPolarity"];
|
|
activeBounds: Rect;
|
|
};
|
|
placement: {
|
|
artboardId: string;
|
|
layerName: string;
|
|
transform: Extract<Layer, { type: "image" | "raster" }>["transform"];
|
|
};
|
|
backend: {
|
|
growMaskBy: number;
|
|
maskedContent: GenerateSettings["inpaint"]["maskedContent"];
|
|
maskBlur: number;
|
|
maskFeather: number;
|
|
maskExpand: number;
|
|
cropPadding: number;
|
|
};
|
|
};
|
|
|
|
type InpaintTarget = {
|
|
artboardId: string;
|
|
layer: Extract<Layer, { type: "image" | "raster" }>;
|
|
asset: Asset;
|
|
bounds: Rect;
|
|
maskLayer: Extract<Layer, { type: "image" | "raster" }>;
|
|
maskAsset: Asset;
|
|
maskBounds: Rect;
|
|
};
|
|
|
|
const modelMultiple = 8;
|
|
const minModelSize = 64;
|
|
const maxModelSize = 4096;
|
|
|
|
export async function buildInpaintBundle(document: ImageDocument, selection: SelectionState, settings: GenerateSettings): Promise<InpaintBundle> {
|
|
const target = resolveInpaintTarget(document, selection);
|
|
validateInpaintTarget(target);
|
|
|
|
const width = Math.max(1, Math.round(target.asset.intrinsicSize.w));
|
|
const height = Math.max(1, Math.round(target.asset.intrinsicSize.h));
|
|
const normalizedMask = await createNormalizedMaskSource(target.maskAsset.source, width, height, {
|
|
polarity: settings.inpaint.maskPolarity,
|
|
expand: settings.inpaint.maskExpand,
|
|
feather: settings.inpaint.maskFeather,
|
|
blur: settings.inpaint.maskBlur,
|
|
despeckle: settings.inpaint.maskDespeckle,
|
|
});
|
|
|
|
if (!normalizedMask.bounds) throw new Error("The selected layer mask has no inpaint pixels.");
|
|
|
|
const crop = settings.inpaint.maskedAreaOnly
|
|
? expandRectWithinBounds(normalizedMask.bounds, settings.inpaint.cropPadding, { w: width, h: height }, modelMultiple, minModelSize)
|
|
: { x: 0, y: 0, w: width, h: height };
|
|
const outputWidth = toModelSize(crop.w, "width");
|
|
const outputHeight = toModelSize(crop.h, "height");
|
|
|
|
const inputCanvas = prepareMaskedContentInputCanvas(await loadImageCanvas(target.asset.source, width, height), normalizedMask.values, settings.inpaint.maskedContent);
|
|
const inputImage = cropCanvas(inputCanvas, crop, outputWidth, outputHeight);
|
|
const maskImage = cropMaskValuesToDataUrl(normalizedMask.values, width, height, crop, outputWidth, outputHeight);
|
|
const scaleX = target.bounds.w / width;
|
|
const scaleY = target.bounds.h / height;
|
|
const documentBounds = {
|
|
x: target.bounds.x + crop.x * scaleX,
|
|
y: target.bounds.y + crop.y * scaleY,
|
|
w: outputWidth * scaleX,
|
|
h: outputHeight * scaleY,
|
|
};
|
|
|
|
return {
|
|
inputImage,
|
|
maskImage,
|
|
width: outputWidth,
|
|
height: outputHeight,
|
|
targetLayerId: target.layer.id,
|
|
maskLayerId: target.maskLayer.id,
|
|
sourceAssetId: target.asset.id,
|
|
maskAssetId: target.maskAsset.id,
|
|
crop: {
|
|
assetBounds: crop,
|
|
documentBounds,
|
|
padding: settings.inpaint.cropPadding,
|
|
maskedAreaOnly: settings.inpaint.maskedAreaOnly,
|
|
},
|
|
mask: {
|
|
polarity: settings.inpaint.maskPolarity,
|
|
activeBounds: normalizedMask.bounds,
|
|
},
|
|
placement: {
|
|
artboardId: target.artboardId,
|
|
layerName: `${target.layer.name} inpaint`,
|
|
transform: {
|
|
position: { x: documentBounds.x, y: documentBounds.y },
|
|
scale: { x: documentBounds.w / outputWidth, y: documentBounds.h / outputHeight },
|
|
rotation: target.layer.transform.rotation,
|
|
},
|
|
},
|
|
backend: {
|
|
growMaskBy: settings.inpaint.growMaskBy,
|
|
maskedContent: settings.inpaint.maskedContent,
|
|
maskBlur: settings.inpaint.maskBlur,
|
|
maskFeather: settings.inpaint.maskFeather,
|
|
maskExpand: settings.inpaint.maskExpand,
|
|
cropPadding: settings.inpaint.cropPadding,
|
|
},
|
|
};
|
|
}
|
|
|
|
function resolveInpaintTarget(document: ImageDocument, selection: SelectionState): InpaintTarget {
|
|
if (selection.layerIds.length !== 1 || !selection.layerIds[0]) throw new Error("Select one image or raster layer to inpaint.");
|
|
|
|
const documentIndex = createDocumentReadIndex(document);
|
|
const layerInfo = documentIndex.layerInfoById.get(selection.layerIds[0]);
|
|
if (!layerInfo || layerInfo.layer.type === "group" || layerInfo.layer.type === "adjustment") throw new Error("Select one image or raster layer to inpaint.");
|
|
|
|
const asset = documentIndex.assetById.get(layerInfo.layer.assetId);
|
|
if (!asset) throw new Error("The selected layer is missing its source image.");
|
|
const layerMask = getLayerMask(layerInfo.layer);
|
|
if (!layerMask?.enabled) throw new Error("Add a layer mask before running inpaint.");
|
|
|
|
const maskLayer = documentIndex.layerById.get(layerMask.maskLayerId);
|
|
if (!maskLayer || maskLayer.type === "group" || maskLayer.type === "adjustment") throw new Error("The selected layer mask is missing.");
|
|
|
|
const maskAsset = documentIndex.assetById.get(maskLayer.assetId);
|
|
if (!maskAsset) throw new Error("The selected layer mask is missing its image data.");
|
|
|
|
const bounds = resolveIndexedLayerBounds(documentIndex, layerInfo.layer);
|
|
const maskBounds = resolveIndexedLayerBounds(documentIndex, maskLayer);
|
|
if (!bounds || !maskBounds) throw new Error("Unable to resolve the selected layer and mask bounds.");
|
|
|
|
return { artboardId: layerInfo.artboardId, layer: layerInfo.layer, asset, bounds, maskLayer, maskAsset, maskBounds };
|
|
}
|
|
|
|
function validateInpaintTarget(target: InpaintTarget) {
|
|
if (target.asset.intrinsicSize.w <= 0 || target.asset.intrinsicSize.h <= 0) throw new Error("The selected image has an invalid size.");
|
|
if (target.maskAsset.intrinsicSize.w <= 0 || target.maskAsset.intrinsicSize.h <= 0) throw new Error("The selected mask has an invalid size.");
|
|
if (Math.round(target.asset.intrinsicSize.w) !== Math.round(target.maskAsset.intrinsicSize.w) || Math.round(target.asset.intrinsicSize.h) !== Math.round(target.maskAsset.intrinsicSize.h)) {
|
|
throw new Error("The selected layer and mask image sizes do not match.");
|
|
}
|
|
if (!rectsAligned(target.bounds, target.maskBounds) || Math.abs(target.layer.transform.rotation - target.maskLayer.transform.rotation) > 0.001) {
|
|
throw new Error("The selected layer and mask are not aligned.");
|
|
}
|
|
}
|
|
|
|
function rectsAligned(a: Rect, b: Rect): boolean {
|
|
return Math.abs(a.x - b.x) <= 0.5 && Math.abs(a.y - b.y) <= 0.5 && Math.abs(a.w - b.w) <= 0.5 && Math.abs(a.h - b.h) <= 0.5;
|
|
}
|
|
|
|
function toModelSize(value: number, axis: "width" | "height"): number {
|
|
const rounded = Math.max(minModelSize, Math.ceil(value / modelMultiple) * modelMultiple);
|
|
if (rounded > maxModelSize) throw new Error(`The inpaint ${axis} is too large for the model. Use masked-area inpaint or a smaller source.`);
|
|
return rounded;
|
|
}
|
|
|
|
function prepareMaskedContentInputCanvas(sourceCanvas: HTMLCanvasElement, maskValues: Uint8ClampedArray, mode: GenerateSettings["inpaint"]["maskedContent"]): HTMLCanvasElement {
|
|
if (mode === "neutral" || mode === "originalColor") return sourceCanvas;
|
|
const context = sourceCanvas.getContext("2d");
|
|
if (!context) return sourceCanvas;
|
|
const imageData = context.getImageData(0, 0, sourceCanvas.width, sourceCanvas.height);
|
|
imageData.data.set(applyMaskedContentModeToRgba(imageData.data, sourceCanvas.width, sourceCanvas.height, maskValues, mode));
|
|
context.putImageData(imageData, 0, 0);
|
|
return sourceCanvas;
|
|
}
|
|
|
|
export function applyMaskedContentModeToRgba(data: Uint8ClampedArray, width: number, height: number, maskValues: Uint8ClampedArray, mode: GenerateSettings["inpaint"]["maskedContent"]): Uint8ClampedArray {
|
|
const next = new Uint8ClampedArray(data);
|
|
if (mode === "neutral" || mode === "originalColor") return next;
|
|
|
|
const edgeValues = mode === "edges" ? createEdgeMapValues(data, width, height) : undefined;
|
|
|
|
for (let pixel = 0; pixel < width * height; pixel += 1) {
|
|
const amount = (maskValues[pixel] ?? 0) / 255;
|
|
if (amount <= 0) continue;
|
|
|
|
const index = pixel * 4;
|
|
const red = data[index] ?? 0;
|
|
const green = data[index + 1] ?? 0;
|
|
const blue = data[index + 2] ?? 0;
|
|
const target = edgeValues ? edgeValues[pixel] ?? 128 : luminance(red, green, blue);
|
|
next[index] = blendChannel(red, target, amount);
|
|
next[index + 1] = blendChannel(green, target, amount);
|
|
next[index + 2] = blendChannel(blue, target, amount);
|
|
}
|
|
|
|
return next;
|
|
}
|
|
|
|
function createEdgeMapValues(data: Uint8ClampedArray, width: number, height: number): Uint8ClampedArray {
|
|
const gray = new Float32Array(width * height);
|
|
const edges = new Uint8ClampedArray(width * height);
|
|
|
|
for (let pixel = 0; pixel < width * height; pixel += 1) {
|
|
const index = pixel * 4;
|
|
gray[pixel] = luminance(data[index] ?? 0, data[index + 1] ?? 0, data[index + 2] ?? 0);
|
|
}
|
|
|
|
for (let y = 0; y < height; y += 1) {
|
|
for (let x = 0; x < width; x += 1) {
|
|
const gx =
|
|
-sampleGray(gray, width, height, x - 1, y - 1) +
|
|
sampleGray(gray, width, height, x + 1, y - 1) -
|
|
2 * sampleGray(gray, width, height, x - 1, y) +
|
|
2 * sampleGray(gray, width, height, x + 1, y) -
|
|
sampleGray(gray, width, height, x - 1, y + 1) +
|
|
sampleGray(gray, width, height, x + 1, y + 1);
|
|
const gy =
|
|
-sampleGray(gray, width, height, x - 1, y - 1) -
|
|
2 * sampleGray(gray, width, height, x, y - 1) -
|
|
sampleGray(gray, width, height, x + 1, y - 1) +
|
|
sampleGray(gray, width, height, x - 1, y + 1) +
|
|
2 * sampleGray(gray, width, height, x, y + 1) +
|
|
sampleGray(gray, width, height, x + 1, y + 1);
|
|
const magnitude = Math.hypot(gx, gy);
|
|
edges[y * width + x] = Math.round(clampNumber(128 + Math.max(0, magnitude - 24) * 0.75, 128, 255));
|
|
}
|
|
}
|
|
|
|
return edges;
|
|
}
|
|
|
|
function sampleGray(values: Float32Array, width: number, height: number, x: number, y: number): number {
|
|
const clampedX = Math.max(0, Math.min(width - 1, x));
|
|
const clampedY = Math.max(0, Math.min(height - 1, y));
|
|
return values[clampedY * width + clampedX] ?? 0;
|
|
}
|
|
|
|
function luminance(red: number, green: number, blue: number): number {
|
|
return Math.round(0.2126 * red + 0.7152 * green + 0.0722 * blue);
|
|
}
|
|
|
|
function blendChannel(previous: number, next: number, amount: number): number {
|
|
return Math.round(previous * (1 - amount) + next * amount);
|
|
}
|
|
|
|
function clampNumber(value: number, min: number, max: number): number {
|
|
return Math.max(min, Math.min(max, value));
|
|
}
|