CoolFace
Apppublic

ASesYusuf1/SESA_Audio_Separation

sourceHugging Facemitupdated 6mo agoView on Hugging Face
14likes
ensemble.py315 linesDownload Raw Back to kaggle
1# Kaggle için gerekli kütüphaneleri yükleme2!pip install librosa soundfile psutil tqdm3import os4import sys5import argparse6import numpy as np7import soundfile as sf8import librosa9import psutil10import gc11import traceback12from scipy.signal import stft, istft13from pathlib import Path14import tempfile15import shutil16import json17from tqdm import tqdm18import time19 20class AudioEnsembleEngine:21    def __init__(self):22        self.temp_dir = None23        self.log_file = "/kaggle/working/ensemble_processor.log"24        25    def __enter__(self):26        self.temp_dir = tempfile.mkdtemp(prefix='audio_ensemble_', dir='/kaggle/working/')27        self.setup_logging()28        return self29        30    def __exit__(self, exc_type, exc_val, exc_tb):31        if self.temp_dir and os.path.exists(self.temp_dir):32            shutil.rmtree(self.temp_dir, ignore_errors=True)33    34    def setup_logging(self):35        """Initialize detailed logging system."""36        with open(self.log_file, 'w') as f:37            f.write("Audio Ensemble Processor Log\n")38            f.write("="*50 + "\n")39            f.write(f"System Memory: {psutil.virtual_memory().total/(1024**3):.2f} GB\n")40            f.write(f"Python Version: {sys.version}\n\n")41    42    def log_message(self, message):43        """Log messages with timestamp."""44        with open(self.log_file, 'a') as f:45            f.write(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {message}\n")46    47    def normalize_path(self, path):48        """Handle all path-related issues comprehensively."""49        try:50            # Kaggle'da dosya yolları /kaggle/input/ veya /kaggle/working/ altında51            path = str(Path(path).absolute().resolve())52            53            # Handle problematic characters54            if any(char in path for char in '[]()|&; '):55                base, ext = os.path.splitext(path)56                safe_name = f"{hash(base)}{ext}"57                temp_path = os.path.join(self.temp_dir, safe_name)58                59                if not os.path.exists(temp_path):60                    data, sr = librosa.load(path, sr=None, mono=False)61                    sf.write(temp_path, data.T, sr)62                63                return temp_path64            65            return path66        except Exception as e:67            self.log_message(f"Path normalization failed: {str(e)}")68            return path69    70    def validate_inputs(self, files, method, output_path):71        """Comprehensive input validation with detailed error reporting."""72        errors = []73        valid_methods = [74            'avg_wave', 'median_wave', 'max_wave', 'min_wave',75            'max_fft', 'min_fft', 'median_fft'76        ]77        78        # Method validation79        if method not in valid_methods:80            errors.append(f"Invalid method '{method}'. Available: {valid_methods}")81        82        # File validation83        valid_files = []84        sample_rates = set()85        durations = []86        channels_set = set()87        88        for f in files:89            try:90                f_normalized = self.normalize_path(f)91                92                # Basic checks93                if not os.path.exists(f_normalized):94                    errors.append(f"File not found: {f_normalized}")95                    continue96                97                if os.path.getsize(f_normalized) == 0:98                    errors.append(f"Empty file: {f_normalized}")99                    continue100                101                # Audio file validation102                try:103                    with sf.SoundFile(f_normalized) as sf_file:104                        sr = sf_file.samplerate105                        frames = sf_file.frames106                        channels = sf_file.channels107                except Exception as e:108                    errors.append(f"Invalid audio file {f_normalized}: {str(e)}")109                    continue110                111                # Audio characteristics112                if channels != 2:113                    errors.append(f"File must be stereo (has {channels} channels): {f_normalized}")114                    continue115                116                sample_rates.add(sr)117                durations.append(frames / sr)118                channels_set.add(channels)119                valid_files.append(f_normalized)120                121            except Exception as e:122                errors.append(f"Error processing {f}: {str(e)}")123                continue124        125        # Final checks126        if len(valid_files) < 2:127            errors.append("At least 2 valid files required")128        129        if len(sample_rates) > 1:130            errors.append(f"Sample rate mismatch: {sample_rates}")131        132        if len(channels_set) > 1:133            errors.append(f"Channel count mismatch: {channels_set}")134        135        # Output path validation136        try:137            output_path = self.normalize_path(output_path)138            output_dir = os.path.dirname(output_path) or '/kaggle/working/'139            140            if not os.path.exists(output_dir):141                os.makedirs(output_dir, exist_ok=True)142            143            if not os.access(output_dir, os.W_OK):144                errors.append(f"No write permission for output directory: {output_dir}")145        except Exception as e:146            errors.append(f"Output path error: {str(e)}")147        148        if errors:149            error_msg = "\n".join(errors)150            self.log_message(f"Validation failed:\n{error_msg}")151            raise ValueError(error_msg)152        153        target_sr = sample_rates.pop() if sample_rates else 44100154        return valid_files, target_sr, min(durations) if durations else None155    156    def process_waveform(self, chunks, method, weights=None):157        """All waveform domain processing methods."""158        if method == 'avg_wave':159            if weights is not None:160                return np.average(chunks, axis=0, weights=weights)161            return np.mean(chunks, axis=0)162        elif method == 'median_wave':163            return np.median(chunks, axis=0)164        elif method == 'max_wave':165            return np.max(chunks, axis=0)166        elif method == 'min_wave':167            return np.min(chunks, axis=0)168    169    def process_spectral(self, chunks, method):170        """All frequency domain processing methods."""171        specs = []172        for c in chunks:173            channel_specs = []174            for channel in range(c.shape[0]):175                _, _, Zxx = stft(c[channel], nperseg=1024, noverlap=512)176                channel_specs.append(Zxx)177            specs.append(np.array(channel_specs))178        179        specs = np.array(specs)180        mag = np.abs(specs)181        182        if method == 'max_fft':183            combined_mag = np.max(mag, axis=0)184        elif method == 'min_fft':185            combined_mag = np.min(mag, axis=0)186        elif method == 'median_fft':187            combined_mag = np.median(mag, axis=0)188        189        # Use phase from first file190        combined_spec = combined_mag * np.exp(1j * np.angle(specs[0]))191        192        # ISTFT reconstruction193        reconstructed = np.zeros((combined_spec.shape[0], chunks[0].shape[1]))194        for channel in range(combined_spec.shape[0]):195            _, xrec = istft(combined_spec[channel], nperseg=1024, noverlap=512)196            reconstructed[channel] = xrec[:chunks[0].shape[1]]197        198        return reconstructed199    200    def run_ensemble(self, files, method, output_path, weights=None, buffer_size=32768):201        """Core ensemble processing with maximum robustness."""202        try:203            # Validate and prepare inputs204            valid_files, target_sr, duration = self.validate_inputs(files, method, output_path)205            output_path = self.normalize_path(output_path)206            207            self.log_message(f"Starting ensemble with method: {method}")208            self.log_message(f"Input files: {json.dumps(valid_files, indent=2)}")209            self.log_message(f"Target sample rate: {target_sr}Hz")210            self.log_message(f"Output path: {output_path}")211            212            # Prepare weights213            if weights and len(weights) == len(valid_files):214                weights = np.array(weights, dtype=np.float32)215                weights /= weights.sum()  # Normalize216                self.log_message(f"Using weights: {weights}")217            else:218                weights = None219            220            # Open all files221            readers = []222            try:223                readers = [sf.SoundFile(f) for f in valid_files]224                shortest_frames = min(int(duration * r.samplerate) for r in readers)225                226                # Prepare output227                with sf.SoundFile(output_path, 'w', target_sr, 2, 'PCM_24') as outfile:228                    # Process in chunks with progress bar229                    progress = tqdm(total=shortest_frames, unit='samples', desc='Processing')230                    231                    for pos in range(0, shortest_frames, buffer_size):232                        chunk_size = min(buffer_size, shortest_frames - pos)233                        234                        # Read aligned chunks from all files235                        chunks = []236                        for r in readers:237                            r.seek(pos)238                            data = r.read(chunk_size)239                            if data.size == 0:240                                data = np.zeros((chunk_size, 2))241                            chunks.append(data.T)  # Transpose to (channels, samples)242                        243                        chunks = np.array(chunks)244                        245                        # Process based on method type246                        if method.endswith('_fft'):247                            result = self.process_spectral(chunks, method)248                        else:249                            result = self.process_waveform(chunks, method, weights)250                        251                        # Write output252                        outfile.write(result.T)  # Transpose back to (samples, channels)253                        254                        # Clean up and update progress255                        del chunks, result256                        if pos % (5 * buffer_size) == 0:257                            gc.collect()258                        259                        progress.update(chunk_size)260                    261                    progress.close()262                263                self.log_message(f"Successfully created output: {output_path}")264                print(f"\nEnsemble completed successfully: {output_path}")265                return True266                267            except Exception as e:268                self.log_message(f"Processing error: {str(e)}\n{traceback.format_exc()}")269                raise270            finally:271                for r in readers:272                    try:273                        r.close()274                    except:275                        pass276                277        except Exception as e:278            self.log_message(f"Fatal error: {str(e)}\n{traceback.format_exc()}")279            print(f"\nError during processing: {str(e)}", file=sys.stderr)280            return False281 282def main():283    parser = argparse.ArgumentParser(284        description='Ultimate Audio Ensemble Processor - Supports all ensemble methods',285        formatter_class=argparse.ArgumentDefaultsHelpFormatter286    )287    parser.add_argument('--files', nargs='+', required=True,288                       help='Input audio files (supports special characters)')289    parser.add_argument('--type', required=True,290                       choices=['avg_wave', 'median_wave', 'max_wave', 'min_wave',291                               'max_fft', 'min_fft', 'median_fft'],292                       help='Ensemble method to use')293    parser.add_argument('--weights', nargs='+', type=float,294                       help='Relative weights for each input file')295    parser.add_argument('--output', required=True,296                       help='Output file path')297    parser.add_argument('--buffer', type=int, default=32768,298                       help='Buffer size in samples (larger=faster but uses more memory)')299    300    args = parser.parse_args()301    302    with AudioEnsembleEngine() as engine:303        success = engine.run_ensemble(304            files=args.files,305            method=args.type,306            output_path=args.output,307            weights=args.weights,308            buffer_size=args.buffer309        )310    311    sys.exit(0 if success else 1)312 313if __name__ == "__main__":314    main()315