ai-chatbot/lib/video-canvas/video-audio-separation.ts

export type VideoAudioSplitMode = "voice" | "background" | "audio-video";

export type VideoCanvasAudioSeparationConfig = {
  backgroundGain: number;
  enabled: boolean;
  frameRate: number;
  model: string;
  voiceGain: number;
};

export const defaultVideoCanvasAudioSeparationConfig: VideoCanvasAudioSeparationConfig = {
  backgroundGain: 1.2,
  enabled: true,
  frameRate: 30,
  model: "center-channel-v1",
  voiceGain: 1,
};

function clamp(value: number, min: number, max: number) {
  return Math.min(max, Math.max(min, value));
}

function finiteNumber(value: unknown, fallback: number, min: number, max: number) {
  const parsed = Number(value);
  return Number.isFinite(parsed) ? clamp(parsed, min, max) : fallback;
}

export function normalizeVideoCanvasAudioSeparationConfig(
  value: unknown,
): VideoCanvasAudioSeparationConfig {
  const incoming = value && typeof value === "object" && !Array.isArray(value)
    ? value as Record<string, unknown>
    : {};
  return {
    backgroundGain: finiteNumber(
      incoming.backgroundGain,
      defaultVideoCanvasAudioSeparationConfig.backgroundGain,
      0,
      4,
    ),
    enabled:
      typeof incoming.enabled === "boolean"
        ? incoming.enabled
        : defaultVideoCanvasAudioSeparationConfig.enabled,
    frameRate: Math.round(finiteNumber(
      incoming.frameRate,
      defaultVideoCanvasAudioSeparationConfig.frameRate,
      12,
      60,
    )),
    model:
      typeof incoming.model === "string" && incoming.model.trim()
        ? incoming.model.trim().slice(0, 100)
        : defaultVideoCanvasAudioSeparationConfig.model,
    voiceGain: finiteNumber(
      incoming.voiceGain,
      defaultVideoCanvasAudioSeparationConfig.voiceGain,
      0,
      4,
    ),
  };
}

function clampSample(value: number) {
  return Math.min(1, Math.max(-1, value));
}

export function separateVideoAudioChannels(
  channels: readonly Float32Array[],
  mode: VideoAudioSplitMode,
  configValue: VideoCanvasAudioSeparationConfig,
) {
  const config = normalizeVideoCanvasAudioSeparationConfig(configValue);
  const left = channels[0];
  if (!left) return [];
  const right = channels[1] || left;
  const length = Math.min(left.length, right.length);
  const outputLeft = new Float32Array(length);
  const outputRight = new Float32Array(length);

  for (let index = 0; index < length; index += 1) {
    const leftSample = left[index] || 0;
    const rightSample = right[index] || 0;
    if (mode === "audio-video") {
      outputLeft[index] = leftSample;
      outputRight[index] = rightSample;
      continue;
    }
    if (mode === "voice") {
      const center = clampSample(((leftSample + rightSample) / 2) * config.voiceGain);
      outputLeft[index] = center;
      outputRight[index] = center;
      continue;
    }
    const side = clampSample(((leftSample - rightSample) / 2) * config.backgroundGain);
    outputLeft[index] = side;
    outputRight[index] = -side;
  }
  return [outputLeft, outputRight];
}

export function getVideoAudioSplitResultCopy(mode: VideoAudioSplitMode) {
  if (mode === "voice") {
    return { audioSuffix: "人声", audioSubtitle: "人声分离" };
  }
  if (mode === "background") {
    return { audioSuffix: "背景音", audioSubtitle: "背景音分离" };
  }
  return { audioSuffix: "音轨", audioSubtitle: "音视频分离" };
}

const MP4_CONTAINER_BOXES = new Set([
  "moov",
  "trak",
  "mdia",
  "minf",
  "stbl",
  "edts",
  "dinf",
  "mvex",
]);

function readBoxType(view: DataView, offset: number) {
  return String.fromCharCode(
    view.getUint8(offset),
    view.getUint8(offset + 1),
    view.getUint8(offset + 2),
    view.getUint8(offset + 3),
  );
}

function boxContainsAudioHandler(view: DataView, start: number, end: number): boolean {
  let offset = start;
  while (offset + 8 <= end) {
    const size32 = view.getUint32(offset);
    const type = readBoxType(view, offset + 4);
    let headerSize = 8;
    let size = size32;
    if (size32 === 1) {
      if (offset + 16 > end) return false;
      const high = view.getUint32(offset + 8);
      const low = view.getUint32(offset + 12);
      size = high * 2 ** 32 + low;
      headerSize = 16;
    } else if (size32 === 0) {
      size = end - offset;
    }
    if (!Number.isSafeInteger(size) || size < headerSize || offset + size > end) return false;
    const contentStart = offset + headerSize;
    const boxEnd = offset + size;
    if (type === "hdlr" && contentStart + 12 <= boxEnd) {
      if (readBoxType(view, contentStart + 8) === "soun") return true;
    }
    if (
      MP4_CONTAINER_BOXES.has(type) &&
      boxContainsAudioHandler(view, contentStart, boxEnd)
    ) {
      return true;
    }
    offset = boxEnd;
  }
  return false;
}

export function hasMp4AudioTrack(arrayBuffer: ArrayBuffer) {
  if (arrayBuffer.byteLength < 8) return false;
  return boxContainsAudioHandler(new DataView(arrayBuffer), 0, arrayBuffer.byteLength);
}

type EbmlVint = {
  length: number;
  unknown: boolean;
  value: bigint;
};

function readEbmlVint(view: DataView, offset: number, stripMarker: boolean): EbmlVint | null {
  if (offset >= view.byteLength) return null;
  const first = view.getUint8(offset);
  let marker = 0x80;
  let length = 1;
  while (length <= 8 && (first & marker) === 0) {
    marker >>= 1;
    length += 1;
  }
  if (length > 8 || offset + length > view.byteLength) return null;

  let value = BigInt(stripMarker ? first & (marker - 1) : first);
  for (let index = 1; index < length; index += 1) {
    value = (value << 8n) | BigInt(view.getUint8(offset + index));
  }
  const unknownValue = (1n << BigInt(length * 7)) - 1n;
  return {
    length,
    unknown: stripMarker && value === unknownValue,
    value,
  };
}

const EBML_SEGMENT_ID = 0x18538067n;
const EBML_TRACKS_ID = 0x1654ae6bn;
const EBML_TRACK_ENTRY_ID = 0xaen;
const EBML_TRACK_TYPE_ID = 0x83n;

function ebmlRangeContainsAudioTrack(
  view: DataView,
  start: number,
  end: number,
  insideTrackEntry = false,
): boolean {
  let offset = start;
  while (offset < end) {
    const id = readEbmlVint(view, offset, false);
    if (!id) return false;
    const size = readEbmlVint(view, offset + id.length, true);
    if (!size) return false;
    const contentStart = offset + id.length + size.length;
    const requestedEnd = size.unknown ? end : contentStart + Number(size.value);
    if (!Number.isSafeInteger(requestedEnd) || requestedEnd < contentStart || requestedEnd > end) {
      return false;
    }

    if (insideTrackEntry && id.value === EBML_TRACK_TYPE_ID) {
      let trackType = 0;
      for (let index = contentStart; index < requestedEnd; index += 1) {
        trackType = trackType * 256 + view.getUint8(index);
      }
      if (trackType === 2) return true;
    }

    const isTrackEntry = id.value === EBML_TRACK_ENTRY_ID;
    if (
      (id.value === EBML_SEGMENT_ID ||
        id.value === EBML_TRACKS_ID ||
        isTrackEntry) &&
      ebmlRangeContainsAudioTrack(view, contentStart, requestedEnd, insideTrackEntry || isTrackEntry)
    ) {
      return true;
    }
    offset = requestedEnd;
  }
  return false;
}

export function hasWebmAudioTrack(arrayBuffer: ArrayBuffer) {
  if (arrayBuffer.byteLength < 4) return false;
  const view = new DataView(arrayBuffer);
  return ebmlRangeContainsAudioTrack(view, 0, view.byteLength);
}

export function hasVideoAudioTrack(arrayBuffer: ArrayBuffer) {
  return hasMp4AudioTrack(arrayBuffer) || hasWebmAudioTrack(arrayBuffer);
}