Add conversation image replacement

2026-08-14 12:56:00 +01:00
1 parent 2500f0a23c
commit 8523873831
@@ -0,0 +1,301 @@
+import fs from 'node:fs/promises';
+import path from 'node:path';
+import sharp from 'sharp';
+import { downloadImage } from '../utils/image.js';
+
+const MAX_INPUT_PIXELS = 25_000_000;
+const DETECTION_MAX_DIMENSION = 900;
+
+export interface ImageReplacementResult {
+  bytes: Buffer;
+  mimeType: 'image/jpeg';
+  width: number;
+  height: number;
+}
+
+export type ImageReplacementTarget =
+  | { kind: 'conversation-media'; localPath: string }
+  | { kind: 'remote'; url: string; userAgent: string };
+
+export interface ImageReplacementGenerator {
+  generate(sourcePath: string, target: ImageReplacementTarget): Promise<ImageReplacementResult>;
+}
+
+export class ImageReplacementInputError extends Error {}
+
+type RawImage = { data: Buffer; width: number; height: number; channels: 3 | 4 };
+type Point = { x: number; y: number };
+
+/**
+ * Ports the green-screen replacement prototype into the Node runtime. The
+ * source is an operator-owned exchange-library image; the replacement is a
+ * conversation image selected from authoritative server data.
+ */
+export class GreenScreenImageReplacementService implements ImageReplacementGenerator {
+  constructor(private readonly mediaRoot = path.join(process.cwd(), 'data', 'media')) {}
+
+  async generate(sourcePath: string, target: ImageReplacementTarget): Promise<ImageReplacementResult> {
+    const sourceBytes = await fs.readFile(sourcePath);
+    const targetBytes = target.kind === 'remote'
+      ? (await downloadImage(target.url, target.userAgent)).buffer
+      : await fs.readFile(resolveInside(this.mediaRoot, target.localPath));
+    return replaceGreenScreen(sourceBytes, targetBytes);
+  }
+}
+
+export async function replaceGreenScreen(sourceBytes: Buffer, targetBytes: Buffer): Promise<ImageReplacementResult> {
+  const [source, detection] = await Promise.all([
+    decode(sourceBytes),
+    decode(sourceBytes, DETECTION_MAX_DIMENSION),
+  ]).catch((error: unknown) => {
+    throw new ImageReplacementInputError(`Unable to read the selected image: ${safeError(error)}`);
+  });
+
+  const smallMask = greenMask(detection);
+  const component = largestComponent(smallMask, detection.width, detection.height);
+  if (component.length < Math.max(64, Math.round(detection.width * detection.height * 0.002))) {
+    throw new ImageReplacementInputError('No usable green-screen region was found in the source image');
+  }
+  const smallQuad = componentExtremes(component, detection.width);
+  const scaleX = source.width / detection.width;
+  const scaleY = source.height / detection.height;
+  const quad = smallQuad.map(({ x, y }) => ({ x: x * scaleX, y: y * scaleY }));
+  const screenWidth = Math.max(1, Math.round(Math.max(distance(quad[0], quad[1]), distance(quad[3], quad[2]))));
+  const screenHeight = Math.max(1, Math.round(Math.max(distance(quad[0], quad[3]), distance(quad[1], quad[2]))));
+  const fitted = await fitCover(targetBytes, screenWidth, screenHeight).catch((error: unknown) => {
+    throw new ImageReplacementInputError(`Unable to read the selected replacement image: ${safeError(error)}`);
+  });
+  const homography = solveHomography(quad, screenWidth, screenHeight);
+
+  const rawMask = greenMask(source);
+  const clippedMask = Buffer.alloc(source.width * source.height);
+  let luminanceTotal = 0;
+  let maskedPixels = 0;
+  for (let y = 0; y < source.height; y++) {
+    for (let x = 0; x < source.width; x++) {
+      const index = y * source.width + x;
+      if (rawMask[index] && pointInQuad(x, y, quad)) {
+        clippedMask[index] = 255;
+        const offset = index * source.channels;
+        luminanceTotal += luminance(source.data[offset], source.data[offset + 1], source.data[offset + 2]);
+        maskedPixels++;
+      }
+    }
+  }
+  if (maskedPixels === 0) throw new ImageReplacementInputError('The green-screen region could not be composited');
+  // Expand by one pixel before feathering so antialiased green edge pixels do
+  // not survive as a bright fringe around the insert.
+  const expandedMask = dilateMask(clippedMask, source.width, source.height);
+  const featheredMask = await sharp(expandedMask, { raw: { width: source.width, height: source.height, channels: 1 } })
+    .blur(1.5)
+    .greyscale()
+    .raw()
+    .toBuffer();
+  const referenceLuminance = luminanceTotal / maskedPixels;
+  const output = Buffer.from(source.data);
+  const bounds = quadBounds(quad, source.width, source.height);
+
+  for (let y = bounds.top; y <= bounds.bottom; y++) {
+    for (let x = bounds.left; x <= bounds.right; x++) {
+      const pixel = y * source.width + x;
+      const alpha = featheredMask[pixel] / 255;
+      if (alpha <= 0.002) continue;
+      const mapped = mapPoint(homography, x, y);
+      if (mapped.x < 0 || mapped.y < 0 || mapped.x >= screenWidth || mapped.y >= screenHeight) continue;
+      const sampled = bilinearSample(fitted, mapped.x, mapped.y);
+      const offset = pixel * source.channels;
+      const sourceLuminance = luminance(source.data[offset], source.data[offset + 1], source.data[offset + 2]);
+      const gain = clamp(1 + 0.35 * (sourceLuminance / Math.max(referenceLuminance, 1) - 1), 0.55, 1.65);
+      for (let channel = 0; channel < 3; channel++) {
+        const replacement = clamp(sampled[channel] * gain, 0, 255);
+        output[offset + channel] = Math.round(source.data[offset + channel] * (1 - alpha) + replacement * alpha);
+      }
+    }
+  }
+
+  const bytes = await sharp(output, {
+    raw: { width: source.width, height: source.height, channels: source.channels },
+  }).jpeg({ quality: 92, mozjpeg: true }).toBuffer();
+  return { bytes, mimeType: 'image/jpeg', width: source.width, height: source.height };
+}
+
+async function decode(bytes: Buffer, maxDimension?: number): Promise<RawImage> {
+  let pipeline = sharp(bytes, { limitInputPixels: MAX_INPUT_PIXELS, animated: false }).rotate().removeAlpha().toColourspace('srgb');
+  if (maxDimension) pipeline = pipeline.resize({ width: maxDimension, height: maxDimension, fit: 'inside', withoutEnlargement: true });
+  const decoded = await pipeline.raw().toBuffer({ resolveWithObject: true });
+  if (decoded.info.channels < 3) throw new ImageReplacementInputError('Selected images must contain RGB colour data');
+  return { data: decoded.data, width: decoded.info.width, height: decoded.info.height, channels: decoded.info.channels as 3 | 4 };
+}
+
+async function fitCover(bytes: Buffer, width: number, height: number): Promise<RawImage> {
+  const decoded = await sharp(bytes, { limitInputPixels: MAX_INPUT_PIXELS, animated: false })
+    .rotate().removeAlpha().toColourspace('srgb')
+    .resize(width, height, { fit: 'cover', position: 'centre' })
+    .raw().toBuffer({ resolveWithObject: true });
+  return { data: decoded.data, width, height, channels: decoded.info.channels as 3 | 4 };
+}
+
+function greenMask(image: RawImage): Uint8Array {
+  const mask = new Uint8Array(image.width * image.height);
+  for (let pixel = 0; pixel < mask.length; pixel++) {
+    const offset = pixel * image.channels;
+    const r = image.data[offset] / 255;
+    const g = image.data[offset + 1] / 255;
+    const b = image.data[offset + 2] / 255;
+    const max = Math.max(r, g, b);
+    const min = Math.min(r, g, b);
+    const saturation = max === 0 ? 0 : (max - min) / max;
+    let hue = 0;
+    if (max !== min) {
+      if (max === r) hue = 60 * (((g - b) / (max - min)) % 6);
+      else if (max === g) hue = 60 * ((b - r) / (max - min) + 2);
+      else hue = 60 * ((r - g) / (max - min) + 4);
+      if (hue < 0) hue += 360;
+    }
+    if (hue >= 70 && hue <= 180 && saturation >= 70 / 255 && max >= 50 / 255) mask[pixel] = 1;
+  }
+  return mask;
+}
+
+function largestComponent(mask: Uint8Array, width: number, height: number): number[] {
+  const seen = new Uint8Array(mask.length);
+  const queue = new Int32Array(mask.length);
+  let largest: number[] = [];
+  for (let start = 0; start < mask.length; start++) {
+    if (!mask[start] || seen[start]) continue;
+    let head = 0;
+    let tail = 0;
+    queue[tail++] = start;
+    seen[start] = 1;
+    const component: number[] = [];
+    while (head < tail) {
+      const pixel = queue[head++];
+      component.push(pixel);
+      const x = pixel % width;
+      const y = Math.floor(pixel / width);
+      const neighbours = [pixel - width, pixel + width, pixel - 1, pixel + 1];
+      for (let index = 0; index < neighbours.length; index++) {
+        const neighbour = neighbours[index];
+        if (neighbour < 0 || neighbour >= mask.length) continue;
+        if (index === 2 && x === 0) continue;
+        if (index === 3 && x === width - 1) continue;
+        if ((index === 0 && y === 0) || (index === 1 && y === height - 1)) continue;
+        if (mask[neighbour] && !seen[neighbour]) {
+          seen[neighbour] = 1;
+          queue[tail++] = neighbour;
+        }
+      }
+    }
+    if (component.length > largest.length) largest = component;
+  }
+  return largest;
+}
+
+function dilateMask(mask: Buffer, width: number, height: number): Buffer {
+  const expanded = Buffer.from(mask);
+  for (let y = 1; y < height - 1; y++) {
+    for (let x = 1; x < width - 1; x++) {
+      const pixel = y * width + x;
+      if (!mask[pixel]) continue;
+      for (let offsetY = -1; offsetY <= 1; offsetY++) {
+        for (let offsetX = -1; offsetX <= 1; offsetX++) expanded[pixel + offsetY * width + offsetX] = 255;
+      }
+    }
+  }
+  return expanded;
+}
+
+function componentExtremes(component: number[], width: number): [Point, Point, Point, Point] {
+  const points = component.map((pixel) => ({ x: pixel % width, y: Math.floor(pixel / width) }));
+  return [
+    points.reduce((best, point) => point.x + point.y < best.x + best.y ? point : best),
+    points.reduce((best, point) => point.x - point.y > best.x - best.y ? point : best),
+    points.reduce((best, point) => point.x + point.y > best.x + best.y ? point : best),
+    points.reduce((best, point) => point.y - point.x > best.y - best.x ? point : best),
+  ];
+}
+
+function solveHomography(quad: Point[], width: number, height: number): number[] {
+  const destinations = [{ x: 0, y: 0 }, { x: width - 1, y: 0 }, { x: width - 1, y: height - 1 }, { x: 0, y: height - 1 }];
+  const matrix: number[][] = [];
+  for (let index = 0; index < 4; index++) {
+    const { x, y } = quad[index];
+    const { x: u, y: v } = destinations[index];
+    matrix.push([x, y, 1, 0, 0, 0, -u * x, -u * y, u]);
+    matrix.push([0, 0, 0, x, y, 1, -v * x, -v * y, v]);
+  }
+  for (let column = 0; column < 8; column++) {
+    let pivot = column;
+    for (let row = column + 1; row < 8; row++) if (Math.abs(matrix[row][column]) > Math.abs(matrix[pivot][column])) pivot = row;
+    [matrix[column], matrix[pivot]] = [matrix[pivot], matrix[column]];
+    const divisor = matrix[column][column];
+    if (Math.abs(divisor) < 1e-9) throw new ImageReplacementInputError('The detected green-screen shape is not usable');
+    for (let cell = column; cell < 9; cell++) matrix[column][cell] /= divisor;
+    for (let row = 0; row < 8; row++) {
+      if (row === column) continue;
+      const factor = matrix[row][column];
+      for (let cell = column; cell < 9; cell++) matrix[row][cell] -= factor * matrix[column][cell];
+    }
+  }
+  return matrix.map((row) => row[8]);
+}
+
+function mapPoint(h: number[], x: number, y: number): Point {
+  const denominator = h[6] * x + h[7] * y + 1;
+  return { x: (h[0] * x + h[1] * y + h[2]) / denominator, y: (h[3] * x + h[4] * y + h[5]) / denominator };
+}
+
+function bilinearSample(image: RawImage, x: number, y: number): [number, number, number] {
+  const x0 = Math.floor(x);
+  const y0 = Math.floor(y);
+  const x1 = Math.min(x0 + 1, image.width - 1);
+  const y1 = Math.min(y0 + 1, image.height - 1);
+  const dx = x - x0;
+  const dy = y - y0;
+  const sample = (sx: number, sy: number, channel: number) => image.data[(sy * image.width + sx) * image.channels + channel];
+  return [0, 1, 2].map((channel) => (
+    sample(x0, y0, channel) * (1 - dx) * (1 - dy)
+    + sample(x1, y0, channel) * dx * (1 - dy)
+    + sample(x0, y1, channel) * (1 - dx) * dy
+    + sample(x1, y1, channel) * dx * dy
+  )) as [number, number, number];
+}
+
+function pointInQuad(x: number, y: number, quad: Point[]): boolean {
+  let sign = 0;
+  for (let index = 0; index < 4; index++) {
+    const a = quad[index];
+    const b = quad[(index + 1) % 4];
+    const cross = (b.x - a.x) * (y - a.y) - (b.y - a.y) * (x - a.x);
+    if (Math.abs(cross) < 1e-6) continue;
+    const nextSign = Math.sign(cross);
+    if (sign !== 0 && nextSign !== sign) return false;
+    sign = nextSign;
+  }
+  return true;
+}
+
+function quadBounds(quad: Point[], width: number, height: number) {
+  const xs = quad.map((point) => point.x);
+  const ys = quad.map((point) => point.y);
+  return {
+    left: clamp(Math.floor(Math.min(...xs)) - 3, 0, width - 1),
+    right: clamp(Math.ceil(Math.max(...xs)) + 3, 0, width - 1),
+    top: clamp(Math.floor(Math.min(...ys)) - 3, 0, height - 1),
+    bottom: clamp(Math.ceil(Math.max(...ys)) + 3, 0, height - 1),
+  };
+}
+
+function resolveInside(root: string, relativePath: string): string {
+  const resolvedRoot = path.resolve(root);
+  const resolved = path.resolve(resolvedRoot, relativePath);
+  if (resolved === resolvedRoot || !resolved.startsWith(`${resolvedRoot}${path.sep}`)) {
+    throw new ImageReplacementInputError('Unsafe conversation media path');
+  }
+  return resolved;
+}
+
+function distance(a: Point, b: Point): number { return Math.hypot(a.x - b.x, a.y - b.y); }
+function luminance(r: number, g: number, b: number): number { return 0.2126 * r + 0.7152 * g + 0.0722 * b; }
+function clamp(value: number, minimum: number, maximum: number): number { return Math.min(maximum, Math.max(minimum, value)); }
+function safeError(error: unknown): string { return error instanceof Error ? error.message : String(error); }