import mido import csv from pathlib import Path from collections import defaultdict import re from tqdm import tqdm from joblib import Parallel, delayed import logging import time import psutil import os logging.basicConfig(level=logging.WARN) logger = logging.getLogger(__name__) MAX_LENGTH_SEC = 480 # 8 min MIN_DENSITY = 0.50 # % of track where a note is active MIN_ONSET_FREQ = 0.25 # once every 4 measures DRUM_CHANNEL = 9 MIDI_FOLDER = "lmd_full" OUT_FOLDER = "lmd_full_stems" INSTRUMENT_CSV = "instrument_categories.csv" # Performance optimizations BATCH_SIZE = 1000 # Process files in batches to reduce memory usage PROGRESS_UPDATE_INTERVAL = 100 # Update progress every N files class InstrumentMapper: """Maps MIDI instrument names to stem categories.""" def __init__(self, csv_path=None): self.mapping = {} # Pre-compile regex patterns for better performance self._cleanup_pattern = re.compile(r"\([^)]*\)|\d+|_") if csv_path: # Instrument, Mapped column names self._load_csv(csv_path) else: logger.warning("WARNING: no stem mapping provided, using default") self._load_defaults() # Pre-compile pattern matching for faster lookups self._pattern_cache = {} self._compile_patterns() def _load_csv(self, csv_path): try: with open(csv_path, "r", encoding="utf-8") as file: for row in csv.DictReader(file): self.mapping[row["Instrument"].strip().lower()] = row["Mapped"].strip().lower() except Exception as e: logger.warning(f"Error loading mapping: {e}. Falling back to default") self._load_defaults() def _load_defaults(self): self.mapping = { "piano": "keyboard", "organ": "keyboard", "guitar": "guitar", "bass": "bass", "violin": "strings", "string": "strings", "trumpet": "brass", "trombone": "brass", "flute": "woodwinds", "sax": "woodwinds", "drum": "drums", "percussion": "percussion", "pad": "synth", "lead": "synth", "choir": "vocals", "voice": "vocals", } def _compile_patterns(self): """Pre-compile regex patterns for faster matching""" self.patterns = { "drums": re.compile(r"\b(drum|kit)\b"), "bass": re.compile(r"\bbass\b"), "guitar": re.compile(r"\b(guitar|banjo)\b"), "keyboard": re.compile(r"\b(piano|organ)\b"), "synth": re.compile(r"\b(pad|lead|synth)\b"), "strings": re.compile(r"\b(violin|viola|cello|string)\b"), "brass": re.compile(r"\b(trumpet|brass)\b"), "woodwinds": re.compile(r"\b(flute|sax|clarinet)\b"), "vocals": re.compile(r"\b(choir|voice|vox)\b"), "percussion": re.compile(r"\b(bell|chime|triangle)\b"), } def get_category(self, instrument_name): # Use cache for repeated lookups if instrument_name in self._pattern_cache: return self._pattern_cache[instrument_name] cleaned = self._cleanup_pattern.sub(" ", instrument_name.lower()).strip() # Match with category name if cleaned in self.mapping: result = self.mapping[cleaned] self._pattern_cache[instrument_name] = result return result # Match with MIDI instrument name for key, category in self.mapping.items(): if key in cleaned: self._pattern_cache[instrument_name] = category return category # Pattern matching with pre-compiled regex for category, pattern in self.patterns.items(): if pattern.search(cleaned): self._pattern_cache[instrument_name] = category return category # No category found self._pattern_cache[instrument_name] = "other" return "other" class MidiAnalyzer: """MIDI analyzer with instrument categorization and dual filtering.""" # Pre-define GM instruments as class variable to avoid recreation GM_INSTRUMENTS = [ "Acoustic Grand Piano", "Bright Acoustic Piano", "Electric Grand Piano", "Honky-tonk Piano", "Electric Piano 1", "Electric Piano 2", "Harpsichord", "Clavi", "Celesta", "Glockenspiel", "Music Box", "Vibraphone", "Marimba", "Xylophone", "Tubular Bells", "Dulcimer", "Drawbar Organ", "Percussive Organ", "Rock Organ", "Church Organ", "Reed Organ", "Accordion", "Harmonica", "Tango Accordion", "Acoustic Guitar (nylon)", "Acoustic Guitar (steel)", "Electric Guitar (jazz)", "Electric Guitar (clean)", "Electric Guitar (muted)", "Overdriven Guitar", "Distortion Guitar", "Guitar Harmonics", "Acoustic Bass", "Electric Bass (finger)", "Electric Bass (pick)", "Fretless Bass", "Slap Bass 1", "Slap Bass 2", "Synth Bass 1", "Synth Bass 2", "Violin", "Viola", "Cello", "Contrabass", "Tremolo Strings", "Pizzicato Strings", "Orchestral Harp", "Timpani", "String Ensemble 1", "String Ensemble 2", "Synth Strings 1", "Synth Strings 2", "Choir Aahs", "Voice Oohs", "Synth Voice", "Orchestra Hit", "Trumpet", "Trombone", "Tuba", "Muted Trumpet", "French Horn", "Brass Section", "Synth Brass 1", "Synth Brass 2", "Soprano Sax", "Alto Sax", "Tenor Sax", "Baritone Sax", "Oboe", "English Horn", "Bassoon", "Clarinet", "Piccolo", "Flute", "Recorder", "Pan Flute", "Blown Bottle", "Shakuhachi", "Whistle", "Ocarina", "Lead 1 (square)", "Lead 2 (sawtooth)", "Lead 3 (calliope)", "Lead 4 (chiff)", "Lead 5 (charang)", "Lead 6 (voice)", "Lead 7 (fifths)", "Lead 8 (bass + lead)", "Pad 1 (new age)", "Pad 2 (warm)", "Pad 3 (polysynth)", "Pad 4 (choir)", "Pad 5 (bowed)", "Pad 6 (metallic)", "Pad 7 (halo)", "Pad 8 (sweep)", "FX 1 (rain)", "FX 2 (soundtrack)", "FX 3 (crystal)", "FX 4 (atmosphere)", "FX 5 (brightness)", "FX 6 (goblins)", "FX 7 (echoes)", "FX 8 (sci-fi)", "Sitar", "Banjo", "Shamisen", "Koto", "Kalimba", "Bag pipe", "Fiddle", "Shanai", "Tinkle Bell", "Agogo", "Steel Drums", "Woodblock", "Taiko Drum", "Melodic Tom", "Synth Drum", "Reverse Cymbal", "Guitar Fret Noise", "Breath Noise", "Seashore", "Bird Tweet", "Telephone Ring", "Helicopter", "Applause", "Gunshot", ] def __init__( self, midi_file_path, instrument_mapper, min_density_threshold=0.5, min_onset_frequency=0.125 ): self.path = Path(midi_file_path) self.instrument_mapper = instrument_mapper self.min_density_threshold = min_density_threshold self.min_onset_frequency = min_onset_frequency self._midi_file = None self._bpm = None self._length = None self._categories = defaultdict(list) self._track_stats = defaultdict(dict) self._load_and_analyze() def _load_and_analyze(self): try: self._midi_file = mido.MidiFile(str(self.path)) except Exception as e: raise Exception(f"Failed to load MIDI file: {e}") self._length = self._midi_file.length if self._length > MAX_LENGTH_SEC: raise Exception(f"Skipping long MIDI file, {self._length} sec") self._bpm = self._extract_bpm() self._calculate_channel_stats() self._categorize_instruments() def _extract_bpm(self): for track in self._midi_file.tracks: for msg in track: if msg.type == "set_tempo": return mido.tempo2bpm(msg.tempo) return 120.0 def _calculate_channel_stats(self): ticks_per_beat = self._midi_file.ticks_per_beat bpm = self._bpm total_ticks = int((self._length * bpm * ticks_per_beat) / 60.0) for track_idx, track in enumerate(self._midi_file.tracks): track_time = 0 active_notes = defaultdict( lambda: defaultdict(lambda: None) ) # channel -> {pitch: start_time} channel_notes = defaultdict(list) # channel -> list of (start, end) tuples for msg in track: track_time += msg.time if not hasattr(msg, "channel"): continue channel = msg.channel if msg.type == "note_on" and msg.velocity > 0: # Note onset active_notes[channel][msg.note] = track_time elif msg.type == "note_off" or (msg.type == "note_on" and msg.velocity == 0): # Note offset - find matching note_on if msg.note in active_notes[channel] and active_notes[channel][msg.note] is not None: start_time = active_notes[channel][msg.note] active_notes[channel][msg.note] = None channel_notes[channel].append((start_time, track_time)) # Calculate stats for each channel for channel, periods in channel_notes.items(): onset_count = len(periods) onset_frequency = onset_count / self._length if self._length > 0 else 0.0 # Calculate density from actual note durations if onset_count > 0: merged = self._merge_periods(periods) active_time = sum(end - start for start, end in merged) density = min(active_time / total_ticks, 1.0) if total_ticks > 0 else 0.0 else: density = 0.0 active_time = 0.0 self._track_stats[track_idx][channel] = { "onset_count": onset_count, "onset_frequency": onset_frequency, "density": density, "active_time": active_time, } def _merge_periods(self, periods): """Helper function for calculating held duration density""" if not periods: return [] sorted_periods = sorted(periods) merged = [sorted_periods[0]] for start, end in sorted_periods[1:]: if start <= merged[-1][1]: merged[-1] = (merged[-1][0], max(merged[-1][1], end)) else: merged.append((start, end)) return merged def _categorize_instruments(self): # Find instruments and channels with notes for track_idx, track in enumerate(self._midi_file.tracks): channel_instruments = {} channels_used = set() for msg in track: if msg.type == "program_change": instrument = self._get_instrument_name(msg.program) category = self.instrument_mapper.get_category(instrument) is_drum = msg.channel == DRUM_CHANNEL if is_drum: # override category = "drums" channel_instruments[msg.channel] = { "instrument": instrument, "program": msg.program, "category": category, "is_drum": is_drum, } elif msg.type == "note_on" and msg.velocity > 0: channels_used.add(msg.channel) # Set defaults for channels without program changes for channel in channels_used: if channel not in channel_instruments: if channel == DRUM_CHANNEL: instrument, category = "Drum Kit", "drums" else: instrument, category = "Acoustic Grand Piano", "keyboard" channel_instruments[channel] = { "instrument": instrument, "category": category, "is_drum": channel == DRUM_CHANNEL, } # Filter and categorize for channel in channels_used: stats = self._track_stats[track_idx].get(channel, {}) density = stats.get("density", 0.0) onset_freq = stats.get("onset_frequency", 0.0) if onset_freq >= self.min_onset_frequency: is_drums = channel_instruments[channel]["is_drum"] if density >= self.min_density_threshold or is_drums: if channel in channel_instruments: info = channel_instruments[channel] self._categories[info["category"]].append( { "channel": channel, "instrument": info["instrument"], "is_drum": info["is_drum"], "density": density, "onset_frequency": onset_freq, "onset_count": stats.get("onset_count", 0), } ) def _get_instrument_name(self, program): return ( self.GM_INSTRUMENTS[program] if 0 <= program < len(self.GM_INSTRUMENTS) else f"Unknown ({program})" ) # Public API methods remain the same... @property def bpm(self): return self._bpm @property def length(self): return self._length @property def filename(self): return self.path.name @property def categories(self): return dict(self._categories) def get_category_list(self): return sorted(self._categories.keys()) def save_category_midis(self, output_dir): output_dir = Path(output_dir) output_dir.mkdir(exist_ok=True) saved_files = [] for category, instruments in self._categories.items(): if not instruments: continue # Create new MIDI file new_midi = mido.MidiFile( type=1, ticks_per_beat=self._midi_file.ticks_per_beat, charset=getattr(self._midi_file, "charset", "latin1"), ) # Create track for this category channels = [inst["channel"] for inst in instruments] track = self._create_category_track(channels, category) if track and self._track_has_notes(track): new_midi.tracks.append(track) filename = f"{self.path.stem}_{category}.mid" filepath = output_dir / filename new_midi.save(str(filepath)) saved_files.append(str(filepath)) return saved_files def _create_category_track(self, channels, category_name): track = mido.MidiTrack() track.append(mido.MetaMessage("track_name", name=category_name)) # Channel mapping: drums/percussion to ch9, others sequential channel_map = {} if category_name.lower() in ["drums", "percussion"]: for ch in channels: channel_map[ch] = DRUM_CHANNEL else: new_ch = 0 for ch in channels: if new_ch == DRUM_CHANNEL: # Skip drum channel new_ch += 1 channel_map[ch] = new_ch new_ch = (new_ch + 1) % 16 # Collect and sort messages messages = [] meta_added = set() for orig_track in self._midi_file.tracks: track_time = 0 for msg in orig_track: track_time += msg.time # Meta messages if msg.type in ["set_tempo", "time_signature", "key_signature"]: key = (msg.type, track_time, str(msg.dict())) if key not in meta_added: new_msg = msg.copy() new_msg.time = 0 messages.append((track_time, new_msg)) meta_added.add(key) # Channel messages - keep notes, program changes, and essential controllers only elif hasattr(msg, "channel") and msg.channel in channels: if msg.type in ["note_on", "note_off", "program_change", "control_change"]: # Only keep essential control changes if msg.type == "control_change" and msg.control in [7]: # Volume only new_msg = msg.copy() new_msg.channel = channel_map[msg.channel] new_msg.time = 0 messages.append((track_time, new_msg)) elif msg.type != "control_change": # Keep all non-CC messages (note_on, note_off, program_change) new_msg = msg.copy() new_msg.channel = channel_map[msg.channel] new_msg.time = 0 messages.append((track_time, new_msg)) # Remove duplicates and add to track messages.sort(key=lambda x: x[0]) deduplicated = self._remove_duplicate_notes(messages) last_time = 0 for abs_time, msg in deduplicated: msg.time = abs_time - last_time track.append(msg) last_time = abs_time return track def _remove_duplicate_notes(self, timed_messages): """Remove duplicate note events, keeping only the first occurrence.""" seen_events = set() result = [] for abs_time, msg in timed_messages: # Create unique key for each event type if msg.type in ["note_on", "note_off"]: event_key = ( msg.type, msg.note, msg.channel, abs_time, msg.velocity if msg.type == "note_on" else 0, ) else: # For non-note events, use message content as key event_key = (msg.type, abs_time, str(msg.dict())) if event_key not in seen_events: seen_events.add(event_key) result.append((abs_time, msg)) return result def _track_has_notes(self, track): return any(msg.type == "note_on" and msg.velocity > 0 for msg in track) def __str__(self): categories = ", ".join(self.get_category_list()) return f"MIDI: {self.filename} | BPM: {self.bpm} | Duration: {self.length:.1f}s | Categories: {categories}" def process_midi_file(input_file, instrument_mapper, input_path, output_path): """Process a single MIDI file - optimized for parallel execution""" relative_path = input_file.relative_to(input_path) output_folder = (output_path / relative_path).parent output_folder.mkdir(parents=True, exist_ok=True) try: midi = MidiAnalyzer( input_file, instrument_mapper, min_density_threshold=MIN_DENSITY, min_onset_frequency=MIN_ONSET_FREQ, ) saved_files = midi.save_category_midis(output_folder) return len(saved_files) # Return number of stems created except Exception as e: # Only log errors for debugging if needed # logger.debug(f"Error processing {input_file.name}: {e}") return 0 # Return 0 on failure def process_batch(batch_files, instrument_mapper, input_path, output_path): """Process a batch of MIDI files""" batch_results = [] for midi_file in batch_files: result = process_midi_file(midi_file, instrument_mapper, input_path, output_path) batch_results.append(result) return batch_results def get_optimal_n_jobs(): """Determine optimal number of jobs based on system resources""" cpu_count = psutil.cpu_count(logical=False) # Physical cores memory_gb = psutil.virtual_memory().total / (1024**3) # Conservative estimate: each job might use ~100-200MB # Adjust based on your system's memory max_jobs_by_memory = int(memory_gb * 0.8 / 0.2) # 80% of memory, 200MB per job max_jobs_by_cpu = cpu_count optimal_jobs = min(max_jobs_by_memory, max_jobs_by_cpu) return max(1, optimal_jobs) def main(): """Main processing function with optimizations""" start_time = time.time() # Initialize instrument mapper once instrument_mapper = InstrumentMapper(INSTRUMENT_CSV) # Setup paths input_path = Path(MIDI_FOLDER) output_path = Path(OUT_FOLDER) output_path.mkdir(parents=True, exist_ok=True) # Find all MIDI files midi_files = list(input_path.rglob("*.mid")) + list(input_path.rglob("*.midi")) total_files = len(midi_files) if total_files == 0: print("No MIDI files found!") return # Determine optimal number of jobs n_jobs = get_optimal_n_jobs() print(f"Found {total_files:,} MIDI files to process") print(f"Using {n_jobs} parallel jobs") print(f"System: {psutil.cpu_count()} CPUs, {psutil.virtual_memory().total / (1024**3):.1f}GB RAM") # Process files with better progress tracking try: # Use 'multiprocessing' backend for better CPU utilization with I/O bound tasks # Use 'threading' if you have memory constraints results = Parallel( n_jobs=n_jobs, backend="multiprocessing", # Try 'threading' if you get memory issues verbose=1, # Show progress from joblib )( delayed(process_midi_file)(midi_file, instrument_mapper, input_path, output_path) for midi_file in tqdm(midi_files, desc="Processing MIDI files", unit="files") ) # Calculate statistics total_stems = sum(r for r in results if r is not None) successful_files = sum(1 for r in results if r and r > 0) failed_files = total_files - successful_files # Performance metrics end_time = time.time() total_time = end_time - start_time files_per_second = total_files / total_time # Results summary print(f"\n{'='*60}") print(f"PROCESSING COMPLETE") print(f"{'='*60}") print(f"Total files processed: {total_files:,}") print(f"Successfully processed: {successful_files:,}") print(f"Failed: {failed_files:,}") print(f"Total stems generated: {total_stems:,}") print(f"Processing time: {total_time:.1f} seconds ({total_time/60:.1f} minutes)") print(f"Average speed: {files_per_second:.1f} files/second") if successful_files > 0: print(f"Average stems per successful file: {total_stems/successful_files:.1f}") # Warnings if successful_files == 0: logger.warning( "No files were successfully processed. Check your input directory and file formats." ) elif successful_files < len(midi_files) * 0.5: logger.warning( f"Low success rate: {successful_files}/{len(midi_files)} files processed successfully" ) else: logger.info(f"Processing completed successfully: {successful_files}/{len(midi_files)} files") except KeyboardInterrupt: print("\nProcessing interrupted by user") except Exception as e: print(f"Error during processing: {e}") logger.error(f"Processing failed with error: {e}") if __name__ == "__main__": main()