sound_plt_comp-re.py

#!/usr/bin/env python
import sys
import numpy as np
import matplotlib.pyplot as plt
import wave
import scipy.signal as signal
import math

LIMIT_DB = 20
MAX_DB = 85
DOWN_RATE = 2
SMOOTHING_WINDOW = 1000
UNIT_OF_TIME = "sec" # sec or min


def downsampling(conversion_rate, data, fr):
    decimation_sample_num = conversion_rate - 1
    nyqF = (fr / conversion_rate) / 2.0
    cF = (fr / conversion_rate / 2.0 - 500.) / nyqF
    taps = 511
    b = signal.firwin(taps, cF)
    data = signal.lfilter(b, 1, data)
    down_data = data[::decimation_sample_num + 1]
    return down_data, int(fr / conversion_rate)


def smoothing(input_data, window_size):
    output = np.zeros_like(input_data)
    for i in range(len(input_data)):
        start = max(0, i - window_size)
        end = min(len(input_data), i + window_size + 1)
        output[i] = np.mean(input_data[start:end])
    return output


def calculate_rms(data, N):
    pad = np.zeros(N // 2)
    pad_data = np.concatenate([pad, data, pad])
    rms = np.array([np.sqrt((1 / N) * (np.sum(pad_data[i:i + N])) ** 2)
                    for i in range(len(data))])
    with np.errstate(divide='ignore'):
        db = 20 * np.log10(rms)
    return db


def moving_average_convolve(data, size):
    b = np.ones(size) / size
    data_mean = np.convolve(data, b, mode="same")
    n_conv = math.ceil(size / 2)
    data_mean[0] *= size / n_conv
    for i in range(1, n_conv):
        data_mean[i] *= size / (i + n_conv)
        data_mean[-i] *= size / (i + n_conv - (size % 2))
    return data_mean


def wave_to_db(wav_fname):
    wave_file = wave.open(wav_fname, "rb")
    fr = wave_file.getframerate()
    nframes = wave_file.getnframes()
    data = np.frombuffer(wave_file.readframes(nframes), dtype="int16")
    wave_file.close()

    down_fr = int(fr / (DOWN_RATE * 1000))
    down_data, down_fr = downsampling(down_fr, data, fr)

    N = int(fr / 42)
    db = calculate_rms(down_data, N)
    time = np.arange(0, db.shape[0] / down_fr, 1 / down_fr) / DOWN_RATE

    sm_db = smoothing(db, SMOOTHING_WINDOW)
    return sm_db, time, fr


def db_stats(db, time):
    db_t = [i for i in db if i >= LIMIT_DB]
    db_mean = np.mean(db_t)
    db_max = np.max(db_t)
    db_max_time = time[np.argmax(db)]
    return db_mean, db_max, db_max_time

def moving_average(db, time, fr):
    max_iterations = 2
    for _ in range(max_iterations):
        db = moving_average_convolve(db, fr)
        time = np.linspace(np.min(time), np.max(time), np.size(db))
    return db, time

def plot_waveforms(wav_fname1, wav_fname2, title):
    sm_db1, time1, fr1 = wave_to_db(wav_fname1)
    db_mean1, db_max1, db_max_time1 = db_stats(sm_db1, time1)
    sm_db_a1, time_a1 = moving_average(sm_db1, time1, fr1)

    sm_db2, time2, fr2 = wave_to_db(wav_fname2)
    db_mean2, db_max2, db_max_time2 = db_stats(sm_db2, time2)
    sm_db_a2, time_a2 = moving_average(sm_db2, time2, fr2)

    if UNIT_OF_TIME == 'min':
        time1 = time1 / 60; db_max_time1 = db_max_time1 / 60
        time_a1 = time_a1 / 60
        time2 = time2 / 60; db_max_time2 = db_max_time2 / 60
        time_a2 = time_a2 / 60
    plt.rcParams["font.family"] = "Noto Sans CJK JP"        

    fig = plt.figure(figsize=(16, 8), dpi=100, facecolor='tan', tight_layout=True)
    ax = fig.add_subplot(111, fc='w', xlabel=UNIT_OF_TIME, ylabel='音量(dB)')
    ax.set_title(title + " 音量推移")
    ax.set_ylim(LIMIT_DB, MAX_DB)
    ax.grid()

    ax.plot(time1, sm_db1, 'c', label='No.1 signal')
    ax.plot(time_a1, sm_db_a1, 'b', label='No.1 moving average')
    ax.axhline(y=db_mean1, color='b', linestyle='dashed', linewidth=1)
    ax.text(-6, db_mean1 - 2, "No.1 average:" + str(round(db_mean1, 1)) + \
            "dB", size=10)
    ax.axvline(x=db_max_time1, color='b', linestyle='dashed', linewidth=1)
    ax.text(db_max_time1, db_max1 + 2, "No.1 max:" + str(round(db_max1, 1)) + \
            'dB(' + str(round(db_max_time1, 1)) + UNIT_OF_TIME + ')', size=10)

    ax.plot(time2, sm_db2, 'm', label='No.2 Signal')
    ax.plot(time_a2, sm_db_a2, 'r', label='No.2 moving average')
    ax.axhline(y=db_mean2, color='r', linestyle='dashed', linewidth=1)
    ax.text(-6, db_mean2 + 1, "No.2 average:" + str(round(db_mean2, 1)) + \
            "dB", size=10)
    ax.axvline(x=db_max_time2, color='r', linestyle='dashed', linewidth=1)
    ax.text(db_max_time2, db_max2 + 1, "No.2 max:" + str(round(db_max2, 1)) + \
            'dB(' + str(round(db_max_time2, 1)) + UNIT_OF_TIME + ')', size=10)

    plt.legend(bbox_to_anchor=(1, 1), loc='upper right', borderaxespad=0)
    fig.savefig(title + '-ret.png', facecolor=fig.get_facecolor())

    plt.show()

def main(argv):
    if len(argv) != 4:
        print("Usage: python script.py <wave_file1> <wave_file2> <title>")
        return

    wav_fname1 = argv[1]
    wav_fname2 = argv[2]
    title = argv[3]

    plot_waveforms(wav_fname1, wav_fname2, title)


if __name__ == '__main__':
    main(sys.argv)