AI-Talent-Force/dev_caio

1"""2ShortSmith v2 - Audio Analyzer Module3 4Audio feature extraction and hype scoring using:5- Librosa for basic audio features (MVP)6- Wav2Vec 2.0 for advanced audio understanding (optional)7 8Features extracted:9- RMS energy (volume/loudness)10- Spectral flux (sudden changes, beat drops)11- Spectral centroid (brightness, crowd noise)12- Onset strength (beats, impacts)13- Speech activity detection14"""15 16from pathlib import Path17from typing import List, Optional, Tuple, Dict18from dataclasses import dataclass19import numpy as np20 21from utils.logger import get_logger, LogTimer22from utils.helpers import ModelLoadError, InferenceError, normalize_scores, batch_list23from config import get_config, ModelConfig24 25logger = get_logger("models.audio_analyzer")26 27 28@dataclass29class AudioFeatures:30    """Audio features for a segment of audio."""31    timestamp: float          # Start time in seconds32    duration: float           # Segment duration33    rms_energy: float         # Root mean square energy (0-1)34    spectral_flux: float      # Spectral change rate (0-1)35    spectral_centroid: float  # Frequency centroid (0-1)36    onset_strength: float     # Beat/impact strength (0-1)37    zero_crossing_rate: float # ZCR (speech indicator) (0-1)38 39    # Optional advanced features40    speech_probability: float = 0.0  # From Wav2Vec if available41 42    @property43    def energy_score(self) -> float:44        """Combined energy-based hype indicator."""45        return (self.rms_energy * 0.4 + self.onset_strength * 0.4 +46                self.spectral_flux * 0.2)47 48    @property49    def excitement_score(self) -> float:50        """Overall audio excitement score."""51        return (self.rms_energy * 0.3 + self.spectral_flux * 0.25 +52                self.onset_strength * 0.25 + self.spectral_centroid * 0.2)53 54 55@dataclass56class AudioSegmentScore:57    """Hype score for an audio segment."""58    start_time: float59    end_time: float60    score: float              # Overall hype score (0-1)61    features: AudioFeatures   # Underlying features62 63    @property64    def duration(self) -> float:65        return self.end_time - self.start_time66 67 68class AudioAnalyzer:69    """70    Audio analysis for hype detection.71 72    Uses Librosa for feature extraction and optionally Wav2Vec 2.073    for advanced semantic understanding.74    """75 76    def __init__(77        self,78        config: Optional[ModelConfig] = None,79        use_advanced: Optional[bool] = None,80    ):81        """82        Initialize audio analyzer.83 84        Args:85            config: Model configuration (uses default if None)86            use_advanced: Override config to use Wav2Vec 2.087 88        Raises:89            ImportError: If librosa is not installed90        """91        self.config = config or get_config().model92        self.use_advanced = use_advanced if use_advanced is not None else self.config.use_advanced_audio93 94        self._librosa = None95        self._wav2vec_model = None96        self._wav2vec_processor = None97 98        # Initialize librosa (required)99        self._init_librosa()100 101        # Initialize Wav2Vec if requested102        if self.use_advanced:103            self._init_wav2vec()104 105        logger.info(f"AudioAnalyzer initialized (advanced={self.use_advanced})")106 107    def _init_librosa(self) -> None:108        """Initialize librosa library."""109        try:110            import librosa111            self._librosa = librosa112        except ImportError as e:113            raise ImportError(114                "Librosa is required for audio analysis. "115                "Install with: pip install librosa"116            ) from e117 118    def _init_wav2vec(self) -> None:119        """Initialize Wav2Vec 2.0 model."""120        try:121            import torch122            from transformers import Wav2Vec2Processor, Wav2Vec2Model123 124            logger.info("Loading Wav2Vec 2.0 model...")125 126            self._wav2vec_processor = Wav2Vec2Processor.from_pretrained(127                self.config.audio_model_id128            )129            self._wav2vec_model = Wav2Vec2Model.from_pretrained(130                self.config.audio_model_id131            )132 133            # Move to device134            device = self.config.device135            if device == "cuda":136                import torch137                if torch.cuda.is_available():138                    self._wav2vec_model = self._wav2vec_model.cuda()139 140            self._wav2vec_model.eval()141            logger.info("Wav2Vec 2.0 model loaded successfully")142 143        except Exception as e:144            logger.warning(f"Failed to load Wav2Vec 2.0, falling back to Librosa only: {e}")145            self.use_advanced = False146 147    def load_audio(148        self,149        audio_path: str | Path,150        sample_rate: int = 22050,151        mono: bool = True,152    ) -> Tuple[np.ndarray, int]:153        """154        Load audio file.155 156        Args:157            audio_path: Path to audio file158            sample_rate: Target sample rate159            mono: Convert to mono if True160 161        Returns:162            Tuple of (audio_array, sample_rate)163 164        Raises:165            InferenceError: If audio loading fails166        """167        try:168            audio, sr = self._librosa.load(169                str(audio_path),170                sr=sample_rate,171                mono=mono,172            )173            logger.debug(f"Loaded audio: {len(audio)/sr:.1f}s at {sr}Hz")174            return audio, sr175 176        except Exception as e:177            raise InferenceError(f"Failed to load audio: {e}") from e178 179    def extract_features(180        self,181        audio: np.ndarray,182        sample_rate: int,183        segment_duration: float = 1.0,184        hop_duration: float = 0.5,185    ) -> List[AudioFeatures]:186        """187        Extract audio features for overlapping segments.188 189        Args:190            audio: Audio array191            sample_rate: Sample rate192            segment_duration: Duration of each segment in seconds193            hop_duration: Hop between segments in seconds194 195        Returns:196            List of AudioFeatures for each segment197        """198        with LogTimer(logger, "Extracting audio features"):199            duration = len(audio) / sample_rate200            segment_samples = int(segment_duration * sample_rate)201            hop_samples = int(hop_duration * sample_rate)202 203            features = []204            position = 0205            timestamp = 0.0206 207            while position + segment_samples <= len(audio):208                segment = audio[position:position + segment_samples]209 210                try:211                    feat = self._extract_segment_features(212                        segment, sample_rate, timestamp, segment_duration213                    )214                    features.append(feat)215                except Exception as e:216                    logger.warning(f"Failed to extract features at {timestamp}s: {e}")217 218                position += hop_samples219                timestamp += hop_duration220 221            logger.info(f"Extracted features for {len(features)} segments")222            return features223 224    def _extract_segment_features(225        self,226        segment: np.ndarray,227        sample_rate: int,228        timestamp: float,229        duration: float,230    ) -> AudioFeatures:231        """Extract features from a single audio segment."""232        librosa = self._librosa233 234        # RMS energy (loudness)235        rms = librosa.feature.rms(y=segment)[0]236        rms_mean = float(np.mean(rms))237 238        # Spectral flux (change rate)239        spec = np.abs(librosa.stft(segment))240        flux = np.mean(np.diff(spec, axis=1) ** 2)241        flux_normalized = min(1.0, flux / 100)  # Normalize242 243        # Spectral centroid (brightness)244        centroid = librosa.feature.spectral_centroid(y=segment, sr=sample_rate)[0]245        centroid_mean = float(np.mean(centroid))246        centroid_normalized = min(1.0, centroid_mean / 8000)  # Normalize247 248        # Onset strength (beats/impacts)249        onset_env = librosa.onset.onset_strength(y=segment, sr=sample_rate)250        onset_mean = float(np.mean(onset_env))251        onset_normalized = min(1.0, onset_mean / 5)  # Normalize252 253        # Zero crossing rate254        zcr = librosa.feature.zero_crossing_rate(segment)[0]255        zcr_mean = float(np.mean(zcr))256 257        return AudioFeatures(258            timestamp=timestamp,259            duration=duration,260            rms_energy=min(1.0, rms_mean * 5),  # Scale up261            spectral_flux=flux_normalized,262            spectral_centroid=centroid_normalized,263            onset_strength=onset_normalized,264            zero_crossing_rate=zcr_mean,265        )266 267    def analyze_file(268        self,269        audio_path: str | Path,270        segment_duration: float = 1.0,271        hop_duration: float = 0.5,272    ) -> List[AudioFeatures]:273        """274        Analyze an audio file and extract features.275 276        Args:277            audio_path: Path to audio file278            segment_duration: Duration of each segment279            hop_duration: Hop between segments280 281        Returns:282            List of AudioFeatures for the file283        """284        audio, sr = self.load_audio(audio_path)285        return self.extract_features(audio, sr, segment_duration, hop_duration)286 287    def compute_hype_scores(288        self,289        features: List[AudioFeatures],290        window_size: int = 5,291    ) -> List[AudioSegmentScore]:292        """293        Compute hype scores from audio features.294 295        Uses a sliding window to smooth scores and identify296        sustained high-energy regions.297 298        Args:299            features: List of AudioFeatures300            window_size: Smoothing window size301 302        Returns:303            List of AudioSegmentScore objects304        """305        if not features:306            return []307 308        with LogTimer(logger, "Computing audio hype scores"):309            # Compute raw excitement scores310            raw_scores = [f.excitement_score for f in features]311 312            # Apply smoothing313            smoothed = self._smooth_scores(raw_scores, window_size)314 315            # Normalize to 0-1316            normalized = normalize_scores(smoothed)317 318            # Create score objects319            scores = []320            for feat, score in zip(features, normalized):321                scores.append(AudioSegmentScore(322                    start_time=feat.timestamp,323                    end_time=feat.timestamp + feat.duration,324                    score=score,325                    features=feat,326                ))327 328            return scores329 330    def _smooth_scores(331        self,332        scores: List[float],333        window_size: int,334    ) -> List[float]:335        """Apply moving average smoothing to scores."""336        if len(scores) < window_size:337            return scores338 339        kernel = np.ones(window_size) / window_size340        padded = np.pad(scores, (window_size // 2, window_size // 2), mode='edge')341        smoothed = np.convolve(padded, kernel, mode='valid')342 343        return smoothed.tolist()344 345    def detect_peaks(346        self,347        scores: List[AudioSegmentScore],348        threshold: float = 0.6,349        min_duration: float = 3.0,350    ) -> List[Tuple[float, float, float]]:351        """352        Detect peak regions in audio hype.353 354        Args:355            scores: List of AudioSegmentScore objects356            threshold: Minimum score to consider a peak357            min_duration: Minimum peak duration in seconds358 359        Returns:360            List of (start_time, end_time, peak_score) tuples361        """362        if not scores:363            return []364 365        peaks = []366        in_peak = False367        peak_start = 0.0368        peak_max = 0.0369 370        for score in scores:371            if score.score >= threshold:372                if not in_peak:373                    in_peak = True374                    peak_start = score.start_time375                    peak_max = score.score376                else:377                    peak_max = max(peak_max, score.score)378            else:379                if in_peak:380                    peak_end = score.start_time381                    if peak_end - peak_start >= min_duration:382                        peaks.append((peak_start, peak_end, peak_max))383                    in_peak = False384 385        # Handle peak at end386        if in_peak:387            peak_end = scores[-1].end_time388            if peak_end - peak_start >= min_duration:389                peaks.append((peak_start, peak_end, peak_max))390 391        logger.info(f"Detected {len(peaks)} audio peaks above threshold {threshold}")392        return peaks393 394    def get_beat_timestamps(395        self,396        audio: np.ndarray,397        sample_rate: int,398    ) -> List[float]:399        """400        Detect beat timestamps in audio.401 402        Args:403            audio: Audio array404            sample_rate: Sample rate405 406        Returns:407            List of beat timestamps in seconds408        """409        try:410            tempo, beats = self._librosa.beat.beat_track(y=audio, sr=sample_rate)411            beat_times = self._librosa.frames_to_time(beats, sr=sample_rate)412            logger.debug(f"Detected {len(beat_times)} beats at {tempo:.1f} BPM")413            return beat_times.tolist()414        except Exception as e:415            logger.warning(f"Beat detection failed: {e}")416            return []417 418    def get_audio_embedding(419        self,420        audio: np.ndarray,421        sample_rate: int = 16000,422    ) -> Optional[np.ndarray]:423        """424        Get Wav2Vec 2.0 embedding for audio segment.425 426        Only available if use_advanced=True.427 428        Args:429            audio: Audio array (should be 16kHz)430            sample_rate: Sample rate431 432        Returns:433            Embedding array or None if not available434        """435        if not self.use_advanced or self._wav2vec_model is None:436            return None437 438        try:439            import torch440 441            # Resample if needed442            if sample_rate != 16000:443                audio = self._librosa.resample(audio, orig_sr=sample_rate, target_sr=16000)444 445            # Process446            inputs = self._wav2vec_processor(447                audio, sampling_rate=16000, return_tensors="pt"448            )449 450            if self.config.device == "cuda" and torch.cuda.is_available():451                inputs = {k: v.cuda() for k, v in inputs.items()}452 453            with torch.no_grad():454                outputs = self._wav2vec_model(**inputs)455                embedding = outputs.last_hidden_state.mean(dim=1).cpu().numpy()456 457            return embedding[0]458 459        except Exception as e:460            logger.warning(f"Wav2Vec embedding extraction failed: {e}")461            return None462 463    def compare_audio_similarity(464        self,465        embedding1: np.ndarray,466        embedding2: np.ndarray,467    ) -> float:468        """469        Compare two audio embeddings using cosine similarity.470 471        Args:472            embedding1: First embedding473            embedding2: Second embedding474 475        Returns:476            Similarity score (0-1)477        """478        norm1 = np.linalg.norm(embedding1)479        norm2 = np.linalg.norm(embedding2)480 481        if norm1 == 0 or norm2 == 0:482            return 0.0483 484        return float(np.dot(embedding1, embedding2) / (norm1 * norm2))485 486 487# Export public interface488__all__ = ["AudioAnalyzer", "AudioFeatures", "AudioSegmentScore"]489