prj-beatrice/dnsmos-torch-native

DNSMOS P.835 and P.808

Microsoft's DNSMOS P.835 and P.808, adapted to use with only torch, torchaudio, and transformers.

import soundfile as sf
import torch
from transformers import AutoModel

device = "cuda" if torch.cuda.is_available() else "cpu"
wav, sr = sf.read("audio.wav", dtype="float32")
wav = torch.from_numpy(wav)
model = AutoModel.from_pretrained(
    "prj-beatrice/dnsmos-torch-native", trust_remote_code=True
).eval().to(device)

with torch.inference_mode():
    scores = model(wav.to(device), sampling_rate=sr)
print(scores.sig, scores.bak, scores.ovrl, scores.p808)

The results match this code:

from dnsmos_local import ComputeScore

# Non-16 kHz input is resampled differently, so exact agreement is not expected.
assert sr == 16_000
reference = ComputeScore(
    "DNSMOS/sig_bak_ovr.onnx", "DNSMOS/model_v8.onnx"
)("audio.wav", 16_000, False)
actual = torch.stack([
    scores.sig[0], scores.bak[0], scores.ovrl[0], scores.p808[0]
]).cpu()
expected = torch.tensor([
    reference[name] for name in ("SIG", "BAK", "OVRL", "P808_MOS")
], dtype=actual.dtype)
torch.testing.assert_close(actual, expected)

Padded batches use a waveform tensor [batch_size, max_length] and an input_lengths tensor [batch_size].

License

The inference code is MIT-licensed, and the converted model weights are licensed under CC BY 4.0; see LICENSE.