0.2_make_audio_features.py

import os
import re
import cv2
import numpy as np
import torch
import librosa
from torchfcpe import spawn_bundled_infer_model

# Directory Configuration
LANDMARK_DIR = "dataset/landmarks"
OUTPUT_DIR = "dataset/audio_features"
# Videos live in the repo root (and optionally dataset/raw_audio or the parent
# dir); search them in that order.
VIDEO_SEARCH_DIRS = [".", "dataset/raw_audio", ".."]

os.makedirs(OUTPUT_DIR, exist_ok=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device} for FCPE extraction")

# bf16 autocast on GPU: ~30% faster on RTX 3050, no F0 accuracy loss vs fp32
# (verified). Keep model weights fp32; autocast casts ops in the forward only.
USE_BF16 = (device == "cuda")
if USE_BF16:
    print("Using bf16 autocast for FCPE")

# Chunked + batched FCPE inference. One infer() call over the full ~2h audio
# OOMs a 4 GB GPU (it stages the whole spectrogram + activations at once).
# Instead: split the clip's audio into CHUNK_SECONDS windows, stack MAX_BATCH
# windows per forward, each window gets output_interp_target_length=
# frames_per_chunk so all batch items share one target length. Concat, trim to V.
# 12 fps -> frames_per_chunk = CHUNK_SECONDS * 12 exactly.
CHUNK_SECONDS = 60
MAX_BATCH = 4  # windows per forward pass; 4x60s fits comfortably in 4 GB at bf16

# Load the FCPE model (it will download pretrained weights automatically on first run)
print("Loading FCPE pitch estimator...")
fcpe_model = spawn_bundled_infer_model(device=device)

# Suffix a landmark stem carries after the video stem, e.g. for
# `XTRlxbOA8Pg_12fps_0001153_to_end` the video stem is
# `XTRlxbOA8Pg_12fps`. Group 1 = start time string (colons stripped), group 2
# = end time string or None for "_to_end".
CLIP_SUFFIX_RE = re.compile(r"_(\d+)_to_(end|\d+)$")


def find_video(video_stem):
    for d in VIDEO_SEARCH_DIRS:
        for ext in (".mp4", ".mkv", ".webm", ".mov", ".avi"):
            p = os.path.join(d, video_stem + ext)
            if os.path.exists(p):
                return p
    return None


def _parse_hhmmss(s: str) -> int:
    """Parse a colon-stripped time string back to seconds.

    Standard 6-char string (HHMMSS, 2 digits each, e.g. "013143" -> 01:31:43)
    is reliable. Non-6-char strings (e.g. the old 7-char "0001153" where the
    user typed minutes as 3 digits) are lossy and get rejected.
    """
    if len(s) == 6:
        h, m, sec = int(s[0:2]), int(s[2:4]), int(s[4:6])
        return h * 3600 + m * 60 + sec
    raise ValueError(f"Time string '{s}' has {len(s)} chars (expected 6). "
                     f"Non-standard padding is ambiguous and can't be parsed.")


def video_frame_bounds(video_path, V, start_str, end_str):
    """Return (start_frame, end_frame) for the landmark clip within the video.

    * "_to_end" -> clip runs to video EOF: start_frame = total - V,
      end_frame = total.  Self-derived from video metadata, no parsing.

    * Explicit start/end -> parse both as HHMMSS (6-char standard), convert
      to frames at 12 fps.  V must match end_frame - start_frame exactly
      (0.1 produces one landmark per source frame, NaN for gaps).
    """
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    cap.release()
    if not (11.9 <= fps <= 12.1):
        raise ValueError(f"{video_path}: fps={fps}, expected ~12")

    if end_str == "end":
        if V > total:
            raise ValueError(f"V={V} > video frames total={total}")
        return total - V, total

    # Explicit end: parse start/end as HHMMSS, convert to frames.
    start_sec = _parse_hhmmss(start_str)
    end_sec = _parse_hhmmss(end_str)
    start_frame = round(start_sec * fps)
    end_frame = round(end_sec * fps)
    expected = end_frame - start_frame
    if V != expected:
        print(f"  Warning: V={V} != computed {expected} (start={start_str} "
              f"end={end_str}). Using computed bounds.")
    return start_frame, end_frame


landmark_files = sorted(os.listdir(LANDMARK_DIR))
processed_count = 0

for lm_file in landmark_files:
    if not lm_file.endswith("_norm_landmarks.npy"):
        continue

    # base = landmark stem, e.g. 'XTRlxbOA8Pg_12fps_0001153_to_end'
    base_name = lm_file.replace("_norm_landmarks.npy", "")
    landmark_path = os.path.join(LANDMARK_DIR, lm_file)

    m = CLIP_SUFFIX_RE.search(base_name)
    if not m:
        print(f"Skipping {base_name}: no '_<time>_to_<time|end>' clip suffix "
              f"(can't derive audio offset).")
        continue
    start_str = m.group(1)
    end_str = m.group(2)
    video_stem = base_name[:m.start()]
    video_path = find_video(video_stem)
    if not video_path:
        print(f"Skipping {base_name}: no video '{video_stem}.*' in "
              f"{VIDEO_SEARCH_DIRS}")
        continue

    print(f"Processing: {base_name}  (video: {os.path.basename(video_path)})")

    output_path = os.path.join(OUTPUT_DIR, f"{base_name}_audio_feats.npy")
    if os.path.exists(output_path):
        print(f"  Skipping (already exists).")
        continue

    # 1. Determine target frame length (V) and the exact video time range the
    #    landmark clip covers, so audio and landmarks stay in sync.
    landmarks = np.load(landmark_path, mmap_mode='r')
    V = landmarks.shape[0]  # Exact target frame count at 12 FPS

    start_frame, end_frame = video_frame_bounds(video_path, V, start_str, end_str)
    start_sec = start_frame / 12.0
    duration_sec = (end_frame - start_frame) / 12.0  # == V / 12
    # Small slack so the melspectrogram frame count lands at >= V (the existing
    # truncate/pad below handles any excess); this is the same hop=1333 drift
    # the original script already tolerated (~1 frame / ~12 min).
    load_duration = duration_sec + 2.0

    print(f"  V={V} frames, clip {start_frame}..{end_frame} "
          f"({start_sec:.1f}s -> {start_sec + duration_sec:.1f}s)")

    # 2. Load the clip's audio segment (16 kHz mono) straight from the video.
    # 0.1 produced one landmark per source frame so audio[start_sec .. +V/12]
    # is exactly what the landmarks were extracted from -- the prior version
    # loaded from t=0, pairing the wrong words to each mouth frame.
    # ponytail: librosa lazy-loads only [offset, offset+duration], it does
    # not decode the whole multi-hour track first.
    y, sr = librosa.load(video_path, sr=16000, mono=True,
                         offset=start_sec, duration=load_duration)

    # 3. Extract 80 Mel-Spectrogram Bins
    # hop_length=1333 aligns exactly with 12 FPS at 16kHz (16000 / 12 = 1333.33)
    mel_spec = librosa.feature.melspectrogram(
        y=y, sr=16000, n_fft=2048, hop_length=1333, n_mels=80
    )
    mel_db = librosa.power_to_db(mel_spec, ref=np.max).T # [Frames, 80]

    # Align Mel-Spectrogram frames to target V
    if mel_db.shape[0] > V:
        mel_db = mel_db[:V, :]
    elif mel_db.shape[0] < V:
        mel_db = np.pad(mel_db, ((0, V - mel_db.shape[0]), (0, 0)), mode='edge')

    # 4. Extract FCPE Pitch (F0) -- chunked + batched.
    # torchfcpe expects [Batch, Samples, Channels]. We split the clip's audio
    # into CHUNK_SECONDS windows and run MAX_BATCH windows per forward, so peak
    # GPU memory is bounded by one window x MAX_BATCH instead of the whole clip.
    sr = 16000
    chunk_samples = CHUNK_SECONDS * sr
    frames_per_chunk = CHUNK_SECONDS * 12  # 12 fps
    y_pad = np.pad(y, (0, (-len(y)) % chunk_samples), mode='constant')
    n_chunks = len(y_pad) // chunk_samples
    chunks = y_pad.reshape(n_chunks, chunk_samples).astype(np.float32)

    f0_parts = []
    for i in range(0, n_chunks, MAX_BATCH):
        batch = torch.from_numpy(chunks[i:i + MAX_BATCH]).to(device)
        batch = batch.unsqueeze(-1)  # [B, chunk_samples, 1]
        with torch.autocast('cuda', dtype=torch.bfloat16, enabled=USE_BF16):
            f0_b = fcpe_model.infer(
                batch,
                sr=sr,
                decoder_mode='local_argmax',
                threshold=0.006,
                f0_min=80,
                f0_max=880,
                interp_uv=False, # Keeps unvoiced segments as 0.0 Hz
                output_interp_target_length=frames_per_chunk, # per-window
            )
        f0_parts.append(f0_b.reshape(-1, 1).cpu().numpy()) # [B*frames_per_chunk, 1]
    f0_np = np.concatenate(f0_parts, axis=0)[:V]  # [V, 1] -- trim trailing pads
    if f0_np.shape[0] < V:  # pad if clip shorter than expected
        f0_np = np.pad(f0_np, ((0, V - f0_np.shape[0]), (0, 0)), mode='edge')

    # 5. Extract RMS Volume (Speech Energy)
    rms = librosa.feature.rms(y=y, frame_length=2048, hop_length=1333).T # [Frames, 1]

    # Align RMS frames to target V
    if rms.shape[0] > V:
        rms = rms[:V, :]
    elif rms.shape[0] < V:
        rms = np.pad(rms, ((0, V - rms.shape[0]), (0, 0)), mode='edge')

    # 6. Concatenate Features into [V, 82] Matrix
    audio_feats = np.concatenate([mel_db, f0_np, rms], axis=-1) # [V, 82]

    # 7. Save output
    # (output_path computed before #1 to support early-skip check above)
    np.save(output_path, audio_feats)
    processed_count += 1

print(f"\nSuccessfully generated {processed_count} files in {OUTPUT_DIR}")