Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import argparse | |
| import json | |
| from pathlib import Path | |
| import torch | |
| from .aasist_loader import build_aasist_backbone | |
| from .dataset import load_audio_fixed | |
| from .explain import build_report | |
| from .labels import ID_TO_ATTACK | |
| from .model_aasistpp import AuralGuardAASISTPP | |
| from .sliding_localization import sliding_window_localization | |
| def load_model(args, device): | |
| backbone = build_aasist_backbone(args.aasist_root, args.aasist_config, device=device) | |
| model = AuralGuardAASISTPP(backbone, feature_dim=args.feature_dim).to(device) | |
| ckpt = torch.load(args.checkpoint, map_location=device) | |
| model.load_state_dict(ckpt["model"] if isinstance(ckpt, dict) and "model" in ckpt else ckpt, strict=True) | |
| model.eval() | |
| return model | |
| def run_auralguard( | |
| audio_path: str, | |
| model, | |
| sample_rate: int = 16000, | |
| duration_sec: float = 4.0, | |
| localization_window_sec: float = 2.0, | |
| localization_hop_sec: float = 1.0, | |
| suspicious_threshold: float = 0.75, | |
| device: str | torch.device = "cpu", | |
| ): | |
| device = torch.device(device) | |
| fixed_wav = load_audio_fixed( | |
| audio_path, sample_rate=sample_rate, num_samples=int(sample_rate * duration_sec), random_crop=False | |
| ).unsqueeze(0).to(device) | |
| outputs = model(fixed_wav) | |
| fake_probability = torch.softmax(outputs["binary_logits"], dim=-1)[0, 1].item() | |
| attack_id = torch.argmax(outputs["attack_logits"], dim=-1)[0].item() | |
| explanation_id = torch.argmax(outputs["explanation_logits"], dim=-1)[0].item() | |
| attack_type = ID_TO_ATTACK[attack_id] | |
| # For timestamps, use the full audio instead of the fixed 4s crop. | |
| full_wav = load_audio_fixed(audio_path, sample_rate=sample_rate, num_samples=None, random_crop=False) | |
| segments = sliding_window_localization( | |
| model=model, | |
| wav=full_wav, | |
| sample_rate=sample_rate, | |
| window_sec=localization_window_sec, | |
| hop_sec=localization_hop_sec, | |
| threshold=suspicious_threshold, | |
| device=device, | |
| ) | |
| return build_report( | |
| fake_probability, | |
| attack_type, | |
| explanation_id, | |
| segments, | |
| suspicious_window_threshold=suspicious_threshold, | |
| ) | |
| def parse_args(): | |
| p = argparse.ArgumentParser(description="Run AuralGuard-AASIST++ inference on one audio file") | |
| p.add_argument("--audio", required=True) | |
| p.add_argument("--checkpoint", required=True) | |
| p.add_argument("--aasist-root", default="external/aasist") | |
| p.add_argument("--aasist-config", default="external/aasist/config/AASIST.conf") | |
| p.add_argument("--sample-rate", type=int, default=16000) | |
| p.add_argument("--duration-sec", type=float, default=4.0) | |
| p.add_argument("--feature-dim", type=int, default=160) | |
| p.add_argument("--window-sec", type=float, default=2.0) | |
| p.add_argument("--hop-sec", type=float, default=1.0) | |
| p.add_argument("--threshold", type=float, default=0.75) | |
| p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") | |
| return p.parse_args() | |
| def main(): | |
| args = parse_args() | |
| device = torch.device(args.device) | |
| model = load_model(args, device) | |
| report = run_auralguard( | |
| args.audio, | |
| model, | |
| sample_rate=args.sample_rate, | |
| duration_sec=args.duration_sec, | |
| localization_window_sec=args.window_sec, | |
| localization_hop_sec=args.hop_sec, | |
| suspicious_threshold=args.threshold, | |
| device=device, | |
| ) | |
| print(json.dumps(report, indent=2)) | |
| if __name__ == "__main__": | |
| main() | |