import argparse import csv import json import os import time from pathlib import Path from typing import Dict, List, Optional os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1") os.environ.setdefault("HF_HUB_OFFLINE", "0") os.environ.setdefault("TRANSFORMERS_OFFLINE", "0") os.environ.setdefault("HF_DATASETS_OFFLINE", "0") import librosa import numpy as np import torch from reasoning_pipeline import aggregate_reasoning_summaries from interface import separate_audio_file, load_wav2vec DEFAULT_INPUT_DIR = Path("input_mixes") DEFAULT_OUTPUT_DIR = Path("separated_audios") DEFAULT_RUNS_DIR = Path("gnn_runs") DEFAULT_PIPELINE_OUTPUT_DIR = Path("pipeline_results") DEFAULT_GNN_CHECKPOINT = "best_audio_separation_model.pt" SUPPORTED_EXTENSIONS = [".wav", ".mp3", ".flac"] FRAME_SECONDS = 0.5 SR = 16000 WINDOW_SECONDS = 5.0 def find_audio_files(directory: Path) -> List[Path]: directory = Path(directory) if not directory.exists(): raise FileNotFoundError(f"Input directory not found: {directory}") files = [] for ext in SUPPORTED_EXTENSIONS: files.extend(sorted(directory.glob(f"*{ext}"))) return sorted(files) def infer_on_audio_file( audio_path: Path, gnn_model, processor, wav2vec_model, run_root: Path, device: torch.device, window_seconds: float = WINDOW_SECONDS, frame_seconds: float = FRAME_SECONDS, ) -> Optional[Dict]: import importlib gnn = importlib.import_module("gnn") build_chain_edge_index = getattr(gnn, "build_chain_edge_index") PatientStateManager = getattr(gnn, "PatientStateManager") torch_geometric_data = importlib.import_module("torch_geometric.data") Data = getattr(torch_geometric_data, "Data") y, _ = librosa.load(str(audio_path), sr=SR, mono=True) audio_duration_s = float(len(y) / SR) frame_len = int(frame_seconds * SR) frames = [] for i in range(0, len(y), frame_len): f = y[i:i + frame_len] if len(f) < frame_len: f = np.pad(f, (0, frame_len - len(f)), mode="constant") frames.append(f.astype(np.float32)) if not frames: return None embeddings = [] batch_size = 8 with torch.no_grad(): for index in range(0, len(frames), batch_size): batch = frames[index:index + batch_size] inputs = processor(batch, sampling_rate=SR, return_tensors="pt", padding=True) input_values = inputs.input_values.to(device) out = wav2vec_model(input_values) emb = out.last_hidden_state.mean(dim=1).cpu().numpy().astype(np.float32) embeddings.append(emb) X = np.concatenate(embeddings, axis=0) if embeddings else np.zeros((0, 768), dtype=np.float32) if X.shape[0] == 0: return None edge_index = build_chain_edge_index(X.shape[0]) if X.shape[0] == 0: return None edge_index = build_chain_edge_index(X.shape[0]) data = Data(x=torch.tensor(X, dtype=torch.float32), edge_index=edge_index) start = time.time() with torch.no_grad(): data = data.to(device) w_logits, c_logits = gnn_model(data) infer_ms = (time.time() - start) * 1000.0 w_prob = float(torch.sigmoid(w_logits).view(-1)[0].cpu().item()) c_prob = float(torch.sigmoid(c_logits).view(-1)[0].cpu().item()) w_pred = int(w_prob >= 0.5) c_pred = int(c_prob >= 0.5) w_unc = c_unc = w_conf = c_conf = None if hasattr(gnn_model, "log_var_wheeze"): w_unc = float(np.exp(0.5 * float(gnn_model.log_var_wheeze.detach().cpu().item()))) w_conf = float(1.0 / (1.0 + w_unc)) if hasattr(gnn_model, "log_var_crackle"): c_unc = float(np.exp(0.5 * float(gnn_model.log_var_crackle.detach().cpu().item()))) c_conf = float(1.0 / (1.0 + c_unc)) breathing_rate_bpm = None if len(y) >= SR: breathing_rate_bpm = float(np.nan) if len(y) == 0 else None from gnn import estimate_breathing_rate_bpm breathing_rate_bpm = estimate_breathing_rate_bpm(y, SR, audio_duration_s) state_mgr = PatientStateManager( ema_alpha=0.12, low_delta=0.08, high_delta=0.20, min_samples_for_baseline=5, force_established_after_s=10.0, ) frames_per_window = max(1, int(round(window_seconds / frame_seconds))) window_rows = [] for w_start in range(0, X.shape[0], frames_per_window): Xw = X[w_start:w_start + frames_per_window] if Xw.shape[0] == 0: continue ew = build_chain_edge_index(Xw.shape[0]) dw = Data(x=torch.tensor(Xw, dtype=torch.float32), edge_index=ew) with torch.no_grad(): dw = dw.to(device) w_l, c_l = gnn_model(dw) w_p = float(torch.sigmoid(w_l).view(-1)[0].cpu().item()) c_p = float(torch.sigmoid(c_l).view(-1)[0].cpu().item()) w_pd = int(w_p >= 0.5) c_pd = int(c_p >= 0.5) start_sec = float(w_start * frame_seconds) end_sec = float(min((w_start + Xw.shape[0]) * frame_seconds, audio_duration_s)) start_sample = int(start_sec * SR) end_sample = int(end_sec * SR) y_window = y[start_sample:end_sample] if end_sample > start_sample else np.array([], dtype=np.float32) breathing_rate_window = None if len(y_window) >= SR: from gnn import estimate_breathing_rate_bpm breathing_rate_window = estimate_breathing_rate_bpm(y_window, SR, max(end_sec - start_sec, 1e-6)) state_out = state_mgr.update_and_get_state(audio_path.stem, w_p, c_p, timestamp=start_sec) window_rows.append({ "start_sec": round(start_sec, 3), "end_sec": round(end_sec, 3), "num_frames": int(Xw.shape[0]), "wheeze_prob": round(w_p, 4), "crackle_prob": round(c_p, 4), "wheeze_pred": w_pd, "crackle_pred": c_pd, "patient_state": state_out.get("overall_state"), "breathing_rate_bpm": None if breathing_rate_window is None else round(breathing_rate_window, 2), "ig_topk": [], "gxi_topk": [], }) run_id = time.strftime("%Y%m%dT%H%M%SZ") + "_" + audio_path.stem run_dir = run_root / run_id run_dir.mkdir(parents=True, exist_ok=True) output = { "request_id": run_id, "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()), "result": { "audio_id": audio_path.stem, "audio_duration_s": round(audio_duration_s, 3), "wheeze": { "probability": round(w_prob, 4), "prediction": bool(w_pred), "confidence": None if w_conf is None else round(w_conf, 4), }, "crackle": { "probability": round(c_prob, 4), "prediction": bool(c_pred), "confidence": None if c_conf is None else round(c_conf, 4), }, "breathing_rate_bpm": None if breathing_rate_bpm is None else round(breathing_rate_bpm, 2), "model_version": str(gnn_model.__class__.__name__), "inference_time_ms": round(float(infer_ms), 2), }, "reasoning": { "thresholds": {"wheeze": 0.5, "crackle": 0.5}, "uncertainty_std": {"wheeze": w_unc, "crackle": c_unc}, "flags": {"short_audio": bool(audio_duration_s < frame_seconds), "near_threshold": bool(abs(w_prob - 0.5) <= 0.02 or abs(c_prob - 0.5) <= 0.02)}, "cumulative_windows": { "window_seconds": float(window_seconds), "frame_seconds": float(frame_seconds), "num_windows": int(len(window_rows)), "frames_per_window": int(frames_per_window), "wheeze_window_positive": int(sum(r["wheeze_pred"] for r in window_rows)), "crackle_window_positive": int(sum(r["crackle_pred"] for r in window_rows)), "wheeze_window_ratio": float(np.mean([r["wheeze_pred"] for r in window_rows])) if window_rows else 0.0, "crackle_window_ratio": float(np.mean([r["crackle_pred"] for r in window_rows])) if window_rows else 0.0, "wheeze_prob_mean_window": float(np.mean([r["wheeze_prob"] for r in window_rows])) if window_rows else 0.0, "crackle_prob_mean_window": float(np.mean([r["crackle_prob"] for r in window_rows])) if window_rows else 0.0, }, }, } with open(run_dir / "result.json", "w", encoding="utf-8") as f: json.dump(output, f, indent=2) with open(run_dir / "reasoning.json", "w", encoding="utf-8") as f: json.dump(output["reasoning"], f, indent=2) with open(run_dir / "window_report.json", "w", encoding="utf-8") as f: json.dump({"audio_id": audio_path.stem, "windows": window_rows}, f, indent=2) with open(run_dir / "window_report.csv", "w", encoding="utf-8", newline="") as f: writer = csv.DictWriter(f, fieldnames=[ "start_sec", "end_sec", "num_frames", "wheeze_prob", "crackle_prob", "wheeze_pred", "crackle_pred", "patient_state", "breathing_rate_bpm", "ig_topk", "gxi_topk", ]) writer.writeheader() for row in window_rows: row_copy = dict(row) row_copy["ig_topk"] = ";".join(map(str, row_copy["ig_topk"])) row_copy["gxi_topk"] = ";".join(map(str, row_copy["gxi_topk"])) writer.writerow(row_copy) output["artifacts"] = { "result_json": "result.json", "reasoning_json": "reasoning.json", "window_report_json": "window_report.json", "window_report_csv": "window_report.csv", "run_dir": str(run_dir), } return output def run_full_pipeline( input_dir: Path, sep_output_dir: Path, run_root: Path, pipeline_output_dir: Path, gnn_checkpoint: str, patient_names: Optional[List[str]] = None, device_str: str = "cuda" if torch.cuda.is_available() else "cpu", ) -> None: device = torch.device(device_str) sep_output_dir.mkdir(parents=True, exist_ok=True) run_root.mkdir(parents=True, exist_ok=True) pipeline_output_dir.mkdir(parents=True, exist_ok=True) mixture_files = find_audio_files(input_dir) if not mixture_files: raise FileNotFoundError(f"No mixture files found in {input_dir}") for mix_file in mixture_files: print(f"Separating mixture: {mix_file.name}") separate_audio_file(str(mix_file), sep_output_dir, patient_names=patient_names) print("Loading GNN and Wav2Vec models...") processor, wav2vec_model = load_wav2vec(device) from gnn import load_gnn_model gnn_model = load_gnn_model(gnn_checkpoint, device) separated_files = find_audio_files(sep_output_dir) if not separated_files: raise FileNotFoundError(f"No separated audio files found in {sep_output_dir}") behavior_results = [] for audio_file in separated_files: print(f"Running GNN inference: {audio_file.name}") result = infer_on_audio_file( audio_file, gnn_model, processor, wav2vec_model, run_root, device, window_seconds=WINDOW_SECONDS, frame_seconds=FRAME_SECONDS, ) if result is not None: behavior_results.append(result) output_path = pipeline_output_dir / "full_pipeline_results.json" with output_path.open("w", encoding="utf-8") as f: json.dump(behavior_results, f, indent=2) print("Aggregating reasoning summaries...") run_dirs = [Path(item["artifacts"]["run_dir"]) for item in behavior_results if item.get("artifacts")] aggregate_reasoning_summaries(run_dirs, pipeline_output_dir) print(f"Saved full pipeline results to: {output_path}") print(f"Saved reasoning summary to: {pipeline_output_dir}") def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Run the full auditory separation + GNN reasoning pipeline.") parser.add_argument("--input-dir", type=Path, default=DEFAULT_INPUT_DIR, help="Folder of mixture audio files.") parser.add_argument("--sep-output-dir", type=Path, default=DEFAULT_OUTPUT_DIR, help="Folder for separated source audio.") parser.add_argument("--run-root", type=Path, default=DEFAULT_RUNS_DIR, help="Base folder for GNN run artifacts.") parser.add_argument("--pipeline-output-dir", type=Path, default=DEFAULT_PIPELINE_OUTPUT_DIR, help="Folder for pipeline summaries.") parser.add_argument("--gnn-checkpoint", type=str, default=DEFAULT_GNN_CHECKPOINT, help="Path to the GNN checkpoint file.") parser.add_argument("--patient-names", nargs="*", default=None, help="Optional names for separated sources.") parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Torch device to use.") return parser.parse_args() def main() -> None: args = parse_args() patient_names = args.patient_names[:3] if args.patient_names else None run_full_pipeline( input_dir=args.input_dir, sep_output_dir=args.sep_output_dir, run_root=args.run_root, pipeline_output_dir=args.pipeline_output_dir, gnn_checkpoint=args.gnn_checkpoint, patient_names=patient_names, device_str=args.device, ) if __name__ == "__main__": main()