Spaces:
Sleeping
Sleeping
File size: 2,078 Bytes
127b976 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | from __future__ import annotations
from typing import List, Dict
import torch
import torch.nn.functional as F
@torch.no_grad()
def sliding_window_localization(
model,
wav: torch.Tensor,
sample_rate: int = 16000,
window_sec: float = 2.0,
hop_sec: float = 1.0,
threshold: float = 0.75,
device: str | torch.device = "cpu",
) -> List[Dict[str, float]]:
"""Localize suspicious regions with overlapping windows.
Args:
model: AuralGuardAASISTPP.
wav: Tensor [samples] or [1, samples].
"""
model.eval()
device = torch.device(device)
if wav.ndim == 2:
wav = wav.squeeze(0)
wav = wav.to(device).float()
window = int(window_sec * sample_rate)
hop = int(hop_sec * sample_rate)
if wav.numel() < window:
wav = F.pad(wav, (0, window - wav.numel()))
suspicious = []
for start in range(0, wav.numel() - window + 1, hop):
end = start + window
chunk = wav[start:end].unsqueeze(0)
outputs = model(chunk)
fake_score = torch.softmax(outputs["binary_logits"], dim=-1)[0, 1].item()
if fake_score >= threshold:
suspicious.append({
"start": round(start / sample_rate, 3),
"end": round(end / sample_rate, 3),
"fake_score": round(fake_score, 4),
})
return merge_overlapping_segments(suspicious)
def merge_overlapping_segments(segments: List[Dict[str, float]]) -> List[Dict[str, float]]:
"""Merge overlapping suspicious windows for cleaner output."""
if not segments:
return []
segments = sorted(segments, key=lambda x: x["start"])
merged = [segments[0].copy()]
for seg in segments[1:]:
last = merged[-1]
if seg["start"] <= last["end"]:
last["end"] = max(last["end"], seg["end"])
last["fake_score"] = round(max(last["fake_score"], seg["fake_score"]), 4)
else:
merged.append(seg.copy())
return merged
|