| import torch |
| from math import gcd |
| from typing import Optional, Union |
| import joblib |
|
|
| import numpy as np |
| from scipy import signal |
|
|
| def load_scaler_joblib(path: str) -> tuple[torch.Tensor, torch.Tensor]: |
| """ |
| Load ecg_scaler.pkl and return center and scale as torch tensors. |
| Args: |
| path: Path to the joblib file. |
| Returns: |
| center: torch.Tensor |
| scale: torch.Tensor |
| """ |
| sc = joblib.load(path) |
| center = torch.from_numpy(sc.mean_.astype(np.float32)) |
| scale = torch.from_numpy(sc.scale_.astype(np.float32)).clamp_min(1e-8) |
| return center, scale |
|
|
| class ECGTransform: |
| """ |
| Unified ECG preprocessing: downsampling and scaling. |
| Usage: |
| transform = ECGTransform(center, scale, src_fs=512, target_fs=100) |
| ecg_out = transform(ecg_in) |
| """ |
| def __init__( |
| self, |
| center: Union[np.ndarray, torch.Tensor], |
| scale: Union[np.ndarray, torch.Tensor], |
| src_fs: int = 100, |
| target_fs: int = 100, |
| band: Optional[tuple[float, float]] = (0.5, 40.0), |
| bp_order: int = 4, |
| axis: int = -1, |
| ) -> None: |
| self.center = torch.as_tensor(center, dtype=torch.float32) |
| self.scale = torch.as_tensor(scale, dtype=torch.float32).clamp_min(1e-8) |
| self.src_fs = src_fs |
| self.target_fs = target_fs |
| self.band = band |
| self.bp_order = bp_order |
| self.axis = axis |
|
|
| def downsample(self, x: np.ndarray) -> np.ndarray: |
| x = np.asarray(x) |
| if self.band is not None: |
| lowcut, highcut = self.band |
| max_high = 0.45 * self.target_fs |
| highcut = min(highcut, max_high) |
| nyq = self.src_fs / 2.0 |
| if lowcut <= 0: |
| wn = highcut / nyq |
| sos = signal.butter(self.bp_order, wn, btype="low", output="sos") |
| else: |
| wn = (lowcut / nyq, highcut / nyq) |
| sos = signal.butter(self.bp_order, wn, btype="band", output="sos") |
| x = signal.sosfiltfilt(sos, x, axis=self.axis) |
| g = gcd(self.src_fs, self.target_fs) |
| up = self.target_fs // g |
| down = self.src_fs // g |
| y = signal.resample_poly(x, up, down, axis=self.axis, window=("kaiser", 5.0), padtype="median") |
| return y |
|
|
| def scale(self, ecg: torch.Tensor) -> torch.Tensor: |
| ecg = ecg.to(torch.float32) |
| ecg = (ecg - self.center[:, None]) / self.scale[:, None] |
| return ecg |
|
|
| def __call__(self, x: np.ndarray) -> torch.Tensor: |
| """ |
| Downsample and scale ECG data. |
| Args: |
| x: np.ndarray, shape (leads, time) |
| Returns: |
| torch.Tensor, shape (leads, time) |
| """ |
| if self.src_fs != self.target_fs: |
| x = self.downsample(x) |
| if not isinstance(x, torch.Tensor): |
| x = torch.from_numpy(x) |
| x = self.scale(x) |
| return x |
