speaker_seperation.py

#!/usr/bin/env python3
"""
YouTube Speaker Diarization Pipeline - Optimized Version

This script downloads a YouTube video, extracts vocals, performs speaker diarization,
and outputs separate audio files for each speaker with parallel processing and memory optimization.
"""

import argparse
import logging
import time
import shutil
from concurrent.futures import ThreadPoolExecutor
from functools import lru_cache
from pathlib import Path
from typing import Dict, List, Tuple

import librosa
import numpy as np
import yt_dlp
from scipy.io.wavfile import write
from simple_diarizer.diarizer import Diarizer
from spleeter.separator import Separator

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('diarization.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

class DiarizationConfig:
    """Configuration for diarization parameters"""
    def __init__(
        self,
        num_speakers: int = 2,
        window: float = 0.5,
        period: float = 0.25,
        model: str = 'ecapa',
        cluster_method: str = 'sc'
    ):
        self.num_speakers = num_speakers
        self.window = window
        self.period = period
        self.model = model
        self.cluster_method = cluster_method

def download_youtube_audio(
    url: str,
    output_dir: Path,
    max_retries: int = 3
) -> Path:
    """
    Download audio from YouTube video with retry logic
    Returns path to downloaded audio file
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / "audio.wav"

    ydl_opts = {
        'format': 'bestaudio/best',
        'outtmpl': str(output_dir / 'temp_audio.%(ext)s'),
        'postprocessors': [{
            'key': 'FFmpegExtractAudio',
            'preferredcodec': 'wav',
            'preferredquality': '192',
        }],
        'retries': max_retries,
        'quiet': True,
    }

    for attempt in range(max_retries):
        try:
            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                info = ydl.extract_info(url, download=True)
                temp_path = Path(ydl.prepare_filename(info)).with_suffix('.wav')
                temp_path.rename(output_path)
                logger.info(f"Downloaded audio to {output_path}")
                return output_path
        except Exception as e:
            logger.warning(f"Attempt {attempt + 1} failed: {str(e)}")
            if attempt == max_retries - 1:
                logger.error("Max retries reached. Download failed.")
                raise
            continue

@lru_cache(maxsize=32)
def load_audio_cached(audio_path: Path) -> Tuple[np.ndarray, int]:
    """Load audio file with caching to avoid repeated disk reads"""
    try:
        signal, fs = librosa.load(str(audio_path), sr=None, mono=True)
        logger.debug(f"Loaded audio from {audio_path}")
        return signal, fs
    except Exception as e:
        logger.error(f"Failed to load audio: {str(e)}")
        raise

def extract_vocals(
    input_path: Path,
    output_dir: Path,
    model: str = 'spleeter:2stems',
    timeout: int = 300,
    poll_interval: float = 0.1
) -> Path:
    """
    Extract vocals using Spleeter with robust file handling
    Returns path to extracted vocals file
    """
    try:
        output_dir.mkdir(parents=True, exist_ok=True)
        temp_dir = output_dir / "temp_spleeter_output"
        
        # Clear any existing temp directory
        if temp_dir.exists():
            import shutil
            shutil.rmtree(temp_dir)
        
        # Use separator with optimized parameters
        separator = Separator(model)
        separator.separate_to_file(
            str(input_path),
            str(output_dir),
            codec='wav',
            bitrate='192k',
            synchronous=False
        )
        
        # Wait for output with timeout
        vocals_path = None
        start_time = time.time()
        
        while time.time() - start_time < timeout:
            # Check for vocals file in expected locations
            possible_paths = [
                output_dir / "temp_spleeter_output" / "vocals.wav",
                output_dir / f"{input_path.stem}" / "vocals.wav",
                output_dir / "vocals.wav"
            ]
            
            for path in possible_paths:
                if path.exists():
                    vocals_path = output_dir / "vocals.wav"
                    path.rename(vocals_path)
                    logger.info(f"Found and moved vocals to {vocals_path}")
                    
                    # Clean up
                    if temp_dir.exists():
                        shutil.rmtree(temp_dir)
                    for other in output_dir.glob('**/*accompaniment.wav'):
                        other.unlink()
                    
                    return vocals_path
            
            time.sleep(poll_interval)
        
        raise TimeoutError(f"Vocal extraction timed out after {timeout} seconds")
        
    except Exception as e:
        logger.error(f"Vocal extraction failed: {str(e)}")
        # Attempt cleanup
        if temp_dir.exists():
            shutil.rmtree(temp_dir)
        raise

def diarize_audio(
    audio_path: Path,
    config: DiarizationConfig,
    output_dir: Path
) -> Tuple[Dict, Path]:
    """
    Perform speaker diarization with optimized configuration
    Returns segments dictionary and RTTM file path
    """
    try:
        output_dir.mkdir(parents=True, exist_ok=True)
        rttm_path = output_dir / "diarization.rttm"

        diarizer = Diarizer(
            embed_model=config.model,
            cluster_method=config.cluster_method,
            window=config.window,
            period=config.period
        )

        logger.info(f"Starting diarization for {audio_path}")
        segments = diarizer.diarize(
            str(audio_path),
            num_speakers=config.num_speakers,
            outfile=str(rttm_path)
        )

        logger.info(f"Diarization completed. Results saved to {rttm_path}")
        return segments, rttm_path
    except Exception as e:
        logger.error(f"Diarization failed: {str(e)}")
        raise

def process_speaker_chunk(args: Tuple[str, List[np.ndarray], int, Path]) -> Path:
    """
    Helper function for parallel processing of speaker chunks
    Returns path to saved speaker audio file
    """
    speaker, chunks, fs, output_dir = args
    try:
        combined = np.concatenate(chunks)
        output_path = output_dir / f"speaker_{speaker}.wav"
        write(output_path, fs, (combined * 32767).astype(np.int16))
        logger.debug(f"Saved speaker {speaker} audio to {output_path}")
        return output_path
    except Exception as e:
        logger.error(f"Failed to process speaker {speaker}: {str(e)}")
        raise

def process_rttm_parallel(
    rttm_path: Path,
    audio_path: Path,
    output_dir: Path,
    max_workers: int = 4
) -> List[Path]:
    """
    Process RTTM file to create separated speaker audio files using parallel processing
    Returns list of paths to speaker audio files
    """
    try:
        signal, fs = load_audio_cached(audio_path)
        speakers = {}
        output_dir.mkdir(parents=True, exist_ok=True)

        # Parse RTTM file
        with open(rttm_path, 'r') as f:
            for line in f:
                parts = line.strip().split()
                start = float(parts[3])
                duration = float(parts[4])
                speaker = parts[7]
                end = start + duration
                
                if speaker not in speakers:
                    speakers[speaker] = []
                
                start_sample = int(start * fs)
                end_sample = int(end * fs)
                speakers[speaker].append(signal[start_sample:end_sample])

        # Process speakers in parallel
        with ThreadPoolExecutor(max_workers=max_workers) as executor:
            args = [(spk, chunks, fs, output_dir) for spk, chunks in speakers.items()]
            results = list(executor.map(process_speaker_chunk, args))

        logger.info(f"Processed {len(results)} speaker files")
        return results
    except Exception as e:
        logger.error(f"RTTM processing failed: {str(e)}")
        raise

def main():
    parser = argparse.ArgumentParser(
        description="Optimized YouTube Speaker Diarization Pipeline",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("url", help="YouTube video URL")
    parser.add_argument(
        "-n", "--num_speakers",
        type=int,
        default=2,
        help="Number of speakers to diarize"
    )
    parser.add_argument(
        "-o", "--output_dir",
        type=Path,
        default="output",
        help="Main output directory"
    )
    parser.add_argument(
        "--window",
        type=float,
        default=0.5,
        help="Window size for diarization in seconds"
    )
    parser.add_argument(
        "--period",
        type=float,
        default=0.25,
        help="Period for diarization in seconds"
    )
    parser.add_argument(
        "--workers",
        type=int,
        default=4,
        help="Number of parallel workers for processing"
    )
    parser.add_argument(
        "--debug",
        action="store_true",
        help="Enable debug logging"
    )
    
    args = parser.parse_args()

    if args.debug:
        logger.setLevel(logging.DEBUG)

    try:
        # Create configuration
        config = DiarizationConfig(
            num_speakers=args.num_speakers,
            window=args.window,
            period=args.period
        )

        # 1. Download audio
        logger.info("Starting YouTube audio download...")
        audio_path = download_youtube_audio(args.url, args.output_dir)
        
        # 2. Extract vocals
        logger.info("Starting vocal extraction...")
        vocals_path = extract_vocals(audio_path, args.output_dir)
        
        # 3. Perform diarization
        logger.info("Starting diarization...")
        segments, rttm_path = diarize_audio(vocals_path, config, args.output_dir)
        
        # 4. Process results in parallel
        logger.info("Processing results with parallel workers...")
        speaker_files = process_rttm_parallel(
            rttm_path,
            vocals_path,
            Path("diarized_audio"),
            max_workers=args.workers
        )
        shutil.rmtree(args.output_dir)
        logger.info(f"Processing completed successfully. Output files: {speaker_files}")

    except Exception as e:
        logger.error(f"Pipeline failed: {str(e)}")
        exit(1)

if __name__ == "__main__":
    main()