Produces record file from WAV format

Introduces the conversion of WAV files to MIT format. Fortunately, doesn't use any outside libraries (even Pythons wave package) since it is simply read using struct.unpack. Further work plans to generate the .dat files. I think the .hea files can already be generated though some information may not appear correct due to the loss of some information while being stored in WAV format, such as only one adc_res and adc_gain value instead of a list provided in EDF format, for example. For this reason, writing test cases for this function may be difficult though I might explore it in the future if time permits and demand commands it.

Author: Lucas-Mc

wfdb/__init__.py

@@ -1,5 +1,5 @@
 from wfdb.io.record import (Record, MultiRecord, rdheader, rdrecord, rdsamp,
-                            wrsamp, dl_database, edf2mit, sampfreq, signame)
+                            wrsamp, dl_database, edf2mit, wav2mit, sampfreq, signame)
 from wfdb.io.annotation import (Annotation, rdann, wrann, show_ann_labels,
                                 show_ann_classes, ann2rr)
 from wfdb.io.download import get_dbs, get_record_list, dl_files, set_db_index_url

wfdb/io/__init__.py

@@ -1,5 +1,5 @@
 from wfdb.io.record import (Record, MultiRecord, rdheader, rdrecord, rdsamp, wrsamp,
-                            dl_database, edf2mit, sampfreq, signame, SIGNAL_CLASSES)
+                            dl_database, edf2mit, wav2mit, sampfreq, signame, SIGNAL_CLASSES)
 from wfdb.io._signal import est_res, wr_dat_file
 from wfdb.io.annotation import (Annotation, rdann, wrann, show_ann_labels,
                                 show_ann_classes, ann2rr)

wfdb/io/record.py

@@ -10,6 +10,7 @@
 import mne
 import math
 import functools
+import struct
 import pdb
 
 from wfdb.io import _header
@@ -1497,6 +1498,168 @@ def edf2mit(record_name, pn_dir=None, delete_file=True, record_only=False):
         pass
 
 
+def wav2mit(record_name, pn_dir=None, delete_file=True, record_only=False):
+    """
+    Convert WAV formatted files to MIT format. See here for more details about
+    the formatting of a WAV file: http://soundfile.sapp.org/doc/WaveFormat/.
+
+    Parameters
+    ----------
+    record_name : str
+        The name of the input WAV record to be read.
+    pn_dir : str, optional
+        Option used to stream data from Physionet. The Physionet
+        database directory from which to find the required record files.
+        eg. For record '100' in 'http://physionet.org/content/mitdb'
+        pn_dir='mitdb'.
+    delete_file : bool, optional
+        Whether to delete the saved WAV file (False) or not (True)
+        after being imported.
+    record_only : bool, optional
+        Whether to only return the record information (True) or not (False).
+        If false, this function will generate both a .dat and .hea file.
+
+    Returns
+    -------
+    record : dict, optional
+        All of the record information needed to generate MIT formatted files.
+        Only returns if 'record_only' is set to True, else generates the
+        corresponding .dat and .hea files. This record file will not match the
+        `rdrecord` output since it will only give us the digital signal for now.
+
+    Examples
+    --------
+    >>> wav_record = wfdb.wav2mit('SC4001E0_PSG.wav', record_only=True)
+
+    """
+    if not record_name.endswith('.wav'):
+        raise Exception('Name of the input file must end in .wav')
+
+    if pn_dir is not None:
+
+        if '.' not in pn_dir:
+            dir_list = pn_dir.split(os.sep)
+            pn_dir = posixpath.join(dir_list[0], get_version(dir_list[0]), *dir_list[1:])
+
+        file_url = posixpath.join(download.PN_INDEX_URL, pn_dir, record_name)
+        # Currently must download file to read it though can give the
+        # user the option to delete it immediately afterwards
+        r = requests.get(file_url, allow_redirects=False)
+        open(record_name, 'wb').write(r.content)
+
+    wave_file = open(record_name, mode='rb')
+    record_name_out = record_name.split(os.sep)[-1].replace('-','_').replace('.wav','')
+
+    chunk_ID = ''.join([s.decode() for s in struct.unpack('>4s', wave_file.read(4))])
+    if chunk_ID != 'RIFF':
+        raise Exception('{} is not a .wav-format file'.format(record_name))
+
+    correct_chunk_size = os.path.getsize(record_name) - 8
+    chunk_size = struct.unpack('<I', wave_file.read(4))[0]
+    if chunk_size != correct_chunk_size:
+        raise Exception('Header chunk has incorrect length (is {} should be {})'.format(chunk_size,correct_chunk_size))
+
+    fmt = struct.unpack('>4s', wave_file.read(4))[0].decode()
+    if fmt != 'WAVE':
+        raise Exception('{} is not a .wav-format file'.format(record_name))
+
+    subchunk1_ID = struct.unpack('>4s', wave_file.read(4))[0].decode()
+    if subchunk1_ID != 'fmt ':
+        raise Exception('Format chunk missing or corrupt')
+
+    subchunk1_size = struct.unpack('<I', wave_file.read(4))[0]
+    audio_format = struct.unpack('<H', wave_file.read(2))[0]
+    if audio_format > 1:
+        print('PCM has compression of {}'.format(audio_format))
+
+    if (subchunk1_size != 16) or (audio_format != 1):
+	    raise Exception('Unsupported format {}'.format(audio_format))
+
+    num_channels = struct.unpack('<H', wave_file.read(2))[0]
+    if num_channels == 1:
+        print('Reading Mono formatted .wav file...')
+    elif num_channels == 2:
+        print('Reading Stereo formatted .wav file...')
+    else:
+        print('Reading {}-channel formatted .wav file...'.format(num_channels))
+
+    sample_rate = struct.unpack('<I', wave_file.read(4))[0]
+    print('Sample rate: {}'.format(sample_rate))
+    byte_rate = struct.unpack('<I', wave_file.read(4))[0]
+    print('Byte rate: {}'.format(byte_rate))
+    block_align = struct.unpack('<H', wave_file.read(2))[0]
+    print('Block align: {}'.format(block_align))
+    bits_per_sample = struct.unpack('<H', wave_file.read(2))[0]
+    print('Bits per sample: {}'.format(bits_per_sample))
+    # I wish this were more precise but unfortunately some information
+    # is lost in .wav files which is needed for these calculations
+    if bits_per_sample <= 8:
+        adc_res = 8
+        adc_gain = 12.5
+    elif bits_per_sample <= 16:
+        adc_res = 16
+        adc_gain = 6400
+    else:
+        raise Exception('Unsupported resolution ({} bits/sample)'.format(bits_per_sample))
+
+    if block_align != (num_channels * int(adc_res / 8)):
+	    raise Exception('Format chunk of {} has incorrect frame length'.format(block_align))
+
+    subchunk2_ID = struct.unpack('>4s', wave_file.read(4))[0].decode()
+    if subchunk2_ID != 'data':
+        raise Exception('Format chunk missing or corrupt')
+
+    correct_subchunk2_size = os.path.getsize(record_name) - 44
+    subchunk2_size = struct.unpack('<I', wave_file.read(4))[0]
+    if subchunk2_size != correct_subchunk2_size:
+        raise Exception('Data chunk has incorrect length.. (is {} should be {})'.format(subchunk2_size, correct_subchunk2_size))
+    sig_len = int(subchunk2_size / block_align)
+
+    sig_data = (np.fromfile(wave_file, dtype=np.int16).reshape((-1,num_channels)) / (2*adc_res)).astype(np.int16)
+
+    init_value = [int(s[0]) for s in np.transpose(sig_data)]
+    checksum = [int(np.sum(v) % 65536) for v in np.transpose(sig_data)]  # not all values correct?
+
+    if pn_dir is not None and delete_file:
+        os.remove(record_name)
+
+    record = Record(
+        record_name = record_name_out,
+        n_sig = num_channels,
+        fs = num_channels * [sample_rate],
+        samps_per_frame = num_channels * [1],
+        counter_freq = None,
+        base_counter = None,
+        sig_len = sig_len,
+        base_time = None,
+        base_date = None,
+        comments = [],
+        sig_name = num_channels * [None],
+        p_signal = None,
+        d_signal = sig_data,
+        e_p_signal = None,
+        e_d_signal = None,
+        file_name = num_channels * [record_name_out + '.dat'],
+        fmt = num_channels * ['16' if (adc_res == 16) else '80'],
+        skew = num_channels * [None],
+        byte_offset = num_channels * [None],
+        adc_gain = num_channels * [adc_gain],
+        baseline = num_channels * [0 if (adc_res == 16) else 128],
+        units = num_channels * [None],
+        adc_res =  num_channels * [adc_res],
+        adc_zero = num_channels * [0 if (adc_res == 16) else 128],
+        init_value = init_value,
+        checksum = checksum,
+        block_size = num_channels * [0]
+    )
+
+    if record_only:
+        return record
+    else:
+        # TODO: Generate the .dat and .hea files
+        pass
+
+
 #------------------------- Reading Records --------------------------- #
 
 
@@ -1626,6 +1789,8 @@ def rdrecord(record_name, sampfrom=0, sampto=None, channels=None,
         parameter is set, this parameter should contain just the base
         record name, and the files fill be searched for remotely.
         Otherwise, the data files will be searched for in the local path.
+        Can also handle .edf and .wav files as long as the extension is
+        provided in the `record_name`.
     sampfrom : int, optional
         The starting sample number to read for all channels.
     sampto : int, 'end', optional
@@ -1713,6 +1878,8 @@ def rdrecord(record_name, sampfrom=0, sampto=None, channels=None,
 
     if record_name.endswith('.edf'):
         record = edf2mit(record_name, pn_dir=pn_dir, record_only=True)
+    elif record_name.endswith('.wav'):
+        record = wav2mit(record_name, pn_dir=pn_dir, record_only=True)
     else:
         record = rdheader(record_name, pn_dir=pn_dir, rd_segments=False)
 
@@ -1785,7 +1952,7 @@ def rdrecord(record_name, sampfrom=0, sampto=None, channels=None,
         if smooth_frames or max([record.samps_per_frame[c] for c in channels]) == 1:
             # Read signals from the associated dat files that contain
             # wanted channels
-            if record_name.endswith('.edf'):
+            if record_name.endswith('.edf') or record_name.endswith('.wav'):
                 record.d_signal = _signal._rd_segment(record.file_name,
                                                       dir_name, pn_dir,
                                                       record.fmt,
@@ -1825,7 +1992,7 @@ def rdrecord(record_name, sampfrom=0, sampto=None, channels=None,
 
         # Return each sample of the signals with multiple samples per frame
         else:
-            if record_name.endswith('.edf'):
+            if record_name.endswith('.edf') or record_name.endswith('.wav'):
                 record.e_d_signal = _signal._rd_segment(record.file_name,
                                                       dir_name, pn_dir,
                                                       record.fmt,