#!/usr/bin/env python3
import numpy as np
import json
import os
import sys
import scipy.io as sio
import wfdb
"""
Written by: Xingyao Wang, Chengyu Liu
School of Instrument Science and Engineering
Southeast University, China
chengyu@seu.edu.cn
"""
R = np.array([[1, -1, -.5], [-2, 1, 0], [-1, 0, 1]])
class RefInfo():
def __init__(self, sample_path):
self.sample_path = sample_path
self.fs, self.len_sig, self.beat_loc, self.af_starts, self.af_ends, self.class_true = self._load_ref()
self.endpoints_true = np.dstack((self.af_starts, self.af_ends))[0, :, :]
# self.endpoints_true = np.concatenate((self.af_starts, self.af_ends), axis=-1)
if self.class_true == 1 or self.class_true == 2:
self.onset_score_range, self.offset_score_range = self._gen_endpoint_score_range()
else:
self.onset_score_range, self.offset_score_range = None, None
def _load_ref(self):
sig, fields = wfdb.rdsamp(self.sample_path)
ann_ref = wfdb.rdann(self.sample_path, 'atr')
fs = fields['fs']
length = len(sig)
sample_descrip = fields['comments']
beat_loc = np.array(ann_ref.sample) # r-peak locations
ann_note = np.array(ann_ref.aux_note) # rhythm change flag
af_start_scripts = np.where((ann_note=='(AFIB') | (ann_note=='(AFL'))[0]
af_end_scripts = np.where(ann_note=='(N')[0]
if 'non atrial fibrillation' in sample_descrip:
class_true = 0
elif 'persistent atrial fibrillation' in sample_descrip:
class_true = 1
elif 'paroxysmal atrial fibrillation' in sample_descrip:
class_true = 2
else:
print('Error: the recording is out of range!')
return -1
return fs, length, beat_loc, af_start_scripts, af_end_scripts, class_true
def _gen_endpoint_score_range(self):
"""
"""
onset_range = np.zeros((self.len_sig, ),dtype=np.float)
offset_range = np.zeros((self.len_sig, ),dtype=np.float)
for i, af_start in enumerate(self.af_starts):
if self.class_true == 2:
if max(af_start-1, 0) == 0:
onset_range[: self.beat_loc[af_start+2]] += 1
elif max(af_start-2, 0) == 0:
onset_range[self.beat_loc[af_start-1]: self.beat_loc[af_start+2]] += 1
onset_range[: self.beat_loc[af_start-1]] += .5
else:
onset_range[self.beat_loc[af_start-1]: self.beat_loc[af_start+2]] += 1
onset_range[self.beat_loc[af_start-2]: self.beat_loc[af_start-1]] += .5
onset_range[self.beat_loc[af_start+2]: self.beat_loc[af_start+3]] += .5
elif self.class_true == 1:
onset_range[: self.beat_loc[af_start+2]] += 1
onset_range[self.beat_loc[af_start+2]: self.beat_loc[af_start+3]] += .5
for i, af_end in enumerate(self.af_ends):
if self.class_true == 2:
if min(af_end+1, len(self.beat_loc)-1) == len(self.beat_loc)-1:
offset_range[self.beat_loc[af_end-2]: ] += 1
elif min(af_end+2, len(self.beat_loc)-1) == len(self.beat_loc)-1:
offset_range[self.beat_loc[af_end-2]: self.beat_loc[af_end+1]] += 1
offset_range[self.beat_loc[af_end+1]: ] += 0.5
else:
offset_range[self.beat_loc[af_end-2]: self.beat_loc[af_end+1]] += 1
offset_range[self.beat_loc[af_end+1]: min(self.beat_loc[af_end+2], self.len_sig-1)] += .5
offset_range[self.beat_loc[af_end-3]: self.beat_loc[af_end-2]] += .5
elif self.class_true == 1:
offset_range[self.beat_loc[af_end-2]: ] += 1
offset_range[self.beat_loc[af_end-3]: self.beat_loc[af_end-2]] += .5
return onset_range, offset_range
def load_ans(ans_file):
endpoints_pred = []
if ans_file.endswith('.json'):
json_file = open(ans_file, "r")
ans_dic = json.load(json_file)
endpoints_pred = np.array(ans_dic['predict_endpoints'])
elif ans_file.endswith('.mat'):
ans_struct = sio.loadmat(ans_file)
endpoints_pred = ans_struct['predict_endpoints']-1
return endpoints_pred
def ue_calculate(endpoints_pred, endpoints_true, onset_score_range, offset_score_range):
score = 0
ma = len(endpoints_true)
mr = len(endpoints_pred)
for [start, end] in endpoints_pred:
score += onset_score_range[int(start)]
score += offset_score_range[int(end)]
score *= (ma / max(ma, mr))
return score
def ur_calculate(class_true, class_pred):
score = R[int(class_true), int(class_pred)]
return score
def score(data_path, ans_path):
# AF burden estimation
SCORE = []
def is_mat_or_json(file):
return (file.endswith('.json')) + (file.endswith('.mat'))
ans_set = filter(is_mat_or_json, os.listdir(ans_path))
# test_set = open(os.path.join(data_path, 'RECORDS'), 'r').read().splitlines()
for i, ans_sample in enumerate(ans_set):
sample_nam = ans_sample.split('.')[0]
sample_path = os.path.join(data_path, sample_nam)
endpoints_pred = load_ans(os.path.join(ans_path, ans_sample))
TrueRef = RefInfo(sample_path)
if len(endpoints_pred) == 0:
class_pred = 0
elif len(endpoints_pred) == 1 and np.diff(endpoints_pred)[-1] == TrueRef.len_sig - 1:
class_pred = 1
else:
class_pred = 2
ur_score = ur_calculate(TrueRef.class_true, class_pred)
if TrueRef.class_true == 1 or TrueRef.class_true == 2:
ue_score = ue_calculate(endpoints_pred, TrueRef.endpoints_true, TrueRef.onset_score_range, TrueRef.offset_score_range)
else:
ue_score = 0
u = ur_score + ue_score
SCORE.append(u)
score_avg = np.mean(SCORE)
return score_avg
if __name__ == '__main__':
TESTSET_PATH = sys.argv[1]
RESULT_PATH = sys.argv[2]
score_avg = score(TESTSET_PATH, RESULT_PATH)
print('AF Endpoints Detection Performance: %0.4f' %score_avg)
with open(os.path.join(RESULT_PATH, 'score.txt'), 'w') as score_file:
print('AF Endpoints Detection Performance: %0.4f' %score_avg, file=score_file)
score_file.close()