""" falcon_viz.py — Time-aligned alignment visualizations for FALCON. Two entry points, both sharing the same per-panel drawers so the styling stays identical: make_alignment_panels(wav, ckpt, out_path, ...) -> ONE tall figure, 5 panels stacked on a shared time axis (for the README / a downloadable overview). make_alignment_separate(wav, ckpt, out_dir, ...) -> a LIST of (path, caption), one wide, full-size figure per panel (for the web app, where a single stacked figure renders too small to read). Panels: waveform, log-mel spectrogram, phoneme posteriors, Soft-DP cost matrix + alignment path, and the contrastive boundary score with its detected peaks. The same predicted boundaries (crimson) and truth boundaries (charcoal) are drawn on every panel so the time alignment is visually verifiable. This module does NOT modify any model/training file. It re-runs the model forward (replicating predict.py:main_predict) and reproduces utils.phoneme_alignment's DP matrix locally (re-using utils.compute_phi_1 / compute_phi_2 unchanged). """ import os from glob import glob import numpy as np import torch import torch.nn.functional as F import torchaudio import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt from matplotlib.gridspec import GridSpec from dataloader import spectral_size from predict import _load_model from utils import ( max_min_norm, get_timit_61_phoneme_mappings, compute_phi_1, compute_phi_2, phoneme_to_idx_MACRO, timit_to_leehon_map_MACRO, timit_leehon_39_phonemes, ) # Hard-coded in the model forward (next_frame_classifier.py) — CNN total stride. # One latent frame ~= len_ratio audio samples ~= 10 ms at 16 kHz. LEN_RATIO = 161.34011627906978 SR = 16000 # Light / pastel palette so the crimson boundary overlays stand out. PRED_COLOR = "crimson" TRUTH_COLOR = "#37474f" # charcoal (distinct from the blue/purple cmaps) WAVE_COLOR = "#3b6fb6" # blue (matches the example waveform figure) SCORE_COLOR = "#3a7ca5" SPEC_CMAP = "viridis" POST_CMAP = "viridis" DP_CMAP = "viridis" # --------------------------------------------------------------------------- # # DP matrix reconstruction (mirrors utils.phoneme_alignment, returns matrix) # # --------------------------------------------------------------------------- # def _build_dp_matrix(p_seq, w_phi, original_lengths, derivative_preds_np, probs_real): """Re-run the exact Soft-DP forward + soft-argmax backtrack from utils.phoneme_alignment, additionally returning the DP matrix so it can be plotted (the stock function only returns the path).""" gamma = 1e-20 T = int(original_lengths[0]) n = len(p_seq) device = derivative_preds_np.device if isinstance(probs_real, np.ndarray): probs_real = torch.tensor(probs_real, device=device) cumsum_probs = torch.cumsum(probs_real, dim=0) phoneme_mappings = { p.lower(): timit_to_leehon_map_MACRO.get(p.lower(), "sil") if p.lower() not in timit_leehon_39_phonemes else p.lower() for p in p_seq } derivatives = torch.cat( [torch.tensor([0], device=device), torch.diff(derivative_preds_np, dim=0)] ) dp_mat = torch.full((n, T, T), float(-1e9), device=device) t_e = torch.arange(T, device=device) dp_mat[0, 0, :] = ( w_phi[0] * compute_phi_1(derivatives, 0, t_e) + w_phi[1] * compute_phi_1(derivatives, 0, t_e) ) for i in range(1, n): p_idx = phoneme_to_idx_MACRO[phoneme_mappings[p_seq[i].lower()]] t_start = torch.arange(T, device=device) t_end = torch.arange(T, device=device) t_start_grid, t_end_grid = torch.meshgrid(t_start, t_end, indexing="ij") valid_mask = t_start_grid < t_end_grid phi1_dev = compute_phi_1(derivatives, t_start_grid, t_end_grid) phi2 = compute_phi_2(cumsum_probs, p_idx, t_start_grid, t_end_grid) total_phi = w_phi[0] * phi1_dev + w_phi[1] * phi2 col_lse = torch.logsumexp(dp_mat[i - 1] / gamma, dim=0) * gamma prev_scores = torch.where( valid_mask, col_lse.unsqueeze(1).expand(T, T), torch.full((T, T), float(-1e9), device=device), ) dp_mat[i] = torch.where( valid_mask, total_phi + prev_scores, torch.full_like(total_phi, float(-1e9)) ) best_start_times = torch.zeros((n), dtype=derivative_preds_np.dtype, device=device) best_prev_t_end = T - 1 for i in range(n): cur_ph = n - 1 - i scores = dp_mat[cur_ph, :, best_prev_t_end] soft_weights = torch.softmax(scores / gamma, dim=0) expected_idx = ( soft_weights * torch.arange(T, device=device, dtype=derivative_preds_np.dtype) ).sum() best_start_times[cur_ph] = expected_idx best_prev_t_end = int(expected_idx.round().item()) dp_to_plot = dp_mat.detach().cpu().max(dim=1)[0].numpy() # (n, T) best_start_frames = best_start_times.detach().cpu().numpy() return dp_to_plot, best_start_frames # --------------------------------------------------------------------------- # # Forward pass + array extraction (replicates predict.py:main_predict) # # --------------------------------------------------------------------------- # def _extract_arrays(wav, ckpt, annotation): model, _peak_params = _load_model(ckpt) model.eval() audio, sr = torchaudio.load(wav) assert sr == SR, "model expects 16 kHz audio" audio = audio[0] audio_len = len(audio) spectral_len = spectral_size(audio_len) len_ratio = audio_len / spectral_len # ~= LEN_RATIO base_dir = os.path.dirname(wav) base_name = os.path.basename(wav).split(".")[0] matches = glob(os.path.join(base_dir, f"{base_name}*.{annotation}")) phn_path = matches[0] if matches else wav.replace("wav", "phn") with open(phn_path, "r") as f: lines = [ln.split() for ln in f.readlines()] truth_secs = [float(ln[1]) / SR for ln in lines][:-1] phonemes = [ln[2].strip() for ln in lines] truth_labels = list(phonemes) length = [audio_len / len_ratio] with torch.no_grad(): preds, original_lengths, probs, frame_labels, seg, total_peaks, w_phi = model( audio.unsqueeze(0), None, [phonemes], length ) # Contrastive / latent boundary score (predict.py 168-172). p = preds[1][0] p = max_min_norm(p) p_np = p.detach().numpy()[0] p_np = p_np - np.median(p_np) # Phoneme posteriors (predict.py 143). probs_real = F.softmax(probs, dim=-1).squeeze(0).detach().numpy() # (T, 39) _, idx_to_phoneme = get_timit_61_phoneme_mappings() phoneme_labels = [idx_to_phoneme[i] for i in range(39)] pred_secs = list(total_peaks[0]) # Phoneme labels at predicted-segment midpoints (segment i = [bound[i], # bound[i+1]] with bounds = predicted boundaries + the utterance end), matching # how the TextGrid assigns each phoneme to a predicted interval. bounds = [float(x) for x in pred_secs] + [audio_len / SR] seg_mids = [(bounds[i] + bounds[i + 1]) / 2.0 for i in range(len(bounds) - 1)] n_lab = min(len(seg_mids), len(phonemes)) label_mids = seg_mids[:n_lab] label_text = phonemes[:n_lab] # DP matrix (reproduce phoneme_alignment locally). w_phi_vec = w_phi.detach() deriv_arg = torch.tensor(p_np, dtype=torch.float32) dp_to_plot, dp_path_frames = _build_dp_matrix( phonemes, w_phi_vec, [int(original_lengths[0])], deriv_arg, probs_real ) return dict( audio=audio.numpy(), sr=sr, len_ratio=len_ratio, duration=audio_len / SR, latent_score=p_np, probs_real=probs_real, phoneme_labels=phoneme_labels, pred_secs=pred_secs, truth_secs=truth_secs, truth_labels=truth_labels, dp_to_plot=dp_to_plot, dp_path_frames=dp_path_frames, phonemes=phonemes, label_mids=label_mids, label_text=label_text, ) # --------------------------------------------------------------------------- # # Shared per-panel drawers # # --------------------------------------------------------------------------- # def _overlay_boundaries(ax, pred_secs, truth_secs, show_truth, label_first=False): """Predicted (crimson dashed) + optional truth (charcoal dotted) lines.""" for j, t in enumerate(truth_secs if show_truth else []): ax.axvline(t, color=TRUTH_COLOR, linestyle=":", linewidth=0.9, alpha=0.6, label="Truth boundary" if (label_first and j == 0) else None, zorder=2) for j, t in enumerate(pred_secs): ax.axvline(t, color=PRED_COLOR, linestyle="--", linewidth=1.1, alpha=0.9, label="Predicted boundary" if (label_first and j == 0) else None, zorder=3) def _annotate_phoneme_tier(ax, d): """Write the input phoneme labels at their predicted-segment midpoints, just below the x-axis (a phoneme tier).""" mids = d.get("label_mids") or [] labels = d.get("label_text") or [] trans = ax.get_xaxis_transform() # x in data coords, y in axes fraction for m, lab in zip(mids, labels): ax.text(m, -0.07, str(lab), transform=trans, rotation=90, ha="center", va="top", fontsize=5.5, color="#333333", clip_on=False) def _panel_waveform(ax, d, show_truth, label_first=True): audio = d["audio"] dur = d["duration"] t = np.linspace(0, dur, num=len(audio)) ax.plot(t, audio, color=WAVE_COLOR, linewidth=0.5) ax.set_ylabel("Amplitude") ax.margins(x=0) ymax = (float(np.abs(audio).max()) or 1.0) * 1.15 ax.set_ylim(-ymax, ymax) _overlay_boundaries(ax, d["pred_secs"], d["truth_secs"], show_truth, label_first=label_first) if label_first: ax.legend(loc="upper right", fontsize=8, framealpha=0.9, ncol=2) def _panel_spectrogram(ax, d, show_truth): audio = d["audio"] dur = d["duration"] mel = torchaudio.transforms.MelSpectrogram( sample_rate=SR, n_fft=400, hop_length=160, n_mels=80 )(torch.tensor(audio).float().unsqueeze(0)) mel_db = torchaudio.transforms.AmplitudeToDB(top_db=80)(mel).squeeze(0).numpy() ax.imshow(mel_db, aspect="auto", origin="lower", extent=[0, dur, 0, SR / 2 / 1000.0], cmap=SPEC_CMAP) ax.set_ylabel("Freq (kHz)") _overlay_boundaries(ax, d["pred_secs"], d["truth_secs"], show_truth) def _panel_posteriors(ax, d, show_truth, colorbar=True): probs_real = d["probs_real"] dur = d["duration"] im = ax.imshow(probs_real.T, aspect="auto", origin="lower", extent=[0, dur, -0.5, 38.5], cmap=POST_CMAP, interpolation="nearest") ax.set_yticks(range(39)) ax.set_yticklabels(d["phoneme_labels"], fontsize=5.5) ax.set_ylabel("Phoneme (LH-39)") if colorbar: ax.figure.colorbar(im, ax=ax, label="P(phoneme)", pad=0.01, fraction=0.025) _overlay_boundaries(ax, d["pred_secs"], d["truth_secs"], show_truth) def _panel_dp(ax, d, show_truth, colorbar=True): dp_to_plot = d["dp_to_plot"] dur = d["duration"] len_ratio = d["len_ratio"] masked = np.ma.masked_where(dp_to_plot <= -1e8, dp_to_plot) cmap = getattr(plt.cm, DP_CMAP).copy() cmap.set_bad(color="white") n_ph = dp_to_plot.shape[0] im = ax.imshow(masked, aspect="auto", origin="lower", extent=[0, dur, -0.5, n_ph - 0.5], cmap=cmap, interpolation="nearest") path_secs = np.asarray(d["dp_path_frames"]) * len_ratio / SR ax.plot(path_secs, np.arange(n_ph), color=PRED_COLOR, marker=".", markersize=4, linewidth=1.2, label="Optimal alignment path") ax.set_ylabel("Phoneme position") if colorbar: ax.figure.colorbar(im, ax=ax, label="DP score", pad=0.01, fraction=0.025) _overlay_boundaries(ax, d["pred_secs"], d["truth_secs"], show_truth) ax.legend(loc="lower right", fontsize=8, framealpha=0.9) def _panel_contrastive(ax, d, show_truth): # Boundary score (red) + its derivative (magenta), as in predict.py's run plot, # but (1) auto-scaled robustly so the per-boundary structure is visible instead # of being squashed by the silence->speech onset spike, and (2) without the # redundant predicted-boundary markers — the predicted boundaries are the red # dashed lines shared across every panel. Ground truth omitted (not available # at inference). x-axis is time (s). s = np.asarray(d["latent_score"], dtype=float) len_ratio = d["len_ratio"] n = len(s) t = np.arange(n) * len_ratio / SR deriv = np.concatenate([[0.0], np.diff(s)]) ax.plot(t, deriv, marker="o", markersize=1.6, linewidth=0.7, alpha=0.8, color="magenta", label="Derivative of latent score") ax.plot(t, s, marker="*", markersize=2.2, linewidth=0.8, color="red", label="Latent score") # Robust symmetric y-limit: zoom to the 96th-percentile magnitude so the small # per-boundary structure fills the panel; the rare large onset spike clips off. mag = np.concatenate([np.abs(s), np.abs(deriv)]) A = max(float(np.percentile(mag, 96)) * 1.5, 0.03) if mag.size else 0.1 ax.set_ylim(-A, A) ax.set_ylabel("Score") ax.margins(x=0) _overlay_boundaries(ax, d["pred_secs"], [], show_truth=False) ax.legend(loc="upper right", fontsize=7, framealpha=0.9, ncol=2) # Order shared by both makers: (key, caption, drawer, separate-figure height). _PANELS = [ ("waveform", "1. Waveform", _panel_waveform, 2.6), ("spectrogram", "2. Log-mel spectrogram", _panel_spectrogram, 2.8), ("posteriors", "3. Phoneme posteriors", _panel_posteriors, 3.6), ("dp", "4. Soft-DP cost matrix + alignment path", _panel_dp, 3.4), ("contrastive", "5. Contrastive boundary score", _panel_contrastive, 2.6), ] # --------------------------------------------------------------------------- # # Combined figure (README) # # --------------------------------------------------------------------------- # def make_alignment_panels(wav, ckpt, out_path, w_phi=0.5, language="english", annotation="phn", show_truth=True): """Build the stacked, time-aligned multi-panel figure and save it to out_path.""" if language != "english": print(f"[falcon_viz] language='{language}' not supported; using english path.") d = _extract_arrays(wav, ckpt, annotation) dur = d["duration"] fig = plt.figure(figsize=(12, 17), dpi=150, constrained_layout=True) gs = GridSpec(5, 1, figure=fig, height_ratios=[1.0, 1.3, 1.7, 1.6, 1.0]) ax_wave = fig.add_subplot(gs[0]) ax_spec = fig.add_subplot(gs[1], sharex=ax_wave) ax_post = fig.add_subplot(gs[2], sharex=ax_wave) ax_dp = fig.add_subplot(gs[3], sharex=ax_wave) ax_score = fig.add_subplot(gs[4], sharex=ax_wave) axes = [ax_wave, ax_spec, ax_post, ax_dp, ax_score] _panel_waveform(ax_wave, d, show_truth, label_first=True) _panel_spectrogram(ax_spec, d, show_truth) _panel_posteriors(ax_post, d, show_truth, colorbar=True) _panel_dp(ax_dp, d, show_truth, colorbar=True) _panel_contrastive(ax_score, d, show_truth) for ax, (_key, caption, _drawer, _h) in zip(axes, _PANELS): ax.set_title(caption, loc="left", fontweight="bold", fontsize=11) ax.tick_params(labelbottom=True) _annotate_phoneme_tier(ax, d) ax_score.set_xlabel("Time (s)", fontsize=12) ax_wave.set_xlim(0, dur) fig.suptitle("FALCON forced-alignment — time-aligned representations", fontsize=14, fontweight="bold") fig.savefig(out_path, dpi=150, bbox_inches="tight") plt.close(fig) print(f"[falcon_viz] saved {out_path}") return out_path # --------------------------------------------------------------------------- # # Separate per-panel figures (web app) # # --------------------------------------------------------------------------- # def make_alignment_separate(wav, ckpt, out_dir, w_phi=0.5, language="english", annotation="phn", show_truth=True): """Build one wide, full-size figure per panel; return [(path, caption), ...].""" if language != "english": print(f"[falcon_viz] language='{language}' not supported; using english path.") d = _extract_arrays(wav, ckpt, annotation) dur = d["duration"] os.makedirs(out_dir, exist_ok=True) out = [] for key, caption, drawer, h in _PANELS: fig, ax = plt.subplots(figsize=(12, h), dpi=130) if key == "waveform": drawer(ax, d, show_truth, label_first=True) else: drawer(ax, d, show_truth) ax.set_title(caption, loc="left", fontweight="bold", fontsize=12) ax.set_xlim(0, dur) _annotate_phoneme_tier(ax, d) ax.set_xlabel("Time (s)", labelpad=26) p = os.path.join(out_dir, f"panel_{key}.png") fig.savefig(p, dpi=130, bbox_inches="tight") plt.close(fig) out.append((p, caption)) return out if __name__ == "__main__": _here = os.path.dirname(os.path.abspath(__file__)) make_alignment_panels( wav=os.path.join(_here, "assets", "fasw0sa2.wav"), ckpt=os.path.join(_here, "pretrained_models", "falcon_timit_english.pt"), out_path=os.path.join(_here, "assets", "example_panels.png"), )