import librosa
import numpy as np
import matplotlib.pyplot as plt
import librosa.display
import soundfile as sf
import os
from pydub import AudioSegment

def preprocess_audio(audio_file):
    sound = AudioSegment.from_mp3(audio_file)
    wav_file = audio_file.rsplit('.',1)[0]+'.wav'
    sound.export(wav_file, format="wav" )
    return wav_file
    
def extract_melody(audio_file,output_folder)：
try:
        # 预处理音频文件
        audio_file = preprocess_audio(audio_file)
    
    # 加载音频文件
    y, sr = librosa.load(audio_file)

    # 使用 HPSS 分离和声和打击乐成分
    y_harmonic, y_percussive = librosa.effects.hpss(y)
    
    # 提取旋律（和声成分）
    y_harmonic, y_percussive = librosa.effects.harmonic(y)
    
    # 计算音高特征（Chroma)
    chroma = librosa.feature.chroma_cqt(y=y_harmonic, sr=sr)
    
    # 显示音高特征
    plt.figure(figsize=(10, 4))
    librosa.display.specshow(chroma, y_axis='chroma', x_axis='time')
    plt.colorbar()
    plt.title('Chromagram')
    plt.tight_layout()
    plt.show()

        # 保存提取的旋律
        base_name = os.path.splitext(os.path.basename(audio_file))[0]
        output_file = os.path.join(output_folder, f'extracted_melody_{base_name}.wav')
        sf.write(output_file, y_harmonic, sr)
        print(f"Extracted melody saved to: {output_file}")
    except Exception as e:
        print(f"An error occurred: {e}")

# 替换为你的音频文件路径
audio_file = 'your_song.mp3'
extract_melody(audio_file)