{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pretty_midi\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "import scipy.io.wavfile\n",
    "import shutil\n",
    "\n",
    "def process_midi_folder(midi_folder, output_folder, base_name, fps=100, sample_rate=44100):\n",
    "    \"\"\"\n",
    "    遍历指定文件夹中的 MIDI 文件，生成可视化图片和音频，并保存到指定文件夹。\n",
    "    \n",
    "    Args:\n",
    "        midi_folder (str): 输入 MIDI 文件夹路径\n",
    "        output_folder (str): 输出文件夹路径\n",
    "        base_name (str): 文件基础名称（不含后缀）\n",
    "        fps (int): 可视化帧率，默认 100\n",
    "        sample_rate (int): 音频采样率，默认 44100 Hz\n",
    "    \"\"\"\n",
    "    # 检查输入 MIDI 文件夹是否存在\n",
    "    if not os.path.exists(midi_folder):\n",
    "        print(f\"MIDI 文件夹 {midi_folder} 不存在！\")\n",
    "        return\n",
    "    \n",
    "    # 创建输出文件夹\n",
    "    os.makedirs(output_folder, exist_ok=True)\n",
    "    \n",
    "    # 遍历 MIDI 文件夹中的所有 .mid 文件\n",
    "    midi_files = [f for f in os.listdir(midi_folder) if f.endswith('.mid') or f.endswith('.midi')]\n",
    "    if not midi_files:\n",
    "        print(f\"文件夹 {midi_folder} 中没有找到 MIDI 文件！\")\n",
    "        return\n",
    "    \n",
    "    for idx, midi_file in enumerate(midi_files):\n",
    "        # 构造输入和输出路径\n",
    "        midi_path = os.path.join(midi_folder, midi_file)\n",
    "        sample_name = f\"{base_name}_{idx}\"  # 例如 sample_0, sample_1\n",
    "        output_midi_path = os.path.join(output_folder, f\"{sample_name}.mid\")\n",
    "        output_image_path = os.path.join(output_folder, f\"{sample_name}.png\")\n",
    "        output_wav_path = os.path.join(output_folder, f\"{sample_name}.wav\")\n",
    "        \n",
    "        print(f\"\\n处理文件: {midi_file}\")\n",
    "        \n",
    "        # 复制 MIDI 文件到输出文件夹\n",
    "        try:\n",
    "            shutil.copy(midi_path, output_midi_path)\n",
    "            print(f\"MIDI 文件已保存到 {output_midi_path}\")\n",
    "        except Exception as e:\n",
    "            print(f\"保存 MIDI 文件时出错: {e}\")\n",
    "            continue\n",
    "        \n",
    "        # 读取 MIDI 文件\n",
    "        try:\n",
    "            midi_data = pretty_midi.PrettyMIDI(midi_path)\n",
    "        except Exception as e:\n",
    "            print(f\"无法读取 MIDI 文件: {e}\")\n",
    "            continue\n",
    "        \n",
    "        # 获取 MIDI 文件总时长（秒）\n",
    "        total_time = midi_data.get_end_time()\n",
    "        if total_time <= 0:\n",
    "            print(\"MIDI 文件中没有音符数据或时长为零！\")\n",
    "            continue\n",
    "        \n",
    "        # --- 可视化部分 ---\n",
    "        num_frames = int(total_time * fps) + 1\n",
    "        time_step = 1.0 / fps\n",
    "        times = np.arange(0, total_time + time_step, time_step)\n",
    "        \n",
    "        piano_roll = np.zeros((128, num_frames), dtype=np.uint8)\n",
    "        \n",
    "        for instrument in midi_data.instruments:\n",
    "            if instrument.is_drum:\n",
    "                continue\n",
    "            for note in instrument.notes:\n",
    "                start_frame = int(note.start / time_step)\n",
    "                end_frame = int(note.end / time_step)\n",
    "                start_frame = min(start_frame, num_frames - 1)\n",
    "                end_frame = min(end_frame, num_frames)\n",
    "                piano_roll[note.pitch, start_frame:end_frame] = note.velocity\n",
    "        \n",
    "        if not piano_roll.any():\n",
    "            print(\"没有找到音符数据！\")\n",
    "            continue\n",
    "        \n",
    "        # 绘制并保存图片\n",
    "        plt.figure(figsize=(12, 6))\n",
    "        plt.imshow(piano_roll, aspect='auto', origin='lower', cmap='Blues',\n",
    "                   interpolation='nearest')\n",
    "        plt.xlabel('Time (s)')\n",
    "        plt.ylabel('Note ')\n",
    "        plt.title(f'MIDI (FPS={fps}): {sample_name}')\n",
    "        plt.colorbar(label='Velocity')\n",
    "        plt.yticks(np.arange(0, 128, 12))\n",
    "        \n",
    "        plt.tight_layout()\n",
    "        try:\n",
    "            plt.savefig(output_image_path)\n",
    "            print(f\"图片已保存到 {output_image_path}\")\n",
    "        except Exception as e:\n",
    "            print(f\"保存图片时出错: {e}\")\n",
    "        finally:\n",
    "            plt.close()  # 关闭图表以释放内存\n",
    "        \n",
    "        # --- 音频保存部分 ---\n",
    "        sf2_path = \"/usr/share/sounds/sf2/FluidR3_GM.sf2\"\n",
    "        if not os.path.exists(sf2_path):\n",
    "            print(f\"SoundFont 文件 {sf2_path} 不存在！\")\n",
    "            continue\n",
    "        \n",
    "        try:\n",
    "            audio = midi_data.fluidsynth(fs=sample_rate, sf2_path=sf2_path)\n",
    "            audio = audio / np.max(np.abs(audio)) if np.max(np.abs(audio)) != 0 else audio\n",
    "            scipy.io.wavfile.write(output_wav_path, sample_rate, audio)\n",
    "            print(f\"音频已保存到 {output_wav_path}，时长 {total_time:.2f} 秒\")\n",
    "        except Exception as e:\n",
    "            print(f\"保存音频时出错: {e}\")\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    # 输入参数\n",
    "    midi_folder = \"/home/zheqid/workspace/musictokenizer/results/maestro/continuation\"  # MIDI 文件夹路径\n",
    "    output_folder = \"trancodec\"  # 输出文件夹路径\n",
    "    base_name = \"sample\"  # 基础文件名\n",
    "    \n",
    "    # 处理 MIDI 文件夹\n",
    "    process_midi_folder(midi_folder, output_folder, base_name, fps=100, sample_rate=44100)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "piano_transcription",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}