Update model and implement abstract class

ECG2HRV.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3669bf4b201fd873f1fb1b48083c41f26c42be00f58d14d234af6b9aac1f6433
+size 39

notebooks/01_Model_Deployment_Research.ipynb

@@ -399,21 +399,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": 1,
    "outputs": [],
    "source": [
     "from huggingface_hub import hf_hub_download\n",
     "import joblib\n",
     "import torch\n",
+    "import numpy as np\n",
     "\n",
-    "from src.model import HR2HRV"
+    "from src.model import ECG2HRV"
    ],
    "metadata": {
-    "collapsed": false,
-    "ExecuteTime": {
-     "end_time": "2024-02-21T11:39:25.775871100Z",
-     "start_time": "2024-02-21T11:39:25.755838Z"
-    }
+    "collapsed": false
    }
   },
   {
@@ -427,11 +424,11 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 27,
+   "execution_count": 2,
    "outputs": [],
    "source": [
     "# Instantiate model\n",
-    "model = HR2HRV()\n",
+    "model = ECG2HRV()\n",
     "# Save\n",
     "joblib.dump(model, \"..\\ECG2HRV.joblib\")\n",
     "# Load in notebook\n",
@@ -440,15 +437,15 @@
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-02-21T12:01:28.600527100Z",
-     "start_time": "2024-02-21T12:01:28.580278200Z"
+     "end_time": "2024-02-21T16:08:51.659030Z",
+     "start_time": "2024-02-21T16:08:51.605730100Z"
     }
    }
   },
   {
    "cell_type": "markdown",
    "source": [
-    "**Test if the model can be loaded from the hub and used**"
+    "**Test the model locally with random ecg**"
    ],
    "metadata": {
     "collapsed": false
@@ -456,38 +453,84 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 3,
    "outputs": [],
    "source": [
-    "# Load from hub\n",
-    "REPO_ID = \"HUBII-Platform/ECG2HRV\"\n",
-    "FILENAME = \"feature-extractor.joblib\"\n",
+    "duration_seconds = 10 # Time duration for ECG signal (in seconds)\n",
+    "sample_rate = 100 # Sample rate (samples per second)\n",
+    "num_samples = duration_seconds * sample_rate # Number of samples\n",
     "\n",
-    "model = joblib.load(\n",
-    "    hf_hub_download(repo_id=REPO_ID, filename=FILENAME)\n",
-    ")"
+    "t = np.linspace(0, duration_seconds, num_samples) # Time array\n",
+    "\n",
+    "# Generate ECG signal (example synthetic data)\n",
+    "ecg_signal = (\n",
+    "    0.2 * np.sin(2 * np.pi * 1 * t) +\n",
+    "    0.5 * np.sin(2 * np.pi * 0.5 * t) -\n",
+    "    0.1 * np.sin(2 * np.pi * 2.5 * t)\n",
+    ")\n",
+    "\n",
+    "# Add some random noise\n",
+    "ecg_signal += np.random.normal(scale=0.1, size=num_samples)"
    ],
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-02-21T11:36:52.302912800Z",
-     "start_time": "2024-02-21T11:36:52.145834500Z"
+     "end_time": "2024-02-21T16:08:51.669938Z",
+     "start_time": "2024-02-21T16:08:51.635032600Z"
     }
    }
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 4,
+   "outputs": [
+    {
+     "data": {
+      "text/plain": "[{'HRV_MeanNN': 413.4782608695652,\n  'HRV_SDNN': 100.97743652790477,\n  'HRV_SDANN1': nan,\n  'HRV_SDNNI1': nan,\n  'HRV_SDANN2': nan,\n  'HRV_SDNNI2': nan,\n  'HRV_SDANN5': nan,\n  'HRV_SDNNI5': nan,\n  'HRV_RMSSD': 92.78518690551262,\n  'HRV_SDSD': 94.96410805236795,\n  'HRV_CVNN': 0.24421462041449105,\n  'HRV_CVSD': 0.22440160870944167,\n  'HRV_MedianNN': 400.0,\n  'HRV_MadNN': 118.60799999999999,\n  'HRV_MCVNN': 0.29651999999999995,\n  'HRV_IQRNN': 150.0,\n  'HRV_SDRMSSD': 1.0882926455785953,\n  'HRV_Prc20NN': 320.0,\n  'HRV_Prc80NN': 490.0,\n  'HRV_pNN50': 52.17391304347826,\n  'HRV_pNN20': 69.56521739130434,\n  'HRV_MinNN': 310.0,\n  'HRV_MaxNN': 640.0,\n  'HRV_HTI': 5.75,\n  'HRV_TINN': 0.0}]"
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model(input_data=ecg_signal, frequency=100.0)"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "ExecuteTime": {
+     "end_time": "2024-02-21T16:08:51.755181400Z",
+     "start_time": "2024-02-21T16:08:51.671014900Z"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "**Test if the model can be loaded from the hub and used**"
+   ],
+   "metadata": {
+    "collapsed": false
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
    "outputs": [],
    "source": [
-    "# Create a example tensor input\n",
-    "tensor = torch.tensor([2.0, 3.0, 4.0])"
+    "# Load from hub\n",
+    "REPO_ID = \"HUBII-Platform/ECG2HRV\"\n",
+    "FILENAME = \"feature-extractor.joblib\"\n",
+    "\n",
+    "model = joblib.load(\n",
+    "    hf_hub_download(repo_id=REPO_ID, filename=FILENAME)\n",
+    ")"
    ],
    "metadata": {
     "collapsed": false,
     "ExecuteTime": {
-     "end_time": "2024-02-21T11:36:55.990475100Z",
-     "start_time": "2024-02-21T11:36:55.989181100Z"
+     "end_time": "2024-02-21T11:36:52.302912800Z",
+     "start_time": "2024-02-21T11:36:52.145834500Z"
     }
    }
   },
@@ -504,9 +547,8 @@
     }
    ],
    "source": [
-    "# Run forward pass\n",
-    "output = model.forward(tensor)\n",
-    "print(output)\n"
+    "# Run model\n",
+    "model(input_data=ecg_signal, frequency=100.0)"
    ],
    "metadata": {
     "collapsed": false,

src/ecg2hrv.py

@@ -0,0 +1,54 @@
+import numpy as np
+import pandas as pd
+import neurokit2 as nk
+import torch
+
+from src.feature_extractor import FeatureExtractor
+
+
+class ECG2HRV(FeatureExtractor):
+    def __init__(self):
+        super().__init__()
+
+    def extract_features(self, ecg, frequency, baseline=None, normalization_method=None):
+        # Ensure the numpy has at least one dimension (i.e. is not a scalar)
+        if ecg.ndim < 1:
+            raise ValueError("Array must have at least one dimension")
+
+        # Preprocess the ecg signal
+        ecg = nk.ecg_clean(ecg_signal=ecg, sampling_rate=frequency, method="pantompkins1985")
+
+        # Compute the HRV features
+        features = self.get_hrv_features(ecg, frequency)
+
+        # Normalize if baseline is available and method is set - TBD
+        if baseline is not None and normalization_method is not None:
+            baseline_features = self.get_hrv_features(baseline)
+            features = self.normalize_features(features, baseline_features)
+
+        return features
+
+    def get_hrv_features(self, ecg, frequency):
+        # Find peaks
+        peaks, info = nk.ecg_peaks(ecg, sampling_rate=frequency, method="pantompkins1985")
+
+        # Compute time domain features
+        hrv_time_features = nk.hrv_time(peaks, sampling_rate=frequency)
+        # Compute frequency domain features
+        # hrv_frequency_features = nk.hrv_frequency(peaks, sampling_rate=fs, method="welch", show=False)
+
+        # Concat features
+        # hrv_features = pd.concat([hrv_time_features, hrv_frequency_features], axis=1)
+        hrv_features = hrv_time_features
+
+        return hrv_features.to_dict(orient="records")
+
+    def normalize_features(self, features, baseline_features, normalization_method=None):
+        if normalization_method == "difference":
+            features = features - baseline_features
+        elif normalization_method == "relative":
+            features = features / baseline_features
+        else:
+            raise ValueError(f"Normalization method {normalization_method} not supported")
+
+        return features

src/feature_extractor.py

@@ -0,0 +1,17 @@
+from abc import ABC, abstractmethod
+
+
+class FeatureExtractor(ABC):
+    def __init__(self):
+        pass
+
+    def __call__(self, input_data, frequency, baseline_data=None, normalization_method=None):
+        return self.extract_features(input_data, frequency, baseline_data, normalization_method)
+
+    @abstractmethod
+    def extract_features(self, input_data, baseline_data, frequency, normalization_method):
+        pass
+
+    @abstractmethod
+    def normalize_features(self, features, baseline_features=None):
+        pass