EEG eye blink normalizer

frontal2standart.py

from pathlib import Path
import numpy as np
import mne

def frontal2standart(X, clear_morgun=False, show=False):
    '''преобразователь фронтальных референтов в "мастоиды"
    
    X - массив (n_sample, 4 chan)
    
    clear_morgun - если True, чистит моргун
    
    show - для отладки рисует наложение
    

    Из EDF-файла:

        twin = 5.0 #s
        A = mne.io.read_raw_edf(u)
        A.pick(range(4))
        A.load_data()
        A.filter(None, 30, filter_length='1s');
        X = A.get_data(tmin=t, tmax=t+twin, units='uV').T

    Работает фактически 1-й компонент
    (1 колонка в W - unmixing)
    (1 ряд в M_ - модицифированная mixing - полностью уходит в лобные, и обнуляется в затылках)
    Разложение подобрано диагонализацией ковариационной матрицы.
    Если при переносе будет работать плохо - надо повторить для новых данных.
    
    '''
    W = np.array([[ 0.58612732,  0.06965951,  0.46354819,  0.07242658],
                  [ 0.45220561,  0.48530181, -0.43015398, -0.04188745],
                  [-0.46949804,  0.22923542,  0.06876178, -0.28639626],
                  [-0.52055864,  0.31274918, -0.01982899,  0.28674805]])
    M_ = np.array([[ 1.        ,  1.        ,  0.        ,  0.        ],
                   [ 0.7814091 ,  0.95039743,  0.92807246,  0.86839795],
                   [ 1.38939577, -0.72303997,  0.73281696,  0.27536562],
                   [ 0.14512948, -0.20708375, -1.80827299,  1.61442051]])

    if clear_morgun:
        M_[0]=0.0
    
    Xs = (X @ W) @ M_
        
    if show:
        # для рисования с наложением поменяйте plotts на подходящую функцию
        plt.figure(figsize=(10,4))
        co = 'g' if clear_morgun else 'b' 
        vivo.plotts(X, 'k', alpha=0.5, scale=50);
        vivo.plotts(Xs, co, scale=50);
        
    return Xs


def convert(u, clear_morgun=False):
    if u.index('_Frontal'):
        usave = u.replace('_Frontal', '_Standart')
    else:
        usave=str(Path(u).with_name('export.edf'))
    
    A = mne.io.read_raw_edf(u)
    A.load_data()
    A.filter(None, 30, filter_length='1s');

    X = A.get_data(picks=A.ch_names[:4], units='uV').T
    Xs = frontal2standart(X, clear_morgun=clear_morgun)

    A._data[:4,:] = Xs.T
    mne.export.export_raw(usave, A, overwrite=True)
    
    return usave


def convert_txt(u, clear_morgun=False):
    if u.index('_Frontal'):
        usave = u.replace('_Frontal', '_Standart')
    else:
        usave=str(Path(u).with_name('export.txt'))
        
    X = np.loadtxt(u)
    
    X[:,:4] = frontal2standart(X[:,:4], clear_morgun=clear_morgun)

    np.savetxt(usave, X, '%g')    
    
    return usave
    
    
if __name__=='__main__':
    import sys
    if len(sys.argv)==2 and (sys.argv[1].lower() in ['--help','-h','/h']):
        print(f'''
Укажите полное имя edf файла в качестве аргумента командной строки
Пример использования:

    {sys.executable}  frontal2standart.py C:\\DATA\\Test_1_2025_01_15_13_40_Raw_Frontal_Ref.edf

''')
        sys.exit()
    
    if len(sys.argv)==1: #GUI
        import win32gui, win32con, os

        file_types = "Text files\0*.txt\0Record\0*.edf\0"
        customfilter = "Other file types\0*.*\0"
        u, customfilter, flags = win32gui.GetOpenFileNameW(
            InitialDir='./',
            Flags=win32con.OFN_ALLOWMULTISELECT | win32con.OFN_EXPLORER |win32con.OFN_NOCHANGEDIR,
            File="выберите файл",
            DefExt="txt",
            Title="GetOpenFileNameW",
            Filter=file_types,
            CustomFilter=customfilter,
            FilterIndex=0,
        )
        if u:
            clear_morgun = False
            convert_ = convert_txt if u.endswith('.txt') else convert
            usave = convert_(u, clear_morgun=clear_morgun)

            # messagebox.showinfo("", "Конвертация завершена!")
            if u.endswith('.edf'):
                print(f'Попробуй: mne browse_raw {usave}')
            else:
                print(f'Проверь {usave}')

        sys.exit()
    
    
    
    u = sys.argv[1]
    clear_morgun = False
    if u.lower() in ['--clear_morgun','-morgun']:
        clear_morgun = True
        u = sys.argv[2]
            
    usave = convert(u, clear_morgun=clear_morgun)

    
    if u.endswith('.edf'):
        print(f'Попробуй: mne browse_raw {usave}')
    else:
        print(f'Проверь {usave}')