utils.py - master - grenki70/siames

utils.py
77 строк2 KB
grenki70
first_commit
28 май 2026, 06:18
28 май 2026, 06:18f3a78a4
import os
import torch
import librosa
import numpy as np
from torchaudio.transforms import MFCC

from config import Config


class AudioProcessor:
    def __init__(self, cfg: Config):
        self.cfg = cfg
        self.mfcc_transform = MFCC(
            sample_rate=cfg.audio.sample_rate,
            n_mfcc=cfg.audio.n_mfcc,
        )

    def process(self, file_path: str, fixed_length: int = None) -> torch.Tensor:
        if fixed_length is None:
            fixed_length = self.cfg.max_audio_len

        try:
            y, _ = librosa.load(file_path, sr=self.cfg.audio.sample_rate)
        except Exception:
            y = np.zeros(self.cfg.audio.sample_rate)

        waveform = torch.from_numpy(y).unsqueeze(0)
        mfcc = self.mfcc_transform(waveform)

        _, n_m, time = mfcc.shape
        if time < fixed_length:
            padding = torch.zeros((1, n_m, fixed_length - time))
            mfcc = torch.cat((mfcc, padding), dim=2)
        else:
            mfcc = mfcc[:, :, :fixed_length]

        return mfcc

    def process_for_inference(self, file_path: str) -> torch.Tensor | None:
        try:
            mfcc = self.process(file_path)
            return mfcc.unsqueeze(0)
        except Exception as e:
            print(f"{file_path}: {e}")
            return None


def calculate_optimal_length(cfg: Config) -> int:
    max_frames = 0
    mfcc_trans = MFCC(sample_rate=cfg.audio.sample_rate, n_mfcc=cfg.audio.n_mfcc)

    all_files = []
    for dirpath, _, filenames in os.walk(cfg.paths.dataset):
        for f in filenames:
            if f.lower().endswith('.wav'):
                all_files.append(os.path.join(dirpath, f))

    check_files = all_files if len(all_files) < 500 else all_files

    for fp in check_files:
        try:
            y, _ = librosa.load(fp, sr=cfg.audio.sample_rate)
            if len(y) == 0:
                continue

            waveform = torch.from_numpy(y).unsqueeze(0)
            frames = mfcc_trans(waveform).shape[2]
            if frames > max_frames:
                max_frames = frames
        except Exception:
            continue

    if max_frames == 0:
        return 320

    optimized = ((max_frames // 8) + 1) * 8
    return optimized