Preprocessing.py

# -*- coding: utf-8 -*-
"""Preprocessing.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1kF8egbS1hChCK-pmvc2AijXghFRRhLa1
"""

### Install waveform-database (WFDB) package. 

conda install -c conda-forge wfdb

# Commented out IPython magic to ensure Python compatibility.
from IPython.display import display
import matplotlib.pyplot as plt
# %matplotlib inline
import numpy as np
import os
import shutil
import posixpath

import wfdb
import pywt
from torch.utils.data import Dataset
import torch

def Download(out_dir):
    """
    Args:
        out_dir (string): Каталог для сохранения датасета
    """
    
    wfdb.dl_database('mitdb', out_dir)
    print('Done!')

#create directory

cwd = os.getcwd()
data_dir = os.path.join(cwd, 'MIT-BIH Arrhythmia Database')

## Download dataset
Download(data_dir)

"""### Попытка решения проблем, возникших на первой версии обработки

При нарезке холтера мы столкнулись с проблемой неоднозначной аннотации: в датасете аннотируются не только R-пики, поэтому нельзя утверждать, что в окне сердцебиения из 300-х точек будет всего-лишь одна аннотация. Более того, код ниже показывает, что это не так 
"""

flag = False
list_of = os.listdir(data_dir)
records = []
for filename in list_of:
    file = os.path.splitext(filename)[0]
    if file in records:
        continue
    else:
        records.append(file)
    
for name in records:
    signal_file = os.path.join(data_dir, name)

    record = wfdb.rdrecord(signal_file, physical=False, channels=[0])
    annotation = wfdb.rdann(signal_file, 'atr')

    for test_ind in annotation.sample:

        #edge case handling 
        if (test_ind - 150) < 0 or (test_ind+149) > record.sig_len:
            continue

        slice_record = wfdb.rdrecord(signal_file, sampfrom=test_ind - 150, 
                                sampto=test_ind+150, channels=[0], physical=False)           
        slice_annotation = wfdb.rdann(signal_file, 'atr', 
                                              sampfrom=test_ind - 150, sampto=test_ind + 150, 
                                              shift_samps=True)

        signal = slice_record.d_signal
        signal = signal.reshape(300)                
        signal = torch.from_numpy(signal)
        label = slice_annotation.symbol

        ##Ignore more than 1 label
        if len(label) > 1:
            wfdb.plot_wfdb(record=slice_record, annotation=slice_annotation, 
                            title='Abnormal signal', time_units='seconds')
            print(label)
            flag = True
            break
    if flag == True:
            break

"""Это вызывает непредсказуемое поведение, поскольку в методе подразумевается, что нарезка ведется по R-пикам. В прошлый раз предлагалось игнорировать все такие окна. Однако это влечет потерю данных в 3 процента, что кажется значительным. 

При дальнейшей классификации меток аннотации мы пользовались таблицей из источника https://github.com/hsd1503/PhysioNet 

В ней присутствуют не все метки, полученные после нарезки. Некоторые метки не лежат ни в одном из классов. Предлагается такие метки игнорировать до нарезки сигнала, поскольку они обозначают шумы, изменения ритма и прочие аннотации, не существенные для нашей классификации.

При таком подходе обработка окна сигнала с несколькими аннотациями становится неотличимой от окна с одной аннотацией. Действительно, центрированная метка в таком окне является "правильной", поэтому мы включаем его аналогично другим сигналам.

Также важен порядок обработки, поскольку меняя его(вначале нарезка, потом обработка шумов) мы не можем осуществить вейвлет преобразование на 6 уровней, поскольку нарезанные сигналы обладают слишком малой длиной. Поэтому сначала стоит убрать шумы. Если это делать отдельно от нарезки(не в одном цикле), то придется после обработки снова сохранять сигнал с аннотациями в wfdb формате. Это вызывает различные баги, связанные с работой библиотечных функций. Поэтому предлагается просто включить обработку шумов сразу в нарезку, получив функцию, которая осуществляет полную предобработку данных.
"""

def Denoise(d_signal):
    """
    Args:
        d_signal (ndarray): массив значений цифрового сырого сигнала
        
    Return:
        denoise_d_signal (darray): массив значений цифрового обработанного сигнала
    """
        
    #Wavelet Decomposion
    coeffs = pywt.wavedec(d_signal, 'db6', level=6)

    # ingor first two levels:
    coeffs[-1] = np.zeros_like(coeffs[-1])
    coeffs[-2] = np.zeros_like(coeffs[-2])

    #Inverse Wavelet Transform
    denoise_d_signal = pywt.waverec(coeffs, 'db6')
    
    return denoise_d_signal

from sklearn.model_selection import train_test_split

def Preprocessing(root_dir, output_dir):
    """
    Args:
        root_dir (string): Каталог с сигналами и аннотациями
        output_dir (string): Каталог для сохранения полученных данных
        
    Return: [data, label]
        data (torch.FloatTensor): Датасет в формате тензора
        label (torch.IntTensor): Мети класса в формате тензора
    """
    
    """
    В методе окно сигнала размера (300) присоединяется ко всему
    датасету методом torch.cat в каждой итерации. По мере увеличения 
    датасета это работает все медленнее. Для повышения скорости
    предлагается сохранять во временный тензор(temp_tensor)  MAX_OUTPUT_SIZE 
    сигналов, а затем присоединять его к формирующемуся датасету.  
    """
    MAX_OUTPUT_SIZE = 5000
    temp_tensor = torch.FloatTensor()
    temp_label = torch.IntTensor() 
    
    X = torch.FloatTensor()
    y = torch.IntTensor() 
    
    """
    encoding:
        1 -- Non-ectopic (N)
        2 -- Supra ventricular ectopic (SVEB)
        3 -- Ventricular ectopic (VEB)
        4 -- Fusion (F)
        5 -- Unknown (Q)
    
    """
       
    label_group_map = {'N':0, 'L':0, 'R':0, 'V':2, '/':4, 
                       'A':1, 'F':3, 'f':4, 'j':1, 'a':1, 
                       'E':2, 'J':1, 'e':1, 'Q':4, 'S':1}    
    valid_labels = label_group_map.keys()
    
    ## Create a list of records
    listdir = os.listdir(root_dir)
    records = []
    for filename in listdir:
        file = os.path.splitext(filename)[0]
        if file in records:
            continue
        else:
            records.append(file)
    
    
    for name in records:
        signal_file = os.path.join(root_dir, name)

        record = wfdb.rdrecord(signal_file, physical=False, channels=[0])
        annotation = wfdb.rdann(signal_file, 'atr')
        
        ##Denoising:
        #lead to one-dimensional array
        d_signal = np.transpose(record.d_signal)[0]
        record.d_signal = Denoise(d_signal)

        for i, test_ind in enumerate(annotation.sample):
            
            #class compliance check
            if annotation.symbol[i] not in valid_labels:
                continue

            #edge case handling 
            if (test_ind - 150) < 0 or (test_ind+149) > record.sig_len:
                continue
                
            #Slicing
            signal = record.d_signal[test_ind - 150:test_ind + 150]
            signal = signal.reshape(1, 300)                
            signal = torch.from_numpy(signal)
            signal = signal.type(torch.FloatTensor)
            str_label = annotation.symbol[i]
            label = label_group_map[str_label]
            label = torch.IntTensor([label])
            
            #Adding
            if temp_tensor.size(0) <= MAX_OUTPUT_SIZE:
                temp_tensor = torch.cat((temp_tensor, signal), 0)               
                temp_label = torch.cat((temp_label, label), 0)
            else:
                X = torch.cat((X, temp_tensor), 0)
                y = torch.cat((y, temp_label), 0)
                temp_tensor = torch.empty(0)
                temp_label = torch.empty(0)
    
    ##SAVING
    y = y.type(torch.LongTensor)
    
    X_train_name = os.path.join(output_dir, 'train_data.pt')
    X_test_name = os.path.join(output_dir, 'test_data.pt')
    y_train_name = os.path.join(output_dir, 'train_label.pt')
    y_test_name = os.path.join(output_dir, 'test_label.pt')
    
    #Splitting train_test 
    train_idx, test_idx = train_test_split(np.arange(y.size(0)), test_size=0.1, 
                                           shuffle=True, stratify=y)
    X_train = X[train_idx]
    y_train = y[train_idx]
    X_test = X[test_idx]
    y_test = y[test_idx]

    torch.save(X_train, X_train_name)
    torch.save(y_train, y_train_name)
    torch.save(X_test, X_test_name)
    torch.save(y_test, y_test_name)
    
    return X, y

## Slicing a dataset
slice_data_dir = os.path.join(cwd, "MIT-BIH Pre-Processed") 
data = torch.FloatTensor()
labels = torch.IntTensor()
data, labels = Preprocessing(data_dir, slice_data_dir)

#Some data analysis

article_data = [107679, 90265, 2503, 7106, 789, 7016]

print('Number of signals {} \t In the article {} \t Lost data {} %'.format(data.size(0), article_data[0], 
                                                                           100 - 100 * data.size(0) / article_data[0]))
counter = [0, 0, 0, 0, 0]
for mark in labels:
    mark = mark.item()
    mark = int(mark)
    counter[mark - 1] += 1
print( """Non-ectopic: {} signals \t In the article {} 
Supra ventricular ectopic: {} signals \t In the article {} 
Ventricular ectopic: {} signals \t In the article {} 
Fusion: {} signals \t In the article {} 
Unknown: {} signals \t In the article {}""".format(
                                                counter[0], article_data[1],
                                                counter[1], article_data[2], 
                                                counter[2], article_data[3],
                                                counter[3], article_data[4],
                                                counter[4], article_data[5]))

"""При данном подходе мы получили на 875 сигналов больше, большинство из которых(примерно 700) прибавились к вентрикулярному ритму. Это кажется полезным, поскольку в этом случае классы становятся сильно сбалансированнее. """