data/data_utils.py

import os
import numpy as np
from scipy.signal import butter, sosfilt, lfilter, iirnotch
from scipy.signal import find_peaks

def butter_bandpass(lowcut, highcut, fs, order=8):
    low = lowcut
    high = highcut
    sos  = butter(order, [low, high], btype='band', fs=fs, output='sos', analog=False)
    return sos 

def butter_highpass(lowcut, fs, order=8):
    low = lowcut
    sos = butter(order, low, btype='highpass', fs=fs, output='sos', analog=False)
    return sos

def butter_lowpass(highcut, fs, order=8):
    high = highcut
    sos = butter(order, high, btype='lowpass', fs=fs, output='sos', analog=False)
    return sos

def butter_bandpass_filter(data, lowcut, highcut, fs, order=8):
    if highcut > fs/2:
        sos = butter_highpass(lowcut, fs, order=order)
    elif lowcut <= 0:
        sos = butter_lowpass(highcut, fs, order=order)
    else:
        sos  = butter_bandpass(lowcut, highcut, fs, order=order)
    #y = lfilter(b, a, data)
    y = sosfilt(sos, data)
    return y

def correct_ecg_inversion(ecg_signal, sampling_rate):
    """
    ECG 신호가 뒤집혔는지 감지하고 올바르게 수정합니다.

    Args:
        ecg_signal (np.array): ECG 신호 데이터.
        sampling_rate (int): 샘플링 주파수 (Hz).

    Returns:
        np.array: 수정된 ECG 신호.
        bool: 뒤집힘이 감지되었는지 여부.
    """
    # QRS군의 R파를 찾기 위해 신호의 양수 피크를 찾음
    # prominence: 피크가 주변보다 얼마나 튀어나왔는지, distance: 피크 간 최소 거리
    positive_peaks, _ = find_peaks(ecg_signal, prominence=0.1, distance=0.3 * sampling_rate)
    
    # QRS군의 S파(뒤집힌 R파)를 찾기 위해 신호를 뒤집어서 피크를 찾음
    negative_peaks, _ = find_peaks(-ecg_signal, prominence=0.1, distance=0.3 * sampling_rate)

    # 피크가 하나도 감지되지 않으면 원본 신호 반환
    if len(positive_peaks) == 0 and len(negative_peaks) == 0:
        return ecg_signal, False

    # 양수 피크와 음수 피크의 평균 진폭 계산
    mean_pos_amp = np.mean(ecg_signal[positive_peaks]) if len(positive_peaks) > 0 else 0
    mean_neg_amp = np.mean(ecg_signal[negative_peaks]) if len(negative_peaks) > 0 else 0

    # 만약 음수 피크의 절대값이 양수 피크보다 크면 뒤집힌 것으로 판단
    is_inverted = abs(mean_neg_amp) > abs(mean_pos_amp)

    if is_inverted:
        # print("💡 ECG 신호가 뒤집힌 것을 감지하여 수정합니다.")
        return ecg_signal * -1, True
    else:
        # print("✅ ECG 신호가 정상 방향입니다.")
        return ecg_signal, False