File size: 3,086 Bytes
6835659 d7c075c c98d24c d7c075c c98d24c 960dff6 d7c075c 960dff6 d7c075c 960dff6 d7c075c 960dff6 d7c075c c98d24c 6835659 fb5a7e6 6835659 d7c075c 6835659 c98d24c d7c075c 6835659 | 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 77 78 79 80 81 82 83 84 85 86 87 88 | from __future__ import annotations
import numpy as np
import torch
import librosa
from transformers import ClapModel, ClapProcessor
class AudioEmbedder:
"""
CLAP-based audio embedder.
Optimized for environmental soundscape semantics.
"""
def __init__(
self,
model_name: str = "laion/clap-htsat-unfused",
device: str = "cpu",
target_sr: int = 48000,
):
self.device = device
self.target_sr = target_sr
self.processor = ClapProcessor.from_pretrained(model_name)
self.model = ClapModel.from_pretrained(model_name)
self.model.to(self.device)
self.model.eval()
def _extract_features(self, output, projection: str) -> torch.Tensor:
"""Extract 1-D projected embedding (512-d) from model output.
Handles both raw tensors and BaseModelOutputWithPooling objects
across different transformers versions.
"""
target_dim = getattr(self.model.config, "projection_dim", 512)
if not isinstance(output, torch.Tensor):
# BaseModelOutputWithPooling — extract pooled features
pooled = output.pooler_output
# Only project if not already at target dim
if pooled.shape[-1] != target_dim:
proj = getattr(self.model, projection, None)
if proj is not None:
pooled = proj(pooled)
output = pooled
if output.dim() == 3:
pooled = output[:, 0, :]
if pooled.shape[-1] != target_dim:
proj = getattr(self.model, projection, None)
if proj is not None:
pooled = proj(pooled)
output = pooled
if output.dim() == 2:
output = output[0]
return output
@torch.no_grad()
def embed(self, audio_path: str) -> np.ndarray:
waveform, _ = librosa.load(audio_path, sr=self.target_sr, mono=True)
# Use 'audio' (newer transformers) with fallback to 'audios' (older)
try:
inputs = self.processor(
audio=waveform,
sampling_rate=self.target_sr,
return_tensors="pt",
).to(self.device)
except TypeError:
inputs = self.processor(
audios=waveform,
sampling_rate=self.target_sr,
return_tensors="pt",
).to(self.device)
outputs = self.model.get_audio_features(**inputs)
emb = self._extract_features(outputs, "audio_projection")
return emb.cpu().numpy().astype("float32")
@torch.no_grad()
def embed_text(self, text: str) -> np.ndarray:
"""Embed text using CLAP's text encoder (for text-audio comparison)."""
inputs = self.processor(
text=[text],
return_tensors="pt",
padding=True,
).to(self.device)
feats = self.model.get_text_features(**inputs)
feats = self._extract_features(feats, "text_projection")
return feats.cpu().numpy().astype("float32")
|