Files
Genarrative/packages/shared/src/assets/chromaKey.ts
2026-04-19 20:33:18 +08:00

479 lines
13 KiB
TypeScript

export type MutableRgbaBuffer = Uint8Array | Uint8ClampedArray;
const SOFT_EDGE_ALPHA_THRESHOLD = 224;
const FOREGROUND_NEIGHBOR_ALPHA_THRESHOLD = 96;
function clamp01(value: number) {
return Math.max(0, Math.min(1, value));
}
function lerp(from: number, to: number, t: number) {
return from + (to - from) * clamp01(t);
}
function computeGreenBackgroundScore(
red: number,
green: number,
blue: number,
alpha: number,
) {
if (alpha === 0) {
return 1;
}
const greenLead = green - Math.max(red, blue);
if (green < 52 || greenLead <= 8) {
return 0;
}
const greenRatio = green / Math.max(1, red + blue);
if (greenRatio <= 0.52) {
return 0;
}
return clamp01(
((green - 52) / 168) * 0.22 +
((greenLead - 8) / 96) * 0.53 +
((greenRatio - 0.52) / 0.82) * 0.25,
);
}
function computeWhiteBackgroundScore(
red: number,
green: number,
blue: number,
alpha: number,
) {
if (alpha === 0) {
return 1;
}
const maxChannel = Math.max(red, green, blue);
const minChannel = Math.min(red, green, blue);
const average = (red + green + blue) / 3;
if (average < 188 || minChannel < 168) {
return 0;
}
const spread = maxChannel - minChannel;
const neutrality = 1 - clamp01((spread - 6) / 34);
const brightness = clamp01((average - 188) / 55);
const floor = clamp01((minChannel - 168) / 60);
return clamp01(neutrality * (brightness * 0.85 + floor * 0.15));
}
function collectForegroundNeighborColor(
pixels: MutableRgbaBuffer,
width: number,
height: number,
x: number,
y: number,
backgroundMask: Uint8Array,
backgroundHints: Float32Array,
) {
let totalWeight = 0;
let totalRed = 0;
let totalGreen = 0;
let totalBlue = 0;
for (let offsetY = -2; offsetY <= 2; offsetY += 1) {
for (let offsetX = -2; offsetX <= 2; offsetX += 1) {
if (offsetX === 0 && offsetY === 0) {
continue;
}
const nextX = x + offsetX;
const nextY = y + offsetY;
if (nextX < 0 || nextX >= width || nextY < 0 || nextY >= height) {
continue;
}
const nextPixelIndex = nextY * width + nextX;
if (backgroundMask[nextPixelIndex]) {
continue;
}
if ((backgroundHints[nextPixelIndex] ?? 0) >= 0.18) {
continue;
}
const nextOffset = nextPixelIndex * 4;
const nextAlpha = pixels[nextOffset + 3] ?? 0;
if (nextAlpha < FOREGROUND_NEIGHBOR_ALPHA_THRESHOLD) {
continue;
}
const distance = Math.abs(offsetX) + Math.abs(offsetY);
const weight =
(nextAlpha / 255) *
(distance <= 1 ? 1.8 : distance === 2 ? 1.2 : 0.7);
totalWeight += weight;
totalRed += (pixels[nextOffset] ?? 0) * weight;
totalGreen += (pixels[nextOffset + 1] ?? 0) * weight;
totalBlue += (pixels[nextOffset + 2] ?? 0) * weight;
}
}
if (totalWeight <= 0) {
return null;
}
return {
red: Math.round(totalRed / totalWeight),
green: Math.round(totalGreen / totalWeight),
blue: Math.round(totalBlue / totalWeight),
};
}
export function removeBackgroundFromRgba(
pixels: MutableRgbaBuffer,
width: number,
height: number,
) {
const pixelCount = width * height;
if (pixelCount <= 0) {
return false;
}
const backgroundMask = new Uint8Array(pixelCount);
const greenScores = new Float32Array(pixelCount);
const whiteScores = new Float32Array(pixelCount);
const backgroundHints = new Float32Array(pixelCount);
const queue: number[] = [];
let queueIndex = 0;
let changed = false;
for (let pixelIndex = 0; pixelIndex < pixelCount; pixelIndex += 1) {
const offset = pixelIndex * 4;
const red = pixels[offset] ?? 0;
const green = pixels[offset + 1] ?? 0;
const blue = pixels[offset + 2] ?? 0;
const alpha = pixels[offset + 3] ?? 0;
const greenScore = computeGreenBackgroundScore(red, green, blue, alpha);
const whiteScore = computeWhiteBackgroundScore(red, green, blue, alpha);
const transparencyHint = clamp01((56 - alpha) / 56) * 0.75;
greenScores[pixelIndex] = greenScore;
whiteScores[pixelIndex] = whiteScore;
backgroundHints[pixelIndex] = Math.max(
greenScore,
whiteScore,
transparencyHint,
);
}
const trySeedBackground = (pixelIndex: number) => {
if (backgroundMask[pixelIndex]) {
return;
}
const offset = pixelIndex * 4;
const alpha = pixels[offset + 3] ?? 0;
const strongCandidate =
alpha < 40 ||
(greenScores[pixelIndex] ?? 0) > 0.12 ||
(whiteScores[pixelIndex] ?? 0) > 0.32;
if (!strongCandidate) {
return;
}
backgroundMask[pixelIndex] = 1;
queue.push(pixelIndex);
};
for (let x = 0; x < width; x += 1) {
trySeedBackground(x);
trySeedBackground((height - 1) * width + x);
}
for (let y = 1; y < height - 1; y += 1) {
trySeedBackground(y * width);
trySeedBackground(y * width + width - 1);
}
while (queueIndex < queue.length) {
const pixelIndex = queue[queueIndex]!;
queueIndex += 1;
const x = pixelIndex % width;
const y = Math.floor(pixelIndex / width);
const neighborIndexes = [
x > 0 ? pixelIndex - 1 : -1,
x + 1 < width ? pixelIndex + 1 : -1,
y > 0 ? pixelIndex - width : -1,
y + 1 < height ? pixelIndex + width : -1,
];
for (const nextPixelIndex of neighborIndexes) {
if (nextPixelIndex < 0 || backgroundMask[nextPixelIndex]) {
continue;
}
const nextOffset = nextPixelIndex * 4;
const nextAlpha = pixels[nextOffset + 3] ?? 0;
const nextGreenScore = greenScores[nextPixelIndex] ?? 0;
const nextWhiteScore = whiteScores[nextPixelIndex] ?? 0;
const nextHint = backgroundHints[nextPixelIndex] ?? 0;
const reachableSoftEdge =
nextHint > 0.08 &&
nextAlpha < SOFT_EDGE_ALPHA_THRESHOLD &&
(nextGreenScore > 0.04 || nextWhiteScore > 0.08 || nextAlpha < 180);
if (
nextAlpha < 40 ||
nextGreenScore > 0.12 ||
nextWhiteScore > 0.32 ||
reachableSoftEdge
) {
backgroundMask[nextPixelIndex] = 1;
queue.push(nextPixelIndex);
}
}
}
for (let iteration = 0; iteration < 2; iteration += 1) {
const expandedMask = new Uint8Array(backgroundMask);
for (let y = 0; y < height; y += 1) {
for (let x = 0; x < width; x += 1) {
const pixelIndex = y * width + x;
if (expandedMask[pixelIndex]) {
continue;
}
const alpha = pixels[pixelIndex * 4 + 3] ?? 0;
const hint = backgroundHints[pixelIndex] ?? 0;
if (alpha >= SOFT_EDGE_ALPHA_THRESHOLD || hint <= 0.06) {
continue;
}
let adjacentBackgroundCount = 0;
for (let offsetY = -1; offsetY <= 1; offsetY += 1) {
for (let offsetX = -1; offsetX <= 1; offsetX += 1) {
if (offsetX === 0 && offsetY === 0) {
continue;
}
const nextX = x + offsetX;
const nextY = y + offsetY;
if (nextX < 0 || nextX >= width || nextY < 0 || nextY >= height) {
continue;
}
if (backgroundMask[nextY * width + nextX]) {
adjacentBackgroundCount += 1;
}
}
}
if (
adjacentBackgroundCount >= 2 ||
(adjacentBackgroundCount >= 1 && hint > 0.18)
) {
expandedMask[pixelIndex] = 1;
}
}
}
backgroundMask.set(expandedMask);
}
for (let y = 0; y < height; y += 1) {
for (let x = 0; x < width; x += 1) {
const pixelIndex = y * width + x;
if (!backgroundMask[pixelIndex]) {
continue;
}
const offset = pixelIndex * 4;
const alpha = pixels[offset + 3] ?? 0;
if (alpha === 0) {
continue;
}
const matteScore = Math.max(
backgroundHints[pixelIndex] ?? 0,
greenScores[pixelIndex] ?? 0,
whiteScores[pixelIndex] ?? 0,
);
let foregroundSupport = 0;
for (let offsetY = -1; offsetY <= 1; offsetY += 1) {
for (let offsetX = -1; offsetX <= 1; offsetX += 1) {
if (offsetX === 0 && offsetY === 0) {
continue;
}
const nextX = x + offsetX;
const nextY = y + offsetY;
if (nextX < 0 || nextX >= width || nextY < 0 || nextY >= height) {
continue;
}
const nextPixelIndex = nextY * width + nextX;
if (backgroundMask[nextPixelIndex]) {
continue;
}
const nextAlpha = pixels[nextPixelIndex * 4 + 3] ?? 0;
if (nextAlpha >= FOREGROUND_NEIGHBOR_ALPHA_THRESHOLD) {
foregroundSupport += 1;
}
}
}
let nextAlpha = alpha;
if (matteScore > 0.9 || foregroundSupport === 0) {
nextAlpha = 0;
} else if (matteScore > 0.72 && foregroundSupport <= 1) {
nextAlpha = Math.min(alpha, Math.round(alpha * 0.08));
} else {
nextAlpha = Math.min(
alpha,
Math.round(alpha * Math.max(0.08, 1 - matteScore * 0.95)),
);
}
if (foregroundSupport >= 3 && matteScore < 0.55) {
nextAlpha = Math.max(nextAlpha, Math.round(alpha * 0.22));
}
if (nextAlpha < 10) {
nextAlpha = 0;
}
if (nextAlpha !== alpha) {
pixels[offset + 3] = nextAlpha;
changed = true;
}
}
}
for (let y = 0; y < height; y += 1) {
for (let x = 0; x < width; x += 1) {
const pixelIndex = y * width + x;
const offset = pixelIndex * 4;
const alpha = pixels[offset + 3] ?? 0;
if (alpha === 0) {
continue;
}
let touchesTransparentEdge = false;
for (let offsetY = -1; offsetY <= 1; offsetY += 1) {
for (let offsetX = -1; offsetX <= 1; offsetX += 1) {
if (offsetX === 0 && offsetY === 0) {
continue;
}
const nextX = x + offsetX;
const nextY = y + offsetY;
if (nextX < 0 || nextX >= width || nextY < 0 || nextY >= height) {
touchesTransparentEdge = true;
continue;
}
const nextPixelIndex = nextY * width + nextX;
if (
backgroundMask[nextPixelIndex] ||
(pixels[nextPixelIndex * 4 + 3] ?? 0) < 16
) {
touchesTransparentEdge = true;
}
}
}
if (!touchesTransparentEdge) {
continue;
}
const greenScore = greenScores[pixelIndex] ?? 0;
const whiteScore = whiteScores[pixelIndex] ?? 0;
const contamination = Math.max(
greenScore,
whiteScore,
backgroundMask[pixelIndex] ? 0.35 : 0,
alpha < 220 ? ((220 - alpha) / 220) * 0.25 : 0,
);
if (contamination < 0.06) {
continue;
}
let red = pixels[offset] ?? 0;
let green = pixels[offset + 1] ?? 0;
let blue = pixels[offset + 2] ?? 0;
const sample = collectForegroundNeighborColor(
pixels,
width,
height,
x,
y,
backgroundMask,
backgroundHints,
);
const blend = clamp01(
Math.max(contamination * 0.82, touchesTransparentEdge ? 0.22 : 0),
);
if (sample) {
red = Math.round(lerp(red, sample.red, blend));
green = Math.round(lerp(green, sample.green, blend));
blue = Math.round(lerp(blue, sample.blue, blend));
if (greenScore > 0.04) {
green = Math.min(green, sample.green + 18);
}
if (whiteScore > 0.1) {
red = Math.min(red, sample.red + 26);
green = Math.min(green, sample.green + 26);
blue = Math.min(blue, sample.blue + 26);
}
} else {
if (greenScore > 0.04) {
green = Math.max(
Math.max(red, blue),
Math.round(green - (green - Math.max(red, blue)) * 0.78),
);
}
if (whiteScore > 0.12) {
const spread = Math.max(red, green, blue) - Math.min(red, green, blue);
if (spread < 20) {
const tonedValue = Math.round(((red + green + blue) / 3) * 0.88);
red = Math.min(red, tonedValue);
green = Math.min(green, tonedValue);
blue = Math.min(blue, tonedValue);
}
}
}
let nextAlpha = alpha;
const edgeFade = Math.max(greenScore * 0.35, whiteScore * 0.28);
if (edgeFade > 0.08) {
nextAlpha = Math.min(alpha, Math.round(alpha * (1 - edgeFade)));
if (nextAlpha < 10) {
nextAlpha = 0;
}
}
if (
red !== (pixels[offset] ?? 0) ||
green !== (pixels[offset + 1] ?? 0) ||
blue !== (pixels[offset + 2] ?? 0) ||
nextAlpha !== alpha
) {
pixels[offset] = red;
pixels[offset + 1] = green;
pixels[offset + 2] = blue;
pixels[offset + 3] = nextAlpha;
changed = true;
}
}
}
return changed;
}