| """Hierarchical audio tagger: frozen Dasheng-0.6B features + an MLP head over a 4,416-node sound taxonomy. |
| |
| pip install torch dasheng soundfile scipy safetensors |
| python inference.py clip1.flac clip2.wav ... |
| |
| Each clip is read from its first 10 s, resampled to 16 kHz mono and zero-padded to 10 s. Dasheng-0.6B frame |
| features are mean-pooled to 1280-d and fed to the head, which has three outputs: 11 top-level groups, 822 second-level |
| nodes and 3,718 classes (nodes at depth <= 3). Class scores use the hierarchical readout |
| logP_L1 + logP_L2 + logP_L3 (each class scored with its own top-level and second-level ancestor). |
| The Dasheng weights are downloaded by the `dasheng` package on first use. |
| """ |
| import json |
| import sys |
| from math import gcd |
|
|
| import numpy as np |
| import soundfile as sf |
| import torch |
| import torch.nn as nn |
| from dasheng import dasheng_06B |
| from safetensors import safe_open |
| from safetensors.torch import load_file |
| from scipy.signal import resample_poly |
|
|
| SR, SECS = 16000, 10 |
|
|
|
|
| def load_audio(path): |
| """First 10 s of `path` as a (160000,) float32 array at 16 kHz mono, zero-padded.""" |
| with sf.SoundFile(path) as f: |
| sr = f.samplerate |
| a = f.read(SECS * sr, dtype="float32", always_2d=True).mean(1) |
| if sr != SR: |
| g = gcd(sr, SR) |
| a = resample_poly(a, SR // g, sr // g).astype(np.float32) |
| x = np.zeros(SECS * SR, np.float32) |
| x[: min(len(a), SECS * SR)] = a[: SECS * SR] |
| return x |
|
|
|
|
| class TaxonomyClassifier(nn.Module): |
| def __init__(self, head_path="head.safetensors"): |
| super().__init__() |
| with safe_open(head_path, "pt") as f: |
| meta = f.metadata() |
| cfg, self.classes = json.loads(meta["config"]), json.loads(meta["classes"]) |
| l1 = [c.split("/")[0] for c in self.classes] |
| l2 = ["/".join(c.split("/")[:2]) for c in self.classes] |
| self.top_names, l2_names = sorted(set(l1)), sorted(set(l2)) |
| self.register_buffer("t1", torch.tensor([self.top_names.index(x) for x in l1]), persistent=False) |
| self.register_buffer("t2", torch.tensor([l2_names.index(x) for x in l2]), persistent=False) |
| h = cfg["hidden"] |
| self.trunk = nn.Sequential(nn.Linear(cfg["in_dim"], h), nn.BatchNorm1d(h), nn.GELU(), nn.Dropout(cfg["dropout"])) |
| self.l3, self.h1, self.h2 = nn.Linear(h, len(self.classes)), nn.Linear(h, len(self.top_names)), nn.Linear(h, len(l2_names)) |
| self.load_state_dict(load_file(head_path)) |
| self.encoder = dasheng_06B() |
| self.eval() |
|
|
| @torch.inference_mode() |
| def forward(self, x): |
| """x: (B, 160000) float32 at 16 kHz -> (top-level probs (B, 11), class probs (B, 3718)).""" |
| xl = torch.full((x.shape[0],), x.shape[1], dtype=torch.long, device=x.device) |
| with torch.autocast(x.device.type, dtype=torch.bfloat16, enabled=x.device.type == "cuda"): |
| feats = self.encoder(x, xl) |
| z = self.trunk(feats.float().mean(1)) |
| o3, o1, o2 = self.l3(z), self.h1(z), self.h2(z) |
| score = o1.log_softmax(1)[:, self.t1] + o2.log_softmax(1)[:, self.t2] + o3.log_softmax(1) |
| return o1.softmax(1), score.softmax(1) |
|
|
|
|
| if __name__ == "__main__": |
| dev = "cuda" if torch.cuda.is_available() else "cpu" |
| model = TaxonomyClassifier().to(dev) |
| paths = sys.argv[1:] |
| top, cls = model(torch.from_numpy(np.stack([load_audio(p) for p in paths])).to(dev)) |
| for path, pt, pc in zip(paths, top.cpu(), cls.cpu()): |
| print(f"{path}\n top-level: {model.top_names[int(pt.argmax())]} ({float(pt.max()):.3f})") |
| for p, i in zip(*pc.topk(5)): |
| print(f" {float(p):.3f} {model.classes[int(i)]}") |
|
|
