#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Audio Feature Extractor Script Main Functions: 1. Extract word-level features from single audio files 2. Use Whisper for speech recognition 3. Use Wav2Vec2 for word-level alignment with Triton optimization 4. Support English audio processing Usage: extractor = AudioFeatureExtractor() features = extractor.extract_features("path/to/audio.wav") Basic Implementation Logic: 1. Load audio file using librosa 2. Use Whisper for speech transcription 3. Use Wav2Vec2 for word-level alignment (Triton optimized) 4. Extract audio features based on word-level timestamps 5. Return feature dictionary with word segments and audio features """ import os import json import warnings import argparse from typing import List, Optional, Dict, Any from dataclasses import dataclass import numpy as np import torch import librosa import soundfile as sf import parselmouth # Check Triton availability try: import triton import triton.language as tl TRITON_AVAILABLE = True except ImportError: TRITON_AVAILABLE = False print("Warning: Triton not available, will use original implementation") # HuggingFace libraries from transformers import ( WhisperProcessor, WhisperForConditionalGeneration, Wav2Vec2Processor, Wav2Vec2ForCTC ) # ===== Configuration Constants ===== # Audio processing parameters SAMPLE_RATE = 16000 # Standard sample rate for Whisper and Wav2Vec2 MAX_DURATION = 30 # Maximum audio segment duration (seconds) # Default model paths for English DEFAULT_WHISPER_MODEL = "openai/whisper-large-v3" DEFAULT_ALIGN_MODEL = "facebook/wav2vec2-large-960h-lv60-self" # ===== Data Structure Definitions ===== @dataclass class WordSegment: """Word-level segment information""" word: str # Word text start: Optional[float] # Start time (seconds) end: Optional[float] # End time (seconds) score: Optional[float] # Confidence score @dataclass class AlignedSegment: """Aligned sentence segment""" text: str # Sentence text start: Optional[float] # Start time (seconds) end: Optional[float] # End time (seconds) words: List[WordSegment] # Word-level information list class AudioFeatureExtractor: """ Audio Feature Extractor Class Main Functions: 1. Load and process single audio files 2. Extract word-level features using Triton optimization 3. Support English audio processing Usage: extractor = AudioFeatureExtractor() features = extractor.extract_features("audio.wav") """ def __init__(self, whisper_model: str = DEFAULT_WHISPER_MODEL, align_model: str = DEFAULT_ALIGN_MODEL, device: str = "auto", merge_threshold: float = 0.5): """ Initialize Audio Feature Extractor Args: whisper_model: Path to Whisper model for speech recognition align_model: Path to Wav2Vec2 model for word alignment device: Computing device ("auto", "cpu", "cuda") merge_threshold: Word merging threshold (seconds) Implementation Logic: 1. Set up device configuration 2. Load Whisper and Wav2Vec2 models 3. Initialize vocabulary for alignment 4. Ensure Triton optimization is available """ self.merge_threshold = merge_threshold # Device selection if device == "auto": self.device = "cuda" if torch.cuda.is_available() else "cpu" else: self.device = device print(f"🚀 Initializing Audio Feature Extractor...") print(f" Device: {self.device}") print(f" Whisper Model: {whisper_model}") print(f" Align Model: {align_model}") print(f" Word Merge Threshold: {merge_threshold}s") print(f" Triton Available: {TRITON_AVAILABLE}") # Ensure Triton is available for optimization if not TRITON_AVAILABLE: raise RuntimeError("Triton is required but not available. Please install triton.") if self.device != "cuda": raise RuntimeError("Triton optimization requires CUDA device.") # Load models self._load_models(whisper_model, align_model) print("✅ Audio Feature Extractor initialized successfully") def _load_models(self, whisper_model: str, align_model: str): """ Load Whisper and Wav2Vec2 models Args: whisper_model: Path to Whisper model align_model: Path to Wav2Vec2 alignment model Implementation Logic: 1. Load Whisper model for speech recognition 2. Load Wav2Vec2 model for word-level alignment 3. Set models to evaluation mode 4. Build character-level vocabulary dictionary """ try: # Load Whisper model print("📥 Loading Whisper model...") self.whisper_processor = WhisperProcessor.from_pretrained(whisper_model) self.whisper_model = WhisperForConditionalGeneration.from_pretrained(whisper_model) self.whisper_model.to(self.device) self.whisper_model.eval() # Load Wav2Vec2 alignment model print(f"📥 Loading Wav2Vec2 alignment model...") self.align_processor = Wav2Vec2Processor.from_pretrained(align_model) self.align_model = Wav2Vec2ForCTC.from_pretrained(align_model) self.align_model.to(self.device) self.align_model.eval() # Build character-level vocabulary dictionary labels = self.align_processor.tokenizer.get_vocab() # Create character to ID mapping, convert all characters to lowercase self.vocab = {char.lower(): code for char, code in labels.items()} self.id_to_token = {v: k for k, v in self.vocab.items()} print("✅ Models loaded successfully") print(f"Vocabulary size: {len(self.vocab)}") except Exception as e: print(f"❌ Model loading failed: {e}") raise def load_audio_file(self, audio_path: str) -> tuple[np.ndarray, int]: """ Load audio file from given path Args: audio_path: Path to audio file (absolute or relative) Returns: tuple: (audio_array, sampling_rate) Implementation Logic: 1. Check if path exists 2. Load audio using librosa 3. Return audio array and sampling rate 4. Handle errors gracefully """ try: # Check if file exists if not os.path.exists(audio_path): raise FileNotFoundError(f"Audio file not found: {audio_path}") # Load audio file audio_array, sampling_rate = librosa.load(audio_path, sr=None) print(f"📁 Loaded audio: {audio_path}") print(f" Duration: {len(audio_array)/sampling_rate:.2f}s") print(f" Sample Rate: {sampling_rate}Hz") return audio_array, sampling_rate except Exception as e: print(f"❌ Failed to load audio file: {audio_path}, Error: {e}") raise def transcribe_audio(self, audio: np.ndarray, sampling_rate: int) -> str: """ Transcribe audio using Whisper model Args: audio: Audio array sampling_rate: Sampling rate Returns: Transcribed text Implementation Logic: 1. Resample audio to 16kHz if needed 2. Preprocess audio for Whisper 3. Generate transcription using Whisper model 4. Return cleaned transcription text """ try: # Resample to 16kHz if sampling_rate != SAMPLE_RATE: audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=SAMPLE_RATE) # Preprocess audio inputs = self.whisper_processor( audio, sampling_rate=SAMPLE_RATE, return_tensors="pt" ) inputs = inputs.to(self.device) # Generate transcription with torch.no_grad(): predicted_ids = self.whisper_model.generate(inputs["input_features"]) transcription = self.whisper_processor.batch_decode( predicted_ids, skip_special_tokens=True )[0] print(f"🎯 Transcription: {transcription}") return transcription.strip() except Exception as e: print(f"❌ Transcription failed: {e}") return "" def get_word_timestamps(self, audio: np.ndarray, text: str) -> List[AlignedSegment]: """ Get word-level timestamps using Wav2Vec2 forced alignment Args: audio: Audio array text: Transcribed text Returns: List of aligned segments with word-level timestamps Implementation Logic: 1. Preprocess text for alignment 2. Use Wav2Vec2 model for CTC alignment 3. Calculate word-level boundaries using Triton optimization 4. Return aligned segments with timestamps """ try: print("🔄 Starting Wav2Vec2 forced alignment...") # Preprocess text clean_transcript = self._preprocess_text(text) if not clean_transcript: print("Warning: Text preprocessing resulted in empty string") return [AlignedSegment( text=text, start=0.0, end=len(audio) / SAMPLE_RATE, words=[WordSegment( word=text, start=0.0, end=len(audio) / SAMPLE_RATE, score=0.0 )] )] # Preprocess audio inputs = self.align_processor( audio, sampling_rate=SAMPLE_RATE, return_tensors="pt" ) inputs = inputs.to(self.device) # Get model output with torch.no_grad(): logits = self.align_model(inputs.input_values).logits # Convert to log probabilities log_probs = torch.log_softmax(logits, dim=-1) emission = log_probs[0] # Remove batch dimension # Perform CTC alignment with Triton optimization aligned_segments = self._ctc_align_triton(emission, clean_transcript, audio) print("✅ Forced alignment completed") return aligned_segments except Exception as e: print(f"❌ Word-level alignment failed: {e}") return [] def _preprocess_text(self, text: str) -> str: """ Preprocess text by removing characters not in vocabulary Args: text: Original text Returns: Cleaned text Implementation Logic: 1. Convert to lowercase 2. Replace spaces with | (Wav2Vec2 convention) 3. Keep only characters in vocabulary 4. Replace unknown characters with wildcards """ # Convert to lowercase text = text.lower().strip() # Replace spaces with | (Wav2Vec2 convention for English) text = text.replace(" ", "|") # Keep only characters in vocabulary clean_chars = [] for char in text: if char in self.vocab: clean_chars.append(char) else: # Replace unknown characters with wildcards clean_chars.append("*") return "".join(clean_chars) def _ctc_align_triton(self, emission: torch.Tensor, transcript: str, audio: np.ndarray) -> List[AlignedSegment]: """ Perform CTC forced alignment using Triton optimization Args: emission: Model output emission probabilities transcript: Cleaned transcript text audio: Original audio array Returns: List of aligned segments Implementation Logic: 1. Convert text to token IDs 2. Build trellis using Triton-optimized kernels 3. Backtrack optimal path 4. Merge repeated characters 5. Generate word alignments with timestamps """ # Convert text to token IDs tokens = [self.vocab.get(char, self.vocab.get("[UNK]", 0)) for char in transcript] # Get blank token ID blank_id = self.vocab.get("[PAD]", 0) if "[PAD]" not in self.vocab: blank_id = self.vocab.get("", 0) # Build trellis using Triton optimization trellis = self._get_trellis(emission, tokens, blank_id) # Backtrack optimal path path = self._backtrack(trellis, emission, tokens, blank_id) if path is None: print("Warning: CTC alignment failed, returning original timestamps") return [AlignedSegment( text=transcript.replace("|", " "), start=0.0, end=len(audio) / SAMPLE_RATE, words=[WordSegment( word=transcript.replace("|", " "), start=0.0, end=len(audio) / SAMPLE_RATE, score=0.0 )] )] # Merge repeated characters char_segments = self._merge_repeats(path, transcript) # Convert to timestamps duration = len(audio) / SAMPLE_RATE time_ratio = duration / (emission.size(0) - 1) # Generate word-level alignments words = self._generate_word_alignments(char_segments, transcript, time_ratio) return [AlignedSegment( text=transcript.replace("|", " "), start=words[0].start if words else 0.0, end=words[-1].end if words else duration, words=words )] @staticmethod @triton.jit def _trellis_row_kernel_optimized( # Pointers trellis_t_ptr, trellis_tm1_ptr, emission_t_ptr, tokens_ptr, # Scalar arguments num_tokens, blank_emit: tl.float32, t, # Tensor strides trellis_stride_n, # Meta-parameters BLOCK_SIZE_N: tl.constexpr, ): """ Triton-optimized kernel for trellis row computation This kernel computes one row of the CTC trellis matrix in parallel, significantly speeding up the forced alignment process. Implementation Logic: 1. Calculate parallel indices for token positions 2. Load previous trellis values (stay and advance paths) 3. Load emission probabilities for current tokens 4. Compute path scores using numerically stable logsumexp 5. Store results back to trellis """ # Calculate parallel indices starting from j=1 pid = tl.program_id(axis=0) offs_n = pid * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + 1 # Ensure indices are within bounds mask = (offs_n < num_tokens) & (offs_n <= t) # Load trellis[t-1, j] (stay path) prev_stay_ptr = trellis_tm1_ptr + offs_n * trellis_stride_n prev_stay_score = tl.load(prev_stay_ptr, mask=mask, other=float('-inf')) # Load trellis[t-1, j-1] (advance path) prev_advance_ptr = trellis_tm1_ptr + (offs_n - 1) * trellis_stride_n prev_advance_score = tl.load(prev_advance_ptr, mask=mask, other=float('-inf')) # Load emission[t, tokens[j]] tokens_j = tl.load(tokens_ptr + offs_n, mask=mask, other=0) emission_token = tl.load(emission_t_ptr + tokens_j, mask=mask, other=float('-inf')) # Calculate path scores stay_score = prev_stay_score + blank_emit advance_score = prev_advance_score + emission_token # Numerically stable logsumexp max_val = tl.maximum(stay_score, advance_score) min_val = tl.minimum(stay_score, advance_score) log_sum = tl.where( max_val > float('-inf'), max_val + tl.log(1.0 + tl.exp(min_val - max_val)), float('-inf') ) # Store results trellis_t_ptr_j = trellis_t_ptr + offs_n * trellis_stride_n tl.store(trellis_t_ptr_j, log_sum, mask=mask) def _get_trellis(self, emission: torch.Tensor, tokens: List[int], blank_id: int) -> torch.Tensor: """ Build CTC alignment trellis using Triton optimization Args: emission: Emission probability matrix [T, V] tokens: Token ID list blank_id: Blank token ID Returns: Trellis matrix [T, N] """ # Use Triton optimized version if available and on CUDA device if TRITON_AVAILABLE and torch.cuda.is_available() and emission.device.type == 'cuda': return self._get_trellis_triton(emission, tokens, blank_id) else: # Fallback to original implementation print("Warning: Triton not enabled or CUDA unavailable, falling back to original CTC alignment implementation") return self._get_trellis_original(emission, tokens, blank_id) def _get_trellis_original(self, emission: torch.Tensor, tokens: List[int], blank_id: int) -> torch.Tensor: """ Original trellis construction implementation (as fallback) """ num_frame = emission.size(0) num_tokens = len(tokens) # Initialize trellis trellis = torch.full((num_frame, num_tokens), float('-inf'), device=emission.device) trellis[0, 0] = emission[0, blank_id] # Fill first row for t in range(1, num_frame): trellis[t, 0] = trellis[t-1, 0] + emission[t, blank_id] # Fill trellis for t in range(1, num_frame): for j in range(1, min(num_tokens, t + 1)): # Stay at current token (insert blank) stay_score = trellis[t-1, j] + emission[t, blank_id] # Advance to next token advance_score = trellis[t-1, j-1] + emission[t, tokens[j]] trellis[t, j] = torch.logsumexp(torch.stack([stay_score, advance_score]), dim=0) return trellis def _get_trellis_triton(self, emission: torch.Tensor, tokens: List[int], blank_id: int) -> torch.Tensor: """ Triton-optimized trellis construction - optimized version """ assert emission.is_cuda, "Input tensor must be on CUDA device" num_frame, vocab_size = emission.size() tokens_tensor = torch.tensor(tokens, device=emission.device, dtype=torch.long) num_tokens = len(tokens_tensor) # --- Optimization 3: Ensure memory contiguity --- # Fix: Remove memory_format parameter, use .contiguous() method to ensure memory contiguity trellis = torch.full((num_frame, num_tokens), float('-inf'), device=emission.device, dtype=torch.float32).contiguous() if num_tokens == 0: return trellis # --- Optimization 4: Use vectorized cumsum to initialize first column --- # Calculate cumulative blank probabilities from t=0 # Note: For consistency with original logic, use emission directly instead of log_softmax trellis[:, 0] = emission[:, blank_id].cumsum(dim=0) # --- Optimization 2: Dynamic block size --- # Adapt to different token sequence lengths, improve GPU utilization BLOCK_SIZE_N = min(1024, triton.next_power_of_2(num_tokens)) if num_tokens > 1 else 1 # Main loop for t in range(1, num_frame): # --- Optimization 1: Scalar broadcasting --- # Pass blank emission as scalar to avoid redundant loading blank_emit = emission[t, blank_id].item() # Launch grid, only compute j > 0 part if num_tokens > 1: grid = lambda meta: (triton.cdiv(num_tokens - 1, meta['BLOCK_SIZE_N']),) self._trellis_row_kernel_optimized[grid]( trellis_t_ptr=trellis[t], trellis_tm1_ptr=trellis[t-1], emission_t_ptr=emission[t], tokens_ptr=tokens_tensor, num_tokens=num_tokens, blank_emit=blank_emit, t=t, trellis_stride_n=trellis.stride(1), BLOCK_SIZE_N=BLOCK_SIZE_N, ) return trellis def _backtrack(self, trellis: torch.Tensor, emission: torch.Tensor, tokens: List[int], blank_id: int) -> Optional[List]: """ Backtrack through trellis to find optimal alignment path Args: trellis: Completed trellis matrix emission: Emission probabilities tokens: Token ID list blank_id: Blank token ID Returns: Optimal path through trellis Implementation Logic: 1. Start from final position in trellis 2. Trace back through highest probability path 3. Record token positions and timestamps 4. Return path for character merging """ # Implementation details would be similar to the original _backtrack method # This is a simplified version for demonstration try: num_frame, num_tokens = trellis.size() # Start from the end t, j = num_frame - 1, num_tokens - 1 path = [] while t >= 0 and j >= 0: path.append((t, j)) if j == 0: t -= 1 elif t == 0: j -= 1 else: # Choose path with higher probability stay_score = trellis[t-1, j] + emission[t, blank_id] advance_score = trellis[t-1, j-1] + emission[t, tokens[j]] if stay_score > advance_score: t -= 1 else: t -= 1 j -= 1 return list(reversed(path)) except Exception as e: print(f"Backtracking failed: {e}") return None def _merge_repeats(self, path: List, transcript: str) -> List: """ Merge repeated characters in alignment path Args: path: Alignment path from backtracking transcript: Original transcript Returns: List of character segments with merged repeats Implementation Logic: 1. Group consecutive identical characters 2. Calculate start and end frames for each character 3. Return character segments for word boundary detection """ if not path: return [] char_segments = [] current_char = None start_frame = None for t, j in path: char = transcript[j] if j < len(transcript) else None if char != current_char: if current_char is not None: char_segments.append({ 'char': current_char, 'start': start_frame, 'end': t - 1 }) current_char = char start_frame = t # Add final character if current_char is not None: char_segments.append({ 'char': current_char, 'start': start_frame, 'end': path[-1][0] }) return char_segments def _generate_word_alignments(self, char_segments: List, transcript: str, time_ratio: float) -> List[WordSegment]: """ Generate word-level alignments from character segments Args: char_segments: Character-level segments transcript: Original transcript time_ratio: Frame to time conversion ratio Returns: List of word segments with timestamps Implementation Logic: 1. Group characters into words using | delimiter 2. Calculate word boundaries from character segments 3. Convert frame indices to time stamps 4. Return word segments with confidence scores """ words = [] current_word = "" word_start = None word_chars = [] for segment in char_segments: char = segment['char'] if char == '|': # Word boundary if current_word and word_chars: # Calculate word timing start_time = word_chars[0]['start'] * time_ratio end_time = word_chars[-1]['end'] * time_ratio words.append(WordSegment( word=current_word, start=start_time, end=end_time, score=1.0 # Simplified confidence score )) current_word = "" word_chars = [] else: current_word += char word_chars.append(segment) # Add final word if current_word and word_chars: start_time = word_chars[0]['start'] * time_ratio end_time = word_chars[-1]['end'] * time_ratio words.append(WordSegment( word=current_word, start=start_time, end=end_time, score=1.0 )) return words def merge_short_words(self, word_segments: List[WordSegment]) -> List[WordSegment]: """ Merge short words with neighboring words Args: word_segments: List of word segments Returns: List of merged word segments Implementation Logic: 1. Identify words shorter than merge threshold 2. Find shortest neighboring word for merging 3. Merge words and update timestamps 4. Repeat until no more merging is needed """ if not word_segments: return [] merged_segments = word_segments.copy() while True: # Find short words short_indices = [] for i, segment in enumerate(merged_segments): if segment.start is not None and segment.end is not None: duration = segment.end - segment.start if duration < self.merge_threshold: short_indices.append(i) if not short_indices: break # Merge shortest word with its shortest neighbor shortest_idx = min(short_indices, key=lambda i: merged_segments[i].end - merged_segments[i].start) neighbor_idx = self._find_shortest_neighbor(merged_segments, shortest_idx) if neighbor_idx is not None: # Merge segments merged_segment = self._merge_two_segments( merged_segments[shortest_idx], merged_segments[neighbor_idx] ) # Remove original segments and insert merged one indices_to_remove = sorted([shortest_idx, neighbor_idx], reverse=True) for idx in indices_to_remove: merged_segments.pop(idx) # Insert merged segment at appropriate position insert_pos = min(shortest_idx, neighbor_idx) merged_segments.insert(insert_pos, merged_segment) else: break return merged_segments def _find_shortest_neighbor(self, segments: List[WordSegment], current_idx: int) -> Optional[int]: """ Find the shortest neighboring word for merging Args: segments: List of word segments current_idx: Index of current word Returns: Index of shortest neighbor, or None if no valid neighbor """ neighbors = [] # Check left neighbor if current_idx > 0: neighbors.append(current_idx - 1) # Check right neighbor if current_idx < len(segments) - 1: neighbors.append(current_idx + 1) if not neighbors: return None # Find shortest neighbor shortest_neighbor = min(neighbors, key=lambda i: segments[i].end - segments[i].start if segments[i].start and segments[i].end else float('inf')) return shortest_neighbor def _merge_two_segments(self, segment1: WordSegment, segment2: WordSegment) -> WordSegment: """ Merge two word segments into one Args: segment1: First word segment segment2: Second word segment Returns: Merged word segment Implementation Logic: 1. Combine word texts with space 2. Use earliest start time 3. Use latest end time 4. Average confidence scores """ # Determine order based on start times if segment1.start <= segment2.start: first, second = segment1, segment2 else: first, second = segment2, segment1 # Merge word texts merged_word = f"{first.word} {second.word}" # Merge timestamps merged_start = first.start merged_end = second.end # Average confidence scores merged_score = (first.score + second.score) / 2 if first.score and second.score else None return WordSegment( word=merged_word, start=merged_start, end=merged_end, score=merged_score ) def extract_audio_features(self, audio: np.ndarray, sampling_rate: int, word_segments: List[WordSegment]) -> Dict[str, Any]: """ Extract comprehensive audio features for word segments (compatible with extractor.py format) Args: audio: Audio array sampling_rate: Sampling rate word_segments: List of word segments with timestamps Returns: Dictionary containing extracted features matching extractor.py format Implementation Logic: 1. Extract features for each word segment 2. Calculate acoustic features (pitch, energy, spectral) 3. Calculate speaking rate and other statistics 4. Return features in extractor.py compatible format """ # Calculate total duration total_duration = len(audio) / sampling_rate # Count original and processed words original_word_count = len(word_segments) processed_word_count = 0 word_features = [] for word_segment in word_segments: if word_segment.start is None or word_segment.end is None: continue # Extract audio segment for this word start_sample = int(word_segment.start * sampling_rate) end_sample = int(word_segment.end * sampling_rate) word_audio = audio[start_sample:end_sample] if len(word_audio) == 0: continue # Extract acoustic features word_feature = self._extract_word_features(word_audio, sampling_rate, word_segment) word_features.append(word_feature) processed_word_count += 1 # Calculate speaking rate (words per minute) speaking_rate = (processed_word_count / total_duration * 60) if total_duration > 0 else 0.0 # Return features in extractor.py compatible format return { "total_duration": total_duration, "speaking_rate": speaking_rate, "original_word_count": original_word_count, "processed_word_count": processed_word_count, "word_features": word_features } def _extract_word_features(self, word_audio: np.ndarray, sampling_rate: int, word_segment: WordSegment) -> Dict[str, Any]: """ Extract acoustic features for a single word (compatible with extractor.py format) Args: word_audio: Audio array for the word sampling_rate: Sampling rate word_segment: Word segment information Returns: Dictionary of acoustic features matching extractor.py format Implementation Logic: 1. Calculate basic timing features 2. Extract pitch features using Parselmouth (pitch_mean, pitch_slope) 3. Calculate energy features (energy_rms, energy_slope) 4. Calculate spectral features (spectral_centroid) 5. Return features in extractor.py compatible format """ # Clean word text, remove special symbols clean_word = word_segment.word # Initialize features with basic information features = { "word": clean_word, "start_time": word_segment.start, "end_time": word_segment.end, "duration": word_segment.end - word_segment.start, "confidence_score": word_segment.score } try: # === Pitch feature extraction === avg_pitch = np.nan pitch_slope = np.nan try: # Create temporary audio file for parselmouth analysis import tempfile with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_audio: sf.write(temp_audio.name, word_audio, sampling_rate) sound = parselmouth.Sound(temp_audio.name) pitch = sound.to_pitch(pitch_floor=50.0, pitch_ceiling=600.0) pitch_times = pitch.xs() pitch_values = pitch.selected_array['frequency'] # Remove unvoiced frames (pitch = 0) pitch_values[pitch_values == 0] = np.nan valid_pitch = pitch_values[~np.isnan(pitch_values)] if len(valid_pitch) > 0: avg_pitch = np.mean(valid_pitch) # Calculate pitch change trend (slope) - safer method if len(valid_pitch) >= 15: # Only calculate slope when there are enough points (using linear regression) duration = features["duration"] time_points = np.linspace(0, duration, len(valid_pitch)) coeffs = np.polyfit(time_points, valid_pitch, 1) pitch_slope = coeffs[0] # Slope else: # Set to NaN when too few data points, safer pitch_slope = np.nan else: # Set to default values when no valid pitch values avg_pitch = np.nan pitch_slope = np.nan # Clean up temporary file try: os.unlink(temp_audio.name) except: pass except Exception as e: print(f"Warning: Pitch extraction failed for word '{clean_word}': {e}") avg_pitch = np.nan pitch_slope = np.nan # === Energy feature extraction === rms_energy = np.sqrt(np.mean(word_audio**2)) # RMS energy try: # 使用已创建的 sound 对象 intensity = sound.to_intensity() intensity_times = intensity.xs() intensity_values = intensity.values[0] # 计算对应时间段的强度值 word_duration = word_segment.end - word_segment.start if len(intensity_values) > 2: start_energy = np.nanmean(intensity_values[:len(intensity_values)//3]) end_energy = np.nanmean(intensity_values[len(intensity_values)*2//3:]) if not (np.isnan(start_energy) or np.isnan(end_energy)): energy_slope = (end_energy - start_energy) / word_duration except Exception as e: print(f"Warning: Energy slope calculation failed for word '{clean_word}': {e}") energy_slope = np.nan # === Spectral feature extraction === try: # Spectral centroid: reflects timbre characteristics segment_length = len(word_audio) if segment_length < 2048: n_fft = 2 ** int(np.log2(segment_length)) n_fft = max(n_fft, 512) # Minimum value is 512 else: n_fft = 2048 # Default value spectral_centroid = librosa.feature.spectral_centroid( y=word_audio, sr=sampling_rate, n_fft=n_fft)[0].mean() except Exception as e: print(f"Warning: Spectral centroid calculation failed for word '{clean_word}': {e}") spectral_centroid = np.nan # Store all features for current word (matching extractor.py format) features.update({ "pitch_mean": float(avg_pitch) if not np.isnan(avg_pitch) else None, "pitch_slope": float(pitch_slope) if not np.isnan(pitch_slope) else None, "energy_rms": float(rms_energy), "energy_slope": float(energy_slope) if not np.isnan(energy_slope) else None, "spectral_centroid": float(spectral_centroid) if not np.isnan(spectral_centroid) else None }) except Exception as e: print(f"Warning: Feature extraction failed for word '{clean_word}': {e}") # Set default values for failed features features.update({ "pitch_mean": None, "pitch_slope": None, "energy_rms": 0.0, "energy_slope": None, "spectral_centroid": None }) return features def extract_features(self, audio_path: str, text: Optional[str] = None, enable_word_merging: bool = True) -> Dict[str, Any]: """ Main method to extract features from audio file Args: audio_path: Path to audio file text: Optional transcription text (if None, will use Whisper) enable_word_merging: Whether to merge short words Returns: Dictionary containing all extracted features Implementation Logic: 1. Load audio file 2. Transcribe audio if text not provided 3. Get word-level timestamps using Triton optimization 4. Optionally merge short words 5. Extract comprehensive audio features 6. Return complete feature dictionary """ try: print(f"🎵 Starting feature extraction for: {audio_path}") # Load audio file audio_array, sampling_rate = self.load_audio_file(audio_path) # Resample to standard rate if needed if sampling_rate != SAMPLE_RATE: audio_array = librosa.resample( audio_array, orig_sr=sampling_rate, target_sr=SAMPLE_RATE ) sampling_rate = SAMPLE_RATE # Transcribe audio if text not provided if text is None: text = self.transcribe_audio(audio_array, sampling_rate) if not text.strip(): return { "error": "Transcription text is empty", "word_features": [], "audio_path": audio_path } # Get word-level timestamps aligned_segments = self.get_word_timestamps(audio_array, text) if not aligned_segments: return { "error": "Word-level alignment failed", "word_features": [], "audio_path": audio_path, "transcribed_text": text } # Collect all word segments all_word_segments = [] for segment in aligned_segments: all_word_segments.extend(segment.words) # Record original word count original_word_count = len(all_word_segments) # Optional word merging if enable_word_merging: all_word_segments = self.merge_short_words(all_word_segments) print(f"📊 Word merging: {original_word_count} → {len(all_word_segments)} words") # Extract audio features features = self.extract_audio_features(audio_array, sampling_rate, all_word_segments) # Format word features using the provided formatting function formatted_word_features = self.format_word_features(features.get("word_features", [])) # Add metadata and formatted features features.update({ "audio_path": audio_path, "transcribed_text": text, "original_word_count": original_word_count, "final_word_count": len(all_word_segments), "word_merging_enabled": enable_word_merging, "triton_optimization": True, "formatted_word_features": formatted_word_features # Add formatted features }) print(f"✅ Feature extraction completed successfully") print(f" Transcribed text: {text}") print(f" Word count: {len(all_word_segments)}") print(f" Total duration: {features['total_duration']:.2f}s") return features except Exception as e: error_msg = f"Feature extraction failed: {str(e)}" print(f"❌ {error_msg}") return { "error": error_msg, "word_features": [], "audio_path": audio_path } def clean_word_text(self, word: str) -> str: """ Clean word text by removing special symbols and formatting Args: word: Original word text Returns: Cleaned word text """ if not word: return "" # Remove common punctuation and special characters import re # Keep only letters, numbers, and basic punctuation cleaned = re.sub(r'[^\w\s\-\']', '', word) # Remove extra whitespace cleaned = cleaned.strip() return cleaned def format_word_features(self, word_features): """ 格式化word_features数据 保留原有的格式化逻辑，处理word_features列表中的每个特征字典 Args: word_features (list): 原始word_features列表 Returns: str: 格式化后的word_features字符串 """ if not isinstance(word_features, list): print(f"Warning: word_features is not a list: {type(word_features)}") return "" formatted_features = [] for feature_dict in word_features: if not isinstance(feature_dict, dict): print(f"Warning: Feature item is not a dict: {type(feature_dict)}") continue # 提取并格式化各个字段 formatted_feature = {} # 处理word字段 word = feature_dict.get('word', '') formatted_feature['word'] = self.clean_word_text(word) # 处理数值字段，按照要求的精度格式化 numeric_fields = { 'pitch_mean': 0, # 仅保留整数 'pitch_slope': 0, # 仅保留整数 'energy_rms': 3, # 保留三位小数 'energy_slope': 0, # 仅保留整数 'spectral_centroid': 0 # 仅保留整数 } for field, decimal_places in numeric_fields.items(): value = feature_dict.get(field) if value is not None: try: if decimal_places == 0: formatted_feature[field] = int(float(value)) else: formatted_feature[field] = round(float(value), decimal_places) except (ValueError, TypeError): print(f"Warning: Failed to format {field}: {value}") formatted_feature[field] = None else: formatted_feature[field] = None formatted_features.append(formatted_feature) # 转换为字符串格式 return str(formatted_features) def extract_json_from_response(response_text: str) -> Dict[str, Any]: """ Extract JSON data from response text Args: response_text: Text containing JSON data Returns: Extracted JSON dictionary Implementation Logic: 1. Try to parse entire text as JSON 2. If that fails, search for JSON blocks 3. Return parsed JSON or error information """ try: # Try to parse entire text as JSON return json.loads(response_text) except json.JSONDecodeError: # Search for JSON blocks in text import re json_pattern = r'\{[^{}]*(?:\{[^{}]*\}[^{}]*)*\}' matches = re.findall(json_pattern, response_text, re.DOTALL) for match in matches: try: return json.loads(match) except json.JSONDecodeError: continue return {"error": "No valid JSON found in response"} def main(): """ Main function for Audio Feature Extractor with command line support Implementation Logic: 1. Parse command line arguments for audio input and output paths 2. Initialize AudioFeatureExtractor with default settings 3. Process specified audio file 4. Extract features using Triton optimization 5. Save results to specified output file """ # Set up command line argument parser parser = argparse.ArgumentParser( description="Audio Feature Extractor - Extract word-level features from audio files", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: python audio_feature_extractor.py --audio input.wav --output features.json python audio_feature_extractor.py -a speech.wav -o output.json python audio_feature_extractor.py --audio /path/to/audio.wav --output /path/to/features.json """ ) parser.add_argument( '--audio', '-a', type=str, required=True, help='Path to input audio file (required)' ) parser.add_argument( '--output', '-o', type=str, required=True, help='Path to output JSON file for extracted features (required)' ) parser.add_argument( '--device', type=str, default="auto", choices=["auto", "cpu", "cuda"], help='Computing device (default: auto)' ) parser.add_argument( '--merge-threshold', type=float, default=0.5, help='Word merging threshold in seconds (default: 0.5)' ) parser.add_argument( '--text', type=str, default=None, help='Pre-transcribed text (optional, if not provided Whisper will transcribe)' ) parser.add_argument( '--no-word-merging', action='store_true', help='Disable word merging for short segments' ) # Parse arguments args = parser.parse_args() print("🎵 Audio Feature Extractor") print("=" * 50) print(f"📁 Input audio: {args.audio}") print(f"💾 Output file: {args.output}") print(f"🖥️ Device: {args.device}") print(f"⏱️ Merge threshold: {args.merge_threshold}s") print(f"📝 Pre-transcribed text: {'Yes' if args.text else 'No (will use Whisper)'}") print(f"🔗 Word merging: {'Disabled' if args.no_word_merging else 'Enabled'}") try: # Check if input audio file exists if not os.path.exists(args.audio): print(f"❌ Audio file not found: {args.audio}") return 1 # Create output directory if it doesn't exist output_dir = os.path.dirname(args.output) if output_dir and not os.path.exists(output_dir): os.makedirs(output_dir, exist_ok=True) print(f"📁 Created output directory: {output_dir}") # Initialize extractor print(f"\n🚀 Initializing Audio Feature Extractor...") extractor = AudioFeatureExtractor( device=args.device, merge_threshold=args.merge_threshold ) # Extract features print(f"\n🔄 Processing audio file: {args.audio}") features = extractor.extract_features( audio_path=args.audio, text=args.text, enable_word_merging=not args.no_word_merging ) # Display results print("\n📊 Extraction Results:") print("-" * 30) if "error" in features: print(f"❌ Error: {features['error']}") return 1 else: print(f"✅ Success!") print(f" Audio file: {features.get('audio_path', 'N/A')}") print(f" Transcribed text: {features.get('transcribed_text', 'N/A')}") print(f" Total duration: {features.get('total_duration', 0):.2f}s") print(f" Word count: {features.get('final_word_count', 0)}") print(f" Triton optimization: {features.get('triton_optimization', False)}") # Display word-level features summary word_features = features.get('word_features', []) if word_features: print(f"\n📝 Word-level Features Summary:") print(f" Total words: {len(word_features)}") if len(word_features) > 0: print(f" First word: '{word_features[0].get('word', 'N/A')}'") print(f" Last word: '{word_features[-1].get('word', 'N/A')}'") # Show first 3 words details if available if len(word_features) >= 3: print(f"\n📝 First 3 words details:") for i, word_feat in enumerate(word_features[:3]): print(f" Word {i+1}: '{word_feat.get('word', 'N/A')}'") print(f" Duration: {word_feat.get('duration', 0):.3f}s") print(f" Pitch mean: {word_feat.get('pitch_mean', 0):.1f}Hz") print(f" RMS energy: {word_feat.get('rms_energy', 0):.4f}") # Save results to specified output file print(f"\n💾 Saving results to: {args.output}") with open(args.output, 'w', encoding='utf-8') as f: json.dump(features, f, indent=2, ensure_ascii=False) print(f"✅ Results saved successfully!") print("\n🎉 Processing completed successfully!") return 0 except Exception as e: print(f"❌ Processing failed: {e}") import traceback traceback.print_exc() return 1 if __name__ == "__main__": main()