Spaces:

vipan-kumar
/

Audio_Deepfake_Detection

Sleeping

Audio_Deepfake_Detection / backend /app /features /audio_preprocessor.py

Initial commit: Audio Deepfake Detector with 8 detectors trained on jay15k

e6a1f55 5 months ago

7.97 kB

"""Audio loading, resampling, and normalisation utilities.
Decode chain (in order):
    1. ``soundfile`` (libsndfile)        — fast path for WAV / FLAC / AIFF
    2. ``ffmpeg`` via imageio-ffmpeg     — handles MP3 / OGG / WebM / M4A / etc.
The ffmpeg fallback is critical on Windows because libsndfile ships without
MP3/OGG/WebM codecs and torchaudio's default backend uses libsndfile too.
Browser MediaRecorder produces WebM/Opus, so without this fallback the
microphone-recording path is completely broken.
"""
from __future__ import annotations

import io
import shutil
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Tuple, Union

import numpy as np
import soundfile as sf
import torch

from app.logging_setup import get_logger

logger = get_logger(__name__)


# --------------------------------------------------------------------------- #
# ffmpeg discovery — prefer the bundled binary, fall back to system PATH.
# --------------------------------------------------------------------------- #

_FFMPEG_EXE: Optional[str] = None


def _resolve_ffmpeg() -> Optional[str]:
    global _FFMPEG_EXE
    if _FFMPEG_EXE is not None:
        return _FFMPEG_EXE
    # Bundled ffmpeg first — works without any system install.
    try:
        import imageio_ffmpeg

        exe = imageio_ffmpeg.get_ffmpeg_exe()
        if exe and Path(exe).exists():
            _FFMPEG_EXE = exe
            logger.info("Using bundled ffmpeg: %s", exe)
            return exe
    except Exception as exc:  # noqa: BLE001
        logger.debug("imageio_ffmpeg unavailable: %s", exc)

    # System ffmpeg as last resort.
    sys_exe = shutil.which("ffmpeg")
    if sys_exe:
        _FFMPEG_EXE = sys_exe
        logger.info("Using system ffmpeg: %s", sys_exe)
        return sys_exe

    logger.warning(
        "No ffmpeg available — only PCM WAV / FLAC / AIFF will decode. "
        "Install imageio-ffmpeg or ffmpeg to enable MP3 / OGG / WebM / M4A."
    )
    return None


# --------------------------------------------------------------------------- #
# Public dataclass
# --------------------------------------------------------------------------- #

@dataclass
class LoadedAudio:
    waveform: torch.Tensor          # [1, T] float32 in ~[-1, 1]
    sample_rate: int                # always == target_sr after preprocessing
    original_sample_rate: int
    original_channels: int
    duration_seconds: float
    file_format: str
    file_size_bytes: int
    decoder: str                    # 'soundfile' | 'ffmpeg' | 'unknown'


# --------------------------------------------------------------------------- #
# Decoders
# --------------------------------------------------------------------------- #

def _decode_with_soundfile(data: Union[bytes, str, Path]) -> Tuple[np.ndarray, int, str]:
    """Try the fast soundfile path. Raises if the format isn't supported."""
    if isinstance(data, (bytes, bytearray)):
        bio = io.BytesIO(data)
        wav, sr = sf.read(bio, dtype="float32", always_2d=True)
        return wav, sr, "wav"
    path = Path(data)
    suffix = path.suffix.lstrip(".").lower() or "unknown"
    wav, sr = sf.read(str(path), dtype="float32", always_2d=True)
    return wav, sr, suffix


def _decode_with_ffmpeg(
    data: Union[bytes, str, Path],
    target_sr: int = 16000,
) -> Tuple[np.ndarray, int, str]:
    """Decode anything ffmpeg understands → mono float32 PCM at target_sr."""
    ffmpeg = _resolve_ffmpeg()
    if ffmpeg is None:
        raise RuntimeError(
            "Audio format not supported by libsndfile and no ffmpeg available. "
            "Install imageio-ffmpeg (pip install imageio-ffmpeg) or ffmpeg."
        )

    is_bytes = isinstance(data, (bytes, bytearray))
    fmt_hint = "compressed"
    if not is_bytes:
        fmt_hint = Path(data).suffix.lstrip(".").lower() or "compressed"

    cmd = [
        ffmpeg,
        "-hide_banner",
        "-loglevel", "error",
        "-nostdin",
        "-i", "pipe:0" if is_bytes else str(data),
        "-vn",                     # discard any video stream
        "-ac", "1",                # mono
        "-ar", str(target_sr),     # resample
        "-f", "wav",
        "-acodec", "pcm_s16le",
        "pipe:1",
    ]

    try:
        proc = subprocess.run(
            cmd,
            input=bytes(data) if is_bytes else None,
            capture_output=True,
            check=False,
            timeout=60,
        )
    except subprocess.TimeoutExpired as exc:
        raise RuntimeError("ffmpeg decode timed out after 60s") from exc

    if proc.returncode != 0:
        msg = proc.stderr.decode(errors="replace").strip() or "unknown ffmpeg error"
        raise RuntimeError(f"ffmpeg decode failed: {msg}")

    bio = io.BytesIO(proc.stdout)
    wav, sr = sf.read(bio, dtype="float32", always_2d=True)
    return wav, sr, fmt_hint


def _read_bytes(
    data: Union[bytes, str, Path],
    target_sr: int = 16000,
) -> Tuple[np.ndarray, int, str, str]:
    """Returns (waveform_np[T,C], sr, format_hint, decoder_used)."""
    # 1) Fast path: soundfile
    try:
        wav, sr, fmt = _decode_with_soundfile(data)
        return wav, sr, fmt, "soundfile"
    except Exception as sf_exc:
        logger.debug("soundfile decode failed: %s; trying ffmpeg.", sf_exc)

    # 2) Fallback: ffmpeg → wav → soundfile
    wav, sr, fmt = _decode_with_ffmpeg(data, target_sr=target_sr)
    return wav, sr, fmt, "ffmpeg"


# --------------------------------------------------------------------------- #
# Public API
# --------------------------------------------------------------------------- #

def preprocess_audio(
    source: Union[bytes, str, Path],
    target_sr: int = 16000,
    max_seconds: float = 30.0,
) -> LoadedAudio:
    """Load → decode → resample → mono → peak-normalize → clip to ``max_seconds``.
    Returns a LoadedAudio with a [1, T] float32 tensor. Raises RuntimeError
    with a clear message on decode failure.
    """
    wav_np, sr, fmt, decoder = _read_bytes(source, target_sr=target_sr)
    file_size = (
        len(source)
        if isinstance(source, (bytes, bytearray))
        else Path(source).stat().st_size
    )
    orig_channels = int(wav_np.shape[1]) if wav_np.ndim == 2 else 1

    # numpy → torch [C, T]
    if wav_np.ndim == 1:
        tensor = torch.from_numpy(wav_np).unsqueeze(0)
    else:
        tensor = torch.from_numpy(wav_np.T)
    tensor = tensor.float()

    # Stereo → mono
    if tensor.shape[0] > 1:
        tensor = tensor.mean(dim=0, keepdim=True)

    # Resample if needed (ffmpeg path is already at target_sr, but be safe).
    if sr != target_sr:
        import torchaudio  # lazy import; not all paths need it

        tensor = torchaudio.functional.resample(tensor, sr, target_sr)

    # Trim to max length
    max_samples = int(max_seconds * target_sr)
    if tensor.shape[1] > max_samples:
        tensor = tensor[:, :max_samples]

    # Peak-normalise (avoid div-by-zero on silence)
    peak = tensor.abs().max()
    if peak > 1e-6:
        tensor = tensor / (peak + 1e-8)

    duration = float(tensor.shape[1]) / target_sr

    return LoadedAudio(
        waveform=tensor.contiguous(),
        sample_rate=target_sr,
        original_sample_rate=sr,
        original_channels=orig_channels,
        duration_seconds=duration,
        file_format=fmt,
        file_size_bytes=int(file_size),
        decoder=decoder,
    )


def waveform_summary(waveform: torch.Tensor, max_points: int = 512) -> list[float]:
    """Down-sample the waveform to ``max_points`` envelope values for the UI."""
    x = waveform.squeeze().detach().cpu().numpy().astype(np.float32)
    if x.size == 0:
        return []
    if x.size <= max_points:
        return x.tolist()
    bucket_size = x.size // max_points
    trimmed = x[: bucket_size * max_points].reshape(max_points, bucket_size)
    env = np.abs(trimmed).max(axis=1)
    return env.tolist()

	"""Audio loading, resampling, and normalisation utilities.

	Decode chain (in order):
	1. ``soundfile`` (libsndfile) — fast path for WAV / FLAC / AIFF
	2. ``ffmpeg`` via imageio-ffmpeg — handles MP3 / OGG / WebM / M4A / etc.

	The ffmpeg fallback is critical on Windows because libsndfile ships without
	MP3/OGG/WebM codecs and torchaudio's default backend uses libsndfile too.
	Browser MediaRecorder produces WebM/Opus, so without this fallback the
	microphone-recording path is completely broken.
	"""
	from __future__ import annotations

	import io
	import shutil
	import subprocess
	from dataclasses import dataclass
	from pathlib import Path
	from typing import Optional, Tuple, Union

	import numpy as np
	import soundfile as sf
	import torch

	from app.logging_setup import get_logger

	logger = get_logger(__name__)


	# --------------------------------------------------------------------------- #
	# ffmpeg discovery — prefer the bundled binary, fall back to system PATH.
	# --------------------------------------------------------------------------- #

	_FFMPEG_EXE: Optional[str] = None


	def _resolve_ffmpeg() -> Optional[str]:
	global _FFMPEG_EXE
	if _FFMPEG_EXE is not None:
	return _FFMPEG_EXE
	# Bundled ffmpeg first — works without any system install.
	try:
	import imageio_ffmpeg

	exe = imageio_ffmpeg.get_ffmpeg_exe()
	if exe and Path(exe).exists():
	_FFMPEG_EXE = exe
	logger.info("Using bundled ffmpeg: %s", exe)
	return exe
	except Exception as exc: # noqa: BLE001
	logger.debug("imageio_ffmpeg unavailable: %s", exc)

	# System ffmpeg as last resort.
	sys_exe = shutil.which("ffmpeg")
	if sys_exe:
	_FFMPEG_EXE = sys_exe
	logger.info("Using system ffmpeg: %s", sys_exe)
	return sys_exe

	logger.warning(
	"No ffmpeg available — only PCM WAV / FLAC / AIFF will decode. "
	"Install imageio-ffmpeg or ffmpeg to enable MP3 / OGG / WebM / M4A."
	)
	return None


	# --------------------------------------------------------------------------- #
	# Public dataclass
	# --------------------------------------------------------------------------- #

	@dataclass
	class LoadedAudio:
	waveform: torch.Tensor # [1, T] float32 in ~[-1, 1]
	sample_rate: int # always == target_sr after preprocessing
	original_sample_rate: int
	original_channels: int
	duration_seconds: float
	file_format: str
	file_size_bytes: int
	decoder: str # 'soundfile' \| 'ffmpeg' \| 'unknown'


	# --------------------------------------------------------------------------- #
	# Decoders
	# --------------------------------------------------------------------------- #

	def _decode_with_soundfile(data: Union[bytes, str, Path]) -> Tuple[np.ndarray, int, str]:
	"""Try the fast soundfile path. Raises if the format isn't supported."""
	if isinstance(data, (bytes, bytearray)):
	bio = io.BytesIO(data)
	wav, sr = sf.read(bio, dtype="float32", always_2d=True)
	return wav, sr, "wav"
	path = Path(data)
	suffix = path.suffix.lstrip(".").lower() or "unknown"
	wav, sr = sf.read(str(path), dtype="float32", always_2d=True)
	return wav, sr, suffix


	def _decode_with_ffmpeg(
	data: Union[bytes, str, Path],
	target_sr: int = 16000,
	) -> Tuple[np.ndarray, int, str]:
	"""Decode anything ffmpeg understands → mono float32 PCM at target_sr."""
	ffmpeg = _resolve_ffmpeg()
	if ffmpeg is None:
	raise RuntimeError(
	"Audio format not supported by libsndfile and no ffmpeg available. "
	"Install imageio-ffmpeg (pip install imageio-ffmpeg) or ffmpeg."
	)

	is_bytes = isinstance(data, (bytes, bytearray))
	fmt_hint = "compressed"
	if not is_bytes:
	fmt_hint = Path(data).suffix.lstrip(".").lower() or "compressed"

	cmd = [
	ffmpeg,
	"-hide_banner",
	"-loglevel", "error",
	"-nostdin",
	"-i", "pipe:0" if is_bytes else str(data),
	"-vn", # discard any video stream
	"-ac", "1", # mono
	"-ar", str(target_sr), # resample
	"-f", "wav",
	"-acodec", "pcm_s16le",
	"pipe:1",
	]

	try:
	proc = subprocess.run(
	cmd,
	input=bytes(data) if is_bytes else None,
	capture_output=True,
	check=False,
	timeout=60,
	)
	except subprocess.TimeoutExpired as exc:
	raise RuntimeError("ffmpeg decode timed out after 60s") from exc

	if proc.returncode != 0:
	msg = proc.stderr.decode(errors="replace").strip() or "unknown ffmpeg error"
	raise RuntimeError(f"ffmpeg decode failed: {msg}")

	bio = io.BytesIO(proc.stdout)
	wav, sr = sf.read(bio, dtype="float32", always_2d=True)
	return wav, sr, fmt_hint


	def _read_bytes(
	data: Union[bytes, str, Path],
	target_sr: int = 16000,
	) -> Tuple[np.ndarray, int, str, str]:
	"""Returns (waveform_np[T,C], sr, format_hint, decoder_used)."""
	# 1) Fast path: soundfile
	try:
	wav, sr, fmt = _decode_with_soundfile(data)
	return wav, sr, fmt, "soundfile"
	except Exception as sf_exc:
	logger.debug("soundfile decode failed: %s; trying ffmpeg.", sf_exc)

	# 2) Fallback: ffmpeg → wav → soundfile
	wav, sr, fmt = _decode_with_ffmpeg(data, target_sr=target_sr)
	return wav, sr, fmt, "ffmpeg"


	# --------------------------------------------------------------------------- #
	# Public API
	# --------------------------------------------------------------------------- #

	def preprocess_audio(
	source: Union[bytes, str, Path],
	target_sr: int = 16000,
	max_seconds: float = 30.0,
	) -> LoadedAudio:
	"""Load → decode → resample → mono → peak-normalize → clip to ``max_seconds``.

	Returns a LoadedAudio with a [1, T] float32 tensor. Raises RuntimeError
	with a clear message on decode failure.
	"""
	wav_np, sr, fmt, decoder = _read_bytes(source, target_sr=target_sr)
	file_size = (
	len(source)
	if isinstance(source, (bytes, bytearray))
	else Path(source).stat().st_size
	)
	orig_channels = int(wav_np.shape[1]) if wav_np.ndim == 2 else 1

	# numpy → torch [C, T]
	if wav_np.ndim == 1:
	tensor = torch.from_numpy(wav_np).unsqueeze(0)
	else:
	tensor = torch.from_numpy(wav_np.T)
	tensor = tensor.float()

	# Stereo → mono
	if tensor.shape[0] > 1:
	tensor = tensor.mean(dim=0, keepdim=True)

	# Resample if needed (ffmpeg path is already at target_sr, but be safe).
	if sr != target_sr:
	import torchaudio # lazy import; not all paths need it

	tensor = torchaudio.functional.resample(tensor, sr, target_sr)

	# Trim to max length
	max_samples = int(max_seconds * target_sr)
	if tensor.shape[1] > max_samples:
	tensor = tensor[:, :max_samples]

	# Peak-normalise (avoid div-by-zero on silence)
	peak = tensor.abs().max()
	if peak > 1e-6:
	tensor = tensor / (peak + 1e-8)

	duration = float(tensor.shape[1]) / target_sr

	return LoadedAudio(
	waveform=tensor.contiguous(),
	sample_rate=target_sr,
	original_sample_rate=sr,
	original_channels=orig_channels,
	duration_seconds=duration,
	file_format=fmt,
	file_size_bytes=int(file_size),
	decoder=decoder,
	)


	def waveform_summary(waveform: torch.Tensor, max_points: int = 512) -> list[float]:
	"""Down-sample the waveform to ``max_points`` envelope values for the UI."""
	x = waveform.squeeze().detach().cpu().numpy().astype(np.float32)
	if x.size == 0:
	return []
	if x.size <= max_points:
	return x.tolist()
	bucket_size = x.size // max_points
	trimmed = x[: bucket_size * max_points].reshape(max_points, bucket_size)
	env = np.abs(trimmed).max(axis=1)
	return env.tolist()