nc-murray/spectrogram-reconstruction
0
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 