diff --git a/tools/waveio.py b/tools/waveio.py new file mode 100755 index 00000000..62b2db53 --- /dev/null +++ b/tools/waveio.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: BSD-3-Clause +# Copyright(c) 2021 Intel Corporation. All rights reserved. + +""" +The WaveIO library is used to read and write standard PCM wave file. +""" + +import struct +import wave +import numpy as np + +def _normalize(data, bits): + signed_max = _get_signed_max_by_bits(bits) + # normalize data with negative max here to avoid the + # normalized negative max < -1.0 + return data / (signed_max + 1) + +def _get_signed_max_by_bits(bits): + return 2 ** (bits - 1) - 1 + +def _quantize(data, bits): + np_dtype = np.int16 if bits == 16 else np.int32 + # The np.astype() behaves like fix() in Matlab, it rounds each element + # in a ndarray to the nearest integer towards zero, eg: + # -2.8 -> -2, -2.1 -> -2, 2.1 -> 2, 2.8 -> 2, + # so add +0.5 to the positive part and -0.5 to the negative part to make + # the arithmetic bahaves like round() in Matlab to ensure correct quantization. + data = data * _get_signed_max_by_bits(bits) + 0.5 * np.sign(data) + return data.astype(np_dtype) + +# reshape linear data to matrix with shape = frames * channels +def _reshape_audio_data(data, channels): + if channels == 1: + return data + frames = data.shape[0] // channels + reshaped_data = np.zeros((frames, channels)) + for i in range(channels): + reshaped_data[:, i] = data[i::channels] + return reshaped_data + +# Wave header: http://soundfile.sapp.org/doc/WaveFormat/ +def _read_wave_file(filepath): + with open(filepath, 'rb') as wfs: + raw_data = wfs.read() + + (audio_format,) = struct.unpack('H', raw_data[20:22]) + if audio_format != 1: + raise Exception("Unknown wave file format, only standard PCM is supported") + + (channels,) = struct.unpack('H', raw_data[22:24]) + (sample_rate,) = struct.unpack('I', raw_data[24:28]) + (block_align,) = struct.unpack('H', raw_data[32:34]) + (sample_bits,) = struct.unpack('H', raw_data[34:36]) + (data_size,) = struct.unpack('I', raw_data[40:44]) + fmt_dict = {(16, 2):'S16_LE', (24, 3):'S24_3LE', (24, 4):'S24_LE', (32, 4):'S32_LE'} + wave_format = fmt_dict[(sample_bits, block_align//channels)] + # raw data starts from 44th byte + wave_data = raw_data[44: 44 + data_size] + return wave_data, channels, sample_rate, wave_format, sample_bits + +def read_wave(filepath='tmp.wav'): + """ + Read wave data from standard PCM wave file to numpy ndarray + + Parameters + ---------- + filepath: Path of standard PCM wave file + + Returns + ---------- + wave_data: tuple contains two elements, normalized data and sample rate + """ + + wave_data_bin, channels, sample_rate, fmt, sample_bits = _read_wave_file(filepath) + + # S24_3LE, append 0x00 to least significant byte + if fmt == 'S24_3LE': + frames = len(wave_data_bin) // 3 + raw_data_arr = [b'\00' + wave_data_bin[3 * i: 3*i + 3] for i in range(frames)] + raw_data_s32 = b''.join(raw_data_arr) + data = np.frombuffer(raw_data_s32, dtype=np.int32) / 256.0 + # S16_LE, S24_LE and S32_LE + else: + np_dtype = np.int16 if sample_bits == 16 else np.int32 + data = np.frombuffer(wave_data_bin, dtype=np_dtype) + + reshaped_data = _reshape_audio_data(data, channels) + normalized_data = _normalize(reshaped_data, sample_bits) + return normalized_data, sample_rate + +def write_wave(filepath, wave_data, sample_rate, fmt): + """ + Write wave data to file + + Parameters + ---------- + filepath: Path of wave file to be written to + wave_data: The data to write to file + sample_rate: Audio sample rate + fmt: Wave file format, supported formats are: S16_LE, S24_3LE, S24_LE and S32_LE + """ + + channels = 1 if len(wave_data.shape) == 1 else wave_data.shape[1] + frames = wave_data.shape[0] + inv_fmt_dict = {'S16_LE':(16, 2), 'S24_3LE': (24, 3), 'S24_LE':(24, 4), 'S32_LE':(32, 4)} + sample_bits, block_align = inv_fmt_dict[fmt] + wave_data = _quantize(wave_data, sample_bits) + wave_data_bin = wave_data.tobytes() + + if fmt == 'S24_3LE': # remove least significant byte + wave_data_bin_arr = [wave_data_bin[4 * i: 4 * i + 3] for i in range(frames * 2)] + wave_data_bin = b''.join(wave_data_bin_arr) + + with wave.open(filepath, 'wb') as wfs: + wave_params = (channels, block_align, sample_rate, frames, 'NONE', 'NONE') + # pylint identifies wfs as a Wave_read object, but it is a Wave_write object + # pylint: disable=E1101 + wfs.setparams(wave_params) + wfs.writeframes(wave_data_bin) + + if fmt == 'S24_LE': # for S24_LE, we need to update sample_bits from 32 to 24 + with open(filepath, 'rb+') as wfs: + sample_bits_bin = np.array([sample_bits]).astype(np.int16).tobytes() + raw_wave = bytearray(wfs.read()) + raw_wave[34:36] = sample_bits_bin + wfs.seek(0, 0) # seek to the beginning of the file + wfs.write(raw_wave) + +if __name__ == '__main__': + import argparse + + def main(): + """ + Entrance of the light-weight format converter + """ + + command_line = parse_command_line() + wave_data, sample_rate = read_wave(command_line.filename) + write_wave(command_line.output, wave_data, sample_rate, command_line.convert_to) + + def parse_command_line(): + """ + This function is used to parse command line parameters + """ + + supported_format = ['S16_LE', 'S24_3LE', 'S24_LE', 'S32_LE'] + parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter, + description='A Light-Weight Wave Format Converter') + argparse.ArgumentParser() + parser.add_argument('filename', type=str, help='Path of input wave file') + parser.add_argument('-t', '--convert_to', type=str, choices=supported_format, + help="Target wave file format") + parser.add_argument('-o', '--output', type=str, help="Path of output wave file") + return parser.parse_args() + + main()