File size: 4,055 Bytes
fb3c11b | 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 75 76 | """PERCH 2 embeddings for audio of any length, using perch_v2_embed.pt2 (needs only torch and huggingface_hub).
from perch2_embed import load, embed
core = load("cpu") # or "cuda", "mps", "auto"
out = embed(core, wave) # mono float waveform at 32 kHz
out["avg"], out["max"], out["temporal"], out["times"]
"""
import math
import torch
REPO, FILE = "DAS-DeepAudioSegmenter/perch2-torch-embed", "perch_v2_embed.pt2"
WINDOW, FRAME = 160_000, 10_240 # 5 s window and 0.32 s frame step, in samples at the 32 kHz model rate
def load(device="cpu", path=None, revision=None):
"""Load the embedding core onto a device ("cpu", "cuda", "mps" or "auto"). Downloads and caches the file unless
`path` points to a local copy."""
from torch.export.passes import move_to_device_pass
if path is None:
from huggingface_hub import hf_hub_download
path = hf_hub_download(REPO, FILE, revision=revision)
if device == "auto":
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
return move_to_device_pass(torch.export.load(path), device).module()
def embed(core, wave, hop_frames=1, max_seconds=None, sample_rate=32_000, batch_size=32):
"""Embed a mono waveform of any length with the perch_v2_embed.pt2 core.
wave 1-D float array or tensor. It is fed to the model unchanged, i.e. treated as 32 kHz audio.
hop_frames window hop in 0.32 s frames, 1-15 (1 = the model's own 0.32 s hop; larger = fewer windows).
max_seconds keep only the first max_seconds of audio (None = everything).
sample_rate rate of `wave`; used only to convert max_seconds and the returned times to seconds.
batch_size 5 s windows per forward pass (memory use).
Returns a dict with avg (1536,), max (1536,), temporal (frames, 1536) (mean over the 4 frequency positions),
spatial (frames, 4, 1536), times (frames,) frame centres in seconds, and n_windows.
"""
if not 1 <= hop_frames <= 15:
raise ValueError("hop_frames must be 1-15: windows must overlap or touch")
wave = torch.as_tensor(wave, dtype=torch.float32).flatten()
if max_seconds is not None:
wave = wave[: int(max_seconds * sample_rate)]
length, hop = len(wave), hop_frames * FRAME
if length == 0:
raise ValueError("empty waveform")
n = math.ceil(max(length - WINDOW, 0) / hop) + 1
windows = torch.nn.functional.pad(wave, (0, (n - 1) * hop + WINDOW - length)).unfold(0, WINDOW, hop)
device = next(core.parameters()).device
with torch.inference_mode():
chunks = torch.cat([core(windows[i:i + batch_size].to(device)).float().cpu() for i in range(0, n, batch_size)])
# Overlap pooling. At hop_frames=1 this uses the full model's rule: windows after the third drop their first
# 2 frames and all but the last two drop their last 2. At larger hops every overlapping frame is averaged.
n_time = chunks.shape[1]
total = (n - 1) * hop_frames + n_time
sums, counts = chunks.new_zeros(total, *chunks.shape[2:]), chunks.new_zeros(total)
maximum = torch.full(chunks.shape[3:], -torch.inf)
trim = 2 if hop_frames == 1 else 0
for i in range(n):
lo = trim if i > 2 else 0
hi = n_time - trim if i < n - 2 else n_time
part, start = chunks[i, lo:hi], i * hop_frames
sums[start + lo:start + hi] += part
counts[start + lo:start + hi] += 1
maximum = torch.maximum(maximum, part.amax(dim=(0, 1)))
avg = sums.sum(dim=(0, 1)) / (counts.sum() * chunks.shape[2])
spatial = sums / counts.clamp_min(1)[:, None, None]
centers = (torch.arange(total) + 0.5) * FRAME
keep = centers < length
if not keep.any(): # audio shorter than half a frame: one frame, centred on the audio
keep[0] = True
centers[0] = length / 2
return {"avg": avg, "max": maximum, "temporal": spatial[keep].mean(dim=1), "spatial": spatial[keep],
"times": centers[keep] / sample_rate, "n_windows": n}
|