Spaces:
Running
Running
| 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, | |
| } | |