Spaces:
Sleeping
Sleeping
| import os | |
| from typing import Dict, Optional, Tuple | |
| from pathlib import Path | |
| import librosa | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| # ============================================================ | |
| # Config inferred from the training pipeline | |
| # ============================================================ | |
| SR = 16000 | |
| TARGET_SECONDS = 1.0 | |
| TARGET_LEN = int(SR * TARGET_SECONDS) | |
| N_MELS = 40 | |
| HOP_LENGTH = int(SR * 0.01) # 10 ms -> 160 | |
| WIN_LENGTH = int(SR * 0.025) # 25 ms -> 400 | |
| N_FFT = 400 | |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| ORDERED_LABELS = [ | |
| "five", "zero", "yes", "seven", "no", "nine", | |
| "down", "one", "go", "two", "stop", "six", | |
| "unknown", "silence", | |
| ] | |
| IDX2LABEL = {i: label for i, label in enumerate(ORDERED_LABELS)} | |
| BASE_DIR = Path(__file__).resolve().parent | |
| BASELINE_CKPT = BASE_DIR / "baseline_cnn14_best.pth" | |
| CRNN_CKPT = BASE_DIR / "crnn_specaugment14_best.pth" | |
| # ============================================================ | |
| # Models copied from the PDF pipeline | |
| # ============================================================ | |
| class CNN2DBaseline(nn.Module): | |
| def __init__(self, num_classes: int, dropout: float = 0.3): | |
| super().__init__() | |
| def block(in_ch: int, out_ch: int) -> nn.Sequential: | |
| return nn.Sequential( | |
| nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False), | |
| nn.BatchNorm2d(out_ch), | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d(2), | |
| ) | |
| self.backbone = nn.Sequential( | |
| block(1, 16), | |
| block(16, 32), | |
| block(32, 64), | |
| block(64, 128), | |
| ) | |
| self.dropout = nn.Dropout(dropout) | |
| self.fc = nn.Linear(128, num_classes) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = self.backbone(x) | |
| x = x.mean(dim=[2, 3]) | |
| x = self.dropout(x) | |
| return self.fc(x) | |
| class CRNN(nn.Module): | |
| def __init__(self, num_classes: int, rnn_hidden: int = 128, dropout: float = 0.2): | |
| super().__init__() | |
| self.cnn = nn.Sequential( | |
| nn.Conv2d(1, 32, 3, padding=1, bias=False), | |
| nn.BatchNorm2d(32), | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d((2, 2)), | |
| nn.Conv2d(32, 64, 3, padding=1, bias=False), | |
| nn.BatchNorm2d(64), | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d((2, 2)), | |
| nn.Conv2d(64, 128, 3, padding=1, bias=False), | |
| nn.BatchNorm2d(128), | |
| nn.ReLU(inplace=True), | |
| nn.MaxPool2d((2, 1)), | |
| ) | |
| self.dropout = nn.Dropout(dropout) | |
| self.bigru = nn.GRU( | |
| input_size=128 * 5, | |
| hidden_size=rnn_hidden, | |
| num_layers=1, | |
| batch_first=True, | |
| bidirectional=True, | |
| ) | |
| self.fc = nn.Linear(rnn_hidden * 2, num_classes) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = self.cnn(x) | |
| b, c, freq, time = x.shape | |
| x = x.permute(0, 3, 1, 2).contiguous().view(b, time, c * freq) | |
| x, _ = self.bigru(x) | |
| x = x.mean(dim=1) | |
| x = self.dropout(x) | |
| return self.fc(x) | |
| def load_model(model_name: str) -> nn.Module: | |
| if model_name == "Baseline CNN2D": | |
| model = CNN2DBaseline(num_classes=len(ORDERED_LABELS), dropout=0.2) | |
| ckpt_path = BASELINE_CKPT | |
| else: | |
| model = CRNN(num_classes=len(ORDERED_LABELS), rnn_hidden=128, dropout=0.3) | |
| ckpt_path = CRNN_CKPT | |
| if not os.path.exists(ckpt_path): | |
| raise FileNotFoundError( | |
| f"Khong tim thay checkpoint: {ckpt_path}. Hay dat file checkpoint cung thu muc voi app.py" | |
| ) | |
| state_dict = torch.load(ckpt_path, map_location="cpu") | |
| model.load_state_dict(state_dict) | |
| model.eval().to(DEVICE) | |
| return model | |
| MODELS: Dict[str, nn.Module] = { | |
| "Baseline CNN2D": load_model("Baseline CNN2D"), | |
| "CRNN + SpecAugment": load_model("CRNN + SpecAugment"), | |
| } | |
| def _to_mono_resampled(audio: np.ndarray, sample_rate: int) -> np.ndarray: | |
| audio = np.asarray(audio) | |
| if audio.ndim == 2: | |
| audio = audio.mean(axis=1) | |
| audio = audio.astype(np.float32) | |
| max_abs = float(np.max(np.abs(audio))) if audio.size else 1.0 | |
| if max_abs > 1.5: | |
| audio = audio / 32768.0 | |
| if sample_rate != SR: | |
| audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=SR) | |
| return audio.astype(np.float32) | |
| def _fix_length(wav: np.ndarray) -> np.ndarray: | |
| if wav.shape[0] < TARGET_LEN: | |
| wav = np.pad(wav, (0, TARGET_LEN - wav.shape[0])) | |
| elif wav.shape[0] > TARGET_LEN: | |
| wav = wav[:TARGET_LEN] | |
| return wav.astype(np.float32) | |
| def _logmel(wav: np.ndarray) -> np.ndarray: | |
| mel = librosa.feature.melspectrogram( | |
| y=wav, | |
| sr=SR, | |
| n_fft=N_FFT, | |
| hop_length=HOP_LENGTH, | |
| win_length=WIN_LENGTH, | |
| n_mels=N_MELS, | |
| center=False, | |
| power=2.0, | |
| ) | |
| logmel = librosa.power_to_db(mel, ref=1.0) | |
| return logmel.astype(np.float32) | |
| def preprocess_numpy_audio(audio: np.ndarray, sample_rate: int) -> torch.Tensor: | |
| wav = _to_mono_resampled(audio, sample_rate) | |
| wav = _fix_length(wav) | |
| logmel = _logmel(wav) | |
| x = torch.from_numpy(logmel).unsqueeze(0).unsqueeze(0) # [1, 1, 40, T] | |
| return x.to(DEVICE) | |
| def run_prediction(audio_input: Optional[Tuple[int, np.ndarray]], model_name: str, top_k: int = 5) -> Tuple[str, float, Dict[str, float]]: | |
| if audio_input is None: | |
| return "", 0.0, {} | |
| sample_rate, audio = audio_input | |
| if audio is None or len(audio) == 0: | |
| return "", 0.0, {} | |
| model = MODELS[model_name] | |
| x = preprocess_numpy_audio(audio, sample_rate) | |
| logits = model(x) | |
| probs = torch.softmax(logits, dim=1).squeeze(0).cpu().numpy() | |
| order = np.argsort(probs)[::-1] | |
| top_k_results = {IDX2LABEL[int(i)]: float(probs[int(i)]) for i in order[:top_k]} | |
| pred_idx = int(order[0]) | |
| pred_label = IDX2LABEL[pred_idx] | |
| pred_conf = float(probs[pred_idx]) | |
| return pred_label, pred_conf, top_k_results | |