batch_process.py

#!/usr/bin/env python3
"""
Batch audio processing pipeline – dostosowane do workflow Roberta (AI Music Creator Aleks).

Funkcje:
 - batch konwersja WAV -> MP3 (ffmpeg)
 - normalizacja RMS (librosa)
 - ekstrakcja stemów przez Spleeter (obsługa: lokalne venv CLI lub Docker)
 - konwersja stemów do MIDI przez Basic Pitch (python API lub CLI fallback)
 - estymacja tempo/key (librosa)
 - zapis spectrogramów PNG (librosa/matplotlib)
 - zapis metadata.json per-track
"""

import argparse
import json
import shutil
import subprocess
from pathlib import Path
from typing import Optional

# audio libs
import numpy as np
import soundfile as sf
import librosa
import librosa.display
import matplotlib
import matplotlib.pyplot as plt
matplotlib.use('GTK3Agg')

# Config
DEFAULT_TARGET_RMS_DB = -20.0
FFMPEG_BIN = shutil.which("ffmpeg") or "ffmpeg"
SPL_VENV_CMD = ["spleeter", "separate", "-p", "spleeter:5stems", "-o"]  # if spleeter installed in venv
SPL_DOCKER_IMAGE = "researchdeezer/spleeter:2.1.1"
DOCKER_BIN = shutil.which("docker") or "docker"

def run_cmd(cmd, check=True):
    print("> RUN:", " ".join(cmd))
    return subprocess.run(cmd, check=check)
    
def ensure_dir(p: Path):
    p.mkdir(parents=True, exist_ok=True)
    
# Audio ops
def load_audio(path: Path):
    y, sr = librosa.load(str(path), sr=None, mono=True)
    return y, sr
    
def write_wav(path: Path, y, sr):
    sf.write(str(path), y, sr)
    
def normalize_rms(y: np.ndarray, target_db: float = DEFAULT_TARGET_RMS_DB):
    target = 10 ** (target_db / 20)
    rms = np.sqrt(np.mean(y ** 2))
    if rms == 0:
        return y
    return y * (target / (rms + 1e-9))
    
def wav_to_mp3(input_wav: Path, output_mp3: Path, bitrate="192k"):
    cmd = [FFMPEG_BIN, "-y", "-i", str(input_wav), "-b:a", bitrate, str(output_mp3)]
    run_cmd(cmd)
    
def save_spectrogram(path: Path, out_png: Path, n_fft=2048, hop_length=512):
    y, sr = load_audio(path)
    S = np.abs(librosa.stft(y, n_fft=n_fft, hop_length=hop_length))
    S_db = librosa.amplitude_to_db(S, ref=np.max)
    try:
        plt.figure(figsize=(8, 3.5))
        librosa.display.specshow(S_db, sr=sr, x_axis="time", y_axis="log", hop_length=hop_length)
        plt.colorbar(format="%+2.0f dB")
        plt.title(path.name)
        plt.tight_layout()
        plt.savefig(str(out_png))
        plt.close()
    except Exception as e:
        print("[exception] ", e)
    
def estimate_key_tempo(path: Path):
    y, sr = load_audio(path)
    tempo, _ = librosa.beat.beat_track(y=y, sr=sr)
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr)
    pitch_class = int(np.argmax(chroma.mean(axis=1)))
    keys = ["C","C#","D","D#","E","F","F#","G","G#","A","A#","B"]
    return {"tempo": float(tempo), "key": keys[pitch_class]}
    
# Spleeter integration: try venv CLI first, else Docker
def spleeter_separate(input_wav: Path, out_dir: Path, use_docker: bool = False):
    ensure_dir(out_dir)
    if not use_docker:
        cmd = SPL_VENV_CMD + [str(out_dir), str(input_wav)]
        run_cmd(cmd)
    else:
        input_parent = input_wav.parent.resolve()
        out_parent = out_dir.resolve()
        container_cmd = [
            DOCKER_BIN, "run", "--rm",
            "-v", f"{input_parent}:/input",
            "-v", f"{out_parent}:/output",
            SPL_DOCKER_IMAGE,
            "separate", "-p", "spleeter:5stems",
            "-o", "/output",
            f"/input/{input_wav.name}"
        ]
        run_cmd(container_cmd)
        
# Basic Pitch conversion with python API fallback to CLI
def basicpitch_convert_to_midi(stem_wav: Path, out_midi: Path):
    """
    Try:
     1) Python API: basic_pitch.inference.predict_and_save_midi(in_wav, out_midi)
     2) Alternative API signatures
     3) CLI 'basic_pitch' or 'basic-pitch' if available
    """
    # Try python API
    try:
        from basic_pitch.inference import predict_and_save_midi  # type: ignore
        try:
            predict_and_save_midi(str(stem_wav), str(out_midi))
            print("BasicPitch (API) saved MIDI:", out_midi)
            return
        except TypeError:
            try:
                # alternative keyword args if present in some versions
                predict_and_save_midi(audio_path=str(stem_wav), midi_path=str(out_midi))
                print("BasicPitch (API alt) saved MIDI:", out_midi)
                return
            except Exception as e:
                print("BasicPitch API call failed (alt):", e)
        except Exception as e:
            print("BasicPitch API call failed:", e)
    except Exception as e:
        print("BasicPitch Python API import failed:", e)
        
    # Fallback: try CLI
    basic_pitch_cli = shutil.which("basic_pitch") or shutil.which("basic-pitch")
    if basic_pitch_cli:
        try:
            run_cmd([basic_pitch_cli, str(stem_wav), str(out_midi)])
            print("BasicPitch (CLI) saved MIDI:", out_midi)
            return
        except Exception as e:
            print("BasicPitch CLI call failed:", e)
            
    raise RuntimeError("BasicPitch not available (python API import failed and CLI not found).")
    
# Main processing per file
def process_file(filepath: Path, outroot: Path, use_spleeter_docker: bool = False,
                 do_basicpitch: bool = True, normalize_db: float = DEFAULT_TARGET_RMS_DB):
    ensure_dir(outroot)
    base = filepath.stem
    workdir = outroot / base
    ensure_dir(workdir)
    
    # 1) load + normalize + save normalized wav
    y, sr = load_audio(filepath)
    y_norm = normalize_rms(y, normalize_db)
    norm_path = workdir / f"{base}_normalized.wav"
    write_wav(norm_path, y_norm, sr)
    print("Normalized saved:", norm_path.name)
    
    # 2) convert to mp3
    mp3_out = workdir / f"{base}.mp3"
    wav_to_mp3(norm_path, mp3_out)
    
    # 3) spectrogram of mix
    save_spectrogram(norm_path, workdir / f"{base}_spectrogram.png")
    
    # 4) metadata
    meta = {"source": filepath.name, "sr": sr, "outputs": {}}
    mix_meta = estimate_key_tempo(norm_path)
    meta["mix"] = mix_meta
    
    # 5) spleeter stems
    spleeter_out_parent = workdir / "spleeter_out"
    try:
        spleeter_separate(norm_path, spleeter_out_parent, use_docker=use_spleeter_docker)
    except Exception as e:
        print("Spleeter failed:", e)
        spleeter_out_parent = None
        
    if spleeter_out_parent:
        candidate = spleeter_out_parent / filepath.stem
        if not candidate.exists():
            candidate = spleeter_out_parent
        stems_meta = {}
        for stem_wav in sorted(candidate.glob("*.wav")):
            stem_name = stem_wav.stem
            stems_meta[stem_name] = {"path": str(stem_wav)}
            try:
                save_spectrogram(stem_wav, workdir / f"{stem_name}_spectrogram.png")
            except Exception as e:
                print("Spectrogram failed for", stem_name, e)
            try:
                kt = estimate_key_tempo(stem_wav)
                stems_meta[stem_name].update(kt)
            except Exception:
                pass
            if do_basicpitch:
                midi_out = workdir / f"{stem_name}.midi"
                try:
                    basicpitch_convert_to_midi(stem_wav, midi_out)
                    stems_meta[stem_name]["midi"] = str(midi_out)
                except Exception as e:
                    print("BasicPitch failed for", stem_name, e)
        meta["stems"] = stems_meta
        
    meta_path = workdir / "metadata.json"
    with open(meta_path, "w", encoding="utf-8") as fh:
        json.dump(meta, fh, indent=2, ensure_ascii=False)
    print("Saved metadata:", meta_path)

# CLI
def main():
    parser = argparse.ArgumentParser(description="Batch audio pipeline (Robert / AI Music Creator Aleks)")
    parser.add_argument("-i", "--input", required=True, help="Input WAV file or folder")
    parser.add_argument("-o", "--output", required=True, help="Output root folder")
    parser.add_argument("--no-basicpitch", action="store_true", help="Disable Basic Pitch MIDI conversion")
    parser.add_argument("--spleeter-docker", action="store_true", help="Run Spleeter via Docker (if not present in venv)")
    parser.add_argument("--target-rms-db", type=float, default=DEFAULT_TARGET_RMS_DB, help="Target RMS dB (default -20)")
    args = parser.parse_args()
    
    inp = Path(args.input)
    out = Path(args.output)
    ensure_dir(out)
    
    files = []
    if inp.is_dir():
        for p in sorted(inp.glob("*.wav")):
            files.append(p)
    elif inp.is_file():
        files.append(inp)
    else:
        print("No input files found.")
        return
        
    for f in files:
        try:
            process_file(f, out, use_spleeter_docker=args.spleeter_docker,
                         do_basicpitch=not args.no_basicpitch,
                         normalize_db=args.target_rms_db)
        except Exception as e:
            print("Error processing", f, e)
        
if __name__ == "__main__":
    main()