Team Ai
Apppublic

nc-murray/spectrogram-reconstruction

sourceHugging Faceupdated 4mo agoView on Hugging Face
0likes
cli.py224 linesDownload Raw Back to specrec
1"""specrec — command-line interface."""2 3import os4import sys5import json6import click7import numpy as np8import soundfile as sf9 10 11@click.group()12def cli():13    """Reconstruct audio from spectrogram images using Griffin-Lim."""14 15 16# ---------------------------------------------------------------------------17# specrec reconstruct18# ---------------------------------------------------------------------------19 20@cli.command()21@click.option("--input",  "image_path",  required=True,  type=click.Path(exists=True), help="Spectrogram image (PNG/JPEG).")22@click.option("--output", "output_path", required=True,  type=click.Path(),            help="Output WAV file.")23@click.option("--colormap", default=None, show_default=True, help="Colormap hint (viridis, jet, magma, …). Auto-detected if omitted.")24@click.option("--freq-max", default=11025.0, show_default=True, help="Highest frequency in image (Hz).")25@click.option("--db-min",   default=-80.0,   show_default=True, help="dB floor used when rendering the image.")26@click.option("--db-max",   default=0.0,     show_default=True, help="dB ceiling used when rendering the image.")27@click.option("--sr",       default=22050,   show_default=True, help="Sample rate (Hz).")28@click.option("--n-fft",    default=2048,    show_default=True, help="FFT size.")29@click.option("--hop",      default=512,     show_default=True, help="Hop length.")30@click.option("--n-iter",   default=60,      show_default=True, help="Griffin-Lim iterations.")31@click.option("--duration", default=None,    type=float,        help="Original audio duration (s). Resizes time axis when provided.")32def reconstruct(image_path, output_path, colormap, freq_max, db_min, db_max,33                sr, n_fft, hop, n_iter, duration):34    """Reconstruct audio from a spectrogram image."""35    from .reconstructor import reconstruct_from_image36 37    click.echo(f"Parsing {image_path} …")38    audio = reconstruct_from_image(39        image_path, output_path,40        colormap=colormap,41        freq_max=freq_max,42        db_min=db_min, db_max=db_max,43        sr=sr, n_fft=n_fft, hop_length=hop,44        n_iter=n_iter, duration=duration,45    )46    click.echo(f"Saved {len(audio)/sr:.2f}s of audio → {output_path}")47 48 49# ---------------------------------------------------------------------------50# specrec test51# ---------------------------------------------------------------------------52 53@cli.command("test")54@click.option("--type", "signal_type", type=click.Choice(["tone", "speech_like"]),55              default="tone", show_default=True, help="Signal type to generate.")56@click.option("--output-dir", default=".", show_default=True, type=click.Path(), help="Directory for output files.")57@click.option("--colormap", default="viridis", show_default=True)58@click.option("--sr", default=22050, show_default=True)59def test_cmd(signal_type, output_dir, colormap, sr):60    """Generate synthetic test audio and its spectrogram image."""61    from .synthesizer import generate_test_tone, generate_test_speech_like, audio_to_spectrogram_image62 63    os.makedirs(output_dir, exist_ok=True)64 65    if signal_type == "tone":66        audio = generate_test_tone(sr=sr)67        stem = "tone_440hz"68    else:69        audio = generate_test_speech_like(sr=sr)70        stem = "speech_like"71 72    wav_path = os.path.join(output_dir, f"{stem}.wav")73    img_path = os.path.join(output_dir, f"{stem}_{colormap}.png")74 75    sf.write(wav_path, audio, sr)76    audio_to_spectrogram_image(audio, sr, img_path, colormap=colormap)77 78    click.echo(f"Audio   → {wav_path}")79    click.echo(f"Image   → {img_path}")80 81 82# ---------------------------------------------------------------------------83# specrec evaluate84# ---------------------------------------------------------------------------85 86@cli.command()87@click.option("--input",     "image_path",     required=True, type=click.Path(exists=True))88@click.option("--reference", "reference_path", required=True, type=click.Path(exists=True),89              help="Original audio WAV for comparison.")90@click.option("--colormap", default=None)91@click.option("--freq-max", default=11025.0, show_default=True)92@click.option("--db-min",   default=-80.0,   show_default=True)93@click.option("--db-max",   default=0.0,     show_default=True)94@click.option("--sr",       default=22050,   show_default=True)95@click.option("--n-fft",    default=2048,    show_default=True)96@click.option("--hop",      default=512,     show_default=True)97@click.option("--n-iter",   default=60,      show_default=True)98def evaluate(image_path, reference_path, colormap, freq_max, db_min, db_max,99             sr, n_fft, hop, n_iter):100    """Reconstruct from an image and compare against a reference WAV."""101    from .reconstructor import reconstruct_from_image102    from .evaluator import compute_spectral_convergence, compute_snr103 104    ref, ref_sr = sf.read(reference_path, dtype="float32")105    if ref.ndim > 1:106        ref = ref[:, 0]107    if ref_sr != sr:108        click.echo(f"Warning: reference SR={ref_sr}, expected {sr}. Proceeding anyway.", err=True)109 110    duration = len(ref) / sr111    audio = reconstruct_from_image(112        image_path, output_path=None,113        colormap=colormap, freq_max=freq_max,114        db_min=db_min, db_max=db_max,115        sr=sr, n_fft=n_fft, hop_length=hop,116        n_iter=n_iter, duration=duration,117    )118 119    sc  = compute_spectral_convergence(ref, audio, sr, n_fft, hop)120    snr = compute_snr(ref, audio)121 122    click.echo(f"Spectral convergence : {sc:.4f}  (lower = better)")123    click.echo(f"SNR                  : {snr:.2f} dB  (higher = better)")124 125 126# ---------------------------------------------------------------------------127# specrec demo128# ---------------------------------------------------------------------------129 130@cli.command()131@click.option("--output-dir", default="examples/synthetic_demo", show_default=True,132              type=click.Path(), help="Directory to write all demo outputs.")133@click.option("--sr", default=22050, show_default=True)134def demo(output_dir, sr):135    """136    Full synthetic round-trip demo.137 138    Generates test signals, renders spectrograms, reconstructs audio,139    runs accuracy metrics, and writes a JSON summary report.140    """141    from .synthesizer import generate_test_tone, generate_test_speech_like, audio_to_spectrogram_image142    from .evaluator import run_roundtrip_accuracy_test, run_colormap_sensitivity_test143    from .visualizer import plot_reconstruction_comparison, plot_accuracy_vs_iterations, plot_colormap_sensitivity144 145    os.makedirs(output_dir, exist_ok=True)146 147    report = {}148 149    for label, audio in [150        ("tone_440hz",   generate_test_tone(sr=sr)),151        ("speech_like",  generate_test_speech_like(sr=sr)),152    ]:153        click.echo(f"\n── {label} ──")154 155        # Save source audio156        wav_path = os.path.join(output_dir, f"{label}_original.wav")157        sf.write(wav_path, audio, sr)158        click.echo(f"  Original audio    → {wav_path}")159 160        # Save spectrogram image (viridis)161        img_path = os.path.join(output_dir, f"{label}_viridis.png")162        audio_to_spectrogram_image(audio, sr, img_path, colormap="viridis")163        click.echo(f"  Spectrogram image → {img_path}")164 165        # Round-trip accuracy across n_iter values166        click.echo("  Running round-trip accuracy test …")167        rt_results = run_roundtrip_accuracy_test(audio, sr, n_iter_values=[10, 30, 60, 100])168 169        click.echo(f"  {'n_iter':>7}  {'spectral_conv':>14}  {'SNR (dB)':>10}")170        click.echo(f"  {'-'*36}")171        for n_iter, m in sorted(rt_results.items()):172            click.echo(f"  {n_iter:>7}  {m['spectral_convergence']:>14.4f}  {m['snr_db']:>10.2f}")173 174        # Save best reconstruction (n_iter=100)175        from .reconstructor import reconstruct_from_image176        recon_path = os.path.join(output_dir, f"{label}_reconstructed.wav")177        recon_audio = reconstruct_from_image(178            img_path, recon_path,179            colormap="viridis", freq_max=sr / 2,180            duration=len(audio) / sr, n_iter=100, sr=sr,181        )182        click.echo(f"  Reconstructed WAV → {recon_path}")183 184        # Comparison plot185        comp_path = os.path.join(output_dir, f"plot_comparison_{label}.png")186        plot_reconstruction_comparison(187            audio, recon_audio, sr,188            title=f"{label.replace('_', ' ').title()} — Original vs Reconstructed",189            output_path=comp_path,190        )191        click.echo(f"  Comparison plot   → {comp_path}")192 193        # Accuracy vs iterations plot194        iter_path = os.path.join(output_dir, f"plot_accuracy_vs_iter_{label}.png")195        plot_accuracy_vs_iterations(196            rt_results,197            title=f"{label.replace('_', ' ').title()} — Accuracy vs Iterations",198            output_path=iter_path,199        )200        click.echo(f"  Accuracy plot     → {iter_path}")201 202        # Colormap sensitivity203        click.echo("  Running colormap sensitivity test …")204        cmap_results = run_colormap_sensitivity_test(audio, sr)205        cmap_path = os.path.join(output_dir, f"plot_colormap_sensitivity_{label}.png")206        plot_colormap_sensitivity(207            cmap_results,208            title=f"{label.replace('_', ' ').title()} — Colormap Sensitivity",209            output_path=cmap_path,210        )211        click.echo(f"  Colormap plot     → {cmap_path}")212 213        report[label] = {214            "roundtrip": {str(k): v for k, v in rt_results.items()},215            "colormap_sensitivity": cmap_results,216        }217 218    # Write JSON summary219    report_path = os.path.join(output_dir, "accuracy_report.json")220    with open(report_path, "w") as f:221        json.dump(report, f, indent=2)222    click.echo(f"\nSummary report → {report_path}")223    click.echo("Demo complete.")224