ALM-backend / main2.py
AIZEN-007's picture
Create main2.py
f491cd5 verified
import os
import tempfile
import numpy as np
import pandas as pd
import torch
import torchaudio
import librosa
import matplotlib.pyplot as plt
import csv
from typing import List, Dict, Tuple, Optional
from scipy.stats import kurtosis, skew
import concurrent.futures
import multiprocessing
from functools import partial
import time
import threading
from queue import Queue
from dotenv import load_dotenv
from groq import Groq
# Import required models
from pyannote.audio import Pipeline
import whisper
from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Processor
from torch_vggish_yamnet import yamnet
from torch_vggish_yamnet.input_proc import WaveformToInput
import warnings
warnings.filterwarnings("ignore")
class UnifiedAudioAnalyzer:
"""
Unified Audio Analysis System combining:
1. Speaker Diarization + Transcription
2. Audio Event Detection (YAMNet)
3. Emotion Recognition + Paralinguistic Features
Enhanced with parallel processing for faster execution
"""
def __init__(self, enable_parallel_processing=True, max_workers=None):
"""Initialize all models and components"""
print("πŸ”„ Initializing Unified Audio Analyzer...")
# Configure device
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {self.device}")
# Parallel processing settings
self.enable_parallel_processing = enable_parallel_processing
self.max_workers = max_workers or max(1, multiprocessing.cpu_count() - 1)
print(f"Parallel processing: {'Enabled' if enable_parallel_processing else 'Disabled'}")
if enable_parallel_processing:
print(f"Max workers: {self.max_workers}")
# Initialize models
self._load_diarization_models()
self._load_emotion_models()
self._load_event_detection_models()
self._load_class_names()
print("βœ… All models loaded successfully!")
def _load_diarization_models(self):
"""Load speaker diarization and transcription models"""
print("Loading speaker diarization and transcription models...")
# Load pyannote diarization pipeline
try:
self.diarization_pipeline = Pipeline.from_pretrained(
"pyannote/speaker-diarization-3.1"
# Uncomment and add your token: use_auth_token="YOUR_HUGGINGFACE_TOKEN"
)
if torch.cuda.is_available():
self.diarization_pipeline = self.diarization_pipeline.to(self.device)
except Exception as e:
print(f"Warning: Could not load diarization model: {e}")
self.diarization_pipeline = None
# Load Whisper transcription model
try:
self.whisper_model = whisper.load_model("base")
except Exception as e:
print(f"Warning: Could not load Whisper model: {e}")
self.whisper_model = None
def _load_emotion_models(self):
"""Load emotion recognition models"""
print("Loading emotion recognition models...")
try:
self.emotion_model = Wav2Vec2ForSequenceClassification.from_pretrained(
"Dpngtm/wav2vec2-emotion-recognition"
)
self.emotion_processor = Wav2Vec2Processor.from_pretrained(
"Dpngtm/wav2vec2-emotion-recognition"
)
self.emotion_labels = ["angry", "calm", "disgust", "fearful", "happy", "neutral", "sad", "surprised"]
except Exception as e:
print(f"Warning: Could not load emotion model: {e}")
self.emotion_model = None
def _load_event_detection_models(self):
"""Load YAMNet for audio event detection"""
print("Loading audio event detection models...")
try:
self.yamnet_model = yamnet.yamnet(pretrained=True)
self.yamnet_model.eval()
self.yamnet_converter = WaveformToInput()
except Exception as e:
print(f"Warning: Could not load YAMNet model: {e}")
self.yamnet_model = None
def _load_class_names(self):
"""Load AudioSet class names for YAMNet from CSV"""
csv_path = "yamnet_class_map.csv"
self.audioset_classes = []
try:
with open(csv_path, "r") as f:
reader = csv.reader(f)
next(reader) # skip header
for row in reader:
self.audioset_classes.append(row[2]) # display_name
except Exception as e:
print(f"Warning: Could not load class names from {csv_path}: {e}")
# Fallback to common AudioSet classes
self.audioset_classes = [
"Speech", "Male speech, man speaking", "Female speech, woman speaking",
"Child speech, kid speaking", "Conversation", "Narration, monologue",
"Babbling", "Speech synthesizer", "Shout", "Bellow", "Whoop", "Yell",
"Children shouting", "Screaming", "Whispering", "Laughter", "Baby laughter",
"Giggle", "Snicker", "Belly laugh", "Chuckle, chortle", "Crying, sobbing",
"Baby cry, infant cry", "Whimper", "Wail, moan", "Sigh", "Singing",
"Choir", "Yodeling", "Chant", "Mantra", "Male singing", "Female singing",
"Child singing", "Synthetic singing", "Rapping", "Humming", "Music",
"Musical instrument", "Piano", "Guitar", "Drum", "Orchestra", "Pop music",
"Rock music", "Jazz", "Classical music", "Electronic music", "Animal",
"Dog", "Cat", "Bird", "Insect", "Vehicle", "Car", "Motorcycle", "Train",
"Aircraft", "Helicopter", "Wind", "Rain", "Thunder", "Water", "Fire",
"Applause", "Crowd", "Footsteps", "Door", "Bell", "Alarm", "Clock"
]
def _transcribe_segment_parallel(self, segment_data):
"""Helper function for parallel transcription of segments"""
segment, sample_rate, speaker, start_time, end_time, whisper_model = segment_data
try:
# Create temporary file for this segment
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
temp_filename = temp_file.name
torchaudio.save(temp_filename, segment, sample_rate)
# Transcribe segment
try:
transcription_result = whisper_model.transcribe(
temp_filename,
language="en",
temperature=0,
no_speech_threshold=0.6
)
segment_text = transcription_result["text"].strip()
if segment_text:
result = {
"speaker": speaker,
"start": round(start_time, 2),
"end": round(end_time, 2),
"duration": round(end_time - start_time, 2),
"text": segment_text,
"confidence": transcription_result.get("language_probability", 0.0)
}
else:
result = None
except Exception as e:
print(f"⚠️ Error transcribing segment: {e}")
result = None
finally:
# Clean up temp file
try:
os.unlink(temp_filename)
except OSError:
pass
return result
except Exception as e:
print(f"⚠️ Error in parallel transcription: {e}")
return None
def transcribe_with_diarization(self, audio_file: str, min_segment_duration: float = 1.0) -> List[Dict]:
"""Perform speaker diarization and transcription (aligned with main.py logic)"""
if self.diarization_pipeline is None or self.whisper_model is None:
print("❌ Diarization or transcription models not available")
return []
print("🎯 Performing speaker diarization and transcription...")
# Perform diarization
diarization_result = self.diarization_pipeline(audio_file,num_speakers=2)
# Load audio
waveform, sample_rate = torchaudio.load(audio_file)
if sample_rate != 16000:
waveform = torchaudio.functional.resample(waveform, sample_rate, 16000)
sample_rate = 16000
results = []
temp_files = []
try:
for turn, _, speaker in diarization_result.itertracks(yield_label=True):
if turn.end - turn.start < min_segment_duration:
continue
# Extract segment
start_sample = int(turn.start * sample_rate)
end_sample = int(turn.end * sample_rate)
segment = waveform[:, start_sample:end_sample]
# Create temporary file for transcription
with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_file:
temp_filename = temp_file.name
temp_files.append(temp_filename)
torchaudio.save(temp_filename, segment, sample_rate)
# Transcribe
try:
transcription_result = self.whisper_model.transcribe(
temp_filename,
language="en",
temperature=0,
no_speech_threshold=0.6
)
segment_text = transcription_result["text"].strip()
if segment_text:
results.append({
"speaker": speaker,
"start": round(turn.start, 2),
"end": round(turn.end, 2),
"duration": round(turn.end - turn.start, 2),
"text": segment_text,
"confidence": transcription_result.get("language_probability", 0.0)
})
except Exception as e:
print(f"⚠️ Error transcribing segment: {e}")
continue
finally:
# Cleanup temp files
for temp_file in temp_files:
try:
os.unlink(temp_file)
except OSError:
pass
return results
def detect_audio_events(self, audio_file: str, top_k: int = 10) -> Dict:
"""Detect audio events using YAMNet"""
if self.yamnet_model is None:
print("❌ YAMNet model not available")
return {}
print("πŸ”Š Detecting audio events...")
try:
# Load and preprocess audio
waveform, sr = torchaudio.load(audio_file)
if sr != 16000:
waveform = torchaudio.functional.resample(waveform, sr, 16000)
# Process through YAMNet
inputs = self.yamnet_converter(waveform, 16000)
with torch.no_grad():
embeddings, logits = self.yamnet_model(inputs)
mean_logits = logits.mean(dim=0)
probs = torch.softmax(mean_logits, dim=-1)
top_probs, top_idx = torch.topk(probs, top_k)
# Format results
events = []
for i in range(top_k):
idx = top_idx[i].item()
prob = top_probs[i].item()
if idx < len(self.audioset_classes):
label = self.audioset_classes[idx]
else:
label = f"Unknown_Class_{idx}"
events.append({
"event": label,
"class_id": idx,
"probability": prob
})
return {
"top_events": events,
"total_classes": len(self.audioset_classes)
}
except Exception as e:
print(f"⚠️ Error in event detection: {e}")
return {}
def _extract_feature_chunk(self, audio_chunk, sr, feature_type):
"""Helper function for parallel feature extraction"""
try:
if feature_type == "mfcc":
mfcc = librosa.feature.mfcc(y=audio_chunk, sr=sr, n_mfcc=13)
features = {}
for i in range(13):
features[f'mfcc_{i+1}_mean'] = float(np.mean(mfcc[i]))
features[f'mfcc_{i+1}_std'] = float(np.std(mfcc[i]))
return features
elif feature_type == "chroma":
chroma = librosa.feature.chroma_stft(y=audio_chunk, sr=sr)
features = {}
for i in range(12):
features[f'chroma_{i+1}_mean'] = float(np.mean(chroma[i]))
return features
elif feature_type == "spectral":
features = {}
features['spectral_centroid_mean'] = float(np.mean(librosa.feature.spectral_centroid(y=audio_chunk, sr=sr)[0]))
features['spectral_rolloff_mean'] = float(np.mean(librosa.feature.spectral_rolloff(y=audio_chunk, sr=sr)[0]))
return features
elif feature_type == "basic":
features = {}
features['rms_energy'] = float(np.mean(librosa.feature.rms(y=audio_chunk)[0]))
features['zero_crossing_rate'] = float(np.mean(librosa.feature.zero_crossing_rate(audio_chunk)[0]))
return features
except Exception as e:
print(f"⚠️ Error extracting {feature_type} features: {e}")
return {}
def extract_paralinguistic_features(self, audio_data, sr):
"""Extract comprehensive paralinguistic features"""
print("🎡 Extracting paralinguistic features...")
features = {}
# Basic properties
features['duration'] = len(audio_data) / sr
features['sample_rate'] = sr
if self.enable_parallel_processing:
print("πŸš€ Using parallel feature extraction...")
# Prepare feature extraction tasks
feature_tasks = [
("mfcc", audio_data, sr),
("chroma", audio_data, sr),
("spectral", audio_data, sr),
("basic", audio_data, sr)
]
# Execute feature extraction in parallel
with concurrent.futures.ThreadPoolExecutor(max_workers=min(4, self.max_workers)) as executor:
future_to_feature = {
executor.submit(self._extract_feature_chunk, audio_chunk, sr, feature_type): feature_type
for feature_type, audio_chunk, sr in feature_tasks
}
for future in concurrent.futures.as_completed(future_to_feature):
feature_result = future.result()
features.update(feature_result)
else:
# Sequential feature extraction (original logic)
# Energy features
features['rms_energy'] = float(np.mean(librosa.feature.rms(y=audio_data)[0]))
features['zero_crossing_rate'] = float(np.mean(librosa.feature.zero_crossing_rate(audio_data)[0]))
# MFCC features
mfcc = librosa.feature.mfcc(y=audio_data, sr=sr, n_mfcc=13)
for i in range(13):
features[f'mfcc_{i+1}_mean'] = float(np.mean(mfcc[i]))
features[f'mfcc_{i+1}_std'] = float(np.std(mfcc[i]))
# Spectral features
features['spectral_centroid_mean'] = float(np.mean(librosa.feature.spectral_centroid(y=audio_data, sr=sr)[0]))
features['spectral_rolloff_mean'] = float(np.mean(librosa.feature.spectral_rolloff(y=audio_data, sr=sr)[0]))
# Chroma features
chroma = librosa.feature.chroma_stft(y=audio_data, sr=sr)
for i in range(12):
features[f'chroma_{i+1}_mean'] = float(np.mean(chroma[i]))
# Pitch features (kept sequential due to complexity)
try:
pitches, magnitudes = librosa.piptrack(y=audio_data, sr=sr, threshold=0.1)
pitch_values = []
for t in range(pitches.shape[1]):
index = magnitudes[:, t].argmax()
pitch = pitches[index, t]
if pitch > 0:
pitch_values.append(pitch)
if pitch_values:
features['pitch_mean'] = float(np.mean(pitch_values))
features['pitch_std'] = float(np.std(pitch_values))
features['pitch_min'] = float(np.min(pitch_values))
features['pitch_max'] = float(np.max(pitch_values))
else:
features.update({'pitch_mean': 0.0, 'pitch_std': 0.0, 'pitch_min': 0.0, 'pitch_max': 0.0})
except:
features.update({'pitch_mean': 0.0, 'pitch_std': 0.0, 'pitch_min': 0.0, 'pitch_max': 0.0})
# Tempo
try:
tempo, _ = librosa.beat.beat_track(y=audio_data, sr=sr)
if isinstance(tempo, np.ndarray):
features['tempo'] = float(tempo.item() if tempo.size == 1 else tempo[0])
else:
features['tempo'] = float(tempo)
except:
features['tempo'] = 0.0
return features
def predict_emotion(self, audio_data, sr):
"""Predict emotion using transformer model"""
if self.emotion_model is None:
return None
print("😊 Predicting emotions...")
try:
# Resample to 16kHz if needed
if sr != 16000:
audio_data = librosa.resample(audio_data, orig_sr=sr, target_sr=16000)
# Process through model
inputs = self.emotion_processor(audio_data, sampling_rate=16000, return_tensors="pt", padding=True)
with torch.no_grad():
outputs = self.emotion_model(**inputs)
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
# Get emotion probabilities
emotion_probs = {}
for i, emotion in enumerate(self.emotion_labels):
emotion_probs[emotion] = predictions[0][i].item()
predicted_emotion = self.emotion_labels[predictions.argmax().item()]
confidence = predictions.max().item()
return {
'predicted_emotion': predicted_emotion,
'confidence': confidence,
'all_emotions': emotion_probs
}
except Exception as e:
print(f"⚠️ Error in emotion prediction: {e}")
return None
def analyze_complete_audio(self, audio_file: str) -> Dict:
"""Perform complete unified audio analysis with parallel processing"""
if not os.path.exists(audio_file):
print(f"❌ Audio file not found: {audio_file}")
return {}
print(f"\nπŸš€ Starting complete analysis of: {audio_file}")
print("="*60)
start_time = time.time()
# Load audio for paralinguistic analysis
try:
audio_data, sr = librosa.load(audio_file, sr=22050)
audio_data, _ = librosa.effects.trim(audio_data, top_db=20)
audio_data = librosa.util.normalize(audio_data)
except Exception as e:
print(f"❌ Error loading audio: {e}")
return {}
if self.enable_parallel_processing:
print("πŸš€ Running analysis components in parallel...")
# Create a queue for results
results_queue = Queue()
# Define analysis functions
def run_diarization():
result = self.transcribe_with_diarization(audio_file)
results_queue.put(('diarization', result))
def run_event_detection():
result = self.detect_audio_events(audio_file)
results_queue.put(('events', result))
def run_feature_extraction():
result = self.extract_paralinguistic_features(audio_data, sr)
results_queue.put(('features', result))
def run_emotion_prediction():
result = self.predict_emotion(audio_data, sr)
results_queue.put(('emotion', result))
# Start threads for parallel execution
threads = [
threading.Thread(target=run_diarization),
threading.Thread(target=run_event_detection),
threading.Thread(target=run_feature_extraction),
threading.Thread(target=run_emotion_prediction)
]
# Start all threads
for thread in threads:
thread.start()
# Wait for all threads to complete
for thread in threads:
thread.join()
# Collect results
analysis_components = {}
while not results_queue.empty():
component, result = results_queue.get()
analysis_components[component] = result
# Assign results
diarization_results = analysis_components.get('diarization', [])
event_results = analysis_components.get('events', {})
paralinguistic_features = analysis_components.get('features', {})
emotion_results = analysis_components.get('emotion', None)
else:
# Sequential processing (original logic)
# 1. Speaker Diarization + Transcription
diarization_results = self.transcribe_with_diarization(audio_file)
# 2. Audio Event Detection
event_results = self.detect_audio_events(audio_file)
# 3. Paralinguistic Features
paralinguistic_features = self.extract_paralinguistic_features(audio_data, sr)
# 4. Emotion Recognition
emotion_results = self.predict_emotion(audio_data, sr)
processing_time = time.time() - start_time
print(f"⏱️ Total processing time: {processing_time:.2f} seconds")
# Combine all results
complete_analysis = {
'file_info': {
'filename': os.path.basename(audio_file),
'filepath': audio_file,
'duration': paralinguistic_features.get('duration', 0),
'sample_rate': paralinguistic_features.get('sample_rate', 0),
'processing_time': processing_time
},
'diarization_transcription': diarization_results,
'audio_events': event_results,
'paralinguistic_features': paralinguistic_features,
'emotion_analysis': emotion_results
}
return complete_analysis
def print_analysis_summary(self, analysis_results: Dict):
"""Print formatted analysis summary"""
if not analysis_results:
print("❌ No analysis results to display")
return
file_info = analysis_results.get('file_info', {})
diarization = analysis_results.get('diarization_transcription', [])
events = analysis_results.get('audio_events', {})
emotion = analysis_results.get('emotion_analysis', {})
print(f"\n{'='*80}")
print("🎯 UNIFIED AUDIO ANALYSIS RESULTS")
print(f"{'='*80}")
# File Information
print(f"πŸ“ File: {file_info.get('filename', 'Unknown')}")
print(f"⏱️ Duration: {file_info.get('duration', 0):.2f} seconds")
print(f"πŸ”Š Sample Rate: {file_info.get('sample_rate', 0)} Hz")
print(f"⚑ Processing Time: {file_info.get('processing_time', 0):.2f} seconds")
# 1. Speaker Diarization Results
print(f"\n{'🎀 SPEAKER DIARIZATION & TRANSCRIPTION'}")
print("-" * 50)
if diarization:
speakers = set(seg['speaker'] for seg in diarization)
print(f"Speakers detected: {len(speakers)}")
print(f"Total segments: {len(diarization)}")
for i, segment in enumerate(diarization, 1):
print(f"{i}. {segment['speaker']} [{segment['start']:.1f}s-{segment['end']:.1f}s]: {segment['text'][:80]}{'...' if len(segment['text']) > 80 else ''}")
else:
print("No diarization results available")
# 2. Audio Event Detection
print(f"\n{'πŸ”Š AUDIO EVENT DETECTION (Top 10)'}")
print("-" * 50)
top_events = events.get('top_events', [])
if top_events:
for i, event in enumerate(top_events[:10], 1):
print(f"{i:2d}. {event['event']:<30} | Probability: {event['probability']:.4f}")
else:
print("No audio events detected")
# 3. Emotion Analysis
print(f"\n{'😊 EMOTION ANALYSIS'}")
print("-" * 30)
if emotion:
print(f"Predicted Emotion: {emotion['predicted_emotion']} (Confidence: {emotion['confidence']:.3f})")
print("\nAll Emotion Probabilities:")
for emo, prob in emotion['all_emotions'].items():
print(f" {emo.capitalize():<12}: {prob:.3f}")
else:
print("No emotion analysis available")
# 4. Key Paralinguistic Features
features = analysis_results.get('paralinguistic_features', {})
if features:
print(f"\n{'🎡 KEY PARALINGUISTIC FEATURES'}")
print("-" * 40)
print(f"RMS Energy: {features.get('rms_energy', 0):.4f}")
print(f"Pitch Mean: {features.get('pitch_mean', 0):.2f} Hz")
print(f"Spectral Centroid: {features.get('spectral_centroid_mean', 0):.2f} Hz")
print(f"Tempo: {features.get('tempo', 0):.2f} BPM")
print(f"Zero Crossing Rate: {features.get('zero_crossing_rate', 0):.4f}")
def save_results_to_csv(self, analysis_results: Dict, output_prefix: str = "unified_analysis"):
"""Save analysis results to CSV files"""
if not analysis_results:
print("❌ No results to save")
return
# Save diarization results
diarization = analysis_results.get('diarization_transcription', [])
if diarization:
df_diarization = pd.DataFrame(diarization)
diarization_file = f"{output_prefix}_diarization.csv"
df_diarization.to_csv(diarization_file, index=False)
print(f"πŸ’Ύ Diarization results saved to: {diarization_file}")
# Save audio events
events = analysis_results.get('audio_events', {}).get('top_events', [])
if events:
df_events = pd.DataFrame(events)
events_file = f"{output_prefix}_audio_events.csv"
df_events.to_csv(events_file, index=False)
print(f"πŸ’Ύ Audio events saved to: {events_file}")
# Save paralinguistic features
features = analysis_results.get('paralinguistic_features', {})
if features:
df_features = pd.DataFrame([features])
features_file = f"{output_prefix}_features.csv"
df_features.to_csv(features_file, index=False)
print(f"πŸ’Ύ Features saved to: {features_file}")
# Save emotion analysis
emotion = analysis_results.get('emotion_analysis', {})
if emotion:
df_emotion = pd.DataFrame([emotion])
emotion_file = f"{output_prefix}_emotion.csv"
df_emotion.to_csv(emotion_file, index=False)
print(f"πŸ’Ύ Emotion analysis saved to: {emotion_file}")
def summarize_audio_analysis_with_llm(analysis_results: dict) -> str:
"""
Send all analysis results to a Groq LLM (gpt-oss-20b) and get a summary
describing relationships between diarization, events, emotion, and features.
Requires GROQ_API_KEY in environment.
"""
# Prepare the prompt
prompt = (
"You are an expert audio scene interpreter. Given the structured audio analysis results, "
"summarize what is happening in plain, natural language, as if explaining the situation to someone. "
"Avoid technical terms, metrics, or probabilities. Instead, combine the speaker's words, background "
"sounds, emotions and other paralingusistic features to infer the most likely real-world context. Keep it short and clear.\n\n"
"Sample input : Recording of a person call reaching an airport (with background noise of airplanes, announcements, and crowd chatter). Sample output : The subway sound and other vehicle sound suggest that person is in Highway, and the aero plane sound indicate nearby Airport, while announcement provide information about the Airplane Schedule, that means person reached in boarding area or into the waiting hall.\n\n"
f"Audio Analysis Results:\n{analysis_results}\n\n"
"Plain Summary:"
)
# Load environment variables
load_dotenv()
api_key = os.getenv("GROQ_API_KEY")
if not api_key:
raise ValueError("GROQ_API_KEY environment variable not set.")
# Initialize Groq client
client = Groq(api_key=api_key)
# Make the API call
response = client.chat.completions.create(
model="openai/gpt-oss-20b",
messages=[
{"role": "system", "content": "You are an expert audio analyst."},
{"role": "user", "content": prompt},
],
)
# Extract summary
summary = response.choices[0].message.content.strip()
return summary
def main():
"""Main function demonstrating usage"""
# Initialize analyzer with parallel processing enabled
analyzer = UnifiedAudioAnalyzer(enable_parallel_processing=True, max_workers=None)
# Specify input audio file
audio_file = "dataset/flight/15.wav" # Update with your audio file path
if os.path.exists(audio_file):
# Perform complete analysis
results = analyzer.analyze_complete_audio(audio_file)
# Print summary
analyzer.print_analysis_summary(results)
# Save results to CSV files
# analyzer.save_results_to_csv(results, "my_audio_analysis")
print(f"\nβœ… Analysis complete! Check CSV files for detailed results.")
summary=summarize_audio_analysis_with_llm(results)
print("\n=== LLM Summary of Audio Analysis ===")
print(summary)
else:
print(f"❌ Audio file not found: {audio_file}")
print("Please update the audio_file path to point to your audio file.")
if __name__ == "__main__":
main()