Upload 54 files

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+csv-file_2021_challenge/collection_of_all_datasets.csv filter=lfs diff=lfs merge=lfs -text
+csv-file_2021_challenge/training_validation_testing/group1/training_group1.csv filter=lfs diff=lfs merge=lfs -text
+csv-file_2021_challenge/training_validation_testing/group2/training_group2.csv filter=lfs diff=lfs merge=lfs -text
+csv-file_2021_challenge/training_validation_testing/group3/training_group3.csv filter=lfs diff=lfs merge=lfs -text
+csv-file_2021_challenge/training_validation_testing/group4/training_group4.csv filter=lfs diff=lfs merge=lfs -text
+csv-file_2021_challenge/training_validation_testing/group5/training_group5.csv filter=lfs diff=lfs merge=lfs -text

csv-file_2021_challenge/collection_of_all_datasets.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:455156a539d4dde56aa9c2d1e5e460a72df64bda32b8ba36a6692837520292ac
+size 12635911

csv-file_2021_challenge/name.csv

@@ -0,0 +1,11 @@
+Name,Target,Label_code,Frequency,Duration,Age,Sex
+JS10647.hea,"[0. 0. 0. 0. 0. 1. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0.
+ 1. 0.]","['426177001', '713427006', '164934002', '39732003']",500.0,10.0,82.0,Male
+JS10648.hea,"[0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0.
+ 0. 0.]","['426177001', '713426002']",500.0,10.0,63.0,Male
+JS10649.hea,"[0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
+ 0. 0.]","['111975006', '713426002', '426761007', '55930002']",500.0,10.0,83.0,Male
+JS10650.hea,"[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.
+ 0. 0.]",['29320008'],500.0,10.0,81.0,Female
+JS10651.hea,"[0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0.
+ 0. 0.]","['233897008', '10370003']",500.0,10.0,63.0,Female

csv-file_2021_challenge/training_validation_testing/group1/training_group1.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd7406d2f748a0a5edd033b389fd967fc9a232d2421ed9be69adcccf964db7a1
+size 10522513

csv-file_2021_challenge/training_validation_testing/group2/training_group2.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db819a51f741bbb0bc88878f4adbec1a16065eaf1cb991758dab8a42aef34fce
+size 10523820

csv-file_2021_challenge/training_validation_testing/group3/training_group3.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:03f7959deaa6cbe865b911c103a1fae71210cce4f5ac99c6920dc2efb4beb18a
+size 10522844

csv-file_2021_challenge/training_validation_testing/group4/training_group4.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:125d50af59c4b1862aefcf389b5a7eda5819bdc2e66ce52f43097f9f60661b25
+size 10525392

csv-file_2021_challenge/training_validation_testing/group5/training_group5.csv

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c6b92de27c915999cf59125328b7f2d73a9b33ef8939f04d942b17d1cb7630a
+size 10522780

ecg_dataset_2020.py

@@ -0,0 +1,397 @@
+import numpy as np
+from tqdm import tqdm
+from helper_code import *
+from scipy import signal
+import pandas as pd
+from skmultilearn.model_selection import iterative_train_test_split
+import random
+import warnings
+import sys
+
+
+def get_nsamp(header):
+    return int(header.split('\n')[0].split(' ')[3])
+
+# Adapted from original scoring function code
+# For each set of equivalent classes, replace each class with the representative class for the set.
+
+
+'''
+the 'output' will be using the zero-padding if the leads are less than 12.(8 leads)
+the 'output_leads' will be [1. 1. 0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
+
+the output will not be changed if the lead is 12
+the 'output_leads' will be [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
+'''
+def mixup(data, mix_data):
+    mix_lambda = np.random.beta(10, 10)# 这里利用beta分布得到一个0到1的值,这里如果可以，尝试别的分布函数
+    if data.shape[1] > mix_data.shape[1]:
+        data = mix_lambda * data[:, :mix_data.shape[1]] + (1 - mix_lambda) * mix_data
+    else:
+        data = mix_lambda * data + (1 - mix_lambda) * mix_data[:, :data.shape[1]]
+    return data, mix_lambda
+
+def cutmix(data, mix_data):
+    cutmix_lambda = np.random.beta(10, 10)
+    
+    if data.shape[1] < mix_data.shape[1]:
+        seq_length = data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    else:
+        seq_length = mix_data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    return data, cutmix_lambda
+
+def cutmix_fix_length(data, mix_data):
+    # cutmix_lambda = np.random.beta(0.2, 0.2)
+    cutmix_lambda = 0.2
+
+    if data.shape[1] < mix_data.shape[1]:
+        seq_length = data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    else:
+        seq_length = mix_data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    return data, cutmix_lambda
+
+def same_shape_mixup(data, mix_data):
+    if data.shape[1] == mix_data.shape[1]:
+        mix_lambda = np.random.beta(10, 10)# 这里利用beta分布得到一个0到1的值,这里如果可以，尝试别的分布函数
+        data = mix_lambda * data + (1 - mix_lambda) * mix_data
+    return data
+
+def expand_leads(recording, input_leads):
+    output = np.zeros((12, recording.shape[1]))
+    # recording.shape[1]: 5000
+    twelve_leads = ('I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6')
+    twelve_leads = [k.lower() for k in twelve_leads]
+    # ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', 'v2', 'v3', 'v4', 'v5', 'v6']
+
+    input_leads = [k.lower() for k in input_leads]
+    # Here we can assume: 
+    # input_leads:I, II, V1, V2, V3, V4, V5, V6, 
+    # so the new input_leads: ['i', 'ii', 'v1', 'v2', 'v3', 'v4', 'v5', 'v6']
+
+    output_leads = np.zeros((12,))
+    # [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
+
+    # idx: [0, 1, 6, 7, 8, 9, 10, 11]
+    for i,k in enumerate(input_leads):
+        idx = twelve_leads.index(k)
+        output[idx,:] = recording[i,:]
+        output_leads[idx] = 1
+
+    return output, output_leads
+
+'''
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 0. 0. 0. 0. 0. 0.]
+
+'''
+
+class lead_exctractor:
+    """
+    used to select specific leads or random choice of configurations
+
+    Twelve leads: I, II, III, aVR, aVL, aVF, V1, V2, V3, V4, V5, V6
+    Eight leads: I, II, V1, V2, V3, V4, V5, V6
+    Six leads: I, II, III, aVR, aVL, aVF
+    Four leads: I, II, III, V2
+    Three leads: I, II, V2
+    Two leads: I, II
+
+    """
+    L2  = np.array([1,1,0,0,0,0,0,0,0,0,0,0])
+    L3  = np.array([1,1,0,0,0,0,0,1,0,0,0,0])
+    L4  = np.array([1,1,1,0,0,0,0,1,0,0,0,0])
+    L6  = np.array([1,1,1,1,1,1,0,0,0,0,0,0])
+    L8  = np.array([1,1,0,0,0,0,1,1,1,1,1,1])
+    L12 = np.array([1,1,1,1,1,1,1,1,1,1,1,1])
+
+    @staticmethod
+    def get (x, num_leads, lead_indicator):
+        if num_leads==None:
+            # random choice output
+            num_leads = random.choice([12,8,6,4,3,2])
+
+        if num_leads==12:
+            # Twelve leads: I, II, III, aVR, aVL, aVF, V1, V2, V3, V4, V5, V6
+            return x, lead_indicator * lead_exctractor.L12
+
+        if num_leads==8:
+            # Six leads:  I, II, V1, V2, V3, V4, V5, V6
+            x = x * lead_exctractor.L8.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L8
+
+        if num_leads==6:
+            # Six leads: I, II, III, aVL, aVR, aVF
+            x = x * lead_exctractor.L6.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L6
+
+        if num_leads==4:
+            # Six leads: I, II, III, V2
+            x = x * lead_exctractor.L4.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L4
+
+        if num_leads==3:
+            # Three leads: I, II, V2
+            x = x * lead_exctractor.L3.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L3
+
+        if num_leads==2:
+            # Two leads: II, V5
+            x = x * lead_exctractor.L2.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L2
+        raise Exception("invalid-leads-number")
+
+class dataset:
+    # #ECG 2021
+    # classes = ['164889003','164890007','6374002','426627000','733534002',
+    #            '713427006','270492004','713426002','39732003','445118002',
+    #            '164947007','251146004','111975006','698252002','426783006',
+    #            '284470004','10370003','365413008','427172004','164917005',
+    #            '47665007','427393009','426177001','427084000','164934002',
+    #            '59931005']
+    
+    # ECG 2020
+
+    
+
+    normal_class = '426783006'
+    
+    def __init__(self, header_files, Mixup = 0, amount = 0, cutMix = 0, Mixup_no_label_interpolate =0, progressive_switch = False):
+        self.files = []
+        self.sample = True
+        self.num_leads = None
+
+        for h in tqdm(header_files):
+            tmp = dict()
+            tmp['header'] = h
+            tmp['record'] = h.replace('.hea','.mat')
+            hdr = load_header(h)
+            tmp['nsamp'] = get_nsamp(hdr)
+            tmp['leads'] = get_leads(hdr)
+            tmp['age'] = get_age(hdr)
+            tmp['sex'] = get_sex(hdr)
+            tmp['dx'] = get_labels(hdr)
+            tmp['fs'] = get_frequency(hdr)
+            # tmp['target'] = np.zeros((26,))
+
+
+            self.files.append(tmp)
+
+        # print("This is the target:", tmp['target'])
+        # set filter parameters
+        self.b, self.a = signal.butter(3, [1 / 250, 47 / 250], 'bandpass')
+
+        self.files = pd.DataFrame(self.files)
+
+        self.Mixup = Mixup
+        self.cutMix = cutMix
+        self.Mixup_no_label_interplate = Mixup_no_label_interpolate
+        self.amount = amount
+        self.progressive = progressive_switch
+
+        self.current_epoch = 0  # Initialize the current epoch
+
+        address_2020 = "./csv-file_2020_challenge/training_validation_testing/group1/train_validation_testing_fold5.csv"
+        self.data_df_2020 = pd.read_csv(address_2020, index_col=0)
+        self.data_df_2020.set_index('Patient', inplace=True)
+        self.classes_2020 = ['270492004', '164889003', '164890007', '426627000', '713427006', 
+                  '713426002', '445118002', '39732003', '164909002', '251146004', 
+                  '698252002', '10370003', '284470004', '427172004', '164947007', 
+                  '111975006', '164917005', '47665007', '59118001', '427393009', 
+                  '426177001', '426783006', '427084000', '63593006', '164934002', 
+                  '59931005', '17338001']
+
+    
+    def set_epoch(self, epoch):
+        self.current_epoch = epoch
+    
+    def summary(self, output):
+        if output=='pandas':
+            # print("@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@")
+            # print("This is the target:",self.files['target'],type(self.files['target']))
+            return pd.Series(np.stack(self.files['target'].to_list(),axis=0).sum(axis=0),index=dataset.classes)
+
+        if output=='numpy':
+            return np.stack(self.files['target'].to_list(),axis=0).sum(axis=0)
+
+    def __len__(self):
+        return len(self.files)
+    
+    '''
+    fs: 500.0   sampling rate
+    target: [0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0.]
+    leads: ('I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6')
+    self.files.iloc[item]['record']: ../../python-classifier-2021-main/training_data/chapman_shaoxing/g1/JS00518.mat
+    data: shape->(12 X 5000)
+    '''
+
+    def __getitem__(self, item):
+
+        fs = self.files.iloc[item]['fs']
+        # target = self.files.iloc[item]['target']
+        header = self.files.iloc[item]['header']
+        name = header.split('/')[-1]
+        row = self.data_df_2020.loc[name[0:-4]]
+        target = row[self.classes_2020].values.astype(np.int_)
+
+        leads = self.files.iloc[item]['leads']
+        data = load_recording(self.files.iloc[item]['record'])
+        
+        # print("This is the target:", target)
+        # Set your threshold
+        
+        threshold = 0.5
+        # print("This is the original target:", target)
+        # print("This is the type of original target:", type(target))
+        
+        # if random.random() < self.Mixup :
+        #     mix_sample_idx = random.randint(0, self.amount-1)
+        #     mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
+        #     mix_target = self.files.iloc[mix_sample_idx]['target']
+        #     data, alpha_value = mixup(data, mix_datum)
+        #     target = alpha_value * target + (1-alpha_value) * mix_target
+        #     # Binarize the labels based on the threshold
+        #     target[target >= threshold] = 1
+        #     target[target < threshold] = 0
+            
+        
+        # if random.random() < self.cutMix:
+        #     mix_sample_idx = random.randint(0, self.amount-1)
+        #     mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
+        #     mix_target = self.files.iloc[mix_sample_idx]['target']
+        #     data, alpha_value = cutmix_fix_length(data, mix_datum)
+        #     target = alpha_value * target + (1-alpha_value) * mix_target
+        #     target[target >= threshold] = 1
+        #     target[target < threshold] = 0
+        
+        # if random.random()< self.Mixup_no_label_interplate:
+        #     mix_sample_idx = random.randint(0, self.amount-1)
+        #     mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
+        #     data= same_shape_mixup(data, mix_datum)
+
+        # proressive_index = 0
+        
+        # # # below is the progessive data augmentation method
+        # # if self.progressive: 
+        # #     if self.current_epoch < 20:
+        # #         if self.current_epoch % 4 == 0 or 3:
+        # #             pass
+        # #         elif self.current_epoch % 4 == 1:
+        # #             proressive_index = 0.2
+        # #         else:
+        # #             proressive_index = 0.5
+
+        # #     elif self.current_epoch < 40:
+        # #         if self.current_epoch % 4 == 0 or 3:
+        # #             pass
+        # #         elif self.current_epoch % 4 == 1:
+        # #             proressive_index = 0.7
+        # #         else:
+        # #             proressive_index = 0.8
+        # #     else:
+        # #         if self.current_epoch % 4 == 0 or 3:
+        # #             pass
+        # #         elif self.current_epoch % 4 == 1 or 2:
+        # #             proressive_index = 1
+                    
+        # # below is the wave-mix data augmentation method
+        # if self.progressive: 
+        #     if self.current_epoch < 5:
+        #         match self.current_epoch:
+        #             case 0:
+        #                 self.progressive=0
+        #             case 1:
+        #                 self.progressive=0.2
+        #             case 2:
+        #                 self.progressive=0.4
+        #             case 3:
+        #                 self.progressive=0.6
+        #             case 4:
+        #                 self.progressive=0.8
+                    
+        #     else :
+        #         proressive_index = 0.8
+            
+        # if random.random() < proressive_index :
+        #     mix_sample_idx = random.randint(0, self.amount-1)
+        #     mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
+        #     mix_target = self.files.iloc[mix_sample_idx]['target']
+        #     data, alpha_value = mixup(data, mix_datum)
+        #     target = alpha_value * target + (1-alpha_value) * mix_target
+        #     # Binarize the labels based on the threshold
+        #     target[target >= threshold] = 1
+        #     target[target < threshold] = 0
+
+
+        # expand to 12 lead setup if original signal has less channels
+        data, lead_indicator = expand_leads(data, input_leads=leads)
+        data = np.nan_to_num(data)
+ 
+
+        # resample to 500hz
+        if fs == float(1000):
+            data = signal.resample_poly(data, up=1, down=2, axis=-1)  # to 500Hz
+            fs = 500
+        elif fs == float(500):
+            pass
+        else:
+            data = signal.resample(data, int(data.shape[1] * 500 / fs), axis=1)
+            fs = 500
+
+        # below is the Butterworth digital and analog filter design.
+        data = signal.filtfilt(self.b, self.a, data)
+
+        '''
+        we filter out the sample which the length is 8192 via the below code.
+        just like the random shift windows.  
+        '''
+        if self.sample:
+            fs = int(fs)
+            # random sample signal if len > 8192 samples
+            if data.shape[-1] >= 8192:
+                idx = data.shape[-1] - 8192-1
+                idx = np.random.randint(idx)
+                data = data[:, idx:idx + 8192]
+
+        '''
+        we can obtain the average and mean for each lead via below code, 90% conclusion.
+        '''
+        mu  = np.nanmean(data, axis=-1, keepdims=True)
+        std = np.nanstd(data, axis=-1, keepdims=True)
+
+        '''
+        Z-Score Normalization
+        '''
+        #std = np.nanstd(data.flatten())
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            data = (data - mu) / std
+        
+        data = np.nan_to_num(data)
+
+        # random choose number of leads to keep
+        data, lead_indicator = lead_exctractor.get(data, self.num_leads, lead_indicator)
+        
+        # print("This is the lead_indicator:(l)", type(lead_indicator), lead_indicator)
+        
+
+        return data, target

ecg_dataset_2021.py

@@ -0,0 +1,422 @@
+import numpy as np
+from tqdm import tqdm
+from helper_code import *
+from scipy import signal
+import pandas as pd
+from skmultilearn.model_selection import iterative_train_test_split
+import random
+import warnings
+import sys
+
+
+def get_nsamp(header):
+    return int(header.split('\n')[0].split(' ')[3])
+
+# Adapted from original scoring function code
+# For each set of equivalent classes, replace each class with the representative class for the set.
+def replace_equivalent_classes(classes, equivalent_classes):
+    for j, x in enumerate(classes):
+        for multiple_classes in equivalent_classes:
+            if x in multiple_classes:
+                classes[j] = multiple_classes[0]  # Use the first class as the representative class.
+    return classes
+
+'''
+the 'output' will be using the zero-padding if the leads are less than 12.(8 leads)
+the 'output_leads' will be [1. 1. 0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
+
+the output will not be changed if the lead is 12
+the 'output_leads' will be [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
+'''
+def mixup(data, mix_data):
+    mix_lambda = np.random.beta(10, 10)# 这里利用beta分布得到一个0到1的值,这里如果可以，尝试别的分布函数
+    if data.shape[1] > mix_data.shape[1]:
+        data = mix_lambda * data[:, :mix_data.shape[1]] + (1 - mix_lambda) * mix_data
+    else:
+        data = mix_lambda * data + (1 - mix_lambda) * mix_data[:, :data.shape[1]]
+    return data, mix_lambda
+
+def cutmix(data, mix_data):
+    cutmix_lambda = np.random.beta(10, 10)
+    
+    if data.shape[1] < mix_data.shape[1]:
+        seq_length = data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    else:
+        seq_length = mix_data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    return data, cutmix_lambda
+
+def cutmix_fix_length(data, mix_data):
+    # cutmix_lambda = np.random.beta(0.2, 0.2)
+    cutmix_lambda = 0.2
+
+    if data.shape[1] < mix_data.shape[1]:
+        seq_length = data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    else:
+        seq_length = mix_data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    return data, cutmix_lambda
+
+def same_shape_mixup(data, mix_data):
+    if data.shape[1] == mix_data.shape[1]:
+        mix_lambda = np.random.beta(10, 10)# 这里利用beta分布得到一个0到1的值,这里如果可以，尝试别的分布函数
+        data = mix_lambda * data + (1 - mix_lambda) * mix_data
+    return data
+
+def expand_leads(recording, input_leads):
+    output = np.zeros((12, recording.shape[1]))
+    # recording.shape[1]: 5000
+    twelve_leads = ('I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6')
+    twelve_leads = [k.lower() for k in twelve_leads]
+    # ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', 'v2', 'v3', 'v4', 'v5', 'v6']
+
+    input_leads = [k.lower() for k in input_leads]
+    # Here we can assume: 
+    # input_leads:I, II, V1, V2, V3, V4, V5, V6, 
+    # so the new input_leads: ['i', 'ii', 'v1', 'v2', 'v3', 'v4', 'v5', 'v6']
+
+    output_leads = np.zeros((12,))
+    # [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
+
+    # idx: [0, 1, 6, 7, 8, 9, 10, 11]
+    for i,k in enumerate(input_leads):
+        idx = twelve_leads.index(k)
+        output[idx,:] = recording[i,:]
+        output_leads[idx] = 1
+
+    return output, output_leads
+
+'''
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 0. 0. 0. 0. 0. 0.]
+
+'''
+
+class lead_exctractor:
+    """
+    used to select specific leads or random choice of configurations
+
+    Twelve leads: I, II, III, aVR, aVL, aVF, V1, V2, V3, V4, V5, V6
+    Eight leads: I, II, V1, V2, V3, V4, V5, V6
+    Six leads: I, II, III, aVR, aVL, aVF
+    Four leads: I, II, III, V2
+    Three leads: I, II, V2
+    Two leads: I, II
+
+    """
+    L2  = np.array([1,1,0,0,0,0,0,0,0,0,0,0])
+    L3  = np.array([1,1,0,0,0,0,0,1,0,0,0,0])
+    L4  = np.array([1,1,1,0,0,0,0,1,0,0,0,0])
+    L6  = np.array([1,1,1,1,1,1,0,0,0,0,0,0])
+    L8  = np.array([1,1,0,0,0,0,1,1,1,1,1,1])
+    L12 = np.array([1,1,1,1,1,1,1,1,1,1,1,1])
+
+    @staticmethod
+    def get (x, num_leads, lead_indicator):
+        if num_leads==None:
+            # random choice output
+            num_leads = random.choice([12,8,6,4,3,2])
+
+        if num_leads==12:
+            # Twelve leads: I, II, III, aVR, aVL, aVF, V1, V2, V3, V4, V5, V6
+            return x, lead_indicator * lead_exctractor.L12
+
+        if num_leads==8:
+            # Six leads:  I, II, V1, V2, V3, V4, V5, V6
+            x = x * lead_exctractor.L8.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L8
+
+        if num_leads==6:
+            # Six leads: I, II, III, aVL, aVR, aVF
+            x = x * lead_exctractor.L6.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L6
+
+        if num_leads==4:
+            # Six leads: I, II, III, V2
+            x = x * lead_exctractor.L4.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L4
+
+        if num_leads==3:
+            # Three leads: I, II, V2
+            x = x * lead_exctractor.L3.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L3
+
+        if num_leads==2:
+            # Two leads: II, V5
+            x = x * lead_exctractor.L2.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L2
+        raise Exception("invalid-leads-number")
+
+class dataset:
+    classes = ['164889003','164890007','6374002','426627000','733534002',
+               '713427006','270492004','713426002','39732003','445118002',
+               '164947007','251146004','111975006','698252002','426783006',
+               '284470004','10370003','365413008','427172004','164917005',
+               '47665007','427393009','426177001','427084000','164934002',
+               '59931005']
+    normal_class = '426783006'
+    equivalent_classes = [['713427006', '59118001'],
+                          ['284470004', '63593006'],
+                          ['427172004', '17338001'],
+                          ['733534002', '164909002']]
+    def __init__(self, header_files, Mixup = 0, amount = 0, cutMix = 0, Mixup_no_label_interpolate =0, progressive_switch = False):
+        self.files = []
+        self.sample = True
+        self.num_leads = None
+
+        for h in tqdm(header_files):
+            tmp = dict()
+            tmp['header'] = h
+            tmp['record'] = h.replace('.hea','.mat')
+            hdr = load_header(h)
+            tmp['nsamp'] = get_nsamp(hdr)
+            tmp['leads'] = get_leads(hdr)
+            tmp['age'] = get_age(hdr)
+            tmp['sex'] = get_sex(hdr)
+            tmp['dx'] = get_labels(hdr)
+            tmp['fs'] = get_frequency(hdr)
+            tmp['target'] = np.zeros((26,))
+            tmp['dx'] = replace_equivalent_classes(tmp['dx'], dataset.equivalent_classes)
+
+            for dx in tmp['dx']:
+                # in SNOMED code is in scored classes
+                if dx in dataset.classes:
+                    idx = dataset.classes.index(dx)
+                    tmp['target'][idx] = 1
+
+            self.files.append(tmp)
+
+        # print("This is the target:", tmp['target'])
+        # set filter parameters
+
+        # Filtering: Data are filtered using a zero-phase method with 3rd order Butterworth bandpass filter 
+        # with frequency band from 1 Hz to 47 Hz.
+        self.b, self.a = signal.butter(3, [1 / 250, 47 / 250], 'bandpass')
+
+        self.files = pd.DataFrame(self.files)
+
+        self.Mixup = Mixup
+        self.cutMix = cutMix
+        self.Mixup_no_label_interplate = Mixup_no_label_interpolate
+        self.amount = amount
+        self.progressive = progressive_switch
+
+        self.current_epoch = 0  # Initialize the current epoch
+
+    def set_epoch(self, epoch):
+        self.current_epoch = epoch
+
+    def train_valid_split(self, test_size):
+        '''
+        test_size: 0.2
+        below is the value of each variable:
+        files[0] <- "/media/jiang/ECG/physionet.org/files/challenge-2021/1.0.3/training/python-classifier-2021-main/training_data/chapman_shaoxing/g1/JS00001.hea" shape -> (999, 1)
+        targets[1] <-[1. 0. 0. ... 0. 1. 0.]  26, shape->(999, 26)
+        x_train[0] <-"/media/jiang/ECG/physionet.org/files/challenge-2021/1.0.3/training/python-classifier-2021-main/training_data/chapman_shaoxing/g1/JS00001.hea" shape -> (799,1)
+        x_valid[0] <-"/media/jiang/ECG/physionet.org/files/challenge-2021/1.0.3/training/python-classifier-2021-main/training_data/chapman_shaoxing/g1/JS00651.hea" shape -> (200,1)
+        '''
+        files = self.files['header'].to_numpy().reshape(-1,1)
+        # print(files[0])
+        targets = np.stack(self.files['target'].to_list(), axis=0)
+        print("This is the targets:", targets)
+
+        x_train, y_train, x_valid, y_valid = iterative_train_test_split(files, targets, test_size=test_size)
+        
+        train = dataset(header_files=x_train[:,0].tolist())
+        
+        train.num_leads=None
+        train.sample=True
+
+        valid = dataset(header_files=x_valid[:,0].tolist())
+        
+        valid.num_leads=12
+        valid.sample=False
+        
+        return train, valid
+
+    def summary(self, output):
+        if output=='pandas':
+            # print("@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@")
+            # print("This is the target:",self.files['target'],type(self.files['target']))
+            return pd.Series(np.stack(self.files['target'].to_list(),axis=0).sum(axis=0),index=dataset.classes)
+
+        if output=='numpy':
+            return np.stack(self.files['target'].to_list(),axis=0).sum(axis=0)
+
+    def Pre_processing(self, fs, leads, data):
+
+        # 1. Expand to a 12-lead setup if the original signal has fewer channels 
+        data, lead_indicator = expand_leads(data, input_leads=leads)
+        data = np.nan_to_num(data)
+ 
+
+        # 2. Resample from 1000 Hz and other rates to 500 Hz
+        if fs == float(1000):
+            data = signal.resample_poly(data, up=1, down=2, axis=-1)  # to 500Hz
+            fs = 500
+        elif fs == float(500):
+            pass
+        else:
+            data = signal.resample(data, int(data.shape[1] * 500 / fs), axis=1)
+            fs = 500
+
+        # 3. Below is the design of the Butterworth digital and analog filter
+        data = signal.filtfilt(self.b, self.a, data)
+            
+        
+        # 4. We compute the average and mean for each lead using the code below; this is Z-score normalization 
+        mu  = np.nanmean(data, axis=-1, keepdims=True)
+        std = np.nanstd(data, axis=-1, keepdims=True)
+        
+        #std = np.nanstd(data.flatten())
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            data = (data - mu) / std
+        
+        data = np.nan_to_num(data)
+
+        # 5. Selection of leads for 12-lead ECGs; the default is 12 leads
+        data, lead_indicator = lead_exctractor.get(data, self.num_leads, lead_indicator)  
+
+
+        # 6. We filter out samples with a length of 8192 using the code below, similar to random shift windows, and zero-padding as well 
+        if self.sample:
+            fs = int(fs)
+            # random sample signal if len > 8192 samples
+            if data.shape[-1] >= 8192:
+                idx = data.shape[-1] - 8192-1
+                idx = np.random.randint(idx)
+                data = data[:, idx:idx + 8192]
+            else:
+                # Apply right zero-padding along the second dimension
+                padding_size = 8192 - data.shape[-1]
+                data = np.pad(data, ((0, 0), (0, padding_size)), mode='constant', constant_values=0) 
+
+        return data
+
+
+    def __len__(self):
+        return len(self.files)
+    
+    '''
+    fs: 500.0   sampling rate
+    target: [0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0.]
+    leads: ('I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6')
+    self.files.iloc[item]['record']: ../../python-classifier-2021-main/training_data/chapman_shaoxing/g1/JS00518.mat
+    data: shape->(12 X 5000)
+    '''
+
+    def __getitem__(self, item):
+
+        fs = self.files.iloc[item]['fs']
+        target = self.files.iloc[item]['target']
+        leads = self.files.iloc[item]['leads']
+        data = load_recording(self.files.iloc[item]['record'])
+        
+        # the phase of pre-processing:
+        data = self.Pre_processing(fs, leads, data)
+
+        '''
+        below are the data augmentation methods
+        '''
+        if random.random() < self.Mixup :
+            mix_sample_idx = random.randint(0, self.amount-1)
+            mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
+            mix_target = self.files.iloc[mix_sample_idx]['target']
+
+            mix_fs = self.files.iloc[mix_sample_idx]['fs']
+            mix_leads = self.files.iloc[mix_sample_idx]['leads']
+            mix_datum = self.Pre_processing(mix_fs, mix_leads, mix_datum)
+
+            data, alpha_value = mixup(data, mix_datum)
+            target = alpha_value * target + (1-alpha_value) * mix_target
+            
+        
+        if random.random() < self.cutMix:
+            mix_sample_idx = random.randint(0, self.amount-1)
+            mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
+            mix_target = self.files.iloc[mix_sample_idx]['target']
+
+            mix_fs = self.files.iloc[mix_sample_idx]['fs']
+            mix_leads = self.files.iloc[mix_sample_idx]['leads']
+            mix_datum = self.Pre_processing(mix_fs, mix_leads, mix_datum)
+
+            data, alpha_value = cutmix(data, mix_datum)
+            target = alpha_value * target + (1-alpha_value) * mix_target
+            
+
+        progressive_index = 0
+        
+        # # below is the progessive data augmentation method
+        # if self.progressive: 
+        #     if self.current_epoch < 20:
+        #         if self.current_epoch % 4 == 0 or 3:
+        #             pass
+        #         elif self.current_epoch % 4 == 1:
+        #             proressive_index = 0.2
+        #         else:
+        #             proressive_index = 0.5
+
+        #     elif self.current_epoch < 40:
+        #         if self.current_epoch % 4 == 0 or 3:
+        #             pass
+        #         elif self.current_epoch % 4 == 1:
+        #             proressive_index = 0.7
+        #         else:
+        #             proressive_index = 0.8
+        #     else:
+        #         if self.current_epoch % 4 == 0 or 3:
+        #             pass
+        #         elif self.current_epoch % 4 == 1 or 2:
+        #             proressive_index = 1
+                    
+        # below is the wave-mix data augmentation method
+        if self.progressive: 
+            if self.current_epoch < 5:
+                match self.current_epoch:
+                    case 0:
+                        progressive_index=0.8
+                    case 1:
+                        progressive_index=0.8
+                    case 2:
+                        progressive_index=0.8
+                    case 3:
+                        progressive_index=0.8
+                    case 4:
+                        progressive_index=0.8
+                    
+            else :
+                progressive_index = 0.8
+            
+        if random.random() < progressive_index :
+            mix_sample_idx = random.randint(0, self.amount-1)
+            mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
+            mix_target = self.files.iloc[mix_sample_idx]['target']
+
+            mix_fs = self.files.iloc[mix_sample_idx]['fs']
+            mix_leads = self.files.iloc[mix_sample_idx]['leads']
+            mix_datum = self.Pre_processing(mix_fs, mix_leads, mix_datum)
+
+            data, alpha_value = mixup(data, mix_datum)
+            target = alpha_value * target + (1-alpha_value) * mix_target
+             
+
+        return data, target

ecg_dataset_2021_DAFirst.py

@@ -0,0 +1,418 @@
+import numpy as np
+from tqdm import tqdm
+from helper_code import *
+from scipy import signal
+import pandas as pd
+from skmultilearn.model_selection import iterative_train_test_split
+import random
+import warnings
+import sys
+
+
+def get_nsamp(header):
+    return int(header.split('\n')[0].split(' ')[3])
+
+# Adapted from original scoring function code
+# For each set of equivalent classes, replace each class with the representative class for the set.
+def replace_equivalent_classes(classes, equivalent_classes):
+    for j, x in enumerate(classes):
+        for multiple_classes in equivalent_classes:
+            if x in multiple_classes:
+                classes[j] = multiple_classes[0]  # Use the first class as the representative class.
+    return classes
+
+'''
+the 'output' will be using the zero-padding if the leads are less than 12.(8 leads)
+the 'output_leads' will be [1. 1. 0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
+
+the output will not be changed if the lead is 12
+the 'output_leads' will be [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
+'''
+def mixup(data, mix_data):
+    mix_lambda = np.random.beta(10, 10)# 这里利用beta分布得到一个0到1的值,这里如果可以，尝试别的分布函数
+    if data.shape[1] > mix_data.shape[1]:
+        data = mix_lambda * data[:, :mix_data.shape[1]] + (1 - mix_lambda) * mix_data
+    else:
+        data = mix_lambda * data + (1 - mix_lambda) * mix_data[:, :data.shape[1]]
+    return data, mix_lambda
+
+def cutmix(data, mix_data):
+    cutmix_lambda = np.random.beta(10, 10)
+    
+    if data.shape[1] < mix_data.shape[1]:
+        seq_length = data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    else:
+        seq_length = mix_data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    return data, cutmix_lambda
+
+def cutmix_fix_length(data, mix_data):
+    # cutmix_lambda = np.random.beta(0.2, 0.2)
+    cutmix_lambda = 0.2
+
+    if data.shape[1] < mix_data.shape[1]:
+        seq_length = data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    else:
+        seq_length = mix_data.shape[1]
+        win_len = int(np.ceil(seq_length * cutmix_lambda))
+        start = np.random.randint(0, seq_length - win_len + 1)
+        end = start + win_len
+        data[:,start:end] = mix_data[:,start:end]
+    return data, cutmix_lambda
+
+def same_shape_mixup(data, mix_data):
+    if data.shape[1] == mix_data.shape[1]:
+        mix_lambda = np.random.beta(10, 10)# 这里利用beta分布得到一个0到1的值,这里如果可以，尝试别的分布函数
+        data = mix_lambda * data + (1 - mix_lambda) * mix_data
+    return data
+
+def expand_leads(recording, input_leads):
+    output = np.zeros((12, recording.shape[1]))
+    # recording.shape[1]: 5000
+    twelve_leads = ('I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6')
+    twelve_leads = [k.lower() for k in twelve_leads]
+    # ['i', 'ii', 'iii', 'avr', 'avl', 'avf', 'v1', 'v2', 'v3', 'v4', 'v5', 'v6']
+
+    input_leads = [k.lower() for k in input_leads]
+    # Here we can assume: 
+    # input_leads:I, II, V1, V2, V3, V4, V5, V6, 
+    # so the new input_leads: ['i', 'ii', 'v1', 'v2', 'v3', 'v4', 'v5', 'v6']
+
+    output_leads = np.zeros((12,))
+    # [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
+
+    # idx: [0, 1, 6, 7, 8, 9, 10, 11]
+    for i,k in enumerate(input_leads):
+        idx = twelve_leads.index(k)
+        output[idx,:] = recording[i,:]
+        output_leads[idx] = 1
+
+    return output, output_leads
+
+'''
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]
+This is the lead_indicator:(l) <class 'numpy.ndarray'> [1. 1. 1. 1. 1. 1. 0. 0. 0. 0. 0. 0.]
+
+'''
+
+class lead_exctractor:
+    """
+    used to select specific leads or random choice of configurations
+
+    Twelve leads: I, II, III, aVR, aVL, aVF, V1, V2, V3, V4, V5, V6
+    Eight leads: I, II, V1, V2, V3, V4, V5, V6
+    Six leads: I, II, III, aVR, aVL, aVF
+    Four leads: I, II, III, V2
+    Three leads: I, II, V2
+    Two leads: I, II
+
+    """
+    L2  = np.array([1,1,0,0,0,0,0,0,0,0,0,0])
+    L3  = np.array([1,1,0,0,0,0,0,1,0,0,0,0])
+    L4  = np.array([1,1,1,0,0,0,0,1,0,0,0,0])
+    L6  = np.array([1,1,1,1,1,1,0,0,0,0,0,0])
+    L8  = np.array([1,1,0,0,0,0,1,1,1,1,1,1])
+    L12 = np.array([1,1,1,1,1,1,1,1,1,1,1,1])
+
+    @staticmethod
+    def get (x, num_leads, lead_indicator):
+        if num_leads==None:
+            # random choice output
+            num_leads = random.choice([12,8,6,4,3,2])
+
+        if num_leads==12:
+            # Twelve leads: I, II, III, aVR, aVL, aVF, V1, V2, V3, V4, V5, V6
+            return x, lead_indicator * lead_exctractor.L12
+
+        if num_leads==8:
+            # Six leads:  I, II, V1, V2, V3, V4, V5, V6
+            x = x * lead_exctractor.L8.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L8
+
+        if num_leads==6:
+            # Six leads: I, II, III, aVL, aVR, aVF
+            x = x * lead_exctractor.L6.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L6
+
+        if num_leads==4:
+            # Six leads: I, II, III, V2
+            x = x * lead_exctractor.L4.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L4
+
+        if num_leads==3:
+            # Three leads: I, II, V2
+            x = x * lead_exctractor.L3.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L3
+
+        if num_leads==2:
+            # Two leads: II, V5
+            x = x * lead_exctractor.L2.reshape(12,1)
+            return x,lead_indicator * lead_exctractor.L2
+        raise Exception("invalid-leads-number")
+
+class dataset:
+    classes = ['164889003','164890007','6374002','426627000','733534002',
+               '713427006','270492004','713426002','39732003','445118002',
+               '164947007','251146004','111975006','698252002','426783006',
+               '284470004','10370003','365413008','427172004','164917005',
+               '47665007','427393009','426177001','427084000','164934002',
+               '59931005']
+    normal_class = '426783006'
+    equivalent_classes = [['713427006', '59118001'],
+                          ['284470004', '63593006'],
+                          ['427172004', '17338001'],
+                          ['733534002', '164909002']]
+    def __init__(self, header_files, Mixup = 0, amount = 0, cutMix = 0, Mixup_no_label_interpolate =0, progressive_switch = False):
+        self.files = []
+        self.sample = True
+        self.num_leads = None
+
+        for h in tqdm(header_files):
+            tmp = dict()
+            tmp['header'] = h
+            tmp['record'] = h.replace('.hea','.mat')
+            hdr = load_header(h)
+            tmp['nsamp'] = get_nsamp(hdr)
+            tmp['leads'] = get_leads(hdr)
+            tmp['age'] = get_age(hdr)
+            tmp['sex'] = get_sex(hdr)
+            tmp['dx'] = get_labels(hdr)
+            tmp['fs'] = get_frequency(hdr)
+            tmp['target'] = np.zeros((26,))
+            tmp['dx'] = replace_equivalent_classes(tmp['dx'], dataset.equivalent_classes)
+
+            for dx in tmp['dx']:
+                # in SNOMED code is in scored classes
+                if dx in dataset.classes:
+                    idx = dataset.classes.index(dx)
+                    tmp['target'][idx] = 1
+
+            self.files.append(tmp)
+
+        # print("This is the target:", tmp['target'])
+        # set filter parameters
+
+        # Filtering: Data are filtered using a zero-phase method with 3rd order Butterworth bandpass filter 
+        # with frequency band from 1 Hz to 47 Hz.
+        self.b, self.a = signal.butter(3, [1 / 250, 47 / 250], 'bandpass')
+
+        self.files = pd.DataFrame(self.files)
+
+        self.Mixup = Mixup
+        self.cutMix = cutMix
+        self.Mixup_no_label_interplate = Mixup_no_label_interpolate
+        self.amount = amount
+        self.progressive = progressive_switch
+
+        self.current_epoch = 0  # Initialize the current epoch
+
+    def set_epoch(self, epoch):
+        self.current_epoch = epoch
+
+    def train_valid_split(self, test_size):
+        '''
+        test_size: 0.2
+        below is the value of each variable:
+        files[0] <- "/media/jiang/ECG/physionet.org/files/challenge-2021/1.0.3/training/python-classifier-2021-main/training_data/chapman_shaoxing/g1/JS00001.hea" shape -> (999, 1)
+        targets[1] <-[1. 0. 0. ... 0. 1. 0.]  26, shape->(999, 26)
+        x_train[0] <-"/media/jiang/ECG/physionet.org/files/challenge-2021/1.0.3/training/python-classifier-2021-main/training_data/chapman_shaoxing/g1/JS00001.hea" shape -> (799,1)
+        x_valid[0] <-"/media/jiang/ECG/physionet.org/files/challenge-2021/1.0.3/training/python-classifier-2021-main/training_data/chapman_shaoxing/g1/JS00651.hea" shape -> (200,1)
+        '''
+        files = self.files['header'].to_numpy().reshape(-1,1)
+        # print(files[0])
+        targets = np.stack(self.files['target'].to_list(), axis=0)
+        print("This is the targets:", targets)
+
+        x_train, y_train, x_valid, y_valid = iterative_train_test_split(files, targets, test_size=test_size)
+        
+        train = dataset(header_files=x_train[:,0].tolist())
+        
+        train.num_leads=None
+        train.sample=True
+
+        valid = dataset(header_files=x_valid[:,0].tolist())
+        
+        valid.num_leads=12
+        valid.sample=False
+        
+        return train, valid
+
+    def summary(self, output):
+        if output=='pandas':
+            # print("@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@")
+            # print("This is the target:",self.files['target'],type(self.files['target']))
+            return pd.Series(np.stack(self.files['target'].to_list(),axis=0).sum(axis=0),index=dataset.classes)
+
+        if output=='numpy':
+            return np.stack(self.files['target'].to_list(),axis=0).sum(axis=0)
+
+    def __len__(self):
+        return len(self.files)
+    
+    '''
+    fs: 500.0   sampling rate
+    target: [0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0. 0. 0. 0. 0. 0. 1. 0. 0. 0.]
+    leads: ('I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6')
+    self.files.iloc[item]['record']: ../../python-classifier-2021-main/training_data/chapman_shaoxing/g1/JS00518.mat
+    data: shape->(12 X 5000)
+    '''
+
+    def __getitem__(self, item):
+
+        fs = self.files.iloc[item]['fs']
+        target = self.files.iloc[item]['target']
+        leads = self.files.iloc[item]['leads']
+        data = load_recording(self.files.iloc[item]['record'])
+        
+        # print("This is the target:", target)
+        # Set your threshold
+        
+        # threshold = 0.5
+        # print("This is the original target:", target)
+        # print("This is the type of original target:", type(target))
+        
+        if random.random() < self.Mixup :
+            mix_sample_idx = random.randint(0, self.amount-1)
+            mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
+            mix_target = self.files.iloc[mix_sample_idx]['target']
+            data, alpha_value = mixup(data, mix_datum)
+            target = alpha_value * target + (1-alpha_value) * mix_target
+            # Binarize the labels based on the threshold
+            # target[target >= threshold] = 1
+            # target[target < threshold] = 0
+            
+        
+        if random.random() < self.cutMix:
+            mix_sample_idx = random.randint(0, self.amount-1)
+            mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
+            mix_target = self.files.iloc[mix_sample_idx]['target']
+            data, alpha_value = cutmix_fix_length(data, mix_datum)
+            target = alpha_value * target + (1-alpha_value) * mix_target
+            # target[target >= threshold] = 1
+            # target[target < threshold] = 0
+        
+        if random.random()< self.Mixup_no_label_interplate:
+            mix_sample_idx = random.randint(0, self.amount-1)
+            mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
+            data= same_shape_mixup(data, mix_datum)
+
+        progressive_index = 0
+        
+        # # below is the progessive data augmentation method
+        # if self.progressive: 
+        #     if self.current_epoch < 20:
+        #         if self.current_epoch % 4 == 0 or 3:
+        #             pass
+        #         elif self.current_epoch % 4 == 1:
+        #             proressive_index = 0.2
+        #         else:
+        #             proressive_index = 0.5
+
+        #     elif self.current_epoch < 40:
+        #         if self.current_epoch % 4 == 0 or 3:
+        #             pass
+        #         elif self.current_epoch % 4 == 1:
+        #             proressive_index = 0.7
+        #         else:
+        #             proressive_index = 0.8
+        #     else:
+        #         if self.current_epoch % 4 == 0 or 3:
+        #             pass
+        #         elif self.current_epoch % 4 == 1 or 2:
+        #             proressive_index = 1
+                    
+        # below is the wave-mix data augmentation method
+        if self.progressive: 
+            if self.current_epoch < 5:
+                match self.current_epoch:
+                    case 0:
+                        progressive_index=0.2
+                    case 1:
+                        progressive_index=0.4
+                    case 2:
+                        progressive_index=0.6
+                    case 3:
+                        progressive_index=0.8
+                    case 4:
+                        progressive_index=0.8
+                    
+            else :
+                progressive_index = 0.8
+            
+        if random.random() < progressive_index :
+            mix_sample_idx = random.randint(0, self.amount-1)
+            mix_datum = load_recording(self.files.iloc[mix_sample_idx]['record'])
+            mix_target = self.files.iloc[mix_sample_idx]['target']
+            data, alpha_value = mixup(data, mix_datum)
+            target = alpha_value * target + (1-alpha_value) * mix_target
+            # Binarize the labels based on the threshold
+            # target[target >= threshold] = 1
+            # target[target < threshold] = 0
+
+
+        # expand to 12 lead setup if original signal has less channels
+        data, lead_indicator = expand_leads(data, input_leads=leads)
+        data = np.nan_to_num(data)
+ 
+
+        # resample to 500hz
+        if fs == float(1000):
+            data = signal.resample_poly(data, up=1, down=2, axis=-1)  # to 500Hz
+            fs = 500
+        elif fs == float(500):
+            pass
+        else:
+            data = signal.resample(data, int(data.shape[1] * 500 / fs), axis=1)
+            fs = 500
+
+        # below is the Butterworth digital and analog filter design.
+        data = signal.filtfilt(self.b, self.a, data)
+
+        '''
+        we filter out the sample which the length is 8192 via the below code.
+        just like the random shift windows.  
+        '''
+        if self.sample:
+            fs = int(fs)
+            # random sample signal if len > 8192 samples
+            if data.shape[-1] >= 8192:
+                idx = data.shape[-1] - 8192-1
+                idx = np.random.randint(idx)
+                data = data[:, idx:idx + 8192]
+
+        '''
+        Z-Score Normalization
+        '''
+        mu  = np.nanmean(data, axis=-1, keepdims=True)
+        std = np.nanstd(data, axis=-1, keepdims=True)
+        
+        
+        #std = np.nanstd(data.flatten())
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            data = (data - mu) / std
+        
+        data = np.nan_to_num(data)
+
+        # random choose number of leads to keep
+        data, lead_indicator = lead_exctractor.get(data, self.num_leads, lead_indicator)
+        
+        # print("This is the lead_indicator:(l)", type(lead_indicator), lead_indicator)
+        
+
+        return data, target

engine_ecg_2020.py

@@ -0,0 +1,309 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+"""
+Train and eval functions used in main.py
+"""
+import math
+import sys
+from typing import Iterable, Optional
+
+import torch
+import torch.nn.functional as F
+
+
+import timm
+from timm.data import Mixup
+from timm.utils import accuracy, ModelEma
+
+from losses import DistillationLoss
+import utils
+
+from sklearn.metrics import average_precision_score, roc_auc_score, f1_score
+from sklearn.metrics import hamming_loss
+from sklearn.metrics import accuracy_score
+
+import numpy as np
+
+from evaluate_12ECG_score import load_weights, compute_challenge_metric
+
+classes = sorted(['270492004', '164889003', '164890007', '426627000', '713427006', 
+                  '713426002', '445118002', '39732003', '164909002', '251146004', 
+                  '698252002', '10370003', '284470004', '427172004', '164947007', 
+                  '111975006', '164917005', '47665007', '59118001', '427393009', 
+                  '426177001', '426783006', '427084000', '63593006', '164934002', 
+                  '59931005', '17338001'])
+
+def normalize_model_outputs(model_outputs):
+    """
+    Normalize model outputs to the range [0, 1].
+    
+    Parameters:
+    model_outputs (numpy.ndarray): The raw output of the deep learning model.
+    
+    Returns:
+    numpy.ndarray: The normalized model outputs.
+    """
+    a = model_outputs.min()
+    b = model_outputs.max()
+    return (model_outputs - a) / (b - a)
+
+weights_file = "./weights_2020.csv"
+normal_class = '426783006'
+weights = load_weights(weights_file, classes)
+
+def train_one_epoch(model: torch.nn.Module, criterion: DistillationLoss,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int, max_norm: float = 0,
+                    model_ema: Optional[ModelEma] = None, mixup_fn: Optional[Mixup] = None,
+                    set_training_mode=True, args = None):
+    
+    model.train(set_training_mode)
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 600
+    
+    
+        
+    output_list = []
+    target_list = []    
+
+    train_auprc_list = []
+
+    loss_value_after_each_epcoh = 0
+    batch_num = 0
+
+    # debug
+    # count = 0
+    # for samples, targets, mix_target, alpha_value in metric_logger.log_every(data_loader, print_freq, header):
+
+    # third_party = 0
+
+
+    for samples, targets in metric_logger.log_every(data_loader, print_freq, header):
+        # count += 1
+        # if count > 20:
+        #     break
+        batch_num += 1
+        # print("code goes here......................*******************************************************")
+
+        samples = samples.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+        # mix_target = mix_target.to(device, non_blocking=True)
+        # alpha_value = alpha_value.to(device, non_blocking=True)
+
+
+        outputs = model(samples.float(), if_random_cls_token_position=args.if_random_cls_token_position, if_random_token_rank=args.if_random_token_rank)
+            
+        # below is original loss            
+        loss = criterion(outputs, targets.float())
+        # # below is the asymmetric loss training
+        # outputs = torch.sigmoid(outputs)
+        # loss = -torch.mean(targets * F.logsigmoid(outputs) + (1 - targets) * F.logsigmoid(-outputs) * 0.1)
+
+        loss_value = loss.item()
+
+        if args.lrschedule == "Noam":
+            optimizer.optimizer.zero_grad()
+        elif args.lrschedule == "CosineAnnealing":
+            optimizer.zero_grad()
+        else:
+            raise ValueError(f"No matching condition for value in the file of engine_ecg_2021.py : {args.lrschedule}")
+
+        loss.backward() # Backward pass: Compute gradient of the loss with respect to model parameters
+
+        # Update parameters/using the Noam
+        optimizer.step() 
+
+        torch.cuda.synchronize()
+
+        metric_logger.update(loss=loss_value)
+
+        if args.lrschedule == "Noam":
+            metric_logger.update(lr=optimizer.optimizer.param_groups[0]["lr"])
+        elif args.lrschedule == "CosineAnnealing":
+            metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+        else:
+            raise ValueError(f"No matching condition for value in the file of engine_ecg_2021.py for updating : {args.lrschedule}")
+
+        target_list.append(targets.data.cpu().numpy())
+        output_list.append(outputs.data.cpu().numpy())
+        loss_value_after_each_epcoh += loss_value
+
+        
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+
+    # below code is for record the AUPRC of training
+    targets_all = np.concatenate(target_list, axis=0)
+    outputs_all = np.concatenate(output_list, axis=0)
+    outputs_all = normalize_model_outputs(outputs_all)
+
+    threshold = 0.5
+    targets_all[targets_all >= threshold] = 1
+    targets_all[targets_all < threshold] = 0
+
+    # When using the function below, y_true must be a binarized value.
+    train_auprc = average_precision_score(y_true = targets_all, y_score = outputs_all)
+    print("This is the training AUPRC:", train_auprc)
+
+    ### The code below is for obtaining the thresholds
+    scores_challengeScore = []
+    scores_F1 = []
+    scores_SubsetAccuracy = []
+    scores_HammingLoss = []
+
+    for thr in np.arange(0., 1., 0.02):
+        outputs_dyn = np.array([[(1 if prob > thr else 0) for prob in probs] for probs in np.array(outputs_all)])
+
+        challenge_value = compute_challenge_metric(weights, targets_all, outputs_dyn, classes, normal_class)
+        scores_challengeScore.append(challenge_value)
+
+        f1 = f1_score(targets_all, outputs_dyn, average='weighted')
+        scores_F1.append(f1)
+
+        subset_accuracy = accuracy_score(targets_all, outputs_dyn)
+        scores_SubsetAccuracy.append(subset_accuracy)
+
+        hamming = hamming_loss(targets_all, outputs_dyn)
+        scores_HammingLoss.append(hamming)
+
+    scores_challengeScore = np.array(scores_challengeScore)
+    scores_F1 = np.array(scores_F1)
+    scores_SubsetAccuracy = np.array(scores_SubsetAccuracy)
+    scores_HammingLoss = np.array(scores_HammingLoss)
+
+    # print("This is the challenge score list from training set:\n", scores)
+
+    # Best thrs for challenge score
+    thrs_CHALL = np.array([np.argmax(scores_challengeScore, axis=0)*0.02])
+    print("This is the best threshold for the challenge score from training set", thrs_CHALL)
+    outputs_best_CHALL = np.array([[(1 if prob > thrs_CHALL else 0) for prob in probs] for probs in np.array(outputs_all)])
+    challenge_value = compute_challenge_metric(weights, targets_all, outputs_best_CHALL, classes, normal_class)
+    print("This is the challenge score from training set:", challenge_value)
+
+    # Best thrs for F1
+    scores_F1 = np.array([np.argmax(scores_F1, axis=0)*0.02])
+    print("This is the best threshold for the F1 from training set", scores_F1)
+    outputs_best_f1 = np.array([[(1 if prob > scores_F1 else 0) for prob in probs] for probs in np.array(outputs_all)])
+    f1 = f1_score(targets_all, outputs_best_f1, average='weighted')
+    print("This is the f1 score from training set:", challenge_value)
+
+    # Best thrs for subset accuracy
+    scores_SubsetAccuracy = np.array([np.argmax(scores_SubsetAccuracy, axis=0)*0.02])
+    print("This is the best threshold for the Subset Accuracy from training set", scores_SubsetAccuracy)
+    outputs_best_SubsetAccuracy = np.array([[(1 if prob > scores_SubsetAccuracy else 0) for prob in probs] for probs in np.array(outputs_all)])
+    subset_accuracy = accuracy_score(targets_all, outputs_best_SubsetAccuracy)
+    print("This is the subset accuracy from training set:", subset_accuracy)
+
+    # Best thrs for hamming loss, here is the loss value from output, we need to find the minimum value.
+    # Determine the optimal threshold for Hamming loss. The loss values are obtained from the output, and the goal is to find the minimum value.
+    scores_HammingLoss = np.array([np.argmin(scores_HammingLoss, axis=0)*0.02])
+    print("This is the best threshold for the Subset Accuracy from training set", scores_HammingLoss)
+    outputs_best_HammingLoss = np.array([[(1 if prob > scores_HammingLoss else 0) for prob in probs] for probs in np.array(outputs_all)])
+    hamming = hamming_loss(targets_all, outputs_best_HammingLoss)
+    print("This is the Hamming from training set:", hamming)
+
+    loss_value_after_each_epcoh /= batch_num
+
+    return train_auprc, loss_value_after_each_epcoh, thrs_CHALL, scores_F1, scores_SubsetAccuracy, scores_HammingLoss
+
+
+
+@torch.no_grad()
+def evaluate(data_loader, model, device):
+    # criterion = torch.nn.CrossEntropyLoss()
+    criterion = torch.nn.BCEWithLogitsLoss()
+
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    header = 'Test:'
+
+    targets = []
+    outputs = []
+    loss_value_after_each_epcoh = 0
+    batch_num = 0
+    # switch to evaluation mode
+    model.eval()
+
+    for images, target in metric_logger.log_every(data_loader, 100, header):
+        
+        batch_num += 1
+        images = images.to(device, non_blocking=True)
+        target = target.to(device, non_blocking=True)
+        # print("NaN in images(input):", torch.isnan(images).any())
+
+        # compute output
+        
+        output = model(images.float())
+            # below is original loss
+        loss = criterion(output, target.float())
+            # below is the logit to proba
+
+            # below is the asymmetric loss testing
+            # output = torch.sigmoid(output)
+            # loss = -torch.mean(target * F.logsigmoid(output) + (1 - target) * F.logsigmoid(-output) * 0.1)
+
+        # print("This is the output:", output.shape,output)
+        # print("This is the target.float():",target.float().shape ,target.float())
+        # print("NaN in output:", torch.isnan(output).any())
+        # print("NaN in target:", torch.isnan(target).any())
+        # print("This is the loss.item():", loss.item())
+        # sys.exit()
+
+        metric_logger.update(loss=loss.item())
+        
+        targets.append(target.data.cpu().numpy())
+        outputs.append(output.data.cpu().numpy())
+    
+        loss_value_after_each_epcoh += loss.item()
+
+    # below code is for record the AUPRC of validation
+    targets = np.concatenate(targets, axis=0)
+    outputs = np.concatenate(outputs, axis=0)
+
+    outputs = normalize_model_outputs(outputs)
+    
+    auprc = average_precision_score(y_true=targets, y_score=outputs)
+    auroc = roc_auc_score(targets, outputs)
+
+    # print("This is the top 3 row:", outputs[0:3])
+    # outputs = np.array([[(1 if prob > 0 else 0) for prob in probs] for probs in np.array(outputs)])
+    # print("This is the top 3 row, after:", outputs[0:3])
+    # for i in range(len(outputs)):
+    #     num_ones = targets[i].count(1)
+    #     third_largest = sorted(outputs[i], reverse=True)[num_ones-1]
+
+    # print(f"Length of targets: {len(targets)}")
+    # print(f"Length of outputs: {len(outputs)}")
+    
+    outputs_0 = np.array([[(1 if prob > 0.5 else 0) for prob in probs] for probs in np.array(outputs)])
+
+    f1 = f1_score(targets, outputs_0, average='weighted')
+    hamming = hamming_loss(targets, outputs_0)
+    subset_accuracy = accuracy_score(targets, outputs_0)
+
+    scores = []
+    for thr in np.arange(0., 1., 0.02):
+        outputs_dyn = np.array([[(1 if prob > thr else 0) for prob in probs] for probs in np.array(outputs)])
+        challenge_value = compute_challenge_metric(weights, targets, outputs_dyn, classes, normal_class)
+        scores.append(challenge_value)
+
+    scores = np.array(scores)
+    print("This is the challenge score list from testing set:\n", scores)
+
+    # Best thrs and preds
+    idxs = np.argmax(scores, axis=0)
+    thrs = np.array([idxs*0.02])
+
+    print("This is the best threshold for the challenge score from testing test", thrs)
+    outputs_best = np.array([[(1 if prob > thrs else 0) for prob in probs] for probs in np.array(outputs)])
+
+    challenge_value = compute_challenge_metric(weights, targets, outputs_best, classes, normal_class)
+    print("This is the challenge score:", challenge_value)
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    loss_value_after_each_epcoh /= batch_num
+
+    return auprc, auroc, f1, hamming, subset_accuracy, loss_value_after_each_epcoh

engine_ecg_2021.py

@@ -0,0 +1,241 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+"""
+Train and eval functions used in main.py
+"""
+import math
+import sys
+from typing import Iterable, Optional
+
+import torch
+
+import timm
+from timm.data import Mixup
+from timm.utils import accuracy, ModelEma
+
+from losses import DistillationLoss
+import utils
+
+from sklearn.metrics import average_precision_score, roc_auc_score, f1_score
+from sklearn.metrics import hamming_loss
+from sklearn.metrics import accuracy_score
+
+import numpy as np
+
+from evaluate_model import load_weights, compute_challenge_metric
+
+
+def normalize_model_outputs(model_outputs):
+    """
+    Normalize model outputs to the range [0, 1].
+    
+    Parameters:
+    model_outputs (numpy.ndarray): The raw output of the deep learning model.
+    
+    Returns:
+    numpy.ndarray: The normalized model outputs.
+    """
+    a = model_outputs.min()
+    b = model_outputs.max()
+    return (model_outputs - a) / (b - a)
+
+
+sinus_rhythm_ID = set(['426783006'])
+
+weights_file = "./weights.csv"
+
+classes, weights = load_weights(weights_file)
+
+def train_one_epoch(model: torch.nn.Module, criterion: DistillationLoss,
+                    data_loader: Iterable, optimizer: torch.optim.Optimizer,
+                    device: torch.device, epoch: int,
+                    set_training_mode=True, args = None):
+    
+    model.train(set_training_mode)
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter('lr', utils.SmoothedValue(window_size=1, fmt='{value:.6f}'))
+    header = 'Epoch: [{}]'.format(epoch)
+    print_freq = 600
+    
+    output_list = []
+    target_list = []    
+
+    loss_value_after_each_epcoh = 0
+    batch_num = 0
+
+    for samples, targets in metric_logger.log_every(data_loader, print_freq, header):        
+        batch_num += 1
+
+        samples = samples.to(device, non_blocking=True)
+        targets = targets.to(device, non_blocking=True)
+        
+        outputs = model(samples.float(), if_random_cls_token_position=args.if_random_cls_token_position, if_random_token_rank=args.if_random_token_rank)
+        loss = criterion(outputs, targets.float())
+            
+        loss_value = loss.item()
+
+        if args.lrschedule == "Noam":
+            optimizer.optimizer.zero_grad()
+        elif args.lrschedule == "CosineAnnealing":
+            optimizer.zero_grad()
+        else:
+            raise ValueError(f"No matching condition for value in the file of engine_ecg_2021.py : {args.lrschedule}")
+        
+        loss.backward() # Backward pass: Compute gradient of the loss with respect to model parameters
+            
+        # Update parameters/using the Noam
+        optimizer.step() 
+
+        torch.cuda.synchronize()
+        
+        metric_logger.update(loss=loss_value)
+        if args.lrschedule == "Noam":
+            metric_logger.update(lr=optimizer.optimizer.param_groups[0]["lr"])
+        elif args.lrschedule == "CosineAnnealing":
+            metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+        else:
+            raise ValueError(f"No matching condition for value in the file of engine_ecg_2021.py for updating : {args.lrschedule}")
+        
+        target_list.append(targets.data.cpu().numpy())
+        output_list.append(outputs.data.cpu().numpy())
+        loss_value_after_each_epcoh += loss_value
+    
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+
+    # below code is for record the AUPRC of training
+    targets_all = np.concatenate(target_list, axis=0)
+    outputs_all = np.concatenate(output_list, axis=0)
+    outputs_all = normalize_model_outputs(outputs_all)
+
+    threshold = 0.5
+    targets_all[targets_all >= threshold] = 1
+    targets_all[targets_all < threshold] = 0
+
+    # When using the function below, y_true must be a binarized value.
+    train_auprc = average_precision_score(y_true = targets_all, y_score = outputs_all)
+    print("This is the training AUPRC:", train_auprc)
+
+    ### The code below is for obtaining the thresholds
+    scores_challengeScore = []
+    scores_F1 = []
+    scores_SubsetAccuracy = []
+    scores_HammingLoss = []
+
+    for thr in np.arange(0., 1., 0.02):
+        outputs_dyn = np.array([[(1 if prob > thr else 0) for prob in probs] for probs in np.array(outputs_all)])
+
+        challenge_value = compute_challenge_metric(weights, targets_all, outputs_dyn, classes, sinus_rhythm_ID)
+        scores_challengeScore.append(challenge_value)
+
+        f1 = f1_score(targets_all, outputs_dyn, average='weighted')
+        scores_F1.append(f1)
+
+        subset_accuracy = accuracy_score(targets_all, outputs_dyn)
+        scores_SubsetAccuracy.append(subset_accuracy)
+
+        hamming = hamming_loss(targets_all, outputs_dyn)
+        scores_HammingLoss.append(hamming)
+
+    scores_challengeScore = np.array(scores_challengeScore)
+    scores_F1 = np.array(scores_F1)
+    scores_SubsetAccuracy = np.array(scores_SubsetAccuracy)
+    scores_HammingLoss = np.array(scores_HammingLoss)
+
+    # print("This is the challenge score list from training set:\n", scores)
+
+    # Best thrs for challenge score
+    thrs_CHALL = np.array([np.argmax(scores_challengeScore, axis=0)*0.02])
+    print("This is the best threshold for the challenge score from training set", thrs_CHALL)
+    outputs_best_CHALL = np.array([[(1 if prob > thrs_CHALL else 0) for prob in probs] for probs in np.array(outputs_all)])
+    challenge_value = compute_challenge_metric(weights, targets_all, outputs_best_CHALL, classes, sinus_rhythm_ID)
+    print("This is the challenge score from training set:", challenge_value)
+
+    # Best thrs for F1
+    scores_F1 = np.array([np.argmax(scores_F1, axis=0)*0.02])
+    print("This is the best threshold for the F1 from training set", scores_F1)
+    outputs_best_f1 = np.array([[(1 if prob > scores_F1 else 0) for prob in probs] for probs in np.array(outputs_all)])
+    f1 = f1_score(targets_all, outputs_best_f1, average='weighted')
+    print("This is the f1 score from training set:", challenge_value)
+
+    # Best thrs for subset accuracy
+    scores_SubsetAccuracy = np.array([np.argmax(scores_SubsetAccuracy, axis=0)*0.02])
+    print("This is the best threshold for the Subset Accuracy from training set", scores_SubsetAccuracy)
+    outputs_best_SubsetAccuracy = np.array([[(1 if prob > scores_SubsetAccuracy else 0) for prob in probs] for probs in np.array(outputs_all)])
+    subset_accuracy = accuracy_score(targets_all, outputs_best_SubsetAccuracy)
+    print("This is the subset accuracy from training set:", subset_accuracy)
+
+    # Best thrs for hamming loss, here is the loss value from output, we need to find the minimum value.
+    # Determine the optimal threshold for Hamming loss. The loss values are obtained from the output, and the goal is to find the minimum value.
+    scores_HammingLoss = np.array([np.argmin(scores_HammingLoss, axis=0)*0.02])
+    print("This is the best threshold for the Subset Accuracy from training set", scores_HammingLoss)
+    outputs_best_HammingLoss = np.array([[(1 if prob > scores_HammingLoss else 0) for prob in probs] for probs in np.array(outputs_all)])
+    hamming = hamming_loss(targets_all, outputs_best_HammingLoss)
+    print("This is the Hamming from training set:", hamming)
+
+    loss_value_after_each_epcoh /= batch_num
+    return train_auprc, loss_value_after_each_epcoh, thrs_CHALL, scores_F1, scores_SubsetAccuracy, scores_HammingLoss
+
+@torch.no_grad()
+def evaluate(data_loader, model, thrs_chall, thrs_F1, thrs_accuracy, thrs_hammingLoss, device):
+    # criterion = torch.nn.CrossEntropyLoss()
+    criterion = torch.nn.BCEWithLogitsLoss()
+
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    header = 'Test:'
+
+    targets = []
+    outputs = []
+    loss_value_after_each_epcoh = 0
+    batch_num = 0
+    # switch to evaluation mode
+    model.eval()
+
+    for images, target in metric_logger.log_every(data_loader, 100, header):
+        
+        batch_num += 1
+        images = images.to(device, non_blocking=True)
+        target = target.to(device, non_blocking=True)
+
+        output = model(images.float())
+
+        loss = criterion(output, target.float())
+        
+        metric_logger.update(loss=loss.item())
+
+        targets.append(target.data.cpu().numpy())
+        outputs.append(output.data.cpu().numpy())
+    
+        loss_value_after_each_epcoh += loss.item()
+
+    # below code is for record the AUPRC of validation
+    targets = np.concatenate(targets, axis=0)
+    outputs = np.concatenate(outputs, axis=0)
+    
+    outputs = normalize_model_outputs(outputs)
+
+    auprc = average_precision_score(y_true=targets, y_score=outputs)
+    auroc = roc_auc_score(targets, outputs)
+
+    # print("This is the top 3 row:", outputs[0:3])
+    outputs_F1 = np.array([[(1 if prob > thrs_F1 else 0) for prob in probs] for probs in np.array(outputs)])
+    f1 = f1_score(targets, outputs_F1, average='weighted')
+
+    outputs_hammingLoss = np.array([[(1 if prob > thrs_hammingLoss else 0) for prob in probs] for probs in np.array(outputs)])
+    hamming = hamming_loss(targets, outputs_hammingLoss)
+
+    outputs_accuracy = np.array([[(1 if prob > thrs_accuracy else 0) for prob in probs] for probs in np.array(outputs)])
+    subset_accuracy = accuracy_score(targets, outputs_accuracy)
+
+    print("This is the best threshold for the challenge score, obtained from the training test, without any testing data leakage:", thrs_chall)
+    outputs_best = np.array([[(1 if prob > thrs_chall else 0) for prob in probs] for probs in np.array(outputs)])
+
+    challenge_value = compute_challenge_metric(weights, targets, outputs_best, classes, sinus_rhythm_ID)
+    print("This is the challenge score:", challenge_value)
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    loss_value_after_each_epcoh /= batch_num
+
+    return auprc, auroc, f1, hamming, subset_accuracy, challenge_value, loss_value_after_each_epcoh

evaluate_12ECG_score.py

@@ -0,0 +1,577 @@
+#!/usr/bin/env python
+
+# This file contains functions for evaluating algorithms for the 2020 PhysioNet/
+# Computing in Cardiology Challenge. You can run it as follows:
+#
+#   python evaluate_12ECG_score.py labels outputs scores.csv
+#
+# where 'labels' is a directory containing files with the labels, 'outputs' is a
+# directory containing files with the outputs from your model, and 'scores.csv'
+# (optional) is a collection of scores for the algorithm outputs.
+#
+# Each file of labels or outputs must have the format described on the Challenge
+# webpage. The scores for the algorithm outputs include the area under the
+# receiver-operating characteristic curve (AUROC), the area under the recall-
+# precision curve (AUPRC), accuracy (fraction of correct recordings), macro F-
+# measure, and the Challenge metric, which assigns different weights to
+# different misclassification errors.
+
+import numpy as np, os, os.path, sys
+
+def evaluate_12ECG_score(label_directory, output_directory):
+    # Define the weights, the SNOMED CT code for the normal class, and equivalent SNOMED CT codes.
+    weights_file = 'weights.csv'
+    normal_class = '426783006'
+    equivalent_classes = [['713427006', '59118001'], ['284470004', '63593006'], ['427172004', '17338001']]
+
+    # Find the label and output files.
+    print('Finding label and output files...')
+    label_files, output_files = find_challenge_files(label_directory, output_directory)
+
+    # Load the labels and outputs.
+    print('Loading labels and outputs...')
+    label_classes, labels = load_labels(label_files, normal_class, equivalent_classes)
+    output_classes, binary_outputs, scalar_outputs = load_outputs(output_files, normal_class, equivalent_classes)
+
+    # Organize/sort the labels and outputs.
+    print('Organizing labels and outputs...')
+    classes, labels, binary_outputs, scalar_outputs = organize_labels_outputs(label_classes, output_classes, labels, binary_outputs, scalar_outputs)
+
+    # Load the weights for the Challenge metric.
+    print('Loading weights...')
+    weights = load_weights(weights_file, classes)
+
+    # Only consider classes that are scored with the Challenge metric.
+    indices = np.any(weights, axis=0) # Find indices of classes in weight matrix.
+    classes = [x for i, x in enumerate(classes) if indices[i]]
+    labels = labels[:, indices]
+    scalar_outputs = scalar_outputs[:, indices]
+    binary_outputs = binary_outputs[:, indices]
+    weights = weights[np.ix_(indices, indices)]
+
+    # Evaluate the model by comparing the labels and outputs.
+    print('Evaluating model...')
+
+    print('- AUROC and AUPRC...')
+    auroc, auprc = compute_auc(labels, scalar_outputs)
+
+    print('- Accuracy...')
+    accuracy = compute_accuracy(labels, binary_outputs)
+
+    print('- F-measure...')
+    f_measure = compute_f_measure(labels, binary_outputs)
+
+    print('- F-beta and G-beta measures...')
+    f_beta_measure, g_beta_measure = compute_beta_measures(labels, binary_outputs, beta=2)
+
+    print('- Challenge metric...')
+    challenge_metric = compute_challenge_metric(weights, labels, binary_outputs, classes, normal_class)
+
+    print('Done.')
+
+    # Return the results.
+    return auroc, auprc, accuracy, f_measure, f_beta_measure, g_beta_measure, challenge_metric
+
+# Check if the input is a number.
+def is_number(x):
+    try:
+        float(x)
+        return True
+    except ValueError:
+        return False
+
+# Find Challenge files.
+def find_challenge_files(label_directory, output_directory):
+    label_files = list()
+    output_files = list()
+    for f in sorted(os.listdir(label_directory)):
+        F = os.path.join(label_directory, f) # Full path for label file
+        if os.path.isfile(F) and F.lower().endswith('.hea') and not f.lower().startswith('.'):
+            root, ext = os.path.splitext(f)
+            g = root + '.csv'
+            G = os.path.join(output_directory, g) # Full path for corresponding output file
+            if os.path.isfile(G):
+                label_files.append(F)
+                output_files.append(G)
+            else:
+                raise IOError('Output file {} not found for label file {}.'.format(g, f))
+
+    if label_files and output_files:
+        return label_files, output_files
+    else:
+        raise IOError('No label or output files found.')
+
+# Load labels from header/label files.
+def load_labels(label_files, normal_class, equivalent_classes_collection):
+    # The labels should have the following form:
+    #
+    # Dx: label_1, label_2, label_3
+    #
+    num_recordings = len(label_files)
+
+    # Load diagnoses.
+    tmp_labels = list()
+    for i in range(num_recordings):
+        with open(label_files[i], 'r') as f:
+            for l in f:
+                if l.startswith('#Dx'):
+                    dxs = set(arr.strip() for arr in l.split(': ')[1].split(','))
+                    tmp_labels.append(dxs)
+
+    # Identify classes.
+    classes = set.union(*map(set, tmp_labels))
+    if normal_class not in classes:
+        classes.add(normal_class)
+        print('- The normal class {} is not one of the label classes, so it has been automatically added, but please check that you chose the correct normal class.'.format(normal_class))
+    classes = sorted(classes)
+    num_classes = len(classes)
+
+    # Use one-hot encoding for labels.
+    labels = np.zeros((num_recordings, num_classes), dtype=np.bool_)
+    for i in range(num_recordings):
+        dxs = tmp_labels[i]
+        for dx in dxs:
+            j = classes.index(dx)
+            labels[i, j] = 1
+
+    # For each set of equivalent class, use only one class as the representative class for the set and discard the other classes in the set.
+    # The label for the representative class is positive if any of the labels in the set is positive.
+    remove_classes = list()
+    remove_indices = list()
+    for equivalent_classes in equivalent_classes_collection:
+        equivalent_classes = [x for x in equivalent_classes if x in classes]
+        if len(equivalent_classes)>1:
+            representative_class = equivalent_classes[0]
+            other_classes = equivalent_classes[1:]
+            equivalent_indices = [classes.index(x) for x in equivalent_classes]
+            representative_index = equivalent_indices[0]
+            other_indices = equivalent_indices[1:]
+
+            labels[:, representative_index] = np.any(labels[:, equivalent_indices], axis=1)
+            remove_classes += other_classes
+            remove_indices += other_indices
+
+    for x in remove_classes:
+        classes.remove(x)
+    labels = np.delete(labels, remove_indices, axis=1)
+
+    # If the labels are negative for all classes, then change the label for the normal class to positive.
+    normal_index = classes.index(normal_class)
+    for i in range(num_recordings):
+        num_positive_classes = np.sum(labels[i, :])
+        if num_positive_classes==0:
+            labels[i, normal_index] = 1
+
+    return classes, labels
+
+# Load outputs from output files.
+def load_outputs(output_files, normal_class, equivalent_classes_collection):
+    # The outputs should have the following form:
+    #
+    # diagnosis_1, diagnosis_2, diagnosis_3
+    #           0,           1,           1
+    #        0.12,        0.34,        0.56
+    #
+    num_recordings = len(output_files)
+
+    tmp_labels = list()
+    tmp_binary_outputs = list()
+    tmp_scalar_outputs = list()
+    for i in range(num_recordings):
+        with open(output_files[i], 'r') as f:
+            for j, l in enumerate(f):
+                arrs = [arr.strip() for arr in l.split(',')]
+                if j==1:
+                    row = arrs
+                    tmp_labels.append(row)
+                elif j==2:
+                    row = list()
+                    for arr in arrs:
+                        number = 1 if arr in ('1', 'True', 'true', 'T', 't') else 0
+                        row.append(number)
+                    tmp_binary_outputs.append(row)
+                elif j==3:
+                    row = list()
+                    for arr in arrs:
+                        number = float(arr) if is_number(arr) else 0
+                        row.append(number)
+                    tmp_scalar_outputs.append(row)
+
+    # Identify classes.
+    classes = set.union(*map(set, tmp_labels))
+    if normal_class not in classes:
+        classes.add(normal_class)
+        print('- The normal class {} is not one of the output classes, so it has been automatically added, but please check that you identified the correct normal class.'.format(normal_class))
+    classes = sorted(classes)
+    num_classes = len(classes)
+
+    # Use one-hot encoding for binary outputs and the same order for scalar outputs.
+    binary_outputs = np.zeros((num_recordings, num_classes), dtype=np.bool_)
+    scalar_outputs = np.zeros((num_recordings, num_classes), dtype=np.float64)
+    for i in range(num_recordings):
+        dxs = tmp_labels[i]
+        for k, dx in enumerate(dxs):
+            j = classes.index(dx)
+            binary_outputs[i, j] = tmp_binary_outputs[i][k]
+            scalar_outputs[i, j] = tmp_scalar_outputs[i][k]
+
+    # For each set of equivalent class, use only one class as the representative class for the set and discard the other classes in the set.
+    # The binary output for the representative class is positive if any of the classes in the set is positive.
+    # The scalar output is the mean of the scalar outputs for the classes in the set.
+    remove_classes = list()
+    remove_indices = list()
+    for equivalent_classes in equivalent_classes_collection:
+        equivalent_classes = [x for x in equivalent_classes if x in classes]
+        if len(equivalent_classes)>1:
+            representative_class = equivalent_classes[0]
+            other_classes = equivalent_classes[1:]
+            equivalent_indices = [classes.index(x) for x in equivalent_classes]
+            representative_index = equivalent_indices[0]
+            other_indices = equivalent_indices[1:]
+
+            binary_outputs[:, representative_index] = np.any(binary_outputs[:, equivalent_indices], axis=1)
+            scalar_outputs[:, representative_index] = np.nanmean(scalar_outputs[:, equivalent_indices], axis=1)
+            remove_classes += other_classes
+            remove_indices += other_indices
+
+    for x in remove_classes:
+        classes.remove(x)
+    binary_outputs = np.delete(binary_outputs, remove_indices, axis=1)
+    scalar_outputs = np.delete(scalar_outputs, remove_indices, axis=1)
+
+    # If any of the outputs is a NaN, then replace it with a zero.
+    binary_outputs[np.isnan(binary_outputs)] = 0
+    scalar_outputs[np.isnan(scalar_outputs)] = 0
+
+    # If the binary outputs are negative for all classes, then change the binary output for the normal class to positive.
+    normal_index = classes.index(normal_class)
+    for i in range(num_recordings):
+        num_positive_classes = np.sum(binary_outputs[i, :])
+        if num_positive_classes==0:
+            binary_outputs[i, normal_index] = 1
+
+    return classes, binary_outputs, scalar_outputs
+
+# Organize labels and outputs.
+def organize_labels_outputs(label_classes, output_classes, tmp_labels, tmp_binary_outputs, tmp_scalar_outputs):
+    # Include all classes from either the labels or the outputs.
+    classes = sorted(set(label_classes) | set(output_classes))
+    num_classes = len(classes)
+
+    # Check that the labels and outputs have the same numbers of recordings.
+    assert(len(tmp_labels)==len(tmp_binary_outputs)==len(tmp_scalar_outputs))
+    num_recordings = len(tmp_labels)
+
+    # Rearrange the columns of the labels and the outputs to be consistent with the order of the classes.
+    labels = np.zeros((num_recordings, num_classes), dtype=np.bool_)
+    for k, dx in enumerate(label_classes):
+        j = classes.index(dx)
+        labels[:, j] = tmp_labels[:, k]
+
+    binary_outputs = np.zeros((num_recordings, num_classes), dtype=np.bool_)
+    scalar_outputs = np.zeros((num_recordings, num_classes), dtype=np.float64)
+    for k, dx in enumerate(output_classes):
+        j = classes.index(dx)
+        binary_outputs[:, j] = tmp_binary_outputs[:, k]
+        scalar_outputs[:, j] = tmp_scalar_outputs[:, k]
+
+    return classes, labels, binary_outputs, scalar_outputs
+
+# Load a table with row and column names.
+def load_table(table_file):
+    # The table should have the following form:
+    #
+    # ,    a,   b,   c
+    # a, 1.2, 2.3, 3.4
+    # b, 4.5, 5.6, 6.7
+    # c, 7.8, 8.9, 9.0
+    #
+    table = list()
+    with open(table_file, 'r') as f:
+        for i, l in enumerate(f):
+            arrs = [arr.strip() for arr in l.split(',')]
+            table.append(arrs)
+
+    # Define the numbers of rows and columns and check for errors.
+    num_rows = len(table)-1
+    if num_rows<1:
+        raise Exception('The table {} is empty.'.format(table_file))
+
+    num_cols = set(len(table[i])-1 for i in range(num_rows))
+    if len(num_cols)!=1:
+        raise Exception('The table {} has rows with different lengths.'.format(table_file))
+    num_cols = min(num_cols)
+    if num_cols<1:
+        raise Exception('The table {} is empty.'.format(table_file))
+
+    # Find the row and column labels.
+    rows = [table[0][j+1] for j in range(num_rows)]
+    cols = [table[i+1][0] for i in range(num_cols)]
+
+    # Find the entries of the table.
+    values = np.zeros((num_rows, num_cols))
+    for i in range(num_rows):
+        for j in range(num_cols):
+            value = table[i+1][j+1]
+            if is_number(value):
+                values[i, j] = float(value)
+            else:
+                values[i, j] = float('nan')
+
+    return rows, cols, values
+
+# Load weights.
+def load_weights(weight_file, classes):
+    # Load the weight matrix.
+    rows, cols, values = load_table(weight_file)
+    assert(rows == cols)
+    num_rows = len(rows)
+
+    # Assign the entries of the weight matrix with rows and columns corresponding to the classes.
+    num_classes = len(classes)
+    weights = np.zeros((num_classes, num_classes), dtype=np.float64)
+    for i, a in enumerate(rows):
+        if a in classes:
+            k = classes.index(a)
+            for j, b in enumerate(rows):
+                if b in classes:
+                    l = classes.index(b)
+                    weights[k, l] = values[i, j]
+
+    return weights
+
+# Compute recording-wise accuracy.
+def compute_accuracy(labels, outputs):
+    num_recordings, num_classes = np.shape(labels)
+
+    num_correct_recordings = 0
+    for i in range(num_recordings):
+        if np.all(labels[i, :]==outputs[i, :]):
+            num_correct_recordings += 1
+
+    return float(num_correct_recordings) / float(num_recordings)
+
+# Compute confusion matrices.
+def compute_confusion_matrices(labels, outputs, normalize=False):
+    # Compute a binary confusion matrix for each class k:
+    #
+    #     [TN_k FN_k]
+    #     [FP_k TP_k]
+    #
+    # If the normalize variable is set to true, then normalize the contributions
+    # to the confusion matrix by the number of labels per recording.
+    num_recordings, num_classes = np.shape(labels)
+
+    if not normalize:
+        A = np.zeros((num_classes, 2, 2))
+        for i in range(num_recordings):
+            for j in range(num_classes):
+                if labels[i, j]==1 and outputs[i, j]==1: # TP
+                    A[j, 1, 1] += 1
+                elif labels[i, j]==0 and outputs[i, j]==1: # FP
+                    A[j, 1, 0] += 1
+                elif labels[i, j]==1 and outputs[i, j]==0: # FN
+                    A[j, 0, 1] += 1
+                elif labels[i, j]==0 and outputs[i, j]==0: # TN
+                    A[j, 0, 0] += 1
+                else: # This condition should not happen.
+                    raise ValueError('Error in computing the confusion matrix.')
+    else:
+        A = np.zeros((num_classes, 2, 2))
+        for i in range(num_recordings):
+            normalization = float(max(np.sum(labels[i, :]), 1))
+            for j in range(num_classes):
+                if labels[i, j]==1 and outputs[i, j]==1: # TP
+                    A[j, 1, 1] += 1.0/normalization
+                elif labels[i, j]==0 and outputs[i, j]==1: # FP
+                    A[j, 1, 0] += 1.0/normalization
+                elif labels[i, j]==1 and outputs[i, j]==0: # FN
+                    A[j, 0, 1] += 1.0/normalization
+                elif labels[i, j]==0 and outputs[i, j]==0: # TN
+                    A[j, 0, 0] += 1.0/normalization
+                else: # This condition should not happen.
+                    raise ValueError('Error in computing the confusion matrix.')
+
+    return A
+
+# Compute macro F-measure.
+def compute_f_measure(labels, outputs):
+    num_recordings, num_classes = np.shape(labels)
+
+    A = compute_confusion_matrices(labels, outputs)
+
+    f_measure = np.zeros(num_classes)
+    for k in range(num_classes):
+        tp, fp, fn, tn = A[k, 1, 1], A[k, 1, 0], A[k, 0, 1], A[k, 0, 0]
+        if 2 * tp + fp + fn:
+            f_measure[k] = float(2 * tp) / float(2 * tp + fp + fn)
+        else:
+            f_measure[k] = float('nan')
+
+    macro_f_measure = np.nanmean(f_measure)
+
+    return macro_f_measure
+
+# Compute F-beta and G-beta measures from the unofficial phase of the Challenge.
+def compute_beta_measures(labels, outputs, beta):
+    num_recordings, num_classes = np.shape(labels)
+
+    A = compute_confusion_matrices(labels, outputs, normalize=True)
+
+    f_beta_measure = np.zeros(num_classes)
+    g_beta_measure = np.zeros(num_classes)
+    for k in range(num_classes):
+        tp, fp, fn, tn = A[k, 1, 1], A[k, 1, 0], A[k, 0, 1], A[k, 0, 0]
+        if (1+beta**2)*tp + fp + beta**2*fn:
+            f_beta_measure[k] = float((1+beta**2)*tp) / float((1+beta**2)*tp + fp + beta**2*fn)
+        else:
+            f_beta_measure[k] = float('nan')
+        if tp + fp + beta*fn:
+            g_beta_measure[k] = float(tp) / float(tp + fp + beta*fn)
+        else:
+            g_beta_measure[k] = float('nan')
+
+    macro_f_beta_measure = np.nanmean(f_beta_measure)
+    macro_g_beta_measure = np.nanmean(g_beta_measure)
+
+    return macro_f_beta_measure, macro_g_beta_measure
+
+# Compute macro AUROC and macro AUPRC.
+def compute_auc(labels, outputs):
+    num_recordings, num_classes = np.shape(labels)
+
+    # Compute and summarize the confusion matrices for each class across at distinct output values.
+    auroc = np.zeros(num_classes)
+    auprc = np.zeros(num_classes)
+
+    for k in range(num_classes):
+        # We only need to compute TPs, FPs, FNs, and TNs at distinct output values.
+        thresholds = np.unique(outputs[:, k])
+        thresholds = np.append(thresholds, thresholds[-1]+1)
+        thresholds = thresholds[::-1]
+        num_thresholds = len(thresholds)
+
+        # Initialize the TPs, FPs, FNs, and TNs.
+        tp = np.zeros(num_thresholds)
+        fp = np.zeros(num_thresholds)
+        fn = np.zeros(num_thresholds)
+        tn = np.zeros(num_thresholds)
+        fn[0] = np.sum(labels[:, k]==1)
+        tn[0] = np.sum(labels[:, k]==0)
+
+        # Find the indices that result in sorted output values.
+        idx = np.argsort(outputs[:, k])[::-1]
+
+        # Compute the TPs, FPs, FNs, and TNs for class k across thresholds.
+        i = 0
+        for j in range(1, num_thresholds):
+            # Initialize TPs, FPs, FNs, and TNs using values at previous threshold.
+            tp[j] = tp[j-1]
+            fp[j] = fp[j-1]
+            fn[j] = fn[j-1]
+            tn[j] = tn[j-1]
+
+            # Update the TPs, FPs, FNs, and TNs at i-th output value.
+            while i < num_recordings and outputs[idx[i], k] >= thresholds[j]:
+                if labels[idx[i], k]:
+                    tp[j] += 1
+                    fn[j] -= 1
+                else:
+                    fp[j] += 1
+                    tn[j] -= 1
+                i += 1
+
+        # Summarize the TPs, FPs, FNs, and TNs for class k.
+        tpr = np.zeros(num_thresholds)
+        tnr = np.zeros(num_thresholds)
+        ppv = np.zeros(num_thresholds)
+        npv = np.zeros(num_thresholds)
+
+        for j in range(num_thresholds):
+            if tp[j] + fn[j]:
+                tpr[j] = float(tp[j]) / float(tp[j] + fn[j])
+            else:
+                tpr[j] = float('nan')
+            if fp[j] + tn[j]:
+                tnr[j] = float(tn[j]) / float(fp[j] + tn[j])
+            else:
+                tnr[j] = float('nan')
+            if tp[j] + fp[j]:
+                ppv[j] = float(tp[j]) / float(tp[j] + fp[j])
+            else:
+                ppv[j] = float('nan')
+
+        # Compute AUROC as the area under a piecewise linear function with TPR/
+        # sensitivity (x-axis) and TNR/specificity (y-axis) and AUPRC as the area
+        # under a piecewise constant with TPR/recall (x-axis) and PPV/precision
+        # (y-axis) for class k.
+        for j in range(num_thresholds-1):
+            auroc[k] += 0.5 * (tpr[j+1] - tpr[j]) * (tnr[j+1] + tnr[j])
+            auprc[k] += (tpr[j+1] - tpr[j]) * ppv[j+1]
+
+    # Compute macro AUROC and macro AUPRC across classes.
+    macro_auroc = np.nanmean(auroc)
+    macro_auprc = np.nanmean(auprc)
+
+    return macro_auroc, macro_auprc
+
+# Compute modified confusion matrix for multi-class, multi-label tasks.
+def compute_modified_confusion_matrix(labels, outputs):
+    # Compute a binary multi-class, multi-label confusion matrix, where the rows
+    # are the labels and the columns are the outputs.
+    num_recordings, num_classes = np.shape(labels)
+    A = np.zeros((num_classes, num_classes))
+
+    # Iterate over all of the recordings.
+    for i in range(num_recordings):
+        # Calculate the number of positive labels and/or outputs.
+        normalization = float(max(np.sum(np.any((labels[i, :], outputs[i, :]), axis=0)), 1))
+        # Iterate over all of the classes.
+        for j in range(num_classes):
+            # Assign full and/or partial credit for each positive class.
+            if labels[i, j]:
+                for k in range(num_classes):
+                    if outputs[i, k]:
+                        A[j, k] += 1.0/normalization
+
+    return A
+
+# Compute the evaluation metric for the Challenge.
+def compute_challenge_metric(weights, labels, outputs, classes, normal_class):
+    num_recordings, num_classes = np.shape(labels)
+    normal_index = classes.index(normal_class)
+
+    # Compute the observed score.
+    A = compute_modified_confusion_matrix(labels, outputs)
+    observed_score = np.nansum(weights * A)
+
+    # Compute the score for the model that always chooses the correct label(s).
+    correct_outputs = labels
+    A = compute_modified_confusion_matrix(labels, correct_outputs)
+    correct_score = np.nansum(weights * A)
+
+    # Compute the score for the model that always chooses the normal class.
+    inactive_outputs = np.zeros((num_recordings, num_classes), dtype=np.bool_)
+    inactive_outputs[:, normal_index] = 1
+    A = compute_modified_confusion_matrix(labels, inactive_outputs)
+    inactive_score = np.nansum(weights * A)
+
+    if correct_score != inactive_score:
+        normalized_score = float(observed_score - inactive_score) / float(correct_score - inactive_score)
+    else:
+        normalized_score = float('nan')
+
+    return normalized_score
+
+if __name__ == '__main__':
+    #auroc, auprc, accuracy, f_measure, f_beta_measure, g_beta_measure, challenge_metric = evaluate_12ECG_score(sys.argv[1], sys.argv[2])
+    lbl_dir = '/home/p2017-999/acs_data/processed_data/physionet2020/jonathan/in'
+    out_dir = '/home/p2017-999/acs_data/processed_data/physionet2020/jonathan/out'
+    auroc, auprc, accuracy, f_measure, f_beta_measure, g_beta_measure, challenge_metric = evaluate_12ECG_score(lbl_dir, out_dir)
+
+    output_string = 'AUROC,AUPRC,Accuracy,F-measure,Fbeta-measure,Gbeta-measure,Challenge metric\n{:.3f},{:.3f},{:.3f},{:.3f},{:.3f},{:.3f},{:.3f}'.format(auroc, auprc, accuracy, f_measure, f_beta_measure, g_beta_measure, challenge_metric)
+    if len(sys.argv) > 3:
+        with open(sys.argv[3], 'w') as f:
+            f.write(output_string)
+    else:
+        print(output_string)

evaluate_model.py

@@ -0,0 +1,434 @@
+#!/usr/bin/env python
+
+# This file contains functions for evaluating algorithms for the 2021 PhysioNet/
+# Computing in Cardiology Challenge. You can run it as follows:
+#
+#   python evaluate_model.py labels outputs scores.csv
+#
+# where 'labels' is a directory containing files with the labels, 'outputs' is a
+# directory containing files with the outputs from your model, and 'scores.csv'
+# (optional) is a collection of scores for the algorithm outputs.
+#
+# Each file of labels or outputs must have the format described on the Challenge
+# webpage. The scores for the algorithm outputs include the area under the
+# receiver-operating characteristic curve (AUROC), the area under the recall-
+# precision curve (AUPRC), accuracy (fraction of correct recordings), macro F-
+# measure, and the Challenge metric, which assigns different weights to
+# different misclassification errors.
+
+import os, os.path, sys, numpy as np
+from helper_code import get_labels, is_finite_number, load_header, load_outputs
+import pandas as pd
+from tabulate import tabulate
+
+def evaluate_model(label_directory, output_directory):
+    # Identify the weights and the SNOMED CT code for the sinus rhythm class.
+    weights_file = 'weights.csv'
+    sinus_rhythm = set(['426783006'])
+
+    # Load the scored classes and the weights for the Challenge metric.
+    print('Loading weights...')
+    classes, weights = load_weights(weights_file)
+
+    # Load the label and output files.
+    print('Loading label and output files...')
+    label_files, output_files = find_challenge_files(label_directory, output_directory)
+    labels = load_labels(label_files, classes)
+    binary_outputs, scalar_outputs = load_classifier_outputs(output_files, classes)
+
+    # Evaluate the model by comparing the labels and outputs.
+    print('Evaluating model...')
+
+    print('- AUROC and AUPRC...')
+    auroc, auprc, auroc_classes, auprc_classes = compute_auc(labels, scalar_outputs)
+
+    print('- Accuracy...')
+    accuracy = compute_accuracy(labels, binary_outputs)
+
+    print('- F-measure...')
+    f_measure, f_measure_classes = compute_f_measure(labels, binary_outputs)
+
+    print('- Challenge metric...')
+    challenge_metric = compute_challenge_metric(weights, labels, binary_outputs, classes, sinus_rhythm)
+
+    print('Done.')
+
+    # Return the results.
+    return classes, auroc, auprc, auroc_classes, auprc_classes, accuracy, f_measure, f_measure_classes, challenge_metric
+
+# Find Challenge files.
+def find_challenge_files(label_directory, output_directory):
+    label_files = list()
+    output_files = list()
+    for label_file in sorted(os.listdir(label_directory)):
+        label_file_path = os.path.join(label_directory, label_file) # Full path for label file
+        if os.path.isfile(label_file_path) and label_file.lower().endswith('.hea') and not label_file.lower().startswith('.'):
+            root, ext = os.path.splitext(label_file)
+            output_file = root + '.csv'
+            output_file_path = os.path.join(output_directory, output_file) # Full path for corresponding output file
+            if os.path.isfile(output_file_path):
+                label_files.append(label_file_path)
+                output_files.append(output_file_path)
+            else:
+                raise IOError('Output file {} not found for label file {}.'.format(output_file, label_file))
+
+    if label_files and output_files:
+        return label_files, output_files
+    else:
+        raise IOError('No label or output files found.')
+
+# Load a table with row and column names.
+def load_table(table_file):
+    # The table should have the following form:
+    #
+    # ,    a,   b,   c
+    # a, 1.2, 2.3, 3.4
+    # b, 4.5, 5.6, 6.7
+    # c, 7.8, 8.9, 9.0
+    #
+    table = list()
+    with open(table_file, 'r') as f:
+        for i, l in enumerate(f):
+            arrs = [arr.strip() for arr in l.split(',')]
+            table.append(arrs)
+
+    # Define the numbers of rows and columns and check for errors.
+    num_rows = len(table)-1
+    if num_rows<1:
+        raise Exception('The table {} is empty.'.format(table_file))
+    row_lengths = set(len(table[i])-1 for i in range(num_rows))
+    if len(row_lengths)!=1:
+        raise Exception('The table {} has rows with different lengths.'.format(table_file))
+    num_cols = min(row_lengths)
+    if num_cols<1:
+        raise Exception('The table {} is empty.'.format(table_file))
+
+    # Find the row and column labels.
+    rows = [table[0][j+1] for j in range(num_rows)]
+    cols = [table[i+1][0] for i in range(num_cols)]
+
+    # Find the entries of the table.
+    values = np.zeros((num_rows, num_cols), dtype=np.float64)
+    for i in range(num_rows):
+        for j in range(num_cols):
+            value = table[i+1][j+1]
+            if is_finite_number(value):
+                values[i, j] = float(value)
+            else:
+                values[i, j] = float('nan')
+
+    return rows, cols, values
+
+# Load weights.
+def load_weights(weight_file):
+    # Load the table with the weight matrix.
+    rows, cols, values = load_table(weight_file)
+
+    # Split the equivalent classes.
+    rows = [set(row.split('|')) for row in rows]
+    cols = [set(col.split('|')) for col in cols]
+    assert(rows == cols)
+
+    # Identify the classes and the weight matrix.
+    classes = rows
+    weights = values
+
+    return classes, weights
+
+# Load labels from header/label files.
+def load_labels(label_files, classes):
+    # The labels should have the following form:
+    #
+    # Dx: label_1, label_2, label_3
+    #
+    num_recordings = len(label_files)
+    num_classes = len(classes)
+
+    # Use one-hot encoding for the labels.
+    labels = np.zeros((num_recordings, num_classes), dtype=np.bool)
+
+    # Iterate over the recordings.
+    for i in range(num_recordings):
+        header = load_header(label_files[i])
+        y = set(get_labels(header))
+        for j, x in enumerate(classes):
+            if x & y:
+                labels[i, j] = 1
+
+    return labels
+
+# Load outputs from output files.
+def load_classifier_outputs(output_files, classes):
+    # The outputs should have the following form:
+    #
+    # #Record ID
+    # diagnosis_1, diagnosis_2, diagnosis_3
+    #           0,           1,           1
+    #        0.12,        0.34,        0.56
+    #
+    num_recordings = len(output_files)
+    num_classes = len(classes)
+
+    # Use one-hot encoding for the outputs.
+    binary_outputs = np.zeros((num_recordings, num_classes), dtype=np.bool)
+    scalar_outputs = np.zeros((num_recordings, num_classes), dtype=np.float64)
+
+    # Iterate over the recordings.
+    for i in range(num_recordings):
+        recording_id, recording_classes, recording_binary_outputs, recording_scalar_outputs = load_outputs(output_files[i])
+
+        # Allow for equivalent classes and sanitize classifier outputs.
+        recording_classes = [set(entry.split('|')) for entry in recording_classes]
+        recording_binary_outputs = [1 if entry in ('1', 'True', 'true', 'T', 't') else 0 for entry in recording_binary_outputs]
+        recording_scalar_outputs = [float(entry) if is_finite_number(entry) else 0 for entry in recording_scalar_outputs]
+
+        # Allow for unordered/reordered and equivalent classes.
+        for j, x in enumerate(classes):
+            binary_values = list()
+            scalar_values = list()
+            for k, y in enumerate(recording_classes):
+                if x & y:
+                    binary_values.append(recording_binary_outputs[k])
+                    scalar_values.append(recording_scalar_outputs[k])
+            if binary_values:
+                binary_outputs[i, j] = any(binary_values) # Define a class as positive if any of the equivalent classes is positive.
+            if scalar_values:
+                scalar_outputs[i, j] = np.mean(scalar_values) # Define the scalar value of a class as the mean value of the scalar values across equivalent classes.
+
+    return binary_outputs, scalar_outputs
+
+# Compute recording-wise accuracy.
+def compute_accuracy(labels, outputs):
+    num_recordings, num_classes = np.shape(labels)
+
+    num_correct_recordings = 0
+    for i in range(num_recordings):
+        if np.all(labels[i, :]==outputs[i, :]):
+            num_correct_recordings += 1
+
+    return float(num_correct_recordings) / float(num_recordings)
+
+# Compute confusion matrices.
+def compute_confusion_matrices(labels, outputs, normalize=False):
+    # Compute a binary confusion matrix for each class k:
+    #
+    #     [TN_k FN_k]
+    #     [FP_k TP_k]
+    #
+    # If the normalize variable is set to true, then normalize the contributions
+    # to the confusion matrix by the number of labels per recording.
+    num_recordings, num_classes = np.shape(labels)
+
+    if not normalize:
+        A = np.zeros((num_classes, 2, 2))
+        for i in range(num_recordings):
+            for j in range(num_classes):
+                if labels[i, j]==1 and outputs[i, j]==1: # TP
+                    A[j, 1, 1] += 1
+                elif labels[i, j]==0 and outputs[i, j]==1: # FP
+                    A[j, 1, 0] += 1
+                elif labels[i, j]==1 and outputs[i, j]==0: # FN
+                    A[j, 0, 1] += 1
+                elif labels[i, j]==0 and outputs[i, j]==0: # TN
+                    A[j, 0, 0] += 1
+                else: # This condition should not happen.
+                    raise ValueError('Error in computing the confusion matrix.')
+    else:
+        A = np.zeros((num_classes, 2, 2))
+        for i in range(num_recordings):
+            normalization = float(max(np.sum(labels[i, :]), 1))
+            for j in range(num_classes):
+                if labels[i, j]==1 and outputs[i, j]==1: # TP
+                    A[j, 1, 1] += 1.0/normalization
+                elif labels[i, j]==0 and outputs[i, j]==1: # FP
+                    A[j, 1, 0] += 1.0/normalization
+                elif labels[i, j]==1 and outputs[i, j]==0: # FN
+                    A[j, 0, 1] += 1.0/normalization
+                elif labels[i, j]==0 and outputs[i, j]==0: # TN
+                    A[j, 0, 0] += 1.0/normalization
+                else: # This condition should not happen.
+                    raise ValueError('Error in computing the confusion matrix.')
+
+    return A
+
+# Compute macro F-measure.
+def compute_f_measure(labels, outputs):
+    num_recordings, num_classes = np.shape(labels)
+
+    A = compute_confusion_matrices(labels, outputs)
+
+    f_measure = np.zeros(num_classes)
+    for k in range(num_classes):
+        tp, fp, fn, tn = A[k, 1, 1], A[k, 1, 0], A[k, 0, 1], A[k, 0, 0]
+        if 2 * tp + fp + fn:
+            f_measure[k] = float(2 * tp) / float(2 * tp + fp + fn)
+        else:
+            f_measure[k] = float('nan')
+
+    if np.any(np.isfinite(f_measure)):
+        macro_f_measure = np.nanmean(f_measure)
+    else:
+        macro_f_measure = float('nan')
+
+    return macro_f_measure, f_measure
+
+# Compute macro AUROC and macro AUPRC.
+def compute_auc(labels, outputs):
+    num_recordings, num_classes = np.shape(labels)
+
+    # Compute and summarize the confusion matrices for each class across at distinct output values.
+    auroc = np.zeros(num_classes)
+    auprc = np.zeros(num_classes)
+
+    for k in range(num_classes):
+        # We only need to compute TPs, FPs, FNs, and TNs at distinct output values.
+        thresholds = np.unique(outputs[:, k])
+        thresholds = np.append(thresholds, thresholds[-1]+1)
+        thresholds = thresholds[::-1]
+        num_thresholds = len(thresholds)
+
+        # Initialize the TPs, FPs, FNs, and TNs.
+        tp = np.zeros(num_thresholds)
+        fp = np.zeros(num_thresholds)
+        fn = np.zeros(num_thresholds)
+        tn = np.zeros(num_thresholds)
+        fn[0] = np.sum(labels[:, k]==1)
+        tn[0] = np.sum(labels[:, k]==0)
+
+        # Find the indices that result in sorted output values.
+        idx = np.argsort(outputs[:, k])[::-1]
+
+        # Compute the TPs, FPs, FNs, and TNs for class k across thresholds.
+        i = 0
+        for j in range(1, num_thresholds):
+            # Initialize TPs, FPs, FNs, and TNs using values at previous threshold.
+            tp[j] = tp[j-1]
+            fp[j] = fp[j-1]
+            fn[j] = fn[j-1]
+            tn[j] = tn[j-1]
+
+            # Update the TPs, FPs, FNs, and TNs at i-th output value.
+            while i < num_recordings and outputs[idx[i], k] >= thresholds[j]:
+                if labels[idx[i], k]:
+                    tp[j] += 1
+                    fn[j] -= 1
+                else:
+                    fp[j] += 1
+                    tn[j] -= 1
+                i += 1
+
+        # Summarize the TPs, FPs, FNs, and TNs for class k.
+        tpr = np.zeros(num_thresholds)
+        tnr = np.zeros(num_thresholds)
+        ppv = np.zeros(num_thresholds)
+        for j in range(num_thresholds):
+            if tp[j] + fn[j]:
+                tpr[j] = float(tp[j]) / float(tp[j] + fn[j])
+            else:
+                tpr[j] = float('nan')
+            if fp[j] + tn[j]:
+                tnr[j] = float(tn[j]) / float(fp[j] + tn[j])
+            else:
+                tnr[j] = float('nan')
+            if tp[j] + fp[j]:
+                ppv[j] = float(tp[j]) / float(tp[j] + fp[j])
+            else:
+                ppv[j] = float('nan')
+
+        # Compute AUROC as the area under a piecewise linear function with TPR/
+        # sensitivity (x-axis) and TNR/specificity (y-axis) and AUPRC as the area
+        # under a piecewise constant with TPR/recall (x-axis) and PPV/precision
+        # (y-axis) for class k.
+        for j in range(num_thresholds-1):
+            auroc[k] += 0.5 * (tpr[j+1] - tpr[j]) * (tnr[j+1] + tnr[j])
+            auprc[k] += (tpr[j+1] - tpr[j]) * ppv[j+1]
+
+    # Compute macro AUROC and macro AUPRC across classes.
+    if np.any(np.isfinite(auroc)):
+        macro_auroc = np.nanmean(auroc)
+    else:
+        macro_auroc = float('nan')
+    if np.any(np.isfinite(auprc)):
+        macro_auprc = np.nanmean(auprc)
+    else:
+        macro_auprc = float('nan')
+
+    return macro_auroc, macro_auprc, auroc, auprc
+
+# Compute a modified confusion matrix for multi-class, multi-label tasks.
+def compute_modified_confusion_matrix(labels, outputs):
+    # Compute a binary multi-class, multi-label confusion matrix, where the rows
+    # are the labels and the columns are the outputs.
+    num_recordings, num_classes = np.shape(labels)
+    A = np.zeros((num_classes, num_classes))
+
+    # Iterate over all of the recordings.
+    for i in range(num_recordings):
+        # Calculate the number of positive labels and/or outputs.
+        normalization = float(max(np.sum(np.any((labels[i, :], outputs[i, :]), axis=0)), 1))
+        # Iterate over all of the classes.
+        for j in range(num_classes):
+            # Assign full and/or partial credit for each positive class.
+            if labels[i, j]:
+                for k in range(num_classes):
+                    if outputs[i, k]:
+                        A[j, k] += 1.0/normalization
+
+    return A
+
+# Compute the evaluation metric for the Challenge.
+def compute_challenge_metric(weights, labels, outputs, classes, sinus_rhythm):
+    num_recordings, num_classes = np.shape(labels)
+    if sinus_rhythm in classes:
+        sinus_rhythm_index = classes.index(sinus_rhythm)
+    else:
+        raise ValueError('The sinus rhythm class is not available.')
+
+    # Compute the observed score.
+    A = compute_modified_confusion_matrix(labels, outputs)
+    observed_score = np.nansum(weights * A)
+
+    # Compute the score for the model that always chooses the correct label(s).
+    correct_outputs = labels
+    A = compute_modified_confusion_matrix(labels, correct_outputs)
+    correct_score = np.nansum(weights * A)
+
+    # Compute the score for the model that always chooses the sinus rhythm class.
+    inactive_outputs = np.zeros((num_recordings, num_classes), dtype=np.bool_)
+    inactive_outputs[:, sinus_rhythm_index] = 1
+    A = compute_modified_confusion_matrix(labels, inactive_outputs)
+    inactive_score = np.nansum(weights * A)
+
+    if correct_score != inactive_score:
+        normalized_score = float(observed_score - inactive_score) / float(correct_score - inactive_score)
+    else:
+        normalized_score = 0.0
+
+    return normalized_score
+
+if __name__ == '__main__':
+    classes, auroc, auprc, auroc_classes, auprc_classes, accuracy, f_measure, f_measure_classes, challenge_metric = evaluate_model(sys.argv[1], sys.argv[2])
+    output_string = 'AUROC,AUPRC,Accuracy,F-measure,Challenge metric\n{:.3f},{:.3f},{:.3f},{:.3f},{:.3f}'.format(auroc, auprc, accuracy, f_measure, challenge_metric)
+    class_output_string = 'Classes,{}\nAUROC,{}\nAUPRC,{}\nF-measure,{}'.format(
+        ','.join('|'.join(sorted(x)) for x in classes),
+        ','.join('{:.3f}'.format(x) for x in auroc_classes),
+        ','.join('{:.3f}'.format(x) for x in auprc_classes),
+        ','.join('{:.3f}'.format(x) for x in f_measure_classes))
+
+    print(output_string)
+    with open('test_outputs/output.txt', 'w') as f:
+        f.write(output_string)
+
+
+    df = pd.DataFrame({'classes':classes,
+                       'auroc_classes':auroc_classes,
+                       'auprc_classes':auprc_classes,
+                       'f_measure_classes':f_measure_classes})
+    # dxname = pd.read_csv('dx_mapping_scored.csv')
+    # dxname = dxname[['Dx','SNOMED CT Code']]
+    # dxname = dxname.rename(columns={'SNOMED CT Code':'classes'})
+    # df = df.merge(dxname,on='classes',how='left')
+    print(tabulate(df, headers='keys', tablefmt='psql'))
+
+    with open('test_outputs/table.txt', 'w') as f:
+        f.write(tabulate(df, headers='keys', tablefmt='psql'))

helper.ipynb

@@ -0,0 +1,182 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### uniform stochastic depth"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "24\n",
+      "25\n",
+      "[0.0, 0.004347826354205608, 0.008695652708411217, 0.013043479062616825, 0.017391305416822433, 0.021739132702350616, 0.02608695812523365, 0.030434783548116684, 0.03478261083364487, 0.03913043811917305, 0.04347826540470123, 0.04782608896493912, 0.052173912525177, 0.056521736085414886, 0.06086956337094307, 0.06521739065647125, 0.06956521421670914, 0.07391304522752762, 0.0782608687877655, 0.08260869979858398, 0.08695652335882187, 0.09130434691905975, 0.09565217792987823, 0.10000000149011612]\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "\n",
+    "drop_path_rate = 0.1\n",
+    "depth = 24\n",
+    "\n",
+    "dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]\n",
+    "\n",
+    "print(len(dpr))\n",
+    "inter_dpr = [0.0] + dpr\n",
+    "print(len(inter_dpr))\n",
+    "print(dpr)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]\n",
+      "24\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "\n",
+    "drop_path_rate = 0\n",
+    "depth = 24\n",
+    "\n",
+    "dpr = [x.item() for x in torch.full((depth,), drop_path_rate)]\n",
+    "print(dpr)\n",
+    "print(len(dpr))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([[[ 0,  1,  2],\n",
+      "         [ 3,  4,  5],\n",
+      "         [ 6,  7,  8],\n",
+      "         [ 9, 10, 11],\n",
+      "         [12, 13, 14]],\n",
+      "\n",
+      "        [[15, 16, 17],\n",
+      "         [18, 19, 20],\n",
+      "         [21, 22, 23],\n",
+      "         [24, 25, 26],\n",
+      "         [27, 28, 29]]])\n",
+      "-------------------------\n",
+      "tensor([[[ 2,  1,  0],\n",
+      "         [ 5,  4,  3],\n",
+      "         [ 8,  7,  6],\n",
+      "         [11, 10,  9],\n",
+      "         [14, 13, 12]],\n",
+      "\n",
+      "        [[17, 16, 15],\n",
+      "         [20, 19, 18],\n",
+      "         [23, 22, 21],\n",
+      "         [26, 25, 24],\n",
+      "         [29, 28, 27]]])\n",
+      "-------------------------\n",
+      "tensor([[[12, 13, 14],\n",
+      "         [ 9, 10, 11],\n",
+      "         [ 6,  7,  8],\n",
+      "         [ 3,  4,  5],\n",
+      "         [ 0,  1,  2]],\n",
+      "\n",
+      "        [[27, 28, 29],\n",
+      "         [24, 25, 26],\n",
+      "         [21, 22, 23],\n",
+      "         [18, 19, 20],\n",
+      "         [15, 16, 17]]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "import torch\n",
+    "\n",
+    "x = torch.arange(30).view(2, 5, 3)\n",
+    "print(x)\n",
+    "print(\"-------------------------\")\n",
+    "\n",
+    "y = x.flip([-1])\n",
+    "print(y)\n",
+    "print(\"-------------------------\")\n",
+    "\n",
+    "z = x.flip([1])\n",
+    "print(z)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[0, 2, 3, 1]\n",
+      "This is the i: 0\n",
+      "This is the idx: 0\n",
+      "This is the i: 1\n",
+      "This is the idx: 2\n",
+      "This is the i: 2\n",
+      "This is the idx: 3\n",
+      "This is the i: 3\n",
+      "This is the idx: 1\n"
+     ]
+    }
+   ],
+   "source": [
+    "import random\n",
+    "\n",
+    "indices = [0, 1, 2, 3]\n",
+    "random.shuffle(indices)\n",
+    "\n",
+    "print(indices)\n",
+    "\n",
+    "for i, idx in enumerate(indices):\n",
+    "    print(\"This is the i:\", i)\n",
+    "    print(\"This is the idx:\", idx)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "mamba",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

helper_code.py

@@ -0,0 +1,241 @@
+#!/usr/bin/env python
+
+# Do *not* edit this script.
+# These are helper functions that you can use with your code.
+
+import os, numpy as np
+
+# Check if a variable is a number or represents a number.
+def is_number(x):
+    try:
+        float(x)
+        return True
+    except (ValueError, TypeError):
+        return False
+
+# Check if a variable is an integer or represents an integer.
+def is_integer(x):
+    if is_number(x):
+        return float(x).is_integer()
+    else:
+        return False
+
+# Check if a variable is a a finite number or represents a finite number.
+def is_finite_number(x):
+    if is_number(x):
+        return np.isfinite(float(x))
+    else:
+        return False
+
+# (Re)sort leads using the standard order of leads for the standard twelve-lead ECG.
+def sort_leads(leads):
+    x = ('I', 'II', 'III', 'aVR', 'aVL', 'aVF', 'V1', 'V2', 'V3', 'V4', 'V5', 'V6')
+    leads = sorted(leads, key=lambda lead: (x.index(lead) if lead in x else len(x) + leads.index(lead)))
+    return tuple(leads)
+
+# Find header and recording files.
+def find_challenge_files(data_directory):
+    header_files = list()
+    recording_files = list()
+    for f in os.listdir(data_directory):
+        root, extension = os.path.splitext(f)
+        if not root.startswith('.') and extension=='.hea':
+            header_file = os.path.join(data_directory, root + '.hea')
+            recording_file = os.path.join(data_directory, root + '.mat')
+            if os.path.isfile(header_file) and os.path.isfile(recording_file):
+                header_files.append(header_file)
+                recording_files.append(recording_file)
+    return header_files, recording_files
+
+# Load header file as a string.
+def load_header(header_file):
+    with open(header_file, 'r') as f:
+        header = f.read()
+    return header
+
+# Load recording file as an array.
+def load_recording(recording_file, header=None, leads=None, key='val'):
+    from scipy.io import loadmat
+    recording = loadmat(recording_file)[key]
+    if header and leads:
+        recording = choose_leads(recording, header, leads)
+    return recording
+
+# Choose leads from the recording file.
+def choose_leads(recording, header, leads):
+    num_leads = len(leads)
+    num_samples = np.shape(recording)[1]
+    chosen_recording = np.zeros((num_leads, num_samples), recording.dtype)
+    available_leads = get_leads(header)
+    for i, lead in enumerate(leads):
+        if lead in available_leads:
+            j = available_leads.index(lead)
+            chosen_recording[i, :] = recording[j, :]
+    return chosen_recording
+
+# Get recording ID.
+def get_recording_id(header):
+    recording_id = None
+    for i, l in enumerate(header.split('\n')):
+        if i==0:
+            try:
+                recording_id = l.split(' ')[0]
+            except:
+                pass
+        else:
+            break
+    return recording_id
+
+# Get leads from header.
+def get_leads(header):
+    leads = list()
+    for i, l in enumerate(header.split('\n')):
+        entries = l.split(' ')
+        if i==0:
+            num_leads = int(entries[1])
+        elif i<=num_leads:
+            leads.append(entries[-1])
+        else:
+            break
+    return tuple(leads)
+
+# Get age from header.
+def get_age(header):
+    age = None
+    for l in header.split('\n'):
+        if l.startswith('# Age'):
+            try:
+                age = float(l.split(': ')[1].strip())
+            except:
+                age = float('nan')
+    return age
+
+# Get sex from header.
+def get_sex(header):
+    sex = None
+    for l in header.split('\n'):
+        if l.startswith('# Sex'):
+            try:
+                sex = l.split(': ')[1].strip()
+            except:
+                pass
+    return sex
+
+# Get frequency from header.
+def get_num_leads(header):
+    num_leads = None
+    for i, l in enumerate(header.split('\n')):
+        if i==0:
+            try:
+                num_samples = float(l.split(' ')[1])
+            except:
+                pass
+        else:
+            break
+    return num_leads
+
+# Get frequency from header.
+def get_frequency(header):
+    frequency = None
+    for i, l in enumerate(header.split('\n')):
+        if i==0:
+            try:
+                frequency = float(l.split(' ')[2])
+            except:
+                pass
+        else:
+            break
+    return frequency
+
+# Get number of samples from header.
+def get_num_samples(header):
+    num_samples = None
+    for i, l in enumerate(header.split('\n')):
+        if i==0:
+            try:
+                num_samples = float(l.split(' ')[3])
+            except:
+                pass
+        else:
+            break
+    return num_samples
+
+# Get analog-to-digital converter (ADC) gains from header.
+def get_adc_gains(header, leads):
+    adc_gains = np.zeros(len(leads))
+    for i, l in enumerate(header.split('\n')):
+        entries = l.split(' ')
+        if i==0:
+            num_leads = int(entries[1])
+        elif i<=num_leads:
+            current_lead = entries[-1]
+            if current_lead in leads:
+                j = leads.index(current_lead)
+                try:
+                    adc_gains[j] = float(entries[2].split('/')[0])
+                except:
+                    pass
+        else:
+            break
+    return adc_gains
+
+# Get baselines from header.
+def get_baselines(header, leads):
+    baselines = np.zeros(len(leads))
+    for i, l in enumerate(header.split('\n')):
+        entries = l.split(' ')
+        if i==0:
+            num_leads = int(entries[1])
+        elif i<=num_leads:
+            current_lead = entries[-1]
+            if current_lead in leads:
+                j = leads.index(current_lead)
+                try:
+                    baselines[j] = float(entries[4].split('/')[0])
+                except:
+                    pass
+        else:
+            break
+    return baselines
+
+# Get labels from header.
+def get_labels(header):
+    labels = list()
+    for l in header.split('\n'):
+        if l.startswith('# Dx'):
+            try:
+                entries = l.split(': ')[1].split(',')
+                for entry in entries:
+                    labels.append(entry.strip())
+            except:
+                pass
+    return labels
+
+# Save outputs from model.
+def save_outputs(output_file, recording_id, classes, labels, probabilities):
+    # Format the model outputs.
+    recording_string = '#{}'.format(recording_id)
+    class_string = ','.join(str(c) for c in classes)
+    label_string = ','.join(str(l) for l in labels)
+    probabilities_string = ','.join(str(p) for p in probabilities)
+    output_string = recording_string + '\n' + class_string + '\n' + label_string + '\n' + probabilities_string + '\n'
+
+    # Save the model outputs.
+    with open(output_file, 'w') as f:
+        f.write(output_string)
+
+# Load outputs from model.
+def load_outputs(output_file):
+    with open(output_file, 'r') as f:
+        for i, l in enumerate(f):
+            if i==0:
+                recording_id = l[1:] if len(l)>1 else None
+            elif i==1:
+                classes = tuple(entry.strip() for entry in l.split(','))
+            elif i==2:
+                labels = tuple(entry.strip() for entry in l.split(','))
+            elif i==3:
+                probabilities = tuple(float(entry) if is_finite_number(entry) else float('nan') for entry in l.split(','))
+            else:
+                break
+    return recording_id, classes, labels, probabilities

losses.py

@@ -0,0 +1,70 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+"""
+Implements the knowledge distillation loss
+"""
+import torch
+from torch.nn import functional as F
+
+
+class DistillationLoss(torch.nn.Module):
+    """
+    This module wraps a standard criterion and adds an extra knowledge distillation loss by
+    taking a teacher model prediction and using it as additional supervision.
+    """
+    def __init__(self, base_criterion: torch.nn.Module, teacher_model: torch.nn.Module,
+                 distillation_type: str, alpha: float, tau: float):
+        super().__init__()
+        self.base_criterion = base_criterion
+        self.teacher_model = teacher_model
+        assert distillation_type in ['none', 'soft', 'hard']
+        self.distillation_type = distillation_type
+        self.alpha = alpha
+        self.tau = tau
+
+    def forward(self, inputs, outputs, labels):
+        """
+        Args:
+            inputs: The original inputs that are feed to the teacher model
+            outputs: the outputs of the model to be trained. It is expected to be
+                either a Tensor, or a Tuple[Tensor, Tensor], with the original output
+                in the first position and the distillation predictions as the second output
+            labels: the labels for the base criterion
+        """
+        outputs_kd = None
+        if not isinstance(outputs, torch.Tensor):
+            # assume that the model outputs a tuple of [outputs, outputs_kd]
+            outputs, outputs_kd = outputs
+        base_loss = self.base_criterion(outputs, labels)
+        if self.distillation_type == 'none':
+            return base_loss
+
+        if outputs_kd is None:
+            raise ValueError("When knowledge distillation is enabled, the model is "
+                             "expected to return a Tuple[Tensor, Tensor] with the output of the "
+                             "class_token and the dist_token")
+        # don't backprop throught the teacher
+        with torch.no_grad():
+            teacher_outputs = self.teacher_model(inputs)
+
+        if self.distillation_type == 'soft':
+            T = self.tau
+            # taken from https://github.com/peterliht/knowledge-distillation-pytorch/blob/master/model/net.py#L100
+            # with slight modifications
+            distillation_loss = F.kl_div(
+                F.log_softmax(outputs_kd / T, dim=1),
+                #We provide the teacher's targets in log probability because we use log_target=True 
+                #(as recommended in pytorch https://github.com/pytorch/pytorch/blob/9324181d0ac7b4f7949a574dbc3e8be30abe7041/torch/nn/functional.py#L2719)
+                #but it is possible to give just the probabilities and set log_target=False. In our experiments we tried both.
+                F.log_softmax(teacher_outputs / T, dim=1),
+                reduction='sum',
+                log_target=True
+            ) * (T * T) / outputs_kd.numel()
+            #We divide by outputs_kd.numel() to have the legacy PyTorch behavior. 
+            #But we also experiments output_kd.size(0) 
+            #see issue 61(https://github.com/facebookresearch/deit/issues/61) for more details
+        elif self.distillation_type == 'hard':
+            distillation_loss = F.cross_entropy(outputs_kd, teacher_outputs.argmax(dim=1))
+
+        loss = base_loss * (1 - self.alpha) + distillation_loss * self.alpha
+        return loss

main_ecg.py

@@ -0,0 +1,509 @@
+import argparse
+from pathlib import Path
+import utils
+import torch
+import numpy as np
+import torch.backends.cudnn as cudnn
+import pandas as pd
+import time
+import datetime
+import os
+
+# mixup
+from timm.data import Mixup
+
+
+# log about
+# import mlflow
+
+# for the challenge 2021
+import ecg_dataset_2021
+import engine_ecg_2021
+
+# for the challenge 2020
+import ecg_dataset_2020
+import engine_ecg_2020
+
+from engine_ecg_2021 import train_one_epoch, evaluate
+
+# for the timm
+from timm.models import create_model
+import timm
+from timm.scheduler.cosine_lr import CosineLRScheduler
+
+import sys
+
+from optimizer import NoamOpt
+
+import models_mamba_ecg
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('DeiT training and evaluation script', add_help=False)
+    parser.add_argument('--batch-size', default=64, type=int)
+    parser.add_argument('--epochs', default=300, type=int)
+
+    # Model parameters
+    parser.add_argument('--model', default='deit_base_patch16_224', type=str, metavar='MODEL', help='Name of model to train')
+
+    # Parameters regarding the Drop
+    parser.add_argument('--drop', type=float, default=0.0, metavar='PCT', help='Dropout rate (default: 0.)')
+    parser.add_argument('--drop-path', type=float, default=0.1, metavar='PCT', help='Drop path rate (default: 0.1)')
+
+    # setting the mode
+    parser.add_argument('--train-mode', action='store_true')
+    parser.add_argument('--no-train-mode', action='store_false', dest='train_mode')
+    parser.set_defaults(train_mode=True)
+    
+    
+    # * data augmentation method
+    parser.add_argument('--mixup', type=float, default=0,
+                        help='mixup alpha, mixup enabled if > 0. (default: 0)')
+    parser.add_argument('--cutmix', type=float, default=0,
+                        help='cutmix alpha, cutmix enabled if > 0. (default: 0)')
+    parser.add_argument('--mixup_no_label', type=float, default=0,
+                        help='this is the mixup but without interpolation of label. (default: 0)')
+    parser.add_argument('--progressive_switch', type=bool, default=False,
+                        help='this is the switch of progressive data augmentation method')
+    
+    # the output directory
+    parser.add_argument('--output_dir', default='', help='path where to save, empty for no saving')
+
+    # the device information
+    parser.add_argument('--device', default='cuda', help='device to use for training / testing')
+
+    parser.add_argument('--seed', default=0, type=int)
+    parser.add_argument('--start_epoch', default=0, type=int, metavar='N', help='start epoch')
+
+
+    parser.add_argument('--num_workers', default=10, type=int)
+    
+    parser.add_argument('--pin-mem', action='store_true', help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
+    parser.add_argument('--no-pin-mem', action='store_false', dest='pin_mem', help='')
+    parser.set_defaults(pin_mem=True)
+    
+    # special parameters for this strategy
+    parser.add_argument('--depth', default=24, type=int)
+    parser.add_argument('--block', default='VisionMamba', choices=['OriginalMamba', 'VisionMamba'], type=str, help='the selection of the block of the paradigm')
+    parser.add_argument('--lead', default='12Lead', choices=['12Lead', 'RandomLead'], type=str, help='the number of Lead')
+    parser.add_argument('--lrschedule', default='CosineAnnealing', choices=['Noam', 'CosineAnnealing'], type=str, help='the selection of learning rate strategy')
+
+    # if use random token position
+    parser.add_argument('--if_random_cls_token_position', action='store_true')
+    parser.add_argument('--no_randargs.if_nan2numom_cls_token_position', action='store_false', dest='if_random_cls_token_position')
+    parser.set_defaults(if_random_cls_token_position=False)    
+
+    # if use random token rank
+    parser.add_argument('--if_random_token_rank', action='store_true')
+    parser.add_argument('--no_random_token_rank', action='store_false', dest='if_random_token_rank')
+    parser.set_defaults(if_random_token_rank=False)
+
+    # using for the journal 
+    parser.add_argument('--fused_add_norm', type=bool, default=True, help='combines the element-wise addition (from residual connections) and normalization (e.g., RMSNorm)')
+    parser.add_argument('--if_divide_out', type=bool, default=True, help='Should the result be divided by two when combining outputs from two directions?')
+    parser.add_argument('--use_middle_cls_token', type=bool, default=True, help='Whether the class is inserted in the middle of sequence')
+
+    # the switch of scenario
+    parser.add_argument('--challenge_scenario', default='2021', choices=[2021, 2020], type=int, help='the scenario of Challenge')
+
+
+    return parser
+
+
+def collate(batch):
+    # Left-zero padding
+    ch = batch[0][0].shape[0]
+    maxL = 8192
+    X = np.zeros((len(batch), ch, maxL))
+    
+    for i in range(len(batch)):
+        X[i, :, -batch[i][0].shape[-1]:] = batch[i][0]
+    
+    t = np.array([b[1] for b in batch])
+    
+    X = torch.from_numpy(X)
+    t = torch.from_numpy(t)
+
+    return X, t
+
+def main(args, data_directory, model_directory, group_number):
+    
+    train_log_fp = open(args.output_dir + '/train_log_group_%d.txt' % group_number, 'a')
+    print(args)
+    train_log_fp.write("The is the configuration: {}\n".format(args))
+
+    
+    device = torch.device(args.device)
+    
+    print("This is the running device", device)
+    train_log_fp.write("This is the running device:{}\n".format(device))
+
+    seed = args.seed + utils.get_rank()
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+    print("This is the random's seed:", seed)
+    train_log_fp.write("This is the random's seed:{}\n".format(seed))
+    
+
+    cudnn.benchmark = True
+
+    
+    dataset_file_all_address = "../../collection_of_all_datasets/"
+    # below is building dataset(train and validation)
+
+    print("This is building the training set:")
+    ############ training area #########################
+    the_training_address = data_directory + "/training_group" + data_directory[-1] + ".csv"
+    df = pd.read_csv(the_training_address)
+    print("Total the {} files will be fed into model for training".format(len(df['Name'])))
+    train_log_fp.write("Total the {} files will be fed into model for training \n".format(len(df['Name'])))
+
+    print("This is the first file of training:",df['Name'][0])
+
+
+    training_header_files=[]
+    
+    for i in range(len(df['Name'])):
+        each_header_file = dataset_file_all_address + df['Name'][i]
+        training_header_files.append(each_header_file)
+    
+    collate_training = collate
+    
+    
+    if args.mixup > 0:
+        print("The mixup is using, and the percentage is:", args.mixup)
+        train_log_fp.write("The mixup is using, and the percentage is:{}\n".format(args.mixup))
+        # collate_training = collate_mixup
+    
+    if args.cutmix > 0:
+        print("The cutmix is using, and the percentage is:", args.cutmix)
+        train_log_fp.write("The cutmix is using, and the percentage is:{}\n".format(args.cutmix))
+    if args.mixup_no_label > 0:
+        print("The mixup is using without label's interpolation , and the percentage is:", args.mixup_no_label)
+        train_log_fp.write("The mixup is using without label's interpolation , and the percentage is:{}\n".format(args.mixup_no_label))
+    if args.progressive_switch:
+        print("The progressive mixup is using.")
+        train_log_fp.write("The progressive mixup is using.")
+
+    if args.challenge_scenario == 2021:
+        train_dataset = ecg_dataset_2021.dataset(training_header_files, Mixup = args.mixup, amount = len(df['Name']), cutMix=args.cutmix, Mixup_no_label_interpolate=args.mixup_no_label, progressive_switch=args.progressive_switch)
+    elif args.challenge_scenario == 2020:
+        train_dataset = ecg_dataset_2020.dataset(training_header_files, Mixup = args.mixup, amount = len(df['Name']), cutMix=args.cutmix, Mixup_no_label_interpolate=args.mixup_no_label, progressive_switch=args.progressive_switch)
+    else:
+        raise ValueError(f"No matching condition for value: {args.challenge_scenario}")
+
+
+    if args.lead == "12Lead":
+        lead_number = 12
+        print("This 12 lead is using.")
+        train_log_fp.write("This 12 lead is using.")
+    elif args.lead == "RandomLead":
+        lead_number = None
+        print("This random lead is using.")
+        train_log_fp.write("This random lead is using.")
+    else:
+        raise ValueError(f"No matching condition for value: {args.lead}")
+    
+    """
+    we filter out the sample which the length is 8192 via the below code.
+    just like the random shift windows.  
+    """
+    train_dataset.num_leads = lead_number
+    train_dataset.sample = True
+    ###################################################
+    print("done")
+    
+
+    print("This is building the validation set:")
+    ############ testing area #########################
+    the_testing_address = data_directory + "/testing_group" + data_directory[-1]+".csv"
+    df = pd.read_csv(the_testing_address)
+    print("Total the {} files will be used as testing".format(len(df['Name'])))
+    train_log_fp.write("Total the {} files will be used as testing \n".format(len(df['Name'])))
+
+    print("This is the first file of testing:",df['Name'][0])
+
+    testing_header_files=[]
+
+    for i in range(len(df['Name'])):
+        each_header_file = dataset_file_all_address + df['Name'][i]
+        testing_header_files.append(each_header_file)
+    
+
+    if args.challenge_scenario == 2021:
+        test_dataset = ecg_dataset_2021.dataset(testing_header_files, Mixup = 0)
+    elif args.challenge_scenario == 2020:
+        test_dataset = ecg_dataset_2020.dataset(testing_header_files, Mixup = 0)
+    else:
+        raise ValueError(f"No matching condition for value: {args.challenge_scenario}")
+
+
+    test_dataset.num_leads = 12
+    test_dataset.sample = True
+    ###################################################
+    print("done")
+
+    sampler_train = torch.utils.data.RandomSampler(train_dataset)
+    sampler_val = torch.utils.data.SequentialSampler(test_dataset)
+
+    data_loader_train = torch.utils.data.DataLoader(
+        train_dataset, sampler=sampler_train,
+        batch_size=args.batch_size,
+        collate_fn=collate_training,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=True,
+    )
+    # Setting `drop_last=True` means that this incomplete batch will be dropped, 
+    # ensuring that all batches fed to the model during training have **the same size.**
+    
+    data_loader_val = torch.utils.data.DataLoader(
+        test_dataset, sampler=sampler_val,
+        batch_size=int(1.5 * args.batch_size),
+        collate_fn=collate,
+        num_workers=args.num_workers,
+        pin_memory=args.pin_mem,
+        drop_last=False
+    )
+
+    if args.challenge_scenario == 2021:
+        args.nb_classes = 26
+    elif args.challenge_scenario == 2020:
+        args.nb_classes = 27
+    else:
+        raise ValueError(f"No matching condition for value: {args.challenge_scenario}")
+
+
+    print(f"Creating model: {args.model}")
+    train_log_fp.write("Creating model:{}\n".format(args.model))
+    model = create_model(
+        args.model,
+        pretrained=False,
+        num_classes=args.nb_classes,
+        drop_rate=args.drop,
+        drop_path_rate=args.drop_path,
+        drop_block_rate=None,
+        block = args.block,
+        depth = args.depth,
+        fused_add_norm = args.fused_add_norm,
+        use_middle_cls_token = args.use_middle_cls_token,
+        img_size=8192
+    )
+
+
+    model.to(device)
+
+
+    n_parameters = sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+    print('number of params:', n_parameters)
+    train_log_fp.write("number of params:{}\n".format(n_parameters))
+    
+    optimizer = torch.optim.Adam(model.parameters(), lr=0.0006, betas=(0.9, 0.98), eps=1e-9)
+
+    if args.lrschedule == "Noam":
+        optimizer = NoamOpt(729, 1, 4000, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))
+        print("The Noam is using as learning rate strategy.")
+        train_log_fp.write("The Noam is using as learning rate strategy.")
+        
+    elif args.lrschedule == "CosineAnnealing":
+        lr_scheduler = CosineLRScheduler(
+        optimizer,
+        t_initial=13,  # Number of epochs after warmup
+        lr_min=1e-6,
+        warmup_lr_init = 1e-5,
+        warmup_t=5,
+        cycle_limit=1,
+        t_in_epochs=True,
+        warmup_prefix=True  # Added to reach initial_lr
+        )
+        print("The cosing annealing is using as learning rate strategy.")
+        train_log_fp.write("The cosing annealing is using as learning rate strategy.")
+    else:
+        raise ValueError(f"No matching condition for value: {args.lrschedule}")
+
+    # below is for the cosine annealing schedule
+    
+    # the loss is set as normal BCE loss.
+    criterion = torch.nn.BCEWithLogitsLoss()
+
+
+    output_dir = Path(args.output_dir)
+
+    print(f"Start training for {args.epochs} epochs")
+    train_log_fp.write(f"Start training for {args.epochs} epochs\n")
+    
+    start_time = time.time()
+
+    max_AUPRC = 0.0
+
+
+    train_auprc_list = []
+    testing_AUPRC_list = []
+    loss_value_after_each_epcoh_list = []
+    loss_value_testing_list = []
+
+    testing_auroc_list = []
+    testing_f1_list = []
+
+    testing_subset_accuracy_list = []
+    hamming_loss_list=[]
+    challenge_score_list = []
+
+
+    jump_count = 0
+
+    thrs_list = []
+    # loss_value_after_each_epcoh_testing_list = []
+    # epoch is started from 0
+    for epoch in range(args.start_epoch, args.epochs):
+        
+        train_log_fp.write("\n")
+        train_log_fp.write("---------------------------------------------------------- \n")
+        train_log_fp.write("---------------------------------------------------------- \n")
+        train_log_fp.write("---------------------------------------------------------- \n")
+        
+        now = datetime.datetime.now()
+        current_time = now.strftime("%H:%M:%S")
+        print("Current Time =", current_time)
+        train_log_fp.write(f"Current Time = {current_time} \n")
+        train_log_fp.write(f"Current epoch: {epoch} \n")
+        
+        train_dataset.set_epoch(epoch)  # Update the current epoch in the dataset
+        
+        if args.lrschedule == "CosineAnnealing":
+            lr_scheduler.step(epoch)
+
+
+        ###### below is for the training ###########
+        if args.challenge_scenario == 2021:
+            train_auprc, loss_value_after_each_epcoh, thrs, scores_F1, scores_SubsetAccuracy, scores_HammingLoss = engine_ecg_2021.train_one_epoch(
+                model, criterion, data_loader_train,
+                optimizer, device, epoch,
+                set_training_mode=args.train_mode,  # keep in eval mode for deit finetuning / train mode for training and deit III finetuning
+                args=args,
+            )
+        elif args.challenge_scenario == 2020: 
+            train_auprc, loss_value_after_each_epcoh, thrs, scores_F1, scores_SubsetAccuracy, scores_HammingLoss = engine_ecg_2020.train_one_epoch(
+                model, criterion, data_loader_train,
+                optimizer, device, epoch,
+                set_training_mode=args.train_mode,  # keep in eval mode for deit finetuning / train mode for training and deit III finetuning
+                args=args,
+            )
+        else:
+            raise ValueError(f"No matching condition for value: {args.challenge_scenario}")
+        
+
+        train_auprc_list.append(round(train_auprc, 4))
+        loss_value_after_each_epcoh_list.append(round(loss_value_after_each_epcoh, 4))
+        print("**************************************************************")
+        print("This is the list of Training of AUPRC:", train_auprc_list)
+        train_log_fp.write(f"This is the list of Training of AUPRC: {train_auprc_list} \n")
+        print("This is the list of Training of loss:", loss_value_after_each_epcoh_list)
+        train_log_fp.write(f"This is the list of Training of loss: {loss_value_after_each_epcoh_list} \n")
+        print("**************************************************************")
+
+
+        ###### below is for the evaluation ###########
+
+        if args.challenge_scenario == 2021:
+            AUPRC, auroc, f1, hamming, subset_accuracy, challenge_score, loss_value_after_each_epoch_testing = engine_ecg_2021.evaluate(data_loader_val, model, thrs, scores_F1, scores_SubsetAccuracy, scores_HammingLoss, device)
+        elif args.challenge_scenario == 2020: 
+            AUPRC, auroc, f1, hamming, subset_accuracy, challenge_score, loss_value_after_each_epoch_testing = engine_ecg_2020.evaluate(data_loader_val, model, thrs, scores_F1, scores_SubsetAccuracy, scores_HammingLoss, device)
+        else:
+            raise ValueError(f"No matching condition for value: {args.challenge_scenario}")
+
+        # print(f"Accuracy of the network on the {len(test_dataset)} test images: {AUPRC:.4f}")
+
+        testing_AUPRC_list.append(round(AUPRC, 4))
+        loss_value_testing_list.append(round(loss_value_after_each_epoch_testing, 4))
+        testing_auroc_list.append(round(auroc, 4))
+        testing_f1_list.append(round(f1, 4))
+        testing_subset_accuracy_list.append((round(subset_accuracy, 4)))  
+        hamming_loss_list.append((round(hamming, 4)))
+        challenge_score_list.append((round(challenge_score, 4)))
+
+
+        # loss_value_after_each_epcoh_testing_list.append(loss)
+        print("**********************************************************")
+        print("This is the list of Testing AUPRC:", testing_AUPRC_list)
+        print("This is the list of Testing loss:", loss_value_testing_list)
+        print("This is the list of Testing auroc:", testing_auroc_list)
+        print("This is the list of Testing f1:", testing_f1_list)
+        print("This is the list of subset accuracy:", testing_subset_accuracy_list)
+        print("This is the list of hamming_loss:", hamming_loss_list)
+        print("This is the list of challenge_score:", challenge_score_list)
+
+        train_log_fp.write("---------------------------------------------------------- \n")
+        train_log_fp.write(f"This is the list of Testing AUPRC: {testing_AUPRC_list} \n")
+        train_log_fp.write(f"This is the list of Testing loss: {loss_value_testing_list} \n")
+        train_log_fp.write(f"This is the list of Testing auroc: {testing_auroc_list} \n")
+        train_log_fp.write(f"This is the list of Testing f1: {testing_f1_list} \n")
+        train_log_fp.write(f"This is the list of Testing subset accuracy: {testing_subset_accuracy_list} \n")
+        train_log_fp.write(f"This is the list of Testing hamming_loss: {hamming_loss_list} \n")
+        train_log_fp.write(f"This is the list of challenge_score: {challenge_score_list} \n")
+        print("**********************************************************")
+        
+        # here I have to put code regarding the auprc
+        if max_AUPRC < AUPRC:
+            max_AUPRC = AUPRC
+            jump_count = 0
+            if args.output_dir:
+                checkpoint_paths = [output_dir / 'best_auprc_checkpoint.pth']
+                for checkpoint_path in checkpoint_paths:
+                    utils.save_on_master({
+                        'model': model.state_dict(),
+                        'optimizer': optimizer.optimizer.state_dict(),
+                        'epoch': epoch,
+                        'args': args,
+                    }, checkpoint_path)
+        else:
+            jump_count = jump_count + 1
+        
+
+        print(f'This is the Max AUPRC: {max_AUPRC:.4f}')
+        train_log_fp.write(f"This is the Max AUPRC: {max_AUPRC:.4f} \n")
+
+        if jump_count > 4:
+            print("This experimental will be finished at epoch:", epoch)
+            train_log_fp.write(f"This experimental will be finished at epoch: {epoch} \n")
+            break
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print('Training time {}'.format(total_time_str))
+     
+    train_log_fp.write('Training time {}\n'.format(total_time_str))
+    train_log_fp.close()
+    old_txt_file_name = args.output_dir + '/train_log_group_%d.txt' % group_number
+    new_txt_file_name = args.output_dir + '/train_log_group_%d_MAX_AUPRC_%.4f_.txt' % (group_number, max_AUPRC)
+    os.rename(old_txt_file_name, new_txt_file_name)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser('DeiT training and evaluation script, but this is for the multiple classification ECG.', parents=[get_args_parser()])
+    args = parser.parse_args()
+    
+    if args.challenge_scenario == 2021: 
+        data_directory = "./csv-file_2021_challenge/training_validation_testing/group"
+    elif args.challenge_scenario == 2020:
+        data_directory = "./csv-file_2020_challenge/training_validation_testing/group"
+    else:
+        raise ValueError(f"No matching condition for value: {args.challenge_scenario}")
+
+    print("This is the scenario:", args.challenge_scenario)
+    print("This is the scenario:", args.challenge_scenario)
+
+    model_directory = "./model/model_group"
+
+    for i in range(1,6):
+        data_directory_x = data_directory + str(i)
+        model_directory_x = model_directory +str(i)
+
+        args.output_dir = f"./output/Scenario_{args.challenge_scenario}_{args.block}_depth_{args.depth}_{args.lead}_batchSize_{args.batch_size}_CutMix(random)_{args.cutmix}_MixUp_{args.mixup}_JTEHM_group{i}"
+
+        if args.output_dir:
+            Path(args.output_dir).mkdir(parents=True, exist_ok=True)
+
+        print("This is the group", i)
+        main(args, data_directory_x, model_directory_x, i)

models_mamba_ecg.py

@@ -0,0 +1,1013 @@
+# Copyright (c) 2015-present, Facebook, Inc.
+# All rights reserved.
+import torch
+import torch.nn as nn
+import torch.nn.functional as functional
+from functools import partial
+from torch import Tensor
+from typing import Optional
+
+from timm.models.vision_transformer import VisionTransformer, _cfg
+# from timm.models.registry import register_model
+# from timm.models.layers import trunc_normal_, lecun_normal_
+
+from timm.models import register_model
+from timm.layers import trunc_normal_, lecun_normal_
+
+from timm.layers import DropPath, to_2tuple
+
+# from timm.models.layers import DropPath, to_2tuple
+from timm.models.vision_transformer import _load_weights
+
+import math
+
+from collections import namedtuple
+
+from mamba_ssm.modules.mamba_simple import Mamba
+from mamba_ssm.utils.generation import GenerationMixin
+from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf
+
+from rope import *
+import random
+import sys
+
+try:
+    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
+except ImportError:
+    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
+
+
+# layer_norm_fn and rms_norm_fn both are normalization method
+
+__all__ = [
+    'vim_tiny_patch16_224', 'vim_small_patch16_224', 'vim_base_patch16_224',
+    'vim_tiny_patch16_384', 'vim_small_patch16_384', 'vim_base_patch16_384',
+]
+
+
+'''
+in the original script(ft-vim-s.sh)
+img_size = 224, 
+patch_size = 16, 
+stride = 8, 
+in_chans = 3, 
+embed_dim = 768
+------------------------------------
+self.img_size: (224, 224)
+self.patch_size: (16, 16)
+self.grid_size: (27, 27)
+self.num_patches: 729
+self.flatten: True
+self.norm: nn.Identity()
+'''
+class PatchEmbed(nn.Module):
+    """ 2D Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, stride=16, in_chans=3, embed_dim=768, norm_layer=None, flatten=True):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.grid_size = ((img_size[0] - patch_size[0]) // stride + 1, (img_size[1] - patch_size[1]) // stride + 1)
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride)
+        # if the norm_layer is not none or null, the self.norm = norm_layer(embed_dim)
+        # otherwise, self.norm = nn.Identity()
+        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+        
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        
+        x = self.proj(x)
+        print("This is the shape after the CNN", x.shape)
+
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+            print("This is the shape after the flatten:", x.shape)
+
+        x = self.norm(x)
+        return x
+
+class PatchEmbed_spectrogram(nn.Module):
+    """ 2D spectrogram to Patch Embedding
+    """
+    def __init__(self, img_size_f = 128, img_size_t = 64, patch_size=6, stride=3, in_chans=12, embed_dim=432, flatten=True):
+        super().__init__()
+        # img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+
+        self.img_size_f = img_size_f
+        self.img_size_t = img_size_t
+
+        self.patch_size = patch_size
+        self.grid_size = ((img_size_f - patch_size[0]) // stride + 1, (img_size_t - patch_size[1]) // stride + 1)
+        self.num_patches = self.grid_size[0] * self.grid_size[1]
+        self.flatten = flatten
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=stride)
+        # if the norm_layer is not none or null, the self.norm = norm_layer(embed_dim)
+        # otherwise, self.norm = nn.Identity()
+        # self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
+        
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert H == self.img_size_f and W == self.img_size_t, \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size_f}*{self.img_size_t})."
+        
+        x = self.proj(x)
+
+        # This is the shape after the CNN torch.Size([1, 432, 41, 20])
+        # print("This is the shape after the CNN", x.shape)
+        if self.flatten:
+            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
+        # This is the shape after the flatten: torch.Size([1, 820, 432])
+        # print("This is the shape after the flatten:", x.shape)
+        return x
+
+
+class CNN_layers(nn.Module):
+
+    def __init__(self, embed_size = 384):
+        super().__init__()
+        self.multiple_cnn = nn.Sequential( 
+            nn.Conv1d(12, 128, kernel_size=14, stride=3, padding=2, bias=False),
+            nn.BatchNorm1d(128),
+            nn.ReLU(inplace=True),
+
+            nn.Conv1d(128, embed_size, kernel_size=15, stride=4, padding=100, bias=False),
+            nn.BatchNorm1d(embed_size),
+            nn.ReLU(inplace=True))
+        
+    def forward(self, x):
+        # print("This is the shape of enconder:(before)", x.shape)
+        # This is the shape of enconder:(before) torch.Size([44, 12, 8192])
+        x = self.multiple_cnn(x)
+        x = x.transpose(1, 2)
+        # print("This is the shape of enconder:(after)", x.shape)
+        # This is the shape of enconder:(after) torch.Size([44, 729, 384])
+
+        # sys.exit()
+        return x 
+
+class CNN_layers_shortcut(nn.Module):
+
+    def __init__(self, embed_size = 384):
+        super().__init__()
+        self.multiple_cnn1 = nn.Sequential( 
+            nn.Conv1d(12, 128, kernel_size=14, stride=3, padding=2, bias=False),
+            nn.BatchNorm1d(128),
+            nn.ReLU(inplace=True),
+
+            nn.Conv1d(128, embed_size, kernel_size=15, stride=4, padding=100, bias=False),
+            nn.BatchNorm1d(embed_size),
+            nn.ReLU(inplace=True))
+
+        self.multiple_cnn2 = nn.Sequential(
+            nn.Conv1d(in_channels=embed_size, out_channels=embed_size, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm1d(embed_size),
+            nn.ReLU(inplace=True),
+
+            nn.Conv1d(in_channels=embed_size, out_channels=embed_size, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm1d(embed_size),
+            nn.ReLU(inplace=True)
+
+            )
+        
+    def forward(self, x):
+        # print("This is the shape of enconder:(before)", x.shape)
+        # This is the shape of enconder:(before) torch.Size([44, 12, 8192])
+        x = self.multiple_cnn1(x)
+
+        shortcut = x
+
+        x = self.multiple_cnn2(x)
+        x = x + shortcut
+
+        x = x.transpose(1, 2)
+        # print("This is the shape of enconder:(after)", x.shape)
+        # This is the shape of enconder:(after) torch.Size([44, 729, 384])
+        # sys.exit()
+        return x 
+
+# class ECG_Patch_embedding(nn.Module):
+
+#     def __init__(self, embed_size = 384):
+#         super().__init__()
+#         self.multiple_cnn = nn.Sequential( 
+#             nn.Conv1d(12, embed_size, kernel_size=16, stride=8),
+#             )
+        
+#     def forward(self, x):
+#         # print("This is the shape of enconder:(before)", x.shape)
+#         # This is the shape of enconder:(before) torch.Size([44, 12, 8192])
+#         x = self.multiple_cnn(x)
+#         x = x.transpose(1, 2)
+#         # print("This is the shape of enconder:(after)", x.shape)
+#         # This is the shape of enconder:(after) torch.Size([44, 1023, 384])
+
+#         # sys.exit()
+#         return x 
+
+# class LeadCombiner(nn.Module):
+#     def __init__(self, lead, out_ch):
+#         super(LeadCombiner, self).__init__()
+#         self.conv2_1 = nn.Conv1d(in_channels=lead * out_ch,
+#                                  out_channels=out_ch,
+#                                  kernel_size=1,
+#                                  bias=False)
+#         self.bn2_1 = nn.BatchNorm1d(out_ch)
+
+#         self.conv2_2 = nn.Conv1d(in_channels=lead * out_ch,
+#                                  out_channels=out_ch,
+#                                  kernel_size=1,
+#                                  bias=False)
+#         self.bn2_2 = nn.BatchNorm1d(out_ch)
+
+#         self.pool1 = nn.AdaptiveMaxPool1d(output_size=1)
+#         self.pool2 = nn.AdaptiveMaxPool1d(output_size=1)
+
+#     def forward(self, x):
+#         # this is the shape of x: torch.Size([78, 128, 12, 128]
+#         x1 = rearrange(x, 'b c l t -> b (c l) t')
+#         x1 = functional.leaky_relu(self.bn2_1(self.conv2_1(x1)))
+
+#         x2 = rearrange(x, 'b c l t -> b (t l) c')
+#         x2 = functional.leaky_relu(self.bn2_2(self.conv2_2(x2)))
+
+#         x1 = functional.dropout(x1, p=0.5, training=self.training)
+#         x2 = functional.dropout(x2, p=0.5, training=self.training)
+
+#         x1 = self.pool1(x1).squeeze(2)
+#         x2 = self.pool2(x2).squeeze(2)
+
+#         x = torch.cat([x1, x2], dim=1)
+#         return x
+
+
+# '''
+# changed the stride and added one more layer
+# '''
+# class CNN_layers(nn.Module):
+
+#     def __init__(self, embed_size = 384):
+#         super().__init__()
+#         self.multiple_cnn = nn.Sequential( 
+#             nn.Conv1d(12, 128, kernel_size=15, stride=2, padding=2, bias=False),
+#             nn.BatchNorm1d(128),
+#             nn.ReLU(inplace=True),
+
+#             nn.Conv1d(128, embed_size, kernel_size=15, stride=2, padding=100, bias=False),
+#             nn.BatchNorm1d(embed_size),
+#             nn.ReLU(inplace=True),
+            
+#             nn.Conv1d(384, embed_size, kernel_size=15, stride=3, padding=100, bias=False),
+#             nn.BatchNorm1d(embed_size),
+#             nn.ReLU(inplace=True)
+#             )
+        
+#     def forward(self, x):
+#         # print("This is the shape of enconder:(before)", x.shape)
+#         # This is the shape of enconder:(before) torch.Size([44, 12, 8192])
+#         x = self.multiple_cnn(x)
+#         x = x.transpose(1, 2)
+        
+
+#         return x
+
+# class D2_CNN_layers(nn.Module):
+#     def __init__(self, embed_size = 384):
+#         super().__init__()
+#         self.multiple_cnn = nn.Sequential(
+#             nn.Conv2d(in_channels=1, out_channels=128, kernel_size=(3, 15), padding=(1, 7), stride=(1, 2), bias=False),
+#             nn.BatchNorm2d(128),
+#             nn.LeakyReLU(inplace=True),
+
+#             nn.Conv2d(in_channels=128, out_channels=384, kernel_size=(3, 15), padding=(1, 7), stride=(1, 2), bias=False),
+#             nn.BatchNorm2d(embed_size),
+#             nn.LeakyReLU(inplace=True),
+
+#             nn.Conv2d(in_channels=384, out_channels=384, kernel_size=(3, 15), padding=(1, 7), stride=(1, 30), bias=False),
+#             nn.BatchNorm2d(embed_size),
+#             nn.LeakyReLU(inplace=True),
+
+#         )
+
+#     def forward(self, x):
+#         x = x.unsqueeze(1)
+#         # print("This is the input shape:", x.shape)
+#         x = self.multiple_cnn(x)
+#         x = torch.flatten(x, start_dim=-2)
+#         x = x.transpose(1, 2)
+#         return x
+
+def broadcat(tensors, dim=-1):
+    num_tensors = len(tensors)
+    shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
+    assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions'
+    shape_len = list(shape_lens)[0]
+    dim = (dim + shape_len) if dim < 0 else dim
+    dims = list(zip(*map(lambda t: list(t.shape), tensors)))
+    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+    assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), 'invalid dimensions for broadcastable concatenation'
+    max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
+    expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
+    expanded_dims.insert(dim, (dim, dims[dim]))
+    expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
+    tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
+    return torch.cat(tensors, dim=dim)
+
+def rotate_half(x):
+    x = rearrange(x, '... (d r) -> ... d r', r=2)
+    x1, x2 = x.unbind(dim=-1)
+    x = torch.stack((-x2, x1), dim=-1)
+    return rearrange(x, '... d r -> ... (d r)')
+
+# Adapted Rotary Embedding for 1D time-series (ECG) with 1024 tokens
+class TimeSeriesRotaryEmbeddingFast(nn.Module):
+    def __init__(
+        self,
+        dim,
+        seq_len=1024,  # Updated to 1024 tokens
+        custom_freqs=None,
+        freqs_for='lang',  # Suitable for 1D sequential data
+        theta=10000,
+        max_freq=10,
+        num_freqs=1,
+    ):
+        super().__init__()
+        if custom_freqs:
+            freqs = custom_freqs
+        elif freqs_for == 'lang':
+            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
+        elif freqs_for == 'pixel':
+            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
+        elif freqs_for == 'constant':
+            freqs = torch.ones(num_freqs).float()
+        else:
+            raise ValueError(f'unknown modality {freqs_for}')
+
+        # 1D sequence for 1024 tokens
+        t = torch.arange(seq_len).float()  # [0, 1, ..., 1023]
+
+        freqs = torch.einsum('..., f -> ... f', t, freqs)
+        freqs = repeat(freqs, '... n -> ... (n r)', r=2)  # Doubles the dimension for rotation
+
+        # Shape: (1024, dim)
+        freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
+        freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
+
+        self.register_buffer("freqs_cos", freqs_cos)  # Shape: (1024, dim)
+        self.register_buffer("freqs_sin", freqs_sin)  # Shape: (1024, dim)
+
+        print('======== shape of rope freq', self.freqs_cos.shape, '========')
+
+    def forward(self, t):
+        # t: (batch_size, seq_len, embed_dim)
+        # Apply rotation to the entire sequence
+        return t * self.freqs_cos + rotate_half(t) * self.freqs_sin
+
+
+
+'''
+dim: 384
+mixer_cls: mixer_cla is an instance of mamba.
+drop_path = 0.
+norm_cls = nn.LayerNorm
+fused_add_norm = True, 
+residual_in_fp32 = True
+'''
+class Block(nn.Module):
+    def __init__(
+        self, dim, mixer_cls, norm_cls=nn.LayerNorm, fused_add_norm=False, residual_in_fp32=False, drop_path=0.,
+    ):
+        """
+        Simple block wrapping a mixer class with LayerNorm/RMSNorm and residual connection"
+
+        This Block has a slightly different structure compared to a regular
+        prenorm Transformer block.
+        The standard block is: LN -> MHA/MLP -> Add.
+        [Ref: https://arxiv.org/abs/2002.04745]
+        Here we have: Add -> LN -> Mixer, returning both
+        the hidden_states (output of the mixer) and the residual.
+        This is purely for performance reasons, as we can fuse add and LayerNorm.
+        The residual needs to be provided (except for the very first block).
+        """
+        super().__init__()
+        
+        self.residual_in_fp32 = residual_in_fp32
+        self.fused_add_norm = fused_add_norm
+        self.mixer = mixer_cls(dim)
+        self.norm = norm_cls(dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        
+        # fused_add_norm true
+        if self.fused_add_norm:
+            assert RMSNorm is not None, "RMSNorm import fails"
+            assert isinstance(self.norm, (nn.LayerNorm, RMSNorm)), "Only LayerNorm and RMSNorm are supported for fused_add_norm"
+
+    def forward(self, hidden_states: Tensor, residual: Optional[Tensor] = None, inference_params=None):
+
+        r"""Pass the input through the encoder layer.
+
+        Args:
+            hidden_states: the sequence to the encoder layer (required).
+            residual: hidden_states = Mixer(LN(residual))
+        """
+        
+        if not self.fused_add_norm:
+            if residual is None:
+                residual = hidden_states
+            else:
+                residual = residual + self.drop_path(hidden_states)
+            
+            hidden_states = self.norm(residual.to(dtype=self.norm.weight.dtype))
+            if self.residual_in_fp32:
+                residual = residual.to(torch.float32)
+        
+        # since the self.fused_add_norm is true, the code is going below
+        # fused_add_norm_fn = layer_norm_fn
+        ###########
+        # hidden_states: Tensor
+        # self.norm.weight = torch.nn.LayerNorm.weight
+        # self.norm.bias = torch.nn.LayerNorm.bias
+        # residual: Optional[Tensor] = None
+        # prenorm=True
+        # residual_in_fp32 = True
+        # eps = torch.nn.LayerNorm.eps
+        
+        else:
+            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm, RMSNorm) else layer_norm_fn
+            if residual is None:
+                hidden_states, residual = fused_add_norm_fn(
+                    hidden_states,
+                    self.norm.weight,
+                    self.norm.bias,
+                    residual=residual,
+                    prenorm=True,
+                    residual_in_fp32=self.residual_in_fp32,
+                    eps=self.norm.eps,
+                )
+            else:
+                hidden_states, residual = fused_add_norm_fn(
+                    self.drop_path(hidden_states),
+                    self.norm.weight,
+                    self.norm.bias,
+                    residual=residual,
+                    prenorm=True,
+                    residual_in_fp32=self.residual_in_fp32,
+                    eps=self.norm.eps,
+                )  
+
+        # inference_params=None
+        hidden_states = self.mixer(hidden_states, inference_params=inference_params)
+        return hidden_states, residual
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        print("the code is going through allocate_inference_cache in the block")
+
+        return self.mixer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+
+'''
+from torch.nn.ModuleList()
+
+embed_dim=384, (embedding dimension)
+device: None 
+dtype: None
+ssm_cfg: None
+norm_epsilon: float = 1e-5
+rms_norm: bool = False
+residual_in_fp32 = True
+fused_add_norm = True
+if_bimamba = False
+bimamba_type = "v2"
+inter_dpr: [0.0, 0.0, 0.004347826354205608, ...,0.09565217792987823, 0.10000000149011612]
+if_devide_out = True
+init_layer_scale = None
+layer_idx = i
+'''
+
+def create_block(
+    d_model,
+    ssm_cfg=None,
+    norm_epsilon=1e-5,
+    drop_path=0.,
+    rms_norm=False,
+    residual_in_fp32=False,
+    fused_add_norm=False,
+    layer_idx=None,
+    device=None,
+    dtype=None,
+    if_bimamba=False,
+    bimamba_type="none",
+    if_devide_out=False,
+    init_layer_scale=None,
+    block_name = "default_value"
+):
+    if if_bimamba:
+        bimamba_type = "v1"
+
+    if ssm_cfg is None:
+        ssm_cfg = {}
+
+    factory_kwargs = {"device": device, "dtype": dtype}
+
+    if block_name == "VisionMamba":
+        mixer_cls = partial(Mamba, layer_idx=layer_idx, bimamba_type=bimamba_type, if_devide_out=if_devide_out, init_layer_scale=init_layer_scale, **ssm_cfg, **factory_kwargs)
+    elif block_name == "OriginalMamba":
+        mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)
+    else:
+        raise ValueError(f"No matching condition for value: {block_name}")
+    
+    
+
+    # rms_norm = False
+    norm_cls = partial(nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs)
+    
+    block = Block(
+        d_model,
+        mixer_cls,
+        norm_cls=norm_cls,
+        drop_path=drop_path,
+        fused_add_norm=fused_add_norm,
+        residual_in_fp32=residual_in_fp32,
+    )
+    block.layer_idx = layer_idx
+    return block
+
+
+# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454
+def _init_weights(
+    module,
+    n_layer,
+    initializer_range=0.02,  # Now only used for embedding layer.
+    rescale_prenorm_residual=True,
+    n_residuals_per_layer=1,  # Change to 2 if we have MLP
+):
+    if isinstance(module, nn.Linear):
+        if module.bias is not None:
+            if not getattr(module.bias, "_no_reinit", False):
+                nn.init.zeros_(module.bias)
+    elif isinstance(module, nn.Embedding):
+        nn.init.normal_(module.weight, std=initializer_range)
+
+    if rescale_prenorm_residual:
+        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:
+        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale
+        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.
+        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/
+        #
+        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py
+        for name, p in module.named_parameters():
+            if name in ["out_proj.weight", "fc2.weight"]:
+                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block
+                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)
+                # We need to reinit p since this code could be called multiple times
+                # Having just p *= scale would repeatedly scale it down
+                nn.init.kaiming_uniform_(p, a=math.sqrt(5))
+                with torch.no_grad():
+                    p /= math.sqrt(n_residuals_per_layer * n_layer)
+
+
+def segm_init_weights(m):
+    if isinstance(m, nn.Linear):
+        trunc_normal_(m.weight, std=0.02)
+        if isinstance(m, nn.Linear) and m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.Conv2d):
+        # NOTE conv was left to pytorch default in my original init
+        lecun_normal_(m.weight)
+        if m.bias is not None:
+            nn.init.zeros_(m.bias)
+    elif isinstance(m, (nn.LayerNorm, nn.GroupNorm, nn.BatchNorm2d)):
+        nn.init.zeros_(m.bias)
+        nn.init.ones_(m.weight)
+
+'''
+below is the 'ft-vim-s.sh':
+
+patch_size=16, (just patch size) 
+stride=8,  (just stride)
+embed_dim=384, (embedding dimension)
+depth=24, (?) the number of block
+rms_norm = True, (?)
+residual_in_fp32 = True, (?) 
+fused_add_norm = True, (?)
+final_pool_type = 'mean', (?)
+if_abs_pos_embed = True, (?)
+if_rope = False, (?)
+if_rope_residual = False, (?) 
+bimamba_type = "v2", (?)
+if_cls_token = True, (?)
+if_devide_out = True, (?)
+use_middle_cls_token = True, 
+**kwargs
+'''
+
+class VisionMamba(nn.Module):
+    def __init__(self, 
+                 img_size=224, 
+                 patch_size=16, 
+                 stride=16,
+                 depth=24, 
+                 embed_dim=192, 
+                 channels=3, 
+                 num_classes=26,
+                 ssm_cfg=None, 
+                 drop_rate=0.,
+                 drop_path_rate=0,
+                 norm_epsilon: float = 1e-5, 
+                 rms_norm: bool = False, 
+                 initializer_cfg=None,
+                 fused_add_norm=False,
+                 residual_in_fp32=False,
+                 device=None,
+                 dtype=None,
+                 ft_seq_len=None,
+                 pt_hw_seq_len=14,
+                 if_bidirectional=False,
+                 final_pool_type='none',
+                 if_abs_pos_embed=False,
+                 if_rope=False,
+                 if_rope_residual=False,
+                 flip_img_sequences_ratio=-1.,
+                 if_bimamba=False,
+                 bimamba_type="none",
+                 if_cls_token=False,
+                 if_devide_out=False,
+                 init_layer_scale=None,
+                 use_double_cls_token=False,
+                 use_middle_cls_token=False,
+                 **kwargs):
+        # print("The program is coming the init")
+        factory_kwargs = {"device": device, "dtype": dtype}
+        # factory_kwargs: {'device': None, 'dtype': None}
+        # add factory_kwargs into kwargs
+        block_name = kwargs.get('block', 'default_value')
+
+        kwargs.update(factory_kwargs) 
+
+        super().__init__()
+        self.residual_in_fp32 = residual_in_fp32
+        self.fused_add_norm = fused_add_norm
+        self.if_bidirectional = if_bidirectional
+        self.final_pool_type = final_pool_type
+        self.if_abs_pos_embed = if_abs_pos_embed
+        self.if_rope = if_rope
+        self.if_rope_residual = if_rope_residual
+        self.flip_img_sequences_ratio = flip_img_sequences_ratio
+        self.if_cls_token = if_cls_token
+        self.use_double_cls_token = use_double_cls_token
+        self.use_middle_cls_token = use_middle_cls_token
+        self.num_tokens = 1 if if_cls_token else 0
+
+        # pretrain parameters
+        self.num_classes = num_classes
+        self.d_model = self.num_features = self.embed_dim = embed_dim  # num_features for consistency with other models
+
+        # self.patch_embed = PatchEmbed(img_size=img_size, patch_size=patch_size, stride=stride, in_chans=channels, embed_dim=embed_dim)
+        # num_patches = self.patch_embed.num_patches
+
+        self.CNN_layers = CNN_layers()
+        num_patches = 729
+
+        # self.ECG_patch_embedding = ECG_Patch_embedding()
+        # num_patches = 1023
+        # self.LC = LeadCombiner(lead=6, out_ch=8)
+
+        # self.CNN_layers = D2_CNN_layers()
+        # num_patches = 828
+
+        # self.CNN_layers = CNN_layers()
+        # num_patches = 775
+        
+        # if_cls_token: True (in the original script)
+        if if_cls_token:
+            # use_double_cls_token: False (in the original script)
+            if use_double_cls_token:
+                self.cls_token_head = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+                self.cls_token_tail = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+                self.num_tokens = 2
+
+            else:
+                self.cls_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
+                # self.num_tokens = 1
+        
+        # if_abs_pos_embed: True (in the original script)
+        if if_abs_pos_embed:
+            self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + self.num_tokens, self.embed_dim))
+            self.pos_drop = nn.Dropout(p=drop_rate)
+
+        # if_rope: False (in the original script)
+
+        if if_rope:
+            half_head_dim = embed_dim // 2
+            self.rope = TimeSeriesRotaryEmbeddingFast(dim=embed_dim, seq_len= num_patches + self.num_tokens)
+        
+        self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
+        # self.head_LC = nn.Linear(16, num_classes)
+        # depth: 24; drop_path_rate: 0.1
+        # TODO: release this comment
+        # dpr: [0.0, 0.004347826354205608, 0.008695652708411217, ..., 0.09130434691905975, 0.09565217792987823, 0.10000000149011612]
+
+        if drop_path_rate == 0:
+            print("This is the drop_path_rate:", drop_path_rate)
+            dpr = [x.item() for x in torch.full((depth,), drop_path_rate)]            
+        else:
+            print("This is the drop_path_rate:", drop_path_rate)
+            print("follow the stochastic depth decay rule")
+            dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+
+
+        # import ipdb;ipdb.set_trace()
+        # inter_dpr: [0.0, 0.0, 0.004347826354205608, ...,0.09565217792987823, 0.10000000149011612]
+        inter_dpr = [0.0] + dpr
+        self.drop_path = DropPath(drop_path_rate) if drop_path_rate > 0. else nn.Identity()
+
+        # transformer blocks
+        # depth: 
+        self.layers = nn.ModuleList(
+            [
+                create_block(
+                    embed_dim,
+                    ssm_cfg=ssm_cfg,
+                    norm_epsilon=norm_epsilon,
+                    rms_norm=rms_norm,
+                    residual_in_fp32=residual_in_fp32,
+                    fused_add_norm=fused_add_norm,
+                    layer_idx=i,
+                    if_bimamba=if_bimamba,
+                    bimamba_type=bimamba_type,
+                    drop_path=inter_dpr[i],
+                    if_devide_out=if_devide_out,
+                    init_layer_scale=init_layer_scale,
+                    block_name = block_name,
+                    **factory_kwargs,
+                )
+                for i in range(depth)
+            ]
+        )
+        
+        # output head
+        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(embed_dim, eps=norm_epsilon, **factory_kwargs)
+
+        # self.pre_logits = nn.Identity()
+
+        # original init
+        # self.patch_embed.apply(segm_init_weights)
+
+        self.CNN_layers.apply(segm_init_weights)
+
+        # self.ECG_patch_embedding.apply(segm_init_weights)
+
+        self.head.apply(segm_init_weights)
+
+        # self.head_LC.apply(segm_init_weights)
+        # self.LC.apply(segm_init_weights)
+
+
+        # if_abs_pos_embed: True (in the original script)
+        if if_abs_pos_embed:
+            trunc_normal_(self.pos_embed, std=.02)
+        # if_cls_token: True (in the original script)
+        if if_cls_token:
+            if use_double_cls_token:
+                trunc_normal_(self.cls_token_head, std=.02)
+                trunc_normal_(self.cls_token_tail, std=.02)
+            # the code is coming here
+            else:
+                trunc_normal_(self.cls_token, std=.02)
+
+        # mamba init
+        self.apply(partial(_init_weights, n_layer=depth, **(initializer_cfg if initializer_cfg is not None else {}),))
+
+
+    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):
+        print("the code is going through allocate_inference_cache in the vision mamba")
+        return {
+            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)
+            for i, layer in enumerate(self.layers)
+        }
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token", "dist_token", "cls_token_head", "cls_token_tail"}
+
+    @torch.jit.ignore()
+    def load_pretrained(self, checkpoint_path, prefix=""):
+        _load_weights(self, checkpoint_path, prefix)
+
+    # x: this is the input 224*224(tensor)
+    # inference_params: None
+    # if_random_cls_token_position：False
+    # if_random_token_rank: False
+
+    def forward_features(self, x, inference_params=None, if_random_cls_token_position=False, if_random_token_rank=False):
+        # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
+        # with slight modifications to add the dist_token
+        
+        # x = self.patch_embed(x)
+
+        # x = self.CNN_layers(x)
+
+        x = self.CNN_layers(x)
+
+        # B: batch size 16
+        # M: 729 the number of patch,(27*27)
+        # D: the hidden state dimension, 384 (small-size variant), this is set by author of Vim, it is 768 in the convential Vit
+        # N: SSM dimension, SSM dimension N to 16.
+        # L: the number of blocks, we set the number of blocks L to 24
+
+        B, M, _ = x.shape
+        
+        # if_cls_token: True (in the original script)
+        if self.if_cls_token:
+
+            # self.use_double_cls_token: False (in the original script)
+            if self.use_double_cls_token:
+                cls_token_head = self.cls_token_head.expand(B, -1, -1)
+                cls_token_tail = self.cls_token_tail.expand(B, -1, -1)
+                token_position = [0, M + 1]
+                x = torch.cat((cls_token_head, x, cls_token_tail), dim=1)
+                M = x.shape[1]
+            else:
+                # self.use_middle_cls_token: True(in the original script)
+                if self.use_middle_cls_token:
+                    cls_token = self.cls_token.expand(B, -1, -1)
+                    token_position = M // 2
+                    # add cls token in the middle
+                    x = torch.cat((x[:, :token_position, :], cls_token, x[:, token_position:, :]), dim=1)
+                elif if_random_cls_token_position:
+                    cls_token = self.cls_token.expand(B, -1, -1)
+                    token_position = random.randint(0, M)
+                    x = torch.cat((x[:, :token_position, :], cls_token, x[:, token_position:, :]), dim=1)
+                    print("token_position: ", token_position)
+                else:
+                    cls_token = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
+                    token_position = 0
+                    x = torch.cat((cls_token, x), dim=1)
+                M = x.shape[1]
+
+        # # if_abs_pos_embed: True (in the original script)
+        if self.if_abs_pos_embed:
+            # if new_grid_size[0] == self.patch_embed.grid_size[0] and new_grid_size[1] == self.patch_embed.grid_size[1]:
+            #     x = x + self.pos_embed
+            # else:
+            #     pos_embed = interpolate_pos_embed_online(
+            #                 self.pos_embed, self.patch_embed.grid_size, new_grid_size,0
+            #             )
+            x = x + self.pos_embed
+            x = self.pos_drop(x)
+
+
+
+        if_flip_img_sequences = False
+        if self.flip_img_sequences_ratio > 0 and (self.flip_img_sequences_ratio - random.random()) > 1e-5:
+            x = x.flip([1])
+            if_flip_img_sequences = True
+
+        # mamba impl
+        # if_bidirectional： false
+        # inference_params： None
+        residual = None
+        hidden_states = x
+        if not self.if_bidirectional:
+            for layer in self.layers:
+                
+                # here is false in the original script
+                if if_flip_img_sequences and self.if_rope:
+                    hidden_states = hidden_states.flip([1])
+                    if residual is not None:
+                        residual = residual.flip([1])
+
+                # rope about, defaule is false
+                if self.if_rope:
+                    hidden_states = self.rope(hidden_states)
+                    if residual is not None and self.if_rope_residual:
+                        residual = self.rope(residual)
+
+                # here is false in the original script
+                if if_flip_img_sequences and self.if_rope:
+                    hidden_states = hidden_states.flip([1])
+                    if residual is not None:
+                        residual = residual.flip([1])
+
+                hidden_states, residual = layer(hidden_states, residual, inference_params=inference_params)
+                # sys.exit()
+
+        else:
+            # get two layers in a single for-loop
+            for i in range(len(self.layers) // 2):
+                if self.if_rope:
+                    hidden_states = self.rope(hidden_states)
+                    if residual is not None and self.if_rope_residual:
+                        residual = self.rope(residual)
+
+                hidden_states_f, residual_f = self.layers[i * 2](
+                    hidden_states, residual, inference_params=inference_params
+                )
+                hidden_states_b, residual_b = self.layers[i * 2 + 1](
+                    hidden_states.flip([1]), None if residual == None else residual.flip([1]), inference_params=inference_params
+                )
+                hidden_states = hidden_states_f + hidden_states_b.flip([1])
+                residual = residual_f + residual_b.flip([1])
+        
+        # fused_add_norm: True
+
+        if not self.fused_add_norm:
+            if residual is None:
+                residual = hidden_states
+            else:
+                residual = residual + self.drop_path(hidden_states)
+            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))
+        else:
+            # Set prenorm = False here since we don't need the residual
+            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn
+
+            hidden_states = fused_add_norm_fn(
+                self.drop_path(hidden_states),
+                self.norm_f.weight,
+                self.norm_f.bias,
+                eps=self.norm_f.eps,
+                residual=residual,
+                prenorm=False,
+                residual_in_fp32=self.residual_in_fp32,
+            )
+
+        # return only cls token if it exists
+        # if_cls_token: True (in the original script)
+        # self.use_middle_cls_token: True
+        if self.if_cls_token:
+            if self.use_double_cls_token:
+                return (hidden_states[:, token_position[0], :] + hidden_states[:, token_position[1], :]) / 2
+            else:
+                if self.use_middle_cls_token:
+                    return hidden_states[:, token_position, :]
+                elif if_random_cls_token_position:
+                    return hidden_states[:, token_position, :]
+                else:
+                    return hidden_states[:, token_position, :]
+
+        # self.final_pol_type = 'mean'
+        if self.final_pool_type == 'none':
+            return hidden_states[:, -1, :]
+        elif self.final_pool_type == 'mean':
+            return hidden_states.mean(dim=1)
+        elif self.final_pool_type == 'max':
+            return hidden_states
+        elif self.final_pool_type == 'all':
+            return hidden_states
+        else:
+            raise NotImplementedError
+
+    def forward(self, x, return_features=False, inference_params=None, if_random_cls_token_position=False, if_random_token_rank=False):
+        x = self.forward_features(x, inference_params, if_random_cls_token_position=if_random_cls_token_position, if_random_token_rank=if_random_token_rank)
+        # batch_number = x.shape[0]
+        # x = x.view(batch_number, 8, 6, 8)
+        # x = self.LC(x)
+        # x = self.head_LC(x)
+        # print("This is the shape of X (after the fully connected layer):", x.shape)
+        # return_features = False 
+        # print("This is the return feature:", return_features)
+        # if return_features:
+        #     return x
+        
+        # print("This is the shape of X (Before the fully connected layer):", x.shape)
+        x = self.head(x)
+        
+        # final_pool_type = 'mean' in original script
+        if self.final_pool_type == 'max':
+            x = x.max(dim=1)[0]
+        # sys.exit()
+        return x
+
+
+
+# below is for the vision in mamba
+@register_model
+def ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, depth=5, fused_add_norm = True, drop_path_rate = 0.1, if_divide_out = True, use_middle_cls_token = True, **kwargs):
+    model = VisionMamba(patch_size=16, stride=8, embed_dim=384, depth=depth, rms_norm=True, residual_in_fp32=True, drop_path_rate = drop_path_rate, fused_add_norm = fused_add_norm, final_pool_type='mean', if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type="v2", if_cls_token=True, if_devide_out=if_divide_out, use_middle_cls_token=use_middle_cls_token, **kwargs)
+    
+    # As a reminder:
+    print("This is whether the fused_add_norm:", fused_add_norm)
+    print("This is whether the if_divide_out:", if_divide_out)
+    print("This is whether the use_middle_cls_token:", use_middle_cls_token)
+
+
+    model.default_cfg = _cfg()
+    
+    return model
+
+# below is for the original mamba and 24 blocks
+# @register_model
+# def ecg_vim_small_patch16_stride8_224_bimambav2_final_pool_mean_abs_pos_embed_with_midclstok_div2(pretrained=False, **kwargs):
+#     model = VisionMamba(patch_size=16, stride=8, embed_dim=384, depth=24, rms_norm=True, residual_in_fp32=True, fused_add_norm=True, final_pool_type='mean', if_abs_pos_embed=True, if_rope=False, if_rope_residual=False, bimamba_type="v2", if_cls_token=True, if_devide_out=True, use_middle_cls_token=True, **kwargs)
+#     model.default_cfg = _cfg()
+    
+#     return model

optimizer.py

@@ -0,0 +1,42 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+
+# from utils import d_model
+
+'''
+optimizer = NoamOpt(d_model, 1, 4000, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))
+model_size: 256
+factor: 1
+warmup: 4000
+optimizer: torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9)
+'''
+class NoamOpt:
+    "Optim wrapper that implements rate."
+    def __init__(self, model_size, factor, warmup, optimizer):
+        self.optimizer = optimizer
+        self._step = 0
+        self.warmup = warmup
+        self.factor = factor
+        self.model_size = model_size
+        self._rate = 0
+        
+    def step(self):
+        "Update parameters and rate"
+        self._step += 1
+        rate = self.rate()
+        # print("This is the rate:", rate)
+        for p in self.optimizer.param_groups:
+            p['lr'] = rate
+        self._rate = rate
+        self.optimizer.step()
+        
+    def rate(self, step = None):
+        "Implement `lrate` above"
+        if step is None:
+            step = self._step
+        return self.factor * (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5)))
+        
+def get_std_opt(model):
+    return NoamOpt(729, 2, 4000, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))

rope.py

@@ -0,0 +1,141 @@
+# --------------------------------------------------------
+# EVA-02: A Visual Representation for Neon Genesis
+# Github source: https://github.com/baaivision/EVA/EVA02
+# Copyright (c) 2023 Beijing Academy of Artificial Intelligence (BAAI)
+# Licensed under The MIT License [see LICENSE for details]
+# By Yuxin Fang
+#
+# Based on https://github.com/lucidrains/rotary-embedding-torch
+# --------------------------------------------------------'
+
+from math import pi
+
+import torch
+from torch import nn
+
+from einops import rearrange, repeat
+
+
+
+def broadcat(tensors, dim = -1):
+    num_tensors = len(tensors)
+    shape_lens = set(list(map(lambda t: len(t.shape), tensors)))
+    assert len(shape_lens) == 1, 'tensors must all have the same number of dimensions'
+    shape_len = list(shape_lens)[0]
+    dim = (dim + shape_len) if dim < 0 else dim
+    dims = list(zip(*map(lambda t: list(t.shape), tensors)))
+    expandable_dims = [(i, val) for i, val in enumerate(dims) if i != dim]
+    assert all([*map(lambda t: len(set(t[1])) <= 2, expandable_dims)]), 'invalid dimensions for broadcastable concatentation'
+    max_dims = list(map(lambda t: (t[0], max(t[1])), expandable_dims))
+    expanded_dims = list(map(lambda t: (t[0], (t[1],) * num_tensors), max_dims))
+    expanded_dims.insert(dim, (dim, dims[dim]))
+    expandable_shapes = list(zip(*map(lambda t: t[1], expanded_dims)))
+    tensors = list(map(lambda t: t[0].expand(*t[1]), zip(tensors, expandable_shapes)))
+    return torch.cat(tensors, dim = dim)
+
+
+
+def rotate_half(x):
+    x = rearrange(x, '... (d r) -> ... d r', r = 2)
+    x1, x2 = x.unbind(dim = -1)
+    x = torch.stack((-x2, x1), dim = -1)
+    return rearrange(x, '... d r -> ... (d r)')
+
+
+
+class VisionRotaryEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim,
+        pt_seq_len,
+        ft_seq_len=None,
+        custom_freqs = None,
+        freqs_for = 'lang',
+        theta = 10000,
+        max_freq = 10,
+        num_freqs = 1,
+    ):
+        super().__init__()
+        if custom_freqs:
+            freqs = custom_freqs
+        elif freqs_for == 'lang':
+            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
+        elif freqs_for == 'pixel':
+            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
+        elif freqs_for == 'constant':
+            freqs = torch.ones(num_freqs).float()
+        else:
+            raise ValueError(f'unknown modality {freqs_for}')
+
+        if ft_seq_len is None: ft_seq_len = pt_seq_len
+        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
+
+        freqs_h = torch.einsum('..., f -> ... f', t, freqs)
+        freqs_h = repeat(freqs_h, '... n -> ... (n r)', r = 2)
+
+        freqs_w = torch.einsum('..., f -> ... f', t, freqs)
+        freqs_w = repeat(freqs_w, '... n -> ... (n r)', r = 2)
+
+        freqs = broadcat((freqs_h[:, None, :], freqs_w[None, :, :]), dim = -1)
+
+        self.register_buffer("freqs_cos", freqs.cos())
+        self.register_buffer("freqs_sin", freqs.sin())
+
+        print('======== shape of rope freq', self.freqs_cos.shape, '========')
+
+    def forward(self, t, start_index = 0):
+        rot_dim = self.freqs_cos.shape[-1]
+        end_index = start_index + rot_dim
+        assert rot_dim <= t.shape[-1], f'feature dimension {t.shape[-1]} is not of sufficient size to rotate in all the positions {rot_dim}'
+        t_left, t, t_right = t[..., :start_index], t[..., start_index:end_index], t[..., end_index:]
+        t = (t * self.freqs_cos) + (rotate_half(t) * self.freqs_sin)
+        return torch.cat((t_left, t, t_right), dim = -1)
+
+
+
+class VisionRotaryEmbeddingFast(nn.Module):
+    def __init__(
+        self,
+        dim,
+        pt_seq_len=16,
+        ft_seq_len=None,
+        custom_freqs = None,
+        freqs_for = 'lang',
+        theta = 10000,
+        max_freq = 10,
+        num_freqs = 1,
+    ):
+        super().__init__()
+        if custom_freqs:
+            freqs = custom_freqs
+        elif freqs_for == 'lang':
+            freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))
+        elif freqs_for == 'pixel':
+            freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi
+        elif freqs_for == 'constant':
+            freqs = torch.ones(num_freqs).float()
+        else:
+            raise ValueError(f'unknown modality {freqs_for}')
+
+        if ft_seq_len is None: ft_seq_len = pt_seq_len
+        t = torch.arange(ft_seq_len) / ft_seq_len * pt_seq_len
+
+        freqs = torch.einsum('..., f -> ... f', t, freqs)
+        freqs = repeat(freqs, '... n -> ... (n r)', r = 2)
+        freqs = broadcat((freqs[:, None, :], freqs[None, :, :]), dim = -1)
+
+        freqs_cos = freqs.cos().view(-1, freqs.shape[-1])
+        freqs_sin = freqs.sin().view(-1, freqs.shape[-1])
+
+        self.register_buffer("freqs_cos", freqs_cos)
+        self.register_buffer("freqs_sin", freqs_sin)
+
+        print('======== shape of rope freq', self.freqs_cos.shape, '========')
+
+    def forward(self, t): 
+        if t.shape[1] % 2 != 0:
+            t_spatial = t[:, 1:, :]
+            t_spatial = t_spatial * self.freqs_cos + rotate_half(t_spatial) * self.freqs_sin
+            return torch.cat((t[:, :1, :], t_spatial), dim=1)
+        else:
+            return  t * self.freqs_cos + rotate_half(t) * self.freqs_sin