File size: 11,979 Bytes
aee40b4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 | import argparse
import sys
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
import torchaudio
from transformers import WavLMModel
# Ensure UTF-8 output encoding for Windows terminal compatibility
if sys.stdout.encoding and sys.stdout.encoding.lower() != 'utf-8':
try:
sys.stdout.reconfigure(encoding='utf-8')
except Exception:
pass
# ============================================================
# CONFIGURATION & CONSTANTS
# ============================================================
WAVLM_MODEL_NAME = "microsoft/wavlm-base-plus"
CHECKPOINT_PATH = Path("checkpoints/best_model.pt")
INPUT_DIR = Path("input")
TARGET_SR = 16000
EMOTION_CLASSES = ["Anger", "Disgust", "Fear", "Happy", "Neutral", "Sad"]
# ============================================================
# MODEL ARCHITECTURE DEFINITION (Self-Contained Single File)
# ============================================================
class BiLSTMFeatureExtractor(nn.Module):
def __init__(self, input_size=768, hidden_size=128, num_layers=1, dropout=0.0):
super(BiLSTMFeatureExtractor, self).__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.num_layers = num_layers
self.bidirectional = True
self.bilstm = nn.LSTM(
input_size=input_size,
hidden_size=hidden_size,
num_layers=num_layers,
batch_first=True,
bidirectional=True,
dropout=dropout if num_layers > 1 else 0.0
)
def forward(self, x, mask=None):
batch_size, seq_len, _ = x.shape
if mask is not None:
lengths = mask.sum(dim=1).cpu()
packed_x = pack_padded_sequence(
x,
lengths,
batch_first=True,
enforce_sorted=False
)
packed_out, (hn, cn) = self.bilstm(packed_x)
out, _ = pad_packed_sequence(
packed_out,
batch_first=True,
total_length=seq_len
)
else:
out, (hn, cn) = self.bilstm(x)
return out
class TemporalAttention(nn.Module):
def __init__(self, input_dim=256):
super(TemporalAttention, self).__init__()
self.input_dim = input_dim
self.w = nn.Linear(input_dim, 1, bias=False)
def forward(self, h, mask=None):
scores = self.w(torch.tanh(h)).squeeze(-1) # [B, T]
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
attn_weights = torch.softmax(scores, dim=1) # [B, T]
context = torch.bmm(attn_weights.unsqueeze(1), h).squeeze(1) # [B, 256]
return context, attn_weights
class BiLSTMAttentionClassifier(nn.Module):
def __init__(self, input_size=768, hidden_size=128, num_classes=6, dropout=0.3):
super(BiLSTMAttentionClassifier, self).__init__()
self.bilstm = BiLSTMFeatureExtractor(
input_size=input_size,
hidden_size=hidden_size,
num_layers=1,
dropout=dropout
)
context_dim = hidden_size * 2 # 256
self.attention = TemporalAttention(input_dim=context_dim)
self.dropout = nn.Dropout(dropout)
self.classifier = nn.Linear(context_dim, num_classes)
def forward(self, x, mask=None):
bilstm_out = self.bilstm(x, mask=mask)
context, attn_weights = self.attention(bilstm_out, mask=mask)
dropped_context = self.dropout(context)
logits = self.classifier(dropped_context)
return logits, attn_weights
# ============================================================
# EXACT AUDIO PREPROCESSING LOGIC (Used During Training)
# ============================================================
def preprocess_audio(audio_path):
"""
Load raw audio file (.wav, .mp3, .flac, .ogg), convert to 16 kHz mono, and peak-normalize amplitude.
"""
audio_path = Path(audio_path)
# 1. Load audio with torchaudio, fallback to soundfile for mp3/flac if needed
try:
waveform, sample_rate = torchaudio.load(str(audio_path))
except Exception:
import soundfile as sf
data, sample_rate = sf.read(str(audio_path))
waveform = torch.tensor(data, dtype=torch.float32)
if waveform.ndim == 1:
waveform = waveform.unsqueeze(0)
elif waveform.ndim == 2:
waveform = waveform.T
# 2. Convert to mono if multi-channel
if waveform.shape[0] > 1:
waveform = waveform.mean(dim=0, keepdim=True)
# 3. Resample to 16 kHz if necessary
if sample_rate != TARGET_SR:
resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=TARGET_SR)
waveform = resampler(waveform)
# 4. Squeeze channel dimension [1, num_samples] -> [num_samples]
waveform = waveform.squeeze(0)
# 5. Amplitude peak normalization
max_val = waveform.abs().max()
if max_val > 0:
waveform = waveform / max_val
return waveform
# ============================================================
# FINAL PREDICTOR PIPELINE CLASS
# ============================================================
class FinalPredictor:
def __init__(self, checkpoint_path=CHECKPOINT_PATH):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("=" * 60)
print("INITIALIZING F1 DRIVER TONE DETECTOR PIPELINE")
print("=" * 60)
print(f"Compute Device : {self.device}")
if self.device.type == "cuda":
print(f"GPU : {torch.cuda.get_device_name(0)}")
print(f"\n[1] Loading WavLM Encoder ({WAVLM_MODEL_NAME})...")
self.wavlm = WavLMModel.from_pretrained(WAVLM_MODEL_NAME).to(self.device)
self.wavlm.eval()
print(f"[2] Loading Trained Downstream Model ({checkpoint_path})...")
if not checkpoint_path.exists():
raise FileNotFoundError(f"Checkpoint file not found at {checkpoint_path}. Please run training first!")
self.classifier = BiLSTMAttentionClassifier(
input_size=768,
hidden_size=128,
num_classes=6
).to(self.device)
checkpoint = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
self.classifier.load_state_dict(checkpoint["model_state_dict"])
self.classifier.eval()
print("Pipeline initialized and ready for inference!")
print("=" * 60)
def predict_single(self, audio_path, chunk_duration=2.5, hop_duration=1.5):
audio_path = Path(audio_path)
if not audio_path.exists():
raise FileNotFoundError(f"Audio file not found: {audio_path}")
# 1. Preprocess raw audio waveform [num_samples]
waveform = preprocess_audio(audio_path)
num_samples = waveform.size(0)
duration_sec = num_samples / float(TARGET_SR)
chunk_samples = int(chunk_duration * TARGET_SR)
hop_samples = int(hop_duration * TARGET_SR)
chunk_probs = []
last_attn_weights = None
with torch.no_grad():
# If audio is long (> 4.0s), use 2.5s sliding window chunking
if duration_sec > 4.0 and num_samples > chunk_samples:
starts = list(range(0, num_samples - chunk_samples + 1, hop_samples))
# Ensure the end of audio is covered
if starts[-1] + chunk_samples < num_samples:
starts.append(num_samples - chunk_samples)
for start in starts:
end = start + chunk_samples
chunk_wave = waveform[start:end].unsqueeze(0).to(self.device)
outputs = self.wavlm(input_values=chunk_wave)
embedding = outputs.last_hidden_state
mask = torch.ones((1, embedding.size(1)), dtype=torch.int64, device=self.device)
logits, attn_weights = self.classifier(embedding, mask=mask)
probs = F.softmax(logits, dim=-1).squeeze(0)
chunk_probs.append(probs)
last_attn_weights = attn_weights.squeeze(0).cpu().numpy()
probabilities = torch.stack(chunk_probs).mean(dim=0)
else:
input_values = waveform.unsqueeze(0).to(self.device)
outputs = self.wavlm(input_values=input_values)
embedding = outputs.last_hidden_state
mask = torch.ones((1, embedding.size(1)), dtype=torch.int64, device=self.device)
logits, attn_weights = self.classifier(embedding, mask=mask)
probabilities = F.softmax(logits, dim=-1).squeeze(0)
last_attn_weights = attn_weights.squeeze(0).cpu().numpy()
pred_id = torch.argmax(probabilities).item()
pred_emotion = EMOTION_CLASSES[pred_id]
confidence = probabilities[pred_id].item() * 100.0
probs_dict = {
EMOTION_CLASSES[i]: probabilities[i].item() * 100.0
for i in range(len(EMOTION_CLASSES))
}
return {
"audio_file": audio_path.name,
"duration_sec": round(duration_sec, 2),
"predicted_emotion": pred_emotion,
"confidence": confidence,
"probabilities": probs_dict,
"attention_weights": last_attn_weights,
"chunks_processed": len(chunk_probs) if chunk_probs else 1
}
# ============================================================
# MAIN CLI DRIVER
# ============================================================
def print_report(result):
print("\n" + "=" * 60)
print(f"PREDICTION REPORT: {result['audio_file']}")
print("=" * 60)
print(f"Predicted Emotion : {result['predicted_emotion']} ({result['confidence']:.2f}% confidence)")
print("-" * 60)
print("EMOTION PROBABILITY BREAKDOWN:")
for emotion, prob in result["probabilities"].items():
bar = "#" * int(prob / 5)
print(f" {emotion:10s} : {prob:6.2f}% | {bar}")
print("=" * 60)
def main():
parser = argparse.ArgumentParser(description="F1 Driver Tone Predictor (Single File Pipeline)")
parser.add_argument("--audio_path", type=str, default=None,
help="Path to specific raw driver audio file (.wav, .mp3, .flac, .ogg)")
args = parser.parse_args()
predictor = FinalPredictor()
if args.audio_path:
target_path = Path(args.audio_path)
result = predictor.predict_single(target_path)
print_report(result)
else:
# Scan input directory
INPUT_DIR.mkdir(parents=True, exist_ok=True)
valid_extensions = {".wav", ".mp3", ".flac", ".ogg", ".m4a"}
audio_files = [
f for f in INPUT_DIR.iterdir()
if f.is_file() and f.suffix.lower() in valid_extensions
]
if not audio_files:
print(f"\n[INFO] No audio files found in input/ directory ({INPUT_DIR.resolve()}).")
print("Usage Options:")
print(" 1. Place .wav or .mp3 files into the 'input/' folder and re-run python finalpredictor.py")
print(" 2. Run with explicit file path: python finalpredictor.py --audio_path <path_to_audio_file>")
else:
print(f"\nFound {len(audio_files)} audio file(s) in {INPUT_DIR}/ directory. Processing...")
for audio_file in audio_files:
result = predictor.predict_single(audio_file)
print_report(result)
if __name__ == "__main__":
main()
|