prepare_data.py

"""풍소음 학습 데이터 준비 스크립트

깨끗한 음성과 풍소음 파일로부터 HDF5 학습 데이터를 생성합니다.
이 스크립트는 로컬에서 실행하여 데이터를 준비한 후 HF Hub에 업로드합니다.

사용법:
    python prepare_data.py --speech-dir ./data/speech \
                           --noise-dir ./data/noise \
                           --output-dir ./data/hdf5 \
                           [--rir-dir ./data/rir] \
                           [--sr 48000]
"""

import argparse
import json
import os
from pathlib import Path

import numpy as np
import soundfile as sf


def find_audio_files(directory, extensions=(".wav", ".flac", ".ogg", ".mp3")):
    """디렉토리에서 오디오 파일 경로 수집"""
    audio_files = []
    for ext in extensions:
        audio_files.extend(Path(directory).rglob(f"*{ext}"))
    return sorted(audio_files)


def split_stereo_to_mono(audio_files, output_dir):
    """Stereo 오디오 파일의 L/R 채널을 각각 mono WAV로 분리.

    DeepFilterNet3은 mono 모델이므로, stereo 파일을 채널별로 분리하여
    학습 데이터를 2배로 확보하는 효과를 얻음.
    이미 mono인 파일은 그대로 복사.

    Returns:
        분리된 mono 파일 경로 리스트
    """
    mono_dir = Path(output_dir) / "mono_split"
    mono_dir.mkdir(parents=True, exist_ok=True)

    mono_files = []
    for f in audio_files:
        info = sf.info(str(f))
        stem = f.stem

        if info.channels >= 2:
            data, sr = sf.read(str(f), always_2d=True)
            # Left channel
            left_path = mono_dir / f"{stem}_L.wav"
            sf.write(str(left_path), data[:, 0], sr, subtype="PCM_16")
            mono_files.append(left_path)
            # Right channel
            right_path = mono_dir / f"{stem}_R.wav"
            sf.write(str(right_path), data[:, 1], sr, subtype="PCM_16")
            mono_files.append(right_path)
        else:
            # Mono 파일은 그대로 복사
            mono_path = mono_dir / f"{stem}.wav"
            data, sr = sf.read(str(f))
            sf.write(str(mono_path), data, sr, subtype="PCM_16")
            mono_files.append(mono_path)

    print(f"  Stereo → Mono 분리 완료: {len(audio_files)}개 → {len(mono_files)}개 파일")
    return sorted(mono_files)


def validate_sample_rate(files, target_sr=48000, max_check=50):
    """샘플링 레이트 검증 (DeepFilterNet3은 48kHz 필수)"""
    issues = []
    for f in files[:max_check]:
        info = sf.info(str(f))
        if info.samplerate != target_sr:
            issues.append((str(f), info.samplerate))
    return issues


def create_file_list(audio_files, output_path):
    """오디오 파일 경로 목록 생성 (HDF5 변환에 사용)"""
    with open(output_path, "w") as f:
        for audio_file in audio_files:
            f.write(str(audio_file.resolve()) + "\n")
    print(f"  파일 목록 생성: {output_path} ({len(audio_files)}개 파일)")
    return output_path


def create_dataset_config(output_dir, has_rir=False):
    """dataset.cfg 생성 (DeepFilterNet 학습에 필요한 JSON 설정)"""
    config = {
        "train": [
            [os.path.join(output_dir, "train_speech.hdf5"), 1.0],
            [os.path.join(output_dir, "train_noise.hdf5"), 1.0],
        ],
        "valid": [
            [os.path.join(output_dir, "valid_speech.hdf5"), 1.0],
            [os.path.join(output_dir, "valid_noise.hdf5"), 1.0],
        ],
        "test": [
            [os.path.join(output_dir, "test_speech.hdf5"), 1.0],
            [os.path.join(output_dir, "test_noise.hdf5"), 1.0],
        ],
    }

    if has_rir:
        config["train"].append([os.path.join(output_dir, "train_rir.hdf5"), 1.0])
        config["valid"].append([os.path.join(output_dir, "valid_rir.hdf5"), 1.0])

    config_path = os.path.join(output_dir, "dataset.cfg")
    with open(config_path, "w") as f:
        json.dump(config, f, indent=2)

    print(f"  dataset.cfg 생성: {config_path}")
    return config_path


def split_files(files, train_ratio=0.8, valid_ratio=0.1):
    """파일을 train/valid/test로 분할"""
    n = len(files)
    n_train = int(n * train_ratio)
    n_valid = int(n * valid_ratio)

    np.random.seed(42)
    indices = np.random.permutation(n)

    train_files = [files[i] for i in indices[:n_train]]
    valid_files = [files[i] for i in indices[n_train:n_train + n_valid]]
    test_files = [files[i] for i in indices[n_train + n_valid:]]

    return train_files, valid_files, test_files


def print_summary(speech_files, noise_files, rir_files=None):
    """데이터 요약 출력"""
    print("\n" + "=" * 50)
    print("데이터 요약")
    print("=" * 50)

    total_speech_dur = 0
    for f in speech_files[:100]:
        info = sf.info(str(f))
        total_speech_dur += info.duration

    est_total = total_speech_dur / min(len(speech_files), 100) * len(speech_files)
    print(f"  음성 파일: {len(speech_files)}개 (예상 총 {est_total/3600:.1f}시간)")

    total_noise_dur = 0
    for f in noise_files[:100]:
        info = sf.info(str(f))
        total_noise_dur += info.duration

    est_total_n = total_noise_dur / min(len(noise_files), 100) * len(noise_files)
    print(f"  노이즈 파일: {len(noise_files)}개 (예상 총 {est_total_n/3600:.1f}시간)")

    if rir_files:
        print(f"  RIR 파일: {len(rir_files)}개")

    print()


def main():
    parser = argparse.ArgumentParser(description="DeepFilterNet3 파인튜닝용 데이터 준비")
    parser.add_argument("--speech-dir", required=True, help="깨끗한 음성 디렉토리")
    parser.add_argument("--noise-dir", required=True, help="풍소음 디렉토리")
    parser.add_argument("--rir-dir", default=None, help="RIR 디렉토리 (선택)")
    parser.add_argument("--output-dir", required=True, help="출력 디렉토리 (HDF5 + 설정파일)")
    parser.add_argument("--sr", type=int, default=48000, help="샘플링 레이트 (기본: 48000)")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    # 1. 오디오 파일 수집
    print("오디오 파일 수집 중...")
    speech_files = find_audio_files(args.speech_dir)
    noise_files = find_audio_files(args.noise_dir)
    rir_files = find_audio_files(args.rir_dir) if args.rir_dir else None

    if not speech_files:
        print(f"오류: {args.speech_dir}에서 음성 파일을 찾을 수 없습니다.")
        return
    if not noise_files:
        print(f"오류: {args.noise_dir}에서 노이즈 파일을 찾을 수 없습니다.")
        return

    # 2. 샘플링 레이트 검증
    print(f"\n샘플링 레이트 검증 ({args.sr}Hz)...")
    sr_issues = validate_sample_rate(speech_files, args.sr)
    sr_issues += validate_sample_rate(noise_files, args.sr)
    if sr_issues:
        print("  경고: 다음 파일들의 샘플링 레이트가 맞지 않습니다:")
        for path, sr in sr_issues[:10]:
            print(f"    {path}: {sr}Hz (필요: {args.sr}Hz)")
        print("  리샘플링이 필요합니다. librosa나 ffmpeg로 48kHz로 변환하세요.")
        print("  예: ffmpeg -i input.wav -ar 48000 output.wav")
        return

    print_summary(speech_files, noise_files, rir_files)

    # 2.5. Stereo → Mono 채널 분리 (DeepFilterNet3은 mono 모델)
    print("\nStereo → Mono 채널 분리 중...")
    print("  Speech 파일:")
    speech_files = split_stereo_to_mono(speech_files, args.output_dir)
    print("  Noise 파일:")
    noise_files = split_stereo_to_mono(noise_files, args.output_dir)
    if rir_files:
        print("  RIR 파일:")
        rir_files = split_stereo_to_mono(rir_files, args.output_dir)

    # 3. Train/Valid/Test 분할
    print("데이터 분할 중 (80/10/10)...")
    train_speech, valid_speech, test_speech = split_files(speech_files)
    train_noise, valid_noise, test_noise = split_files(noise_files)

    print(f"  Speech - train: {len(train_speech)}, valid: {len(valid_speech)}, test: {len(test_speech)}")
    print(f"  Noise  - train: {len(train_noise)}, valid: {len(valid_noise)}, test: {len(test_noise)}")

    # 4. 파일 목록 생성
    print("\n파일 목록 생성 중...")
    lists_dir = os.path.join(args.output_dir, "lists")
    os.makedirs(lists_dir, exist_ok=True)

    for split, files in [("train", train_speech), ("valid", valid_speech), ("test", test_speech)]:
        create_file_list(files, os.path.join(lists_dir, f"{split}_speech.txt"))
    for split, files in [("train", train_noise), ("valid", valid_noise), ("test", test_noise)]:
        create_file_list(files, os.path.join(lists_dir, f"{split}_noise.txt"))

    if rir_files:
        train_rir, valid_rir, _ = split_files(rir_files)
        create_file_list(train_rir, os.path.join(lists_dir, "train_rir.txt"))
        create_file_list(valid_rir, os.path.join(lists_dir, "valid_rir.txt"))

    # 5. dataset.cfg 생성
    print("\ndataset.cfg 생성 중...")
    create_dataset_config(args.output_dir, has_rir=rir_files is not None)

    # 6. HDF5 변환 안내
    print("\n" + "=" * 50)
    print("다음 단계: HDF5 변환")
    print("=" * 50)
    print()
    print("DeepFilterNet의 prepare_data.py 스크립트로 HDF5를 생성하세요:")
    print("(Docker 컨테이너 내부 또는 Linux 환경에서 실행)")
    print()

    dfn_script = "python /opt/DeepFilterNet/DeepFilterNet/df/scripts/prepare_data.py"
    for split in ["train", "valid", "test"]:
        for dtype in ["speech", "noise"]:
            txt_path = os.path.join(lists_dir, f"{split}_{dtype}.txt")
            hdf5_path = os.path.join(args.output_dir, f"{split}_{dtype}.hdf5")
            print(f"  {dfn_script} --sr {args.sr} {dtype} {txt_path} {hdf5_path}")
    print()

    if rir_files:
        for split in ["train", "valid"]:
            txt_path = os.path.join(lists_dir, f"{split}_rir.txt")
            hdf5_path = os.path.join(args.output_dir, f"{split}_rir.hdf5")
            print(f"  {dfn_script} --sr {args.sr} rir {txt_path} {hdf5_path}")
        print()

    print("HDF5 생성 후 HF Hub에 업로드:")
    print(f"  huggingface-cli upload YOUR_USERNAME/wind-noise-data {args.output_dir} --repo-type dataset")
    print()
    print("완료!")


if __name__ == "__main__":
    main()