Spaces:
Running
Running
File size: 3,303 Bytes
41e6846 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | from __future__ import annotations
import logging
import numpy as np
from config import MIN_ECG_SECONDS
LOGGER = logging.getLogger(__name__)
def process_ecg_signal(ecg_signal: np.ndarray, sampling_rate: int) -> dict[str, object]:
"""Clean ECG, detect R peaks, and compute RR intervals."""
signal = np.asarray(ecg_signal, dtype=float).reshape(-1)
min_samples = int(MIN_ECG_SECONDS * sampling_rate)
if signal.size < min_samples:
return _error(f"ECG signal is too short. Need at least {MIN_ECG_SECONDS} seconds.")
if not np.isfinite(signal).all():
signal = signal[np.isfinite(signal)]
if signal.size < min_samples:
return _error("ECG signal has too few finite samples after cleaning invalid values.")
try:
cleaned_ecg, rpeaks = _process_with_neurokit(signal, sampling_rate)
except Exception:
LOGGER.warning("neurokit2 ECG processing unavailable; falling back to scipy/manual peak detection.")
cleaned_ecg, rpeaks = _process_with_fallback(signal, sampling_rate)
rpeaks = np.asarray(rpeaks, dtype=int)
if rpeaks.size < 3:
return _error("Too few R peaks were detected for HR/HRV estimation.")
rr_intervals_ms = np.diff(rpeaks) / sampling_rate * 1000.0
if rr_intervals_ms.size == 0 or float(np.mean(rr_intervals_ms)) <= 0:
return _error("Unable to compute valid RR intervals.")
heart_rate = 60000.0 / float(np.mean(rr_intervals_ms))
return {
"ok": True,
"heart_rate": round(heart_rate, 2),
"rr_intervals_ms": rr_intervals_ms.round(2).tolist(),
"rpeaks": rpeaks.tolist(),
"cleaned_ecg": np.asarray(cleaned_ecg, dtype=float).round(6).tolist(),
"message": "success",
}
def _process_with_neurokit(signal: np.ndarray, sampling_rate: int) -> tuple[np.ndarray, np.ndarray]:
import neurokit2 as nk
cleaned = nk.ecg_clean(signal, sampling_rate=sampling_rate)
_, info = nk.ecg_peaks(cleaned, sampling_rate=sampling_rate)
return np.asarray(cleaned), np.asarray(info.get("ECG_R_Peaks", []), dtype=int)
def _process_with_fallback(signal: np.ndarray, sampling_rate: int) -> tuple[np.ndarray, np.ndarray]:
centered = signal - np.median(signal)
try:
from scipy.signal import find_peaks
min_distance = int(0.45 * sampling_rate)
threshold = np.percentile(centered, 90)
peaks, _ = find_peaks(centered, distance=min_distance, height=threshold)
return centered, peaks
except Exception:
threshold = np.percentile(centered, 95)
candidate_indices = np.where(centered >= threshold)[0]
peaks = []
last_peak = -int(0.45 * sampling_rate)
for idx in candidate_indices:
if idx - last_peak >= int(0.45 * sampling_rate):
window_end = min(idx + int(0.2 * sampling_rate), len(centered))
local_idx = idx + int(np.argmax(centered[idx:window_end]))
peaks.append(local_idx)
last_peak = local_idx
return centered, np.asarray(peaks, dtype=int)
def _error(message: str) -> dict[str, object]:
return {
"ok": False,
"heart_rate": None,
"rr_intervals_ms": [],
"rpeaks": [],
"cleaned_ecg": [],
"message": message,
}
|