Janady13/Jennys-Drumste / drum_pattern_extractor.py

Source identity

Field	Recorded value
Source type	GitHub source
Original record	Janady13/Jennys-Drumste / drum_pattern_extractor.py
Date	Not specified in source
Repository	Janady13/Jennys-Drumste
Revision	c601bf1648fe71058835589331c7f7e669940b8b
Branch	main
Path or archive member	drum_pattern_extractor.py
Original SHA-256	b6d521c6c531a4b37bae20e34a9d861667785ef7a4da6d73233282879d731895
Git blob SHA	62af68b831d1aa620611262af50e40abfaa8f8b4
Reading-copy SHA-256	b6d521c6c531a4b37bae20e34a9d861667785ef7a4da6d73233282879d731895
Original source ↗
Every recorded occurrence

These 3 records contain identical original source bytes. Their paths and revisions remain part of the archive.
Original record	Branch or archive	Revision / identifier	Original location
Janady13/Jennys-Drumste / drum_pattern_extractor.py	main	c601bf1648fe71058835589331c7f7e669940b8b	Open original
Janady13/Jennys-Drumste / drum_pattern_extractor.py	claude/complete-integration-feature-011CUkTzyyk8bKmXaRXJAR61	01c36ac1a180916e11864b4573b8c1999e2c5db3	Open original
Janady13/Jennys-Drumste / drum_pattern_extractor.py	claude/jenny-drumsets-research-MofSb	1f946b9e18e552f80f4b2ffd7f59c01d0126aa53	Open original
Recovered source transcript

Download complete source text ↓321 lines · 11.6 KB
"""
drum_pattern_extractor.py - Extract drum patterns from YouTube videos
This script analyzes audio to detect drum hits and create patterns for DrumMaster Pro
"""

import os
import json
import numpy as np
from typing import Dict, List, Tuple
import librosa
import librosa.display
from pytube import YouTube
from pydub import AudioSegment
import matplotlib.pyplot as plt
from scipy.signal import find_peaks
import warnings
warnings.filterwarnings('ignore')

class DrumPatternExtractor:
    """Extract drum patterns from audio files or YouTube videos"""
    
    def __init__(self):
        self.sample_rate = 22050
        self.hop_length = 512
        self.drum_frequencies = {
            'kick': (20, 100),      # Bass drum frequency range
            'snare': (200, 400),    # Snare drum frequency range
            'hihat': (6000, 12000), # Hi-hat frequency range
            'crash': (400, 4000),   # Crash cymbal range
            'ride': (300, 3000)     # Ride cymbal range
        }
        
    def download_youtube_audio(self, url: str, output_path: str = 'temp_audio.mp3') -> str:
        """Download audio from YouTube video"""
        try:
            yt = YouTube(url)
            audio_stream = yt.streams.filter(only_audio=True).first()
            
            if audio_stream:
                print(f"Downloading: {yt.title}")
                audio_file = audio_stream.download(filename='temp_audio')
                
                # Convert to mp3
                audio = AudioSegment.from_file(audio_file)
                audio.export(output_path, format='mp3')
                
                # Clean up temp file
                if os.path.exists(audio_file):
                    os.remove(audio_file)
                    
                return output_path
            else:
                raise Exception("No audio stream found")
                
        except Exception as e:
            print(f"Error downloading YouTube audio: {e}")
            return None
    
    def load_audio(self, file_path: str) -> Tuple[np.ndarray, int]:
        """Load audio file and return audio data and sample rate"""
        y, sr = librosa.load(file_path, sr=self.sample_rate)
        return y, sr
    
    def detect_tempo(self, y: np.ndarray) -> float:
        """Detect BPM of the audio"""
        tempo, beats = librosa.beat.beat_track(y=y, sr=self.sample_rate)
        return tempo
    
    def extract_drum_hits(self, y: np.ndarray) -> Dict[str, List[float]]:
        """Extract drum hit timings from audio"""
        drum_hits = {}
        
        # Percussive/Harmonic separation
        y_harmonic, y_percussive = librosa.effects.hpss(y)
        
        for drum_type, (low_freq, high_freq) in self.drum_frequencies.items():
            # Apply bandpass filter for specific drum frequency
            drum_hits[drum_type] = self._detect_hits_in_frequency_band(
                y_percussive, low_freq, high_freq
            )
        
        return drum_hits
    
    def _detect_hits_in_frequency_band(self, y: np.ndarray, 
                                      low_freq: float, 
                                      high_freq: float) -> List[float]:
        """Detect hits within a specific frequency band"""
        # Compute spectrogram
        D = librosa.stft(y, hop_length=self.hop_length)
        freqs = librosa.fft_frequencies(sr=self.sample_rate)
        
        # Find frequency bins
        low_bin = np.argmin(np.abs(freqs - low_freq))
        high_bin = np.argmin(np.abs(freqs - high_freq))
        
        # Extract energy in frequency band
        band_energy = np.sum(np.abs(D[low_bin:high_bin, :]), axis=0)
        
        # Detect peaks (drum hits)
        peaks, properties = find_peaks(
            band_energy, 
            height=np.mean(band_energy) + np.std(band_energy),
            distance=int(0.1 * self.sample_rate / self.hop_length)  # Min 100ms between hits
        )
        
        # Convert to timestamps
        times = librosa.frames_to_time(peaks, sr=self.sample_rate, hop_length=self.hop_length)
        
        return times.tolist()
    
    def quantize_to_grid(self, hit_times: List[float], 
                        tempo: float, 
                        bars: int = 1, 
                        subdivision: int = 16) -> str:
        """Quantize hit times to a rhythmic grid"""
        beat_duration = 60.0 / tempo  # Duration of one beat in seconds
        sixteenth_duration = beat_duration / 4  # Duration of 16th note
        
        # Create pattern string
        pattern_length = bars * subdivision
        pattern = ['-'] * pattern_length
        
        # Map hits to grid positions
        for hit_time in hit_times:
            # Find nearest grid position
            grid_position = round(hit_time / sixteenth_duration) % pattern_length
            if 0 <= grid_position < pattern_length:
                pattern[grid_position] = '1'
        
        return ''.join(pattern)
    
    def analyze_song_structure(self, y: np.ndarray) -> Dict:
        """Analyze song structure (verse, chorus, etc.)"""
        # Use spectral features to detect song sections
        chroma = librosa.feature.chroma_cqt(y=y, sr=self.sample_rate)
        
        # Compute self-similarity matrix
        similarity = librosa.segment.recurrence_matrix(chroma, mode='affinity')
        
        # Detect segment boundaries
        boundaries = self._detect_boundaries(similarity)
        
        sections = []
        for i in range(len(boundaries) - 1):
            start_time = librosa.frames_to_time(boundaries[i], sr=self.sample_rate)
            end_time = librosa.frames_to_time(boundaries[i + 1], sr=self.sample_rate)
            sections.append({
                'start': start_time,
                'end': end_time,
                'type': self._classify_section(chroma[:, boundaries[i]:boundaries[i + 1]])
            })
        
        return sections
    
    def _detect_boundaries(self, similarity_matrix: np.ndarray) -> List[int]:
        """Detect section boundaries from similarity matrix"""
        # Simple boundary detection using diagonal sum
        diagonal_sum = np.sum(similarity, axis=0)
        
        # Find valleys (potential boundaries)
        peaks, _ = find_peaks(-diagonal_sum, distance=50)
        
        # Add start and end
        boundaries = [0] + peaks.tolist() + [similarity_matrix.shape[0] - 1]
        
        return boundaries
    
    def _classify_section(self, chroma_section: np.ndarray) -> str:
        """Classify section type based on harmonic content"""
        # Simple classification based on energy
        energy = np.mean(chroma_section)
        
        if energy > 0.5:
            return 'chorus'
        elif energy > 0.3:
            return 'verse'
        else:
            return 'intro/outro'
    
    def create_drummaster_pattern(self, youtube_url: str, 
                                 output_file: str = 'pattern.json') -> Dict:
        """Complete pipeline to create DrumMaster Pro pattern from YouTube URL"""
        print(f"Processing: {youtube_url}")
        
        # Download audio
        audio_file = self.download_youtube_audio(youtube_url)
        if not audio_file:
            return None
        
        # Load audio
        y, sr = self.load_audio(audio_file)
        
        # Detect tempo
        tempo = self.detect_tempo(y)
        print(f"Detected BPM: {tempo:.1f}")
        
        # Extract drum hits
        drum_hits = self.extract_drum_hits(y)
        
        # Convert to patterns
        patterns = {}
        for drum_type, hits in drum_hits.items():
            pattern = self.quantize_to_grid(hits[:100], tempo)  # First 100 hits
            patterns[drum_type] = pattern
            print(f"{drum_type}: {pattern}")
        
        # Analyze structure
        sections = self.analyze_song_structure(y)
        
        # Create DrumMaster Pro format
        drummaster_data = {
            'title': YouTube(youtube_url).title,
            'youtube_url': youtube_url,
            'youtube_id': youtube_url.split('v=')[1].split('&')[0],
            'bpm': round(tempo),
            'patterns': patterns,
            'sections': sections,
            'difficulty': self._estimate_difficulty(patterns)
        }
        
        # Save to file
        with open(output_file, 'w') as f:
            json.dump(drummaster_data, f, indent=2)
        
        # Clean up
        if os.path.exists(audio_file):
            os.remove(audio_file)
        
        print(f"Pattern saved to {output_file}")
        return drummaster_data
    
    def _estimate_difficulty(self, patterns: Dict[str, str]) -> str:
        """Estimate difficulty based on pattern complexity"""
        complexity = 0
        
        for pattern in patterns.values():
            # Count hits
            hits = pattern.count('1')
            complexity += hits
            
            # Check for complex patterns
            if '111' in pattern:  # Rapid hits
                complexity += 5
            if pattern.count('1') > 8:  # Many hits
                complexity += 3
        
        if complexity < 20:
            return 'easy'
        elif complexity < 40:
            return 'medium'
        else:
            return 'hard'
    
    def visualize_pattern(self, pattern_data: Dict):
        """Visualize the extracted drum pattern"""
        fig, axes = plt.subplots(5, 1, figsize=(15, 8), sharex=True)
        
        drums = ['kick', 'snare', 'hihat', 'crash', 'ride']
        colors = ['red', 'blue', 'green', 'orange', 'purple']
        
        for idx, (drum, ax) in enumerate(zip(drums, axes)):
            pattern = pattern_data['patterns'].get(drum, '-' * 16)
            
            # Convert pattern to binary array
            binary = [1 if c == '1' else 0 for c in pattern]
            
            # Plot
            ax.bar(range(len(binary)), binary, color=colors[idx], alpha=0.7)
            ax.set_ylabel(drum.capitalize())
            ax.set_ylim(0, 1.5)
            ax.set_yticks([])
            ax.grid(True, alpha=0.3)
        
        axes[-1].set_xlabel('Beat (16th notes)')
        plt.suptitle(f"{pattern_data['title']} - {pattern_data['bpm']} BPM")
        plt.tight_layout()
        plt.show()

def main():
    """Example usage"""
    extractor = DrumPatternExtractor()
    
    # Example YouTube URLs - Popular songs with clear drum patterns
    example_songs = [
        "https://www.youtube.com/watch?v=0J2QdDbelmY",  # Seven Nation Army
        "https://www.youtube.com/watch?v=Zi_XLOBDo_Y",  # Billie Jean
        "https://www.youtube.com/watch?v=hTWKbfoikeg",  # Smells Like Teen Spirit
    ]
    
    # Process a song
    youtube_url = input("Enter YouTube URL (or press Enter for example): ")
    if not youtube_url:
        youtube_url = example_songs[0]
    
    pattern_data = extractor.create_drummaster_pattern(youtube_url)
    
    if pattern_data:
        # Visualize the pattern
        extractor.visualize_pattern(pattern_data)
        
        # Print JavaScript format for DrumMaster Pro
        print("\n// Add this to SONG_DATABASE in drum-song-mode.html:")
        print(f"""{{
    id: '{pattern_data['title'].lower().replace(' ', '-')}',
    title: '{pattern_data['title']}',
    artist: 'Extracted',
    youtubeId: '{pattern_data['youtube_id']}',
    bpm: {pattern_data['bpm']},
    difficulty: '{pattern_data['difficulty']}',
    pattern: {{
        kick:   '{pattern_data['patterns'].get('kick', '----------------')}',
        snare:  '{pattern_data['patterns'].get('snare', '----------------')}',
        hihat:  '{pattern_data['patterns'].get('hihat', '----------------')}',
        crash:  '{pattern_data['patterns'].get('crash', '----------------')}',
        ride:   '{pattern_data['patterns'].get('ride', '----------------')}'
    }}
}}""")

if __name__ == "__main__":
    main()