liammatt5/GLAM_Web_App

1import argparse2import csv3import json4import os5import time6from pathlib import Path7from typing import Dict, List, Optional8 9os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")10os.environ.setdefault("HF_HUB_OFFLINE", "0")11os.environ.setdefault("TRANSFORMERS_OFFLINE", "0")12os.environ.setdefault("HF_DATASETS_OFFLINE", "0")13 14import librosa15import numpy as np16import torch17from reasoning_pipeline import aggregate_reasoning_summaries18from interface import separate_audio_file, load_wav2vec19 20DEFAULT_INPUT_DIR = Path("input_mixes")21DEFAULT_OUTPUT_DIR = Path("separated_audios")22DEFAULT_RUNS_DIR = Path("gnn_runs")23DEFAULT_PIPELINE_OUTPUT_DIR = Path("pipeline_results")24DEFAULT_GNN_CHECKPOINT = "best_audio_separation_model.pt"25SUPPORTED_EXTENSIONS = [".wav", ".mp3", ".flac"]26FRAME_SECONDS = 0.527SR = 1600028WINDOW_SECONDS = 5.029 30 31def find_audio_files(directory: Path) -> List[Path]:32    directory = Path(directory)33    if not directory.exists():34        raise FileNotFoundError(f"Input directory not found: {directory}")35    files = []36    for ext in SUPPORTED_EXTENSIONS:37        files.extend(sorted(directory.glob(f"*{ext}")))38    return sorted(files)39 40 41def infer_on_audio_file(42    audio_path: Path,43    gnn_model,44    processor,45    wav2vec_model,46    run_root: Path,47    device: torch.device,48    window_seconds: float = WINDOW_SECONDS,49    frame_seconds: float = FRAME_SECONDS,50) -> Optional[Dict]:51    import importlib52 53    gnn = importlib.import_module("gnn")54    build_chain_edge_index = getattr(gnn, "build_chain_edge_index")55    PatientStateManager = getattr(gnn, "PatientStateManager")56    torch_geometric_data = importlib.import_module("torch_geometric.data")57    Data = getattr(torch_geometric_data, "Data")58 59    y, _ = librosa.load(str(audio_path), sr=SR, mono=True)60    audio_duration_s = float(len(y) / SR)61    frame_len = int(frame_seconds * SR)62 63    frames = []64    for i in range(0, len(y), frame_len):65        f = y[i:i + frame_len]66        if len(f) < frame_len:67            f = np.pad(f, (0, frame_len - len(f)), mode="constant")68        frames.append(f.astype(np.float32))69 70    if not frames:71        return None72 73    embeddings = []74    batch_size = 875    with torch.no_grad():76        for index in range(0, len(frames), batch_size):77            batch = frames[index:index + batch_size]78            inputs = processor(batch, sampling_rate=SR, return_tensors="pt", padding=True)79            input_values = inputs.input_values.to(device)80            out = wav2vec_model(input_values)81            emb = out.last_hidden_state.mean(dim=1).cpu().numpy().astype(np.float32)82            embeddings.append(emb)83 84    X = np.concatenate(embeddings, axis=0) if embeddings else np.zeros((0, 768), dtype=np.float32)85    if X.shape[0] == 0:86        return None87 88    edge_index = build_chain_edge_index(X.shape[0])89    if X.shape[0] == 0:90        return None91 92    edge_index = build_chain_edge_index(X.shape[0])93    data = Data(x=torch.tensor(X, dtype=torch.float32), edge_index=edge_index)94 95    start = time.time()96    with torch.no_grad():97        data = data.to(device)98        w_logits, c_logits = gnn_model(data)99    infer_ms = (time.time() - start) * 1000.0100 101    w_prob = float(torch.sigmoid(w_logits).view(-1)[0].cpu().item())102    c_prob = float(torch.sigmoid(c_logits).view(-1)[0].cpu().item())103    w_pred = int(w_prob >= 0.5)104    c_pred = int(c_prob >= 0.5)105 106    w_unc = c_unc = w_conf = c_conf = None107    if hasattr(gnn_model, "log_var_wheeze"):108        w_unc = float(np.exp(0.5 * float(gnn_model.log_var_wheeze.detach().cpu().item())))109        w_conf = float(1.0 / (1.0 + w_unc))110    if hasattr(gnn_model, "log_var_crackle"):111        c_unc = float(np.exp(0.5 * float(gnn_model.log_var_crackle.detach().cpu().item())))112        c_conf = float(1.0 / (1.0 + c_unc))113 114    breathing_rate_bpm = None115    if len(y) >= SR:116        breathing_rate_bpm = float(np.nan) if len(y) == 0 else None117        from gnn import estimate_breathing_rate_bpm118 119        breathing_rate_bpm = estimate_breathing_rate_bpm(y, SR, audio_duration_s)120 121    state_mgr = PatientStateManager(122        ema_alpha=0.12,123        low_delta=0.08,124        high_delta=0.20,125        min_samples_for_baseline=5,126        force_established_after_s=10.0,127    )128 129    frames_per_window = max(1, int(round(window_seconds / frame_seconds)))130    window_rows = []131    for w_start in range(0, X.shape[0], frames_per_window):132        Xw = X[w_start:w_start + frames_per_window]133        if Xw.shape[0] == 0:134            continue135        ew = build_chain_edge_index(Xw.shape[0])136        dw = Data(x=torch.tensor(Xw, dtype=torch.float32), edge_index=ew)137 138        with torch.no_grad():139            dw = dw.to(device)140            w_l, c_l = gnn_model(dw)141 142        w_p = float(torch.sigmoid(w_l).view(-1)[0].cpu().item())143        c_p = float(torch.sigmoid(c_l).view(-1)[0].cpu().item())144        w_pd = int(w_p >= 0.5)145        c_pd = int(c_p >= 0.5)146 147        start_sec = float(w_start * frame_seconds)148        end_sec = float(min((w_start + Xw.shape[0]) * frame_seconds, audio_duration_s))149        start_sample = int(start_sec * SR)150        end_sample = int(end_sec * SR)151        y_window = y[start_sample:end_sample] if end_sample > start_sample else np.array([], dtype=np.float32)152        breathing_rate_window = None153        if len(y_window) >= SR:154            from gnn import estimate_breathing_rate_bpm155 156            breathing_rate_window = estimate_breathing_rate_bpm(y_window, SR, max(end_sec - start_sec, 1e-6))157 158        state_out = state_mgr.update_and_get_state(audio_path.stem, w_p, c_p, timestamp=start_sec)159 160        window_rows.append({161            "start_sec": round(start_sec, 3),162            "end_sec": round(end_sec, 3),163            "num_frames": int(Xw.shape[0]),164            "wheeze_prob": round(w_p, 4),165            "crackle_prob": round(c_p, 4),166            "wheeze_pred": w_pd,167            "crackle_pred": c_pd,168            "patient_state": state_out.get("overall_state"),169            "breathing_rate_bpm": None if breathing_rate_window is None else round(breathing_rate_window, 2),170            "ig_topk": [],171            "gxi_topk": [],172        })173 174    run_id = time.strftime("%Y%m%dT%H%M%SZ") + "_" + audio_path.stem175    run_dir = run_root / run_id176    run_dir.mkdir(parents=True, exist_ok=True)177 178    output = {179        "request_id": run_id,180        "timestamp": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),181        "result": {182            "audio_id": audio_path.stem,183            "audio_duration_s": round(audio_duration_s, 3),184            "wheeze": {185                "probability": round(w_prob, 4),186                "prediction": bool(w_pred),187                "confidence": None if w_conf is None else round(w_conf, 4),188            },189            "crackle": {190                "probability": round(c_prob, 4),191                "prediction": bool(c_pred),192                "confidence": None if c_conf is None else round(c_conf, 4),193            },194            "breathing_rate_bpm": None if breathing_rate_bpm is None else round(breathing_rate_bpm, 2),195            "model_version": str(gnn_model.__class__.__name__),196            "inference_time_ms": round(float(infer_ms), 2),197        },198        "reasoning": {199            "thresholds": {"wheeze": 0.5, "crackle": 0.5},200            "uncertainty_std": {"wheeze": w_unc, "crackle": c_unc},201            "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)},202            "cumulative_windows": {203                "window_seconds": float(window_seconds),204                "frame_seconds": float(frame_seconds),205                "num_windows": int(len(window_rows)),206                "frames_per_window": int(frames_per_window),207                "wheeze_window_positive": int(sum(r["wheeze_pred"] for r in window_rows)),208                "crackle_window_positive": int(sum(r["crackle_pred"] for r in window_rows)),209                "wheeze_window_ratio": float(np.mean([r["wheeze_pred"] for r in window_rows])) if window_rows else 0.0,210                "crackle_window_ratio": float(np.mean([r["crackle_pred"] for r in window_rows])) if window_rows else 0.0,211                "wheeze_prob_mean_window": float(np.mean([r["wheeze_prob"] for r in window_rows])) if window_rows else 0.0,212                "crackle_prob_mean_window": float(np.mean([r["crackle_prob"] for r in window_rows])) if window_rows else 0.0,213            },214        },215    }216 217    with open(run_dir / "result.json", "w", encoding="utf-8") as f:218        json.dump(output, f, indent=2)219    with open(run_dir / "reasoning.json", "w", encoding="utf-8") as f:220        json.dump(output["reasoning"], f, indent=2)221    with open(run_dir / "window_report.json", "w", encoding="utf-8") as f:222        json.dump({"audio_id": audio_path.stem, "windows": window_rows}, f, indent=2)223    with open(run_dir / "window_report.csv", "w", encoding="utf-8", newline="") as f:224        writer = csv.DictWriter(f, fieldnames=[225            "start_sec",226            "end_sec",227            "num_frames",228            "wheeze_prob",229            "crackle_prob",230            "wheeze_pred",231            "crackle_pred",232            "patient_state",233            "breathing_rate_bpm",234            "ig_topk",235            "gxi_topk",236        ])237        writer.writeheader()238        for row in window_rows:239            row_copy = dict(row)240            row_copy["ig_topk"] = ";".join(map(str, row_copy["ig_topk"]))241            row_copy["gxi_topk"] = ";".join(map(str, row_copy["gxi_topk"]))242            writer.writerow(row_copy)243 244    output["artifacts"] = {245        "result_json": "result.json",246        "reasoning_json": "reasoning.json",247        "window_report_json": "window_report.json",248        "window_report_csv": "window_report.csv",249        "run_dir": str(run_dir),250    }251    return output252 253 254def run_full_pipeline(255    input_dir: Path,256    sep_output_dir: Path,257    run_root: Path,258    pipeline_output_dir: Path,259    gnn_checkpoint: str,260    patient_names: Optional[List[str]] = None,261    device_str: str = "cuda" if torch.cuda.is_available() else "cpu",262) -> None:263    device = torch.device(device_str)264    sep_output_dir.mkdir(parents=True, exist_ok=True)265    run_root.mkdir(parents=True, exist_ok=True)266    pipeline_output_dir.mkdir(parents=True, exist_ok=True)267 268    mixture_files = find_audio_files(input_dir)269    if not mixture_files:270        raise FileNotFoundError(f"No mixture files found in {input_dir}")271 272    for mix_file in mixture_files:273        print(f"Separating mixture: {mix_file.name}")274        separate_audio_file(str(mix_file), sep_output_dir, patient_names=patient_names)275 276    print("Loading GNN and Wav2Vec models...")277    processor, wav2vec_model = load_wav2vec(device)278    from gnn import load_gnn_model279 280    gnn_model = load_gnn_model(gnn_checkpoint, device)281 282    separated_files = find_audio_files(sep_output_dir)283    if not separated_files:284        raise FileNotFoundError(f"No separated audio files found in {sep_output_dir}")285 286    behavior_results = []287    for audio_file in separated_files:288        print(f"Running GNN inference: {audio_file.name}")289        result = infer_on_audio_file(290            audio_file,291            gnn_model,292            processor,293            wav2vec_model,294            run_root,295            device,296            window_seconds=WINDOW_SECONDS,297            frame_seconds=FRAME_SECONDS,298        )299        if result is not None:300            behavior_results.append(result)301 302    output_path = pipeline_output_dir / "full_pipeline_results.json"303    with output_path.open("w", encoding="utf-8") as f:304        json.dump(behavior_results, f, indent=2)305 306    print("Aggregating reasoning summaries...")307    run_dirs = [Path(item["artifacts"]["run_dir"]) for item in behavior_results if item.get("artifacts")]308    aggregate_reasoning_summaries(run_dirs, pipeline_output_dir)309 310    print(f"Saved full pipeline results to: {output_path}")311    print(f"Saved reasoning summary to: {pipeline_output_dir}")312 313 314def parse_args() -> argparse.Namespace:315    parser = argparse.ArgumentParser(description="Run the full auditory separation + GNN reasoning pipeline.")316    parser.add_argument("--input-dir", type=Path, default=DEFAULT_INPUT_DIR, help="Folder of mixture audio files.")317    parser.add_argument("--sep-output-dir", type=Path, default=DEFAULT_OUTPUT_DIR, help="Folder for separated source audio.")318    parser.add_argument("--run-root", type=Path, default=DEFAULT_RUNS_DIR, help="Base folder for GNN run artifacts.")319    parser.add_argument("--pipeline-output-dir", type=Path, default=DEFAULT_PIPELINE_OUTPUT_DIR, help="Folder for pipeline summaries.")320    parser.add_argument("--gnn-checkpoint", type=str, default=DEFAULT_GNN_CHECKPOINT, help="Path to the GNN checkpoint file.")321    parser.add_argument("--patient-names", nargs="*", default=None, help="Optional names for separated sources.")322    parser.add_argument("--device", type=str, default="cuda" if torch.cuda.is_available() else "cpu", help="Torch device to use.")323    return parser.parse_args()324 325 326def main() -> None:327    args = parse_args()328    patient_names = args.patient_names[:3] if args.patient_names else None329    run_full_pipeline(330        input_dir=args.input_dir,331        sep_output_dir=args.sep_output_dir,332        run_root=args.run_root,333        pipeline_output_dir=args.pipeline_output_dir,334        gnn_checkpoint=args.gnn_checkpoint,335        patient_names=patient_names,336        device_str=args.device,337    )338 339 340if __name__ == "__main__":341    main()342