Skip to content

Commit f1c0355

Browse files
Lucas-McLucas-Mc
committed
Adds mrgann function
1 parent a2defa6 commit f1c0355

File tree

3 files changed

+257
-10
lines changed

3 files changed

+257
-10
lines changed

‎wfdb/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,8 @@
22
wrsamp, dl_database, edf2mit, mit2edf, wav2mit, mit2wav,
33
wfdb2mat, csv2mit, sampfreq, signame, wfdbdesc, wfdbtime)
44
from wfdb.io.annotation import (Annotation, rdann, wrann, show_ann_labels,
5-
show_ann_classes, ann2rr, rr2ann, csv2ann, rdedfann)
5+
show_ann_classes, ann2rr, rr2ann, csv2ann,
6+
rdedfann, mrgann)
67
from wfdb.io.download import get_dbs, get_record_list, dl_files, set_db_index_url
78
from wfdb.plot.plot import plot_items, plot_wfdb, plot_all_records
89

‎wfdb/io/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
csv2mit, sampfreq, signame, wfdbdesc, wfdbtime, SIGNAL_CLASSES)
44
from wfdb.io._signal import est_res, wr_dat_file
55
from wfdb.io.annotation import (Annotation, rdann, wrann, show_ann_labels,
6-
show_ann_classes, ann2rr, rr2ann, csv2ann, rdedfann)
6+
show_ann_classes, ann2rr, rr2ann, csv2ann,
7+
rdedfann, mrgann)
78
from wfdb.io.download import get_dbs, get_record_list, dl_files, set_db_index_url
89
from wfdb.io.tff import rdtff

‎wfdb/io/annotation.py‎

Lines changed: 253 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
import posixpath
77
import pdb
88
import struct
9+
import sys
910

1011
from wfdb.io import download
1112
from wfdb.io import _header
@@ -2570,20 +2571,22 @@ def rdedfann(record_name, pn_dir=None, delete_file=True, info_only=True,
25702571
by the original WFDB package. Must not be True if `record_only` is
25712572
True.
25722573
record_only : bool, optional
2573-
Whether to only return the record information (True) or not (False).
2574-
If False, this function will generate both a .dat and .hea file. Must
2575-
not be True if `info_only` is True.
2574+
Whether to only return the annotation information (True) or not
2575+
(False). If False, this function will generate a WFDB-formatted
2576+
annotation file. If True, it will return the object returned if that
2577+
file were read with `rdann`. Must not be True if `info_only` is True.
25762578
verbose : bool, optional
25772579
Whether to print all the information read about the file (True) or
25782580
not (False).
25792581
25802582
Returns
25812583
-------
2582-
record : dict, optional
2583-
All of the record information needed to generate MIT formatted files.
2584-
Only returns if 'record_only' is set to True, else generates the
2585-
corresponding .dat and .hea files. This record file will not match the
2586-
`rdrecord` output since it will only give us the digital signal for now.
2584+
N/A : dict, Annotation, optional
2585+
If 'info_only' is set to True, return all of the annotation
2586+
information needed to generate WFDB-formatted annotation files.
2587+
If 'record_only' is set to True, return the WFDB-formatted annotation
2588+
object generated by the `rdann` output. If none are set to True, write
2589+
the WFDB-formatted annotation file.
25872590
25882591
Notes
25892592
-----
@@ -2707,6 +2710,248 @@ def rdedfann(record_name, pn_dir=None, delete_file=True, info_only=True,
27072710
fs=fs)
27082711

27092712

2713+
def mrgann(ann_file1, ann_file2, out_file_name='merged_ann.atr', start_ann=0,
2714+
end_ann='e', merge_method='combine', record_only=True, verbose=False):
2715+
"""
2716+
This function reads a pair of annotation files (specified by `ann_file1`
2717+
and `ann_file2`) for the specified record and writes a third annotation
2718+
file (specified by `out_file_name`) for the same record. The header (.hea)
2719+
file should be included in the same directory as each annotation file so
2720+
that the sampling rate can be read. Typical applications of `mrgann`
2721+
include combining annotation files that apply to different signals within
2722+
a multi-signal record, and replacing a segment of an annotation file with
2723+
annotations from another file. For example, setting 'merge_method' to
2724+
'combine' will simply blindly merge the annotation files for the specified
2725+
'start_ann' and 'end_ann' range while setting 'merge_method' to 'replace1'
2726+
will replace the contents of the first file with the second in that
2727+
specified range. Setting 'merge_method' to 'replace2' will replace the
2728+
contents of the second file with the first in that specified range.
2729+
2730+
Parameters
2731+
----------
2732+
ann_file1 : string
2733+
The file path of the first annotation file (with extension included).
2734+
ann_file2 : string
2735+
The file path of the second annotation file (with extension included).
2736+
out_file_name : string
2737+
The name of the output file name (with extension included). The
2738+
default is 'merged_ann.atr'.
2739+
start_ann : float, int, string, optional
2740+
The location (sample, time, etc.) to start the annotation filtering.
2741+
If float, it will be interpreted as time in seconds. If int, it will
2742+
be interpreted as sample number. If string, it will be interpreted
2743+
as time formatted in HH:MM:SS format (the same as that in `wfdbtime`).
2744+
The default is 0 to represent sample number 0. A value of 0.0 would
2745+
represent 0 seconds instead.
2746+
end_ann : float, int, string, optional
2747+
The location (sample, time, etc.) to stop the annotation filtering.
2748+
If float, it will be interpreted as time in seconds. If int, it will
2749+
be interpreted as sample number. If string, it will be interpreted
2750+
as time formatted in HH:MM:SS format (the same as that in `wfdbtime`).
2751+
The default is 'e' to represent the end of the annotation.
2752+
merge_method : string, optional
2753+
The method used to merge the two annotation files. The default is
2754+
'combine' which simply combines the two files along every attribute;
2755+
duplicates will be preserved. The other options are 'replace1' which
2756+
replaces attributes of the first annotation file with attributes of
2757+
the second for the desired time range and 'replace2' which does the
2758+
same thing except switched (first file replaces second).
2759+
record_only : bool, optional
2760+
Whether to only return the annotation information (True) or not
2761+
(False). If False, this function will generate a WFDB-formatted
2762+
annotation file. If True, it will return the object returned if that
2763+
file was read with `rdann`.
2764+
verbose : bool, optional
2765+
Whether to print all the information read about each annotation file
2766+
and the methodology for merging them (True) or not (False).
2767+
2768+
Returns
2769+
-------
2770+
N/A : Annotation, optional
2771+
If 'record_only' is set to True, then return the new WFDB-formatted
2772+
annotation object which is the same as generated by the `rdann`
2773+
output. Else, create the WFDB-formatted annotation file.
2774+
2775+
"""
2776+
ann1 = rdann(ann_file1.split('.')[0], ann_file1.split('.')[1])
2777+
ann2 = rdann(ann_file2.split('.')[0], ann_file2.split('.')[1])
2778+
if ann1.fs != ann2.fs:
2779+
raise Exception('Annotation sample rates do not match up: samples '
2780+
'can be aligned but final sample rate can not be '
2781+
'determined')
2782+
2783+
if start_ann == 'e':
2784+
raise Exception('Start time can not be set to the end of the record')
2785+
if end_ann == 0:
2786+
raise Exception('End time can not be set to the start of the record')
2787+
2788+
samples = []
2789+
for i,time in enumerate([start_ann, end_ann]):
2790+
if time == 'e':
2791+
# End of annotation, set end sample to largest int, roughly
2792+
sample = sys.maxsize
2793+
else:
2794+
if type(time) is int:
2795+
# Sample number
2796+
sample = time
2797+
elif type(time) is float:
2798+
# Time in seconds
2799+
sample = int(time * ann1.fs)
2800+
else:
2801+
# HH:MM:SS format, loosely
2802+
time_split = [t if t != '' else '0' for t in time.split(':')]
2803+
if len(time_split) == 1:
2804+
seconds = float(time)%60
2805+
minutes = int(float(time)//60)
2806+
hours = int(float(time)//60//60)
2807+
elif len(time_split) == 2:
2808+
seconds = float(time_split[1])
2809+
minutes = int(time_split[0])
2810+
hours = 0
2811+
elif len(time_split) == 3:
2812+
seconds = float(time_split[2])
2813+
minutes = int(time_split[1])
2814+
hours = int(time_split[0])
2815+
if seconds >= 60:
2816+
raise Exception('Seconds not in correct format')
2817+
if minutes >= 60:
2818+
raise Exception('Minutes not in correct format')
2819+
total_seconds = hours*60*60 + minutes*60 + seconds
2820+
if (i == 1) and (total_seconds == 0):
2821+
raise Exception('End time can not be set to the start of '
2822+
'the record')
2823+
sample = int(total_seconds * ann1.fs)
2824+
if sample > max([max(ann1.sample), max(ann2.sample)]):
2825+
if i == 0:
2826+
raise Exception('Start time can not be set to the '
2827+
'end of the record')
2828+
else:
2829+
print("'end_ann' greater than the highest "
2830+
"annotation... reverting to the highest "
2831+
"annotation")
2832+
samples.append(sample)
2833+
start_sample = samples[0]
2834+
end_sample = samples[1]
2835+
if verbose:
2836+
print(f'Start sample: {start_sample}, end sample: {end_sample}')
2837+
2838+
if merge_method == 'combine':
2839+
if verbose:
2840+
print('Combining the two files together')
2841+
# The sample should never be empty but others can (though they
2842+
# shouldn't be)
2843+
both_sample = np.concatenate([ann1.sample, ann2.sample]).astype(np.int64)
2844+
# Generate a list of sorted indices then sort the array
2845+
sort_indices = np.argsort(both_sample)
2846+
both_sample = np.sort(both_sample)
2847+
# Find where to filter the array
2848+
sample_range = ((both_sample >= start_sample) &
2849+
(both_sample <= end_sample))
2850+
index_range = np.where(sample_range)[0]
2851+
both_sample = both_sample[sample_range]
2852+
# Combine both annotation attributes
2853+
ann_attr = {}
2854+
blank_array = np.array([], dtype=np.int64)
2855+
for cat in ['chan', 'num', 'subtype', 'label_store', 'symbol',
2856+
'aux_note']:
2857+
ann1_cat = ann1.__dict__[cat]
2858+
ann2_cat = ann2.__dict__[cat]
2859+
if cat in ['symbol', 'aux_note']:
2860+
ann1_cat = ann1_cat if ann1_cat is not None else []
2861+
ann2_cat = ann2_cat if ann2_cat is not None else []
2862+
temp_cat = ann1_cat
2863+
temp_cat.extend(ann2_cat)
2864+
if len(temp_cat) == 0:
2865+
ann_attr[cat] = None
2866+
else:
2867+
temp_cat = [temp_cat[i] for i in sort_indices]
2868+
ann_attr[cat] = [temp_cat[i] for i in index_range]
2869+
else:
2870+
ann1_cat = ann1_cat if ann1_cat is not None else blank_array
2871+
ann2_cat = ann2_cat if ann2_cat is not None else blank_array
2872+
temp_cat = np.concatenate([ann1_cat, ann2_cat]).astype(np.int64)
2873+
if temp_cat.shape[0] == 0:
2874+
ann_attr[cat] = None
2875+
else:
2876+
temp_cat = np.array([temp_cat[i] for i in sort_indices])
2877+
ann_attr[cat] = np.array([temp_cat[i] for i in index_range])
2878+
2879+
elif (merge_method == 'replace1') or (merge_method == 'replace2'):
2880+
if merge_method == 'replace1':
2881+
if verbose:
2882+
print('Replacing the contents of the first file with the '
2883+
'contents of the second')
2884+
keep_ann = ann2
2885+
remove_ann = ann1
2886+
elif merge_method == 'replace2':
2887+
if verbose:
2888+
print('Replacing the contents of the second file with the '
2889+
'contents of the first')
2890+
keep_ann = ann1
2891+
remove_ann = ann2
2892+
# Find where to filter the first array
2893+
keep_sample_range = ((keep_ann.sample >= start_sample) &
2894+
(keep_ann.sample <= end_sample))
2895+
keep_index_range = np.where(keep_sample_range)[0]
2896+
# Find where to filter the second array
2897+
remove_sample_range = ((remove_ann.sample < start_sample) |
2898+
(remove_ann.sample > end_sample))
2899+
remove_index_range = np.where(remove_sample_range)[0]
2900+
# The sample should never be empty but others can (though they
2901+
# shouldn't be)
2902+
keep_ann_sample = keep_ann.sample[keep_index_range]
2903+
remove_ann_sample = remove_ann.sample[remove_index_range]
2904+
both_sample = np.concatenate([keep_ann_sample, remove_ann_sample]).astype(np.int64)
2905+
# Generate a list of sorted indices then sort the array
2906+
sort_indices = np.argsort(both_sample)
2907+
both_sample = np.sort(both_sample)
2908+
# Combine both annotation attributes
2909+
ann_attr = {}
2910+
blank_array = np.array([], dtype=np.int64)
2911+
for cat in ['chan', 'num', 'subtype', 'label_store', 'symbol',
2912+
'aux_note']:
2913+
keep_cat = keep_ann.__dict__[cat]
2914+
remove_cat = remove_ann.__dict__[cat]
2915+
if cat in ['symbol', 'aux_note']:
2916+
keep_cat = [keep_cat[i] for i in keep_index_range] if keep_cat is not None else []
2917+
remove_cat = [remove_cat[i] for i in remove_index_range] if remove_cat is not None else []
2918+
temp_cat = keep_cat
2919+
temp_cat.extend(remove_cat)
2920+
if len(temp_cat) == 0:
2921+
ann_attr[cat] = None
2922+
else:
2923+
ann_attr[cat] = [temp_cat[i] for i in sort_indices]
2924+
else:
2925+
keep_cat = np.array([keep_cat[i] for i in keep_index_range]) if keep_cat is not None else blank_array
2926+
remove_cat = np.array([remove_cat[i] for i in remove_index_range]) if remove_cat is not None else blank_array
2927+
temp_cat = np.concatenate([keep_cat, remove_cat]).astype(np.int64)
2928+
if temp_cat.shape[0] == 0:
2929+
ann_attr[cat] = None
2930+
else:
2931+
ann_attr[cat] = np.array([temp_cat[i] for i in sort_indices])
2932+
else:
2933+
raise Exception("Invalid value for 'merge_method': options are "
2934+
"'combine', 'replace1', and 'replace2'")
2935+
2936+
if record_only:
2937+
if verbose:
2938+
print('Returning Annotation object')
2939+
return Annotation(record_name=out_file_name.split('.')[0],
2940+
extension=out_file_name.split('.')[1],
2941+
sample=both_sample, symbol=ann_attr['symbol'],
2942+
subtype=ann_attr['subtype'], chan=ann_attr['chan'],
2943+
num=ann_attr['num'], aux_note=ann_attr['aux_note'],
2944+
label_store=ann_attr['label_store'], fs=ann1.fs)
2945+
else:
2946+
if verbose:
2947+
print(f'Creating annotation file called: {out_file_name}')
2948+
wrann(out_file_name.split('.')[0], out_file_name.split('.')[1],
2949+
sample=both_sample, symbol=ann_attr['symbol'],
2950+
subtype=ann_attr['subtype'], chan=ann_attr['chan'],
2951+
num=ann_attr['num'], aux_note=ann_attr['aux_note'],
2952+
label_store=ann_attr['label_store'], fs=ann1.fs)
2953+
2954+
27102955
def _format_ann_from_df(df_in):
27112956
"""
27122957
Parameters

0 commit comments

Comments
 (0)