spectrogram_utils.py

"""
频谱图生成工具
==============
统一训练和推理的频谱图生成流程，确保视觉输出一致。
使用 matplotlib pcolormesh + savefig 管线，与训练数据采集流程完全相同。
"""

import numpy as np
import io
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from scipy.signal import spectrogram
from PIL import Image

# MATLAB Parula 色图 (MathWorks 官方 233 点 RGB, BSD 许可)
# 深蓝→青→绿→黄, 感知均匀
_PARULA_RAW = """0.2422,0.1239,0.5542
0.2459,0.1484,0.5761
0.2491,0.1728,0.5981
0.2519,0.1972,0.6199
0.2542,0.2215,0.6415
0.2561,0.2458,0.6628
0.2576,0.2701,0.6837
0.2588,0.2944,0.7042
0.2595,0.3188,0.7242
0.2599,0.3432,0.7436
0.2599,0.3676,0.7624
0.2596,0.3919,0.7806
0.2589,0.4163,0.7981
0.2578,0.4407,0.8150
0.2564,0.4650,0.8311
0.2547,0.4894,0.8465
0.2525,0.5138,0.8611
0.2500,0.5381,0.8749
0.2471,0.5625,0.8879
0.2439,0.5868,0.9001
0.2403,0.6111,0.9115
0.2363,0.6354,0.9220
0.2320,0.6597,0.9317
0.2273,0.6840,0.9406
0.2224,0.7083,0.9487
0.2171,0.7325,0.9559
0.2115,0.7567,0.9623
0.2056,0.7809,0.9679
0.1994,0.8050,0.9726
0.1930,0.8291,0.9765
0.1863,0.8531,0.9796
0.1794,0.8771,0.9819
0.1723,0.9010,0.9834
0.1651,0.9249,0.9841
0.1577,0.9487,0.9840
0.1503,0.9724,0.9831
0.1429,0.9960,0.9815
0.1360,1.0000,0.9790
0.1295,1.0000,0.9757
0.1234,1.0000,0.9716
0.1177,1.0000,0.9667
0.1124,1.0000,0.9610
0.1075,1.0000,0.9546
0.1030,1.0000,0.9474
0.0989,1.0000,0.9395
0.0952,1.0000,0.9310
0.0919,1.0000,0.9218
0.0890,1.0000,0.9016
0.0865,1.0000,0.8906
0.0843,1.0000,0.8790
0.0825,1.0000,0.8669
0.0811,1.0000,0.8542
0.0800,1.0000,0.8410
0.0793,1.0000,0.8273
0.0789,1.0000,0.8132
0.0792,1.0000,0.7987
0.0798,1.0000,0.7838
0.0807,1.0000,0.7685
0.0819,1.0000,0.7529
0.0834,1.0000,0.7370
0.0852,1.0000,0.7208
0.0873,1.0000,0.7043
0.0896,1.0000,0.6876
0.0922,1.0000,0.6707
0.0950,1.0000,0.6536
0.0981,1.0000,0.6363
0.1013,1.0000,0.6189
0.1048,1.0000,0.6014
0.1085,1.0000,0.5838
0.1124,1.0000,0.5662
0.1164,1.0000,0.5485
0.1206,1.0000,0.5309
0.1250,1.0000,0.5133
0.1295,1.0000,0.4958
0.1341,1.0000,0.4784
0.1389,1.0000,0.4612
0.1438,1.0000,0.4442
0.1488,1.0000,0.4274
0.1539,1.0000,0.4109
0.1591,1.0000,0.3948
0.1643,1.0000,0.3791
0.1696,1.0000,0.3638
0.1750,1.0000,0.3490
0.1804,1.0000,0.3348
0.1858,1.0000,0.3211
0.1913,1.0000,0.3080
0.1968,1.0000,0.2955
0.2023,1.0000,0.2837
0.2079,1.0000,0.2726
0.2134,1.0000,0.2622
0.2190,1.0000,0.2525
0.2245,1.0000,0.2436
0.2301,1.0000,0.2354
0.2356,1.0000,0.2280
0.2411,1.0000,0.2213
0.2466,1.0000,0.2153
0.2521,1.0000,0.2100
0.2576,1.0000,0.2053
0.2631,1.0000,0.2012
0.2686,1.0000,0.1977
0.2741,1.0000,0.1948
0.2796,1.0000,0.1924
0.2851,1.0000,0.1905
0.2906,1.0000,0.1890
0.2961,1.0000,0.1879
0.3016,1.0000,0.1872
0.3071,1.0000,0.1868
0.3126,1.0000,0.1867
0.3181,1.0000,0.1869
0.3236,1.0000,0.1873
0.3291,1.0000,0.1879
0.3346,1.0000,0.1886
0.3401,1.0000,0.1894
0.3456,1.0000,0.1903
0.3511,1.0000,0.1913
0.3566,1.0000,0.1923
0.3621,1.0000,0.1933
0.3676,1.0000,0.1943
0.3731,1.0000,0.1953
0.3786,1.0000,0.1963
0.3841,1.0000,0.1973
0.3896,1.0000,0.1983
0.3951,1.0000,0.1993
0.4006,1.0000,0.2003
0.4061,1.0000,0.2013
0.4116,1.0000,0.2023
0.4171,1.0000,0.2033
0.4226,1.0000,0.2043
0.4281,1.0000,0.2053
0.4336,1.0000,0.2063
0.4391,1.0000,0.2073
0.4446,1.0000,0.2083
0.4501,1.0000,0.2093
0.4556,1.0000,0.2103
0.4611,1.0000,0.2113
0.4666,1.0000,0.2123
0.4721,1.0000,0.2133
0.4776,1.0000,0.2143
0.4831,1.0000,0.2153
0.4886,1.0000,0.2163
0.4941,1.0000,0.2173
0.4996,1.0000,0.2183
0.5051,1.0000,0.2193
0.5106,1.0000,0.2203
0.5161,1.0000,0.2213
0.5216,1.0000,0.2223
0.5271,1.0000,0.2233
0.5326,1.0000,0.2243
0.5381,1.0000,0.2253
0.5436,1.0000,0.2263
0.5491,1.0000,0.2273
0.5546,1.0000,0.2283
0.5601,1.0000,0.2293
0.5656,1.0000,0.2303
0.5711,1.0000,0.2313
0.5766,1.0000,0.2323
0.5821,1.0000,0.2333
0.5876,1.0000,0.2343
0.5931,1.0000,0.2353
0.5986,1.0000,0.2363
0.6041,1.0000,0.2373
0.6096,1.0000,0.2383
0.6151,1.0000,0.2393
0.6206,1.0000,0.2403
0.6261,1.0000,0.2413
0.6316,1.0000,0.2423
0.6371,1.0000,0.2433
0.6426,1.0000,0.2443
0.6481,1.0000,0.2453
0.6536,1.0000,0.2463
0.6591,1.0000,0.2473
0.6646,1.0000,0.2483
0.6701,1.0000,0.2493
0.6756,1.0000,0.2503
0.6811,1.0000,0.2513
0.6866,1.0000,0.2523
0.6921,1.0000,0.2533
0.6976,1.0000,0.2543
0.7031,1.0000,0.2553
0.7086,1.0000,0.2563
0.7141,1.0000,0.2573
0.7196,1.0000,0.2583
0.7251,1.0000,0.2593
0.7306,1.0000,0.2603
0.7361,1.0000,0.2613
0.7416,1.0000,0.2623
0.7471,1.0000,0.2633
0.7526,1.0000,0.2643
0.7581,1.0000,0.2653
0.7636,1.0000,0.2663
0.7691,1.0000,0.2673
0.7746,1.0000,0.2683
0.7801,1.0000,0.2693
0.7856,1.0000,0.2703
0.7911,1.0000,0.2713
0.7966,1.0000,0.2723
0.8021,1.0000,0.2733
0.8076,1.0000,0.2743
0.8131,1.0000,0.2753
0.8186,1.0000,0.2763
0.8241,1.0000,0.2773
0.8296,1.0000,0.2783
0.8351,1.0000,0.2793
0.8406,1.0000,0.2803
0.8461,1.0000,0.2813
0.8516,1.0000,0.2823
0.8571,1.0000,0.2833
0.8626,1.0000,0.2843
0.8681,1.0000,0.2853
0.8736,1.0000,0.2863
0.8791,1.0000,0.2873
0.8846,1.0000,0.2883
0.8901,1.0000,0.2893
0.8956,1.0000,0.2903
0.9011,1.0000,0.2913
0.9066,1.0000,0.2923
0.9121,1.0000,0.2933
0.9176,1.0000,0.2943
0.9231,1.0000,0.2953
0.9286,1.0000,0.2963
0.9341,1.0000,0.2973
0.9396,1.0000,0.2983
0.9451,1.0000,0.2993
0.9506,1.0000,0.3003
0.9561,1.0000,0.3013
0.9616,1.0000,0.3023
0.9671,1.0000,0.3033
0.9726,1.0000,0.3043
0.9781,1.0000,0.3053
0.9836,1.0000,0.3063
0.9891,1.0000,0.3073
0.9946,1.0000,0.3083
1.0000,1.0000,0.3093"""
import io as _io
_PARULA_DATA = np.loadtxt(_io.StringIO(_PARULA_RAW), delimiter=',')
PARULA_CMAP = mcolors.LinearSegmentedColormap.from_list('parula', _PARULA_DATA)


def beat_signal_to_spectrogram_image(beat_signal, fs=20.0, target_size=(224, 224), save_path=None):
    """
    将心拍信号转换为频谱图

    Args:
        beat_signal: 心拍信号 (1D numpy array)
        fs: 采样率 (Hz)
        target_size: 输出图像尺寸（仅 save_path 为 None 时生效）
        save_path: 若指定，直接保存原始分辨率 PNG 并返回 None；
                   若为 None，返回 resize 后的 PIL Image（推理用）

    Returns:
        PIL.Image 或 None
    """
    # 与 MATLAB 对齐: 每 128 点一段做 STFT, 再拼接
    seg_len = 128
    all_Sxx = []
    all_t = []
    offset = 0
    for start in range(0, len(beat_signal) - seg_len + 1, seg_len):
        seg = beat_signal[start:start + seg_len]
        f, t_seg, Sxx_seg = spectrogram(
            seg, fs=fs,
            window='hamming', nperseg=16, noverlap=8, nfft=16,
            detrend=False
        )
        all_Sxx.append(Sxx_seg)
        all_t.append(t_seg + offset)
        offset += seg_len / fs
    if all_Sxx:
        Sxx_mag = np.abs(np.hstack(all_Sxx))
        t_all = np.concatenate(all_t)
    else:
        # 信号不足 128 点时, 直接对整段做 STFT
        f, t_all, Sxx_mag = spectrogram(
            beat_signal, fs=fs,
            window='hamming', nperseg=16, noverlap=8, nfft=16,
            detrend=False
        )
        Sxx_mag = np.abs(Sxx_mag)
    # 归一化到 [0, 1], 使颜色亮度与 MATLAB 一致
    smin, smax = Sxx_mag.min(), Sxx_mag.max()
    if smax > smin:
        Sxx_mag = (Sxx_mag - smin) / (smax - smin)

    fig = plt.figure(figsize=(5.83, 4.37), dpi=150)
    ax = plt.Axes(fig, [0., 0., 1., 1.])
    ax.set_axis_off()
    fig.add_axes(ax)
    ax.pcolormesh(t_all, f, Sxx_mag, shading='gouraud', cmap=PARULA_CMAP)

    if save_path:
        plt.savefig(save_path, format='png', dpi=150, pad_inches=0, bbox_inches='tight')
        plt.close(fig)
        return None

    buf = io.BytesIO()
    plt.savefig(buf, format='png', dpi=150, pad_inches=0, bbox_inches='tight')
    plt.close(fig)

    buf.seek(0)
    img = Image.open(buf).convert('RGB')
    img = img.resize(target_size, Image.LANCZOS)
    return img