import os os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # Suppress TensorFlow logs (must be set before importing TensorFlow) import tensorflow as tf tf.get_logger().setLevel('ERROR') # Suppress TensorFlow ERROR logs import warnings warnings.filterwarnings("ignore") # Suppress all warnings import argparse from functools import reduce from typing import List, Tuple import shutil import librosa import numpy as np from matplotlib import pyplot as plt from pydub import AudioSegment from pydub.silence import detect_nonsilent from pytube import YouTube from sklearn.preprocessing import StandardScaler import shutil import streamlit as st import tempfile # Constants SR = 12000 HOP_LENGTH = 128 MAX_FRAMES = 300 MAX_METERS = 201 N_FEATURES = 15 MODEL_PATH = "models/CRNN/best_model_V3.h5" AUDIO_TEMP_PATH = "output/temp" def extract_audio(url): try: yt = YouTube(url) video_title = yt.title audio_stream = yt.streams.filter(only_audio=True).first() if audio_stream: temp_dir = tempfile.mkdtemp() out_file = audio_stream.download(temp_dir) base, _ = os.path.splitext(out_file) audio_file = base + '.mp3' if os.path.exists(audio_file): os.remove(audio_file) os.rename(out_file, audio_file) return audio_file, video_title, temp_dir else: st.error("No audio stream found") return None, None, None except Exception as e: st.error(f"An error occurred: {e}") return None, None, None def strip_silence(audio_path): """Removes silent parts from an audio file.""" sound = AudioSegment.from_file(audio_path) nonsilent_ranges = detect_nonsilent( sound, min_silence_len=500, silence_thresh=-50) stripped = reduce(lambda acc, val: acc + sound[val[0]:val[1]], nonsilent_ranges, AudioSegment.empty()) stripped.export(audio_path, format='mp3') class AudioFeature: """Class for extracting and processing audio features.""" def __init__(self, audio_path, sr=SR, hop_length=HOP_LENGTH): self.audio_path = audio_path self.beats = None self.chroma_acts = None self.chromagram = None self.combined_features = None self.hop_length = hop_length self.key, self.mode = None, None self.mel_acts = None self.melspectrogram = None self.meter_grid = None self.mfccs = None self.mfcc_acts = None self.n_frames = None self.onset_env = None self.rms = None self.spectrogram = None self.sr = sr self.tempo = None self.tempogram = None self.tempogram_acts = None self.time_signature = 4 self.y = None self.y_harm, self.y_perc = None, None def detect_key(self, chroma_vals: np.ndarray) -> Tuple[str, str]: """Detect the key and mode (major or minor) of the audio segment.""" note_names = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B'] major_profile = np.array( [6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88]) minor_profile = np.array( [6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17]) major_profile /= np.linalg.norm(major_profile) minor_profile /= np.linalg.norm(minor_profile) major_correlations = [np.corrcoef(chroma_vals, np.roll(major_profile, i))[ 0, 1] for i in range(12)] minor_correlations = [np.corrcoef(chroma_vals, np.roll(minor_profile, i))[ 0, 1] for i in range(12)] max_major_idx = np.argmax(major_correlations) max_minor_idx = np.argmax(minor_correlations) self.mode = 'major' if major_correlations[max_major_idx] > minor_correlations[max_minor_idx] else 'minor' self.key = note_names[max_major_idx if self.mode == 'major' else max_minor_idx] return self.key, self.mode def calculate_ki_chroma(self, waveform: np.ndarray, sr: int, hop_length: int) -> np.ndarray: """Calculate a normalized, key-invariant chromagram for the given audio waveform.""" chromagram = librosa.feature.chroma_cqt( y=waveform, sr=sr, hop_length=hop_length, bins_per_octave=24) chromagram = (chromagram - chromagram.min()) / \ (chromagram.max() - chromagram.min()) chroma_vals = np.sum(chromagram, axis=1) key, mode = self.detect_key(chroma_vals) key_idx = ['C', 'C#', 'D', 'D#', 'E', 'F', 'F#', 'G', 'G#', 'A', 'A#', 'B'].index(key) shift_amount = -key_idx if mode == 'major' else -(key_idx + 3) % 12 return librosa.util.normalize(np.roll(chromagram, shift_amount, axis=0), axis=1) def extract_features(self): """Extract various audio features from the loaded audio.""" self.y, self.sr = librosa.load(self.audio_path, sr=self.sr) self.y_harm, self.y_perc = librosa.effects.hpss(self.y) self.spectrogram, _ = librosa.magphase( librosa.stft(self.y, hop_length=self.hop_length)) self.rms = librosa.feature.rms( S=self.spectrogram, hop_length=self.hop_length).astype(np.float32) self.melspectrogram = librosa.feature.melspectrogram( y=self.y, sr=self.sr, n_mels=128, hop_length=self.hop_length).astype(np.float32) self.mel_acts = librosa.decompose.decompose( self.melspectrogram, n_components=3, sort=True)[1].astype(np.float32) self.chromagram = self.calculate_ki_chroma( self.y_harm, self.sr, self.hop_length).astype(np.float32) self.chroma_acts = librosa.decompose.decompose( self.chromagram, n_components=4, sort=True)[1].astype(np.float32) self.onset_env = librosa.onset.onset_strength( y=self.y_perc, sr=self.sr, hop_length=self.hop_length) self.tempogram = np.clip(librosa.feature.tempogram( onset_envelope=self.onset_env, sr=self.sr, hop_length=self.hop_length), 0, None) self.tempogram_acts = librosa.decompose.decompose( self.tempogram, n_components=3, sort=True)[1] self.mfccs = librosa.feature.mfcc( y=self.y, sr=self.sr, n_mfcc=20, hop_length=self.hop_length) self.mfccs += abs(np.min(self.mfccs)) self.mfcc_acts = librosa.decompose.decompose( self.mfccs, n_components=4, sort=True)[1].astype(np.float32) features = [self.rms, self.mel_acts, self.chroma_acts, self.tempogram_acts, self.mfcc_acts] feature_names = ['rms', 'mel_acts', 'chroma_acts', 'tempogram_acts', 'mfcc_acts'] dims = {name: feature.shape[0] for feature, name in zip(features, feature_names)} total_inv_dim = sum(1 / dim for dim in dims.values()) weights = {name: 1 / (dims[name] * total_inv_dim) for name in feature_names} std_weighted_features = [StandardScaler().fit_transform(feature.T).T * weights[name] for feature, name in zip(features, feature_names)] self.combined_features = np.concatenate( std_weighted_features, axis=0).T.astype(np.float32) self.n_frames = len(self.combined_features) def create_meter_grid(self): """Create a grid based on the meter of the song, using tempo and beats.""" self.tempo, self.beats = librosa.beat.beat_track( onset_envelope=self.onset_env, sr=self.sr, hop_length=self.hop_length) self.tempo = self.tempo * 2 if self.tempo < 70 else self.tempo / \ 2 if self.tempo > 140 else self.tempo self.meter_grid = self._create_meter_grid() return self.meter_grid def _create_meter_grid(self) -> np.ndarray: """ Helper function to create a meter grid for the song, extrapolating both forwards and backwards from an anchor frame. Returns: - np.ndarray: The meter grid. """ seconds_per_beat = 60 / self.tempo beat_interval = int(librosa.time_to_frames( seconds_per_beat, sr=self.sr, hop_length=self.hop_length)) # Find the best matching start beat based on the tempo and existing beats best_match_start = max((1 - abs(np.mean(self.beats[i:i+3]) - beat_interval) / beat_interval, self.beats[i]) for i in range(len(self.beats) - 2))[1] anchor_frame = best_match_start if best_match_start > 0.95 else self.beats[0] first_beat_time = librosa.frames_to_time( anchor_frame, sr=self.sr, hop_length=self.hop_length) # Calculate the number of beats forward and backward time_duration = librosa.frames_to_time( self.n_frames, sr=self.sr, hop_length=self.hop_length) num_beats_forward = int( (time_duration - first_beat_time) / seconds_per_beat) num_beats_backward = int(first_beat_time / seconds_per_beat) + 1 # Create beat times forward and backward beat_times_forward = first_beat_time + \ np.arange(num_beats_forward) * seconds_per_beat beat_times_backward = first_beat_time - \ np.arange(1, num_beats_backward) * seconds_per_beat # Combine and sort the beat times beat_grid = np.concatenate( (np.array([0.0]), beat_times_backward[::-1], beat_times_forward)) meter_indices = np.arange(0, len(beat_grid), self.time_signature) meter_grid = beat_grid[meter_indices] # Ensure the meter grid starts at 0 and ends at frame_duration if meter_grid[0] != 0.0: meter_grid = np.insert(meter_grid, 0, 0.0) meter_grid = librosa.time_to_frames( meter_grid, sr=self.sr, hop_length=self.hop_length) if meter_grid[-1] != self.n_frames: meter_grid = np.append(meter_grid, self.n_frames) return meter_grid def segment_data_meters(data: np.ndarray, meter_grid: List[int]) -> List[np.ndarray]: """ Divide song data into segments based on measure grid frames. Parameters: - data (np.ndarray): The song data to be segmented. - meter_grid (List[int]): The grid indicating the start of each measure. Returns: - List[np.ndarray]: A list of song data segments. """ meter_segments = [data[s:e] for s, e in zip(meter_grid[:-1], meter_grid[1:])] meter_segments = [segment.astype(np.float32) for segment in meter_segments] return meter_segments def positional_encoding(position: int, d_model: int) -> np.ndarray: """ Generate a positional encoding for a given position and model dimension. Parameters: - position (int): The position for which to generate the encoding. - d_model (int): The dimension of the model. Returns: - np.ndarray: The positional encoding. """ angle_rads = np.arange(position)[:, np.newaxis] / np.power( 10000, (2 * (np.arange(d_model)[np.newaxis, :] // 2)) / np.float32(d_model)) return np.concatenate([np.sin(angle_rads[:, 0::2]), np.cos(angle_rads[:, 1::2])], axis=-1) def apply_hierarchical_positional_encoding(segments: List[np.ndarray]) -> List[np.ndarray]: """ Apply positional encoding at the meter and frame levels to a list of segments. Parameters: - segments (List[np.ndarray]): The list of segments to encode. Returns: - List[np.ndarray]: The list of segments with applied positional encoding. """ n_features = segments[0].shape[1] measure_level_encodings = positional_encoding(len(segments), n_features) return [ seg + positional_encoding(len(seg), n_features) + measure_level_encodings[i] for i, seg in enumerate(segments) ] def pad_song(encoded_segments: List[np.ndarray], max_frames: int = MAX_FRAMES, max_meters: int = MAX_METERS, n_features: int = N_FEATURES) -> np.ndarray: """ Pad or truncate the encoded segments to have the specified max_frames and max_meters dimensions. Parameters: - encoded_segments (List[np.ndarray]): The encoded segments to pad or truncate. - max_frames (int): The maximum number of frames per segment. - max_meters (int): The maximum number of meters. - n_features (int): The number of features per frame. Returns: - np.ndarray: The padded or truncated song. """ padded_meters = [ np.pad(meter[:max_frames], ((0, max(0, max_frames - meter.shape[0])), (0, 0)), 'constant', constant_values=0) for meter in encoded_segments ] padding_meter = np.zeros((max_frames, n_features)) padded_song = np.array( padded_meters[:max_meters] + [padding_meter] * max(0, max_meters - len(padded_meters))) return padded_song def process_audio(audio_path, trim_silence=True, sr=SR, hop_length=HOP_LENGTH): """ Process an audio file, extracting features and applying positional encoding. Parameters: - audio_path (str): The path to the audio file. - trim_silence (bool): Whether to trim silence from the audio. - sr (int): The sample rate to use when loading the audio. - hop_length (int): The hop length to use for feature extraction. Returns: - Tuple[np.ndarray, AudioFeature]: The processed audio and its features. """ if trim_silence: strip_silence(audio_path) audio_features = AudioFeature( audio_path=audio_path, sr=sr, hop_length=hop_length) audio_features.extract_features() audio_features.create_meter_grid() audio_segments = segment_data_meters( audio_features.combined_features, audio_features.meter_grid) encoded_audio_segments = apply_hierarchical_positional_encoding( audio_segments) processed_audio = np.expand_dims(pad_song(encoded_audio_segments), axis=0) return processed_audio, audio_features def load_model(model_path=MODEL_PATH): # Placeholder functions for loading the model def custom_binary_crossentropy(y_true, y_pred): return y_pred def custom_accuracy(y_true, y_pred): return y_pred custom_objects = { 'custom_binary_crossentropy': custom_binary_crossentropy, 'custom_accuracy': custom_accuracy } model = tf.keras.models.load_model(model_path, custom_objects=custom_objects) return model def smooth_predictions(data: np.ndarray) -> np.ndarray: """ Smooth predictions by correcting isolated mispredictions and removing short sequences of 1s. This function applies a smoothing algorithm to correct isolated zeros and ones in a sequence of binary predictions. It also removes isolated sequences of 1s that are shorter than 5. Parameters: - data (np.ndarray): Array of binary predictions. Returns: - np.ndarray: Smoothed array of binary predictions. """ if not isinstance(data, np.ndarray): data = np.array(data) # First pass: Correct isolated 0's data_first_pass = data.copy() for i in range(1, len(data) - 1): if data[i] == 0 and data[i - 1] == 1 and data[i + 1] == 1: data_first_pass[i] = 1 # Second pass: Correct isolated 1's corrected_data = data_first_pass.copy() for i in range(1, len(data_first_pass) - 1): if data_first_pass[i] == 1 and data_first_pass[i - 1] == 0 and data_first_pass[i + 1] == 0: corrected_data[i] = 0 # Third pass: Remove short sequences of 1s (less than 5) smoothed_data = corrected_data.copy() sequence_start = None for i in range(len(corrected_data)): if corrected_data[i] == 1: if sequence_start is None: sequence_start = i else: if sequence_start is not None: sequence_length = i - sequence_start if sequence_length < 5: smoothed_data[sequence_start:i] = 0 sequence_start = None return smoothed_data def make_predictions(model, processed_audio, audio_features, url, video_name): """ Generate predictions from the model and process them to binary and smoothed predictions. Parameters: - model: The loaded model for making predictions. - processed_audio: The audio data that has been processed for prediction. - audio_features: Audio features object containing necessary metadata like meter grid. - url (str): YouTube URL of the audio file. - video_name (str): Name of the video. Returns: - np.ndarray: The smoothed binary predictions. """ predictions = model.predict(processed_audio)[0] binary_predictions = np.round( predictions[:(len(audio_features.meter_grid) - 1)]).flatten() smoothed_predictions = smooth_predictions(binary_predictions) meter_grid_times = librosa.frames_to_time( audio_features.meter_grid, sr=audio_features.sr, hop_length=audio_features.hop_length) chorus_start_times = [meter_grid_times[i] for i in range(len( smoothed_predictions)) if smoothed_predictions[i] == 1 and (i == 0 or smoothed_predictions[i - 1] == 0)] chorus_end_times = [meter_grid_times[i + 1] for i in range(len( smoothed_predictions)) if smoothed_predictions[i] == 1 and (i == len(smoothed_predictions) - 1 or smoothed_predictions[i + 1] == 0)] st.write(f"**Video Title:** {video_name}") st.write(f"**Number of choruses identified:** {len(chorus_start_times)}") for start_time, end_time in zip(chorus_start_times, chorus_end_times): link = f"{url}&t={int(start_time)}s" st.write(f"Chorus from {start_time:.2f}s to {end_time:.2f}s: [Link]({link})") if len(chorus_start_times) == 0: st.write("No choruses identified.") return smoothed_predictions def plot_meter_lines(ax: plt.Axes, meter_grid_times: np.ndarray) -> None: """Draw meter grid lines on the plot.""" for time in meter_grid_times: ax.axvline(x=time, color='grey', linestyle='--', linewidth=1, alpha=0.6) def plot_predictions(audio_features, predictions): meter_grid_times = librosa.frames_to_time( audio_features.meter_grid, sr=audio_features.sr, hop_length=audio_features.hop_length) fig, ax = plt.subplots(figsize=(12.5, 3), dpi=96) # Display harmonic and percussive components without adding them to the legend librosa.display.waveshow(audio_features.y_harm, sr=audio_features.sr, alpha=0.8, ax=ax, color='deepskyblue') librosa.display.waveshow(audio_features.y_perc, sr=audio_features.sr, alpha=0.7, ax=ax, color='plum') plot_meter_lines(ax, meter_grid_times) for i, prediction in enumerate(predictions): start_time = meter_grid_times[i] end_time = meter_grid_times[i + 1] if i < len( meter_grid_times) - 1 else len(audio_features.y) / audio_features.sr if prediction == 1: ax.axvspan(start_time, end_time, color='green', alpha=0.3, label='Predicted Chorus' if i == 0 else None) ax.set_xlim([0, len(audio_features.y) / audio_features.sr]) ax.set_ylabel('Amplitude') audio_file_name = os.path.basename(audio_features.audio_path) ax.set_title( f'Chorus Predictions for {os.path.splitext(audio_file_name)[0]}') # Add a green square patch to represent "Chorus" in the legend chorus_patch = plt.Rectangle((0, 0), 1, 1, fc='green', alpha=0.3) handles, labels = ax.get_legend_handles_labels() handles.append(chorus_patch) labels.append('Chorus') ax.legend(handles=handles, labels=labels) # Set x-tick labels every 10 seconds in single-digit minutes format duration = len(audio_features.y) / audio_features.sr xticks = np.arange(0, duration, 10) xlabels = [f"{int(tick // 60)}:{int(tick % 60):02d}" for tick in xticks] ax.set_xticks(xticks) ax.set_xticklabels(xlabels) plt.tight_layout() st.pyplot(plt) def main(): st.title("Chorus Finder") st.write("Upload a YouTube URL to find the chorus in the song.") url = st.text_input("YouTube URL") if st.button("Find Chorus"): if url: audio_file, video_title, temp_dir = extract_audio(url) if audio_file: strip_silence(audio_file) processed_audio, audio_features = process_audio(audio_path=audio_file) model = load_model() smoothed_predictions = make_predictions(model, processed_audio, audio_features, url, video_title) plot_predictions(audio_features=audio_features, predictions=smoothed_predictions) shutil.rmtree(temp_dir) else: st.error("Please enter a valid YouTube URL") if __name__ == "__main__": main()