| 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, |
| ) |
|
|
