|
|
import argparse |
|
|
from typing import List, Union |
|
|
|
|
|
import hydra |
|
|
import soundfile |
|
|
import torch |
|
|
from omegaconf import DictConfig, ListConfig, OmegaConf |
|
|
|
|
|
|
|
|
class SpecScaler(torch.nn.Module): |
|
|
def forward(self, x: torch.Tensor) -> torch.Tensor: |
|
|
return torch.log(x.clamp_(1e-9, 1e9)) |
|
|
|
|
|
|
|
|
class GigaAMEmo(torch.nn.Module): |
|
|
def __init__(self, conf: Union[DictConfig, ListConfig]): |
|
|
super().__init__() |
|
|
self.id2name = conf.id2name |
|
|
self.feature_extractor = hydra.utils.instantiate(conf.feature_extractor) |
|
|
self.conformer = hydra.utils.instantiate(conf.encoder) |
|
|
self.linear_head = hydra.utils.instantiate(conf.classification_head) |
|
|
|
|
|
@property |
|
|
def device(self): |
|
|
return next(self.parameters()).device |
|
|
|
|
|
def forward(self, features, features_length=None): |
|
|
if features.dim() == 2: |
|
|
features = features.unsqueeze(0) |
|
|
if not features_length: |
|
|
features_length = torch.ones(features.shape[0], device=self.device) * features.shape[-1] |
|
|
encoded, _ = self.conformer(audio_signal=features, length=features_length) |
|
|
encoded_pooled = torch.nn.functional.avg_pool1d( |
|
|
encoded, kernel_size=encoded.shape[-1] |
|
|
).squeeze(-1) |
|
|
|
|
|
logits = self.linear_head(encoded_pooled) |
|
|
return logits |
|
|
|
|
|
def get_probs(self, audio_path: str) -> List[List[float]]: |
|
|
audio_signal, _ = soundfile.read(audio_path, dtype="float32") |
|
|
audio_tensor = torch.tensor(audio_signal).float().to(self.device) |
|
|
features = self.feature_extractor(audio_tensor) |
|
|
logits = self.forward(features) |
|
|
probs = torch.nn.functional.softmax(logits, dim=1).detach().tolist() |
|
|
return probs |
|
|
|
|
|
|
|
|
def _parse_args(): |
|
|
parser = argparse.ArgumentParser( |
|
|
description="Run inference using GigaAM-Emo checkpoint" |
|
|
) |
|
|
parser.add_argument("--model_config", help="Path to GigaAM-Emo config file (.yaml)") |
|
|
parser.add_argument( |
|
|
"--model_weights", help="Path to GigaAM-Emo checkpoint file (.ckpt)" |
|
|
) |
|
|
parser.add_argument("--audio_path", help="Path to audio signal") |
|
|
parser.add_argument("--device", help="Device: cpu / cuda") |
|
|
return parser.parse_args() |
|
|
|
|
|
|
|
|
def main(model_config: str, model_weights: str, device: str, audio_path: str): |
|
|
conf = OmegaConf.load(model_config) |
|
|
model = GigaAMEmo(conf) |
|
|
ckpt = torch.load(model_weights, map_location="cpu") |
|
|
model.load_state_dict(ckpt, strict=False) |
|
|
model = model.to(device) |
|
|
model.eval() |
|
|
with torch.no_grad(): |
|
|
probs = model.get_probs(audio_path)[0] |
|
|
print(", ".join([f"{model.id2name[i]}: {p:.3f}" for i, p in enumerate(probs)])) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
args = _parse_args() |
|
|
main( |
|
|
model_config=args.model_config, |
|
|
model_weights=args.model_weights, |
|
|
device=args.device, |
|
|
audio_path=args.audio_path, |
|
|
) |
|
|
|
