Merge pull request #22 from AmbiqAI/feat/2023q3

Version 1.1.0 Release
AmbiqAI · Aug 11, 2023 · fd663d9 · fd663d9
2 parents 13b09d3 + 408b5da
commit fd663d9
Show file tree

Hide file tree

Showing 107 changed files with 24,909 additions and 26,218 deletions.
diff --git a/.devcontainer/devcontainer.json b/.devcontainer/devcontainer.json
@@ -1,6 +1,6 @@
 {
     "name": "Python 3",
-    "image": "mcr.microsoft.com/devcontainers/python:0-3.10",
+    "image": "mcr.microsoft.com/devcontainers/python:1-3.11",
 
     // Features: https://containers.dev/features.
     "features": {
@@ -24,7 +24,9 @@
         }
     },
 
-    "runArgs": [ "--gpus", "all"],
+    "hostRequirements": {
+        "gpu": "optional"
+    },
 
     "mounts": [
         "source=${localEnv:HOME}${localEnv:USERPROFILE}/tensorflow_datasets,target=/home/vscode/tensorflow_datasets,type=bind,consistency=cached"

diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -13,7 +13,7 @@ jobs:
         run: pipx install poetry==1.3.2
       - uses: actions/setup-python@v4
         with:
-          python-version: '3.10'
+          python-version: '3.11'
           cache: 'poetry'
       - run: poetry install
       - run: poetry run task lint

diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml
@@ -16,14 +16,15 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v3
-      - name: Install poetry
+      - name: Install Poetry
         run: pipx install poetry==1.3.2
-      - name: Install heartkit
+      - name: Install Python
         uses: actions/setup-python@v4
         with:
-          python-version: '3.10'
+          python-version: '3.11'
           cache: 'poetry'
-      - run: |
+      - name: Build Docs
+        run: |
           git config --global --add safe.directory "$GITHUB_WORKSPACE"
           git config --global user.email "adam.page@ambiq.com"
           git remote rm origin

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -10,11 +10,12 @@ repos:
     rev: 23.3.0
     hooks:
       - id: black
-        language_version: python3
+        language_version: python3.11
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v4.4.0
     hooks:
       - id: check-added-large-files
+        args: ["--maxkb=800"]
       - id: check-toml
       - id: check-yaml
         args: [--unsafe]

diff --git a/README.md b/README.md
@@ -22,7 +22,7 @@ The architecture consists of an **ECG segmentation** model followed by three ups
 
 ## Requirements
 
-* [Python 3.10+](https://www.python.org)
+* [Python 3.11+](https://www.python.org)
 * [Poetry 1.2.1+](https://python-poetry.org/docs/#installation)
 
 The following are also required to compile/flash the binary for the EVB demo:
@@ -99,7 +99,7 @@ heartkit --task arrhythmia --mode train --config ./configs/train-arrhythmia-mode
 The `evaluate` command will evaluate the performance of the model on the reserved test set. A confidence threshold can also be set such that a label is only assigned when the model's probability is greater than the threshold; otherwise, a label of inconclusive will be assigned.
 
 ```bash
-heartkit --task arrhythmia --mode evaluate --config ./configs/test-arrhythmia-model.json
+heartkit --task arrhythmia --mode evaluate --config ./configs/evaluate-arrhythmia-model.json
 ```
 
 #### __4. Export Model__
@@ -133,19 +133,18 @@ HeartKit leverages a multi-head network- a backbone segmentation model followed
 
 HeartKit leverages several open-source datasets for training each of the HeartKit models. Additionally, HeartKit contains a customizable synthetic 12-lead ECG generator. Check out the [Datasets Guide](./datasets.md) to learn more about the datasets used along with their corresponding licenses and limitations.
 
-## Results Summary
+## Results
 
 The following table provides the latest performance and accuracy results of all models when running on Apollo4 Plus EVB.
 
-| Task           | Params   | FLOPS   | Metric      |
-| -------------- | -------- | ------- | ----------- |
-| Segmentation   | 105K     | 19.3M   | IOU=85.3%   |
-| Arrhythmia     | 76K      | 7.2M    | F1=99.4%    |
-| Beat           | 79K      | 1.6M    | F1=91.6%    |
-| HRV            | N/A      | N/A     | N/A         |
+| Task           | Params   | FLOPS   | Metric     | Cycles/Inf | Time/Inf   |
+| -------------- | -------- | ------- | ---------- | ---------- | ---------- |
+| Segmentation   | 33K      | 6.5M    | 87.0% IOU  | 531ms      | 102M       |
+| Arrhythmia     | 50K      | 3.6M    | 99.0% F1   | 465ms      | 89M        |
+| Beat           | 73K      | 2.2M    | 91.5% F1   | 241ms      | 46M        |
 
 
-## Reference Papers
+## References
 
 * [ECG Heartbeat classification using deep transfer learning with Convolutional Neural Network and STFT technique](https://arxiv.org/abs/2206.14200)
 * [Classification of ECG based on Hybrid Features using CNNs for Wearable Applications](https://arxiv.org/pdf/2206.07648.pdf)

diff --git a/configs/arrhythmia-demo.json b/configs/arrhythmia-demo.json
@@ -10,5 +10,5 @@
     "vid_pid": "51966:16385",
     "baudrate": "115200",
     "data_parallelism": 1,
-    "datasets": ["icentia11k"]
+    "dataset": "icentia11k"
 }
diff --git a/configs/evaluate-arrhythmia-model.json b/configs/evaluate-arrhythmia-model.json
@@ -1,11 +1,12 @@
 {
     "job_dir": "./results/arrhythmia",
     "ds_path": "./datasets",
-    "sampling_rate": 250,
-    "frame_size": 1000,
-    "samples_per_patient": [100, 800, 400],
+    "sampling_rate": 200,
+    "frame_size": 800,
+    "samples_per_patient": [50, 400, 200],
     "test_patients": 1000,
     "test_size": 100000,
     "model_file": "./results/arrhythmia/model.tf",
-    "threshold": 0.95
+    "model_file_rmt": "wandb://ambiq/model-registry/heartkit-arrhythmia:v0",
+    "threshold": 0.75
 }
diff --git a/configs/evaluate-beat-model.json b/configs/evaluate-beat-model.json
@@ -1,10 +1,12 @@
 {
     "job_dir": "./results/beat",
     "ds_path": "./datasets",
-    "sampling_rate": 250,
-    "frame_size": 200,
-    "samples_per_patient": [200, 1000, 1000],
+    "sampling_rate": 200,
+    "frame_size": 160,
+    "samples_per_patient": [50, 100, 100],
     "test_patients": 1000,
-    "test_size": 100000,
-    "model_file": "./results/beat/model.tf"
+    "test_size": 50000,
+    "model_file": "./results/beat/model.tf",
+    "model_file_rmt": "wandb://ambiq/model-registry/heartkit-beat:v1",
+    "threshold": 0.50
 }
diff --git a/configs/evaluate-segmentation-model.json b/configs/evaluate-segmentation-model.json
@@ -1,9 +1,14 @@
 {
     "job_dir": "./results/segmentation",
     "ds_path": "./datasets",
-    "sampling_rate": 250,
-    "frame_size": 624,
-    "samples_per_patient": 200,
-    "test_size": 4000,
-    "model_file": "./results/segmentation/model.tf"
+    "sampling_rate": 200,
+    "frame_size": 512,
+    "samples_per_patient": 100,
+    "test_size": 10000,
+    "num_pts": 600,
+    "use_logits": false,
+    "datasets": ["synthetic", "ludb"],
+    "model_file": "./results/segmentation/model.tf",
+    "model_file_rmt": "wandb://ambiq/model-registry/heartkit-segmentation:v1",
+    "threshold": 0.50
 }
diff --git a/configs/export-arrhythmia-model.json b/configs/export-arrhythmia-model.json
@@ -1,13 +1,17 @@
 {
     "job_dir": "./results/arrhythmia",
     "ds_path": "./datasets",
-    "sampling_rate": 250,
-    "frame_size": 1000,
-    "samples_per_patient": [100, 500, 100],
-    "model_file_rmt": "wandb://ambiq/model-registry/ecg-arrhythmia-afib:v2",
+    "sampling_rate": 200,
+    "frame_size": 800,
+    "samples_per_patient": [5, 40, 20],
+    "test_patients": 1000,
+    "test_size": 10000,
     "model_file": "./results/arrhythmia/model.tf",
+    "model_file_rmt": "wandb://ambiq/model-registry/heartkit-arrhythmia:v0",
     "quantization": true,
-    "threshold": 0.95,
+    "use_logits": false,
+    "threshold": 0.80,
+    "val_acc_threshold": 0.98,
     "tflm_var_name": "g_arrhythmia_model",
     "tflm_file": "./evb/src/arrhythmia_model_buffer.h"
 }
diff --git a/configs/export-beat-model.json b/configs/export-beat-model.json
@@ -1,11 +1,17 @@
 {
     "job_dir": "./results/beat",
     "ds_path": "./datasets",
-    "sampling_rate": 250,
-    "frame_size": 200,
-    "samples_per_patient": [25, 125, 125],
+    "sampling_rate": 200,
+    "frame_size": 160,
+    "samples_per_patient": [10, 20, 20],
+    "test_patients": 1000,
+    "test_size": 10000,
     "model_file": "./results/beat/model.tf",
+    "model_file_rmt": "wandb://ambiq/model-registry/heartkit-beat:v1",
     "quantization": true,
+    "use_logits": false,
+    "threshold": 0.50,
+    "val_acc_threshold": 0.98,
     "tflm_var_name": "g_beat_model",
     "tflm_file": "./evb/src/beat_model_buffer.h"
 }
diff --git a/configs/export-segmentation-model.json b/configs/export-segmentation-model.json
@@ -1,11 +1,19 @@
 {
     "job_dir": "./results/segmentation",
     "ds_path": "./datasets",
-    "sampling_rate": 250,
-    "frame_size": 624,
-    "samples_per_patient": 200,
+    "sampling_rate": 200,
+    "frame_size": 512,
+    "samples_per_patient": 100,
+    "num_pts": 400,
+    "test_size": 8000,
+    "datasets": ["synthetic", "ludb"],
     "model_file": "./results/segmentation/model.tf",
+    "model_file_rmt": "wandb://ambiq/model-registry/heartkit-segmentation:v1",
     "quantization": true,
+    "use_logits": false,
+    "threshold": 0.50,
+    "val_acc_threshold": 0.98,
     "tflm_var_name": "g_segmentation_model",
     "tflm_file": "./evb/src/segmentation_model_buffer.h"
+
 }
diff --git a/configs/finetune-segmentation-model.json b/configs/finetune-segmentation-model.json
diff --git a/configs/heartkit-demo.json b/configs/heartkit-demo.json
@@ -7,10 +7,11 @@
     "arrhythmia_model": "./results/arrhythmia/model.tf",
     "segmentation_model": "./results/segmentation/model.tf",
     "beat_model": "./results/beat/model.tf",
-    "samples_per_patient": 1,
-    "frame_size": 2500,
+    "samples_per_patient": 5,
+    "sampling_rate": 200,
+    "frame_size": 2000,
     "vid_pid": "51966:16385",
     "baudrate": "115200",
     "data_parallelism": 1,
-    "datasets": ["icentia11k"]
+    "dataset": "ludb"
 }
diff --git a/configs/pretrain-segmentation-model.json b/configs/pretrain-segmentation-model.json
@@ -0,0 +1,57 @@
+{
+    "job_dir": "./results/segmentation-pre",
+    "ds_path": "./datasets",
+    "sampling_rate": 200,
+    "frame_size": 512,
+    "samples_per_patient": 100,
+    "val_samples_per_patient": 100,
+    "val_patients": 0.10,
+    "batch_size": 512,
+    "buffer_size": 25000,
+    "epochs": 80,
+    "steps_per_epoch": 100,
+    "val_metric": "loss",
+    "datasets": ["ludb", "synthetic"],
+    "lr_rate": 5e-3,
+    "num_pts": 600,
+    "quantization": false,
+    "augmentations": [
+        {
+            "name": "baseline_wander",
+            "args": {
+                "amplitude": [0.5, 1.5],
+                "frequency": [0.4, 0.5]
+            }
+        },
+        {
+            "name": "motion_noise",
+            "args": {
+                "amplitude": [0.5, 1.5],
+                "frequency": [0.4, 0.7]
+            }
+        },
+        {
+            "name": "burst_noise",
+            "args": {
+                "burst_number": [2, 10],
+                "amplitude": [0.5, 1.5],
+                "frequency": [40, 100]
+            }
+        },
+        {
+            "name": "powerline_noise",
+            "args": {
+                "amplitude": [0.005, 0.01],
+                "frequency": [50, 60]
+            }
+        },
+        {
+            "name": "noise_sources",
+            "args": {
+                "num_sources": [1, 8],
+                "amplitude": [0.05, 0.25],
+                "frequency": [10, 40]
+            }
+        }
+    ]
+}
diff --git a/configs/train-arrhythmia-model.json b/configs/train-arrhythmia-model.json
@@ -1,18 +1,58 @@
 {
     "job_dir": "./results/arrhythmia",
     "ds_path": "./datasets",
-    "sampling_rate": 250,
-    "frame_size": 1000,
-    "samples_per_patient": [100, 800, 400],
-    "val_samples_per_patient": [100, 800, 400],
+    "sampling_rate": 200,
+    "frame_size": 800,
+    "samples_per_patient": [100, 800, 800],
+    "val_samples_per_patient": [100, 800, 800],
     "train_patients": 10000,
-    "val_file": "./results/arrhythmia-10000pt-4s.pkl",
-    "val_patients": 0.20,
-    "val_size": 100000,
-    "batch_size": 6144,
+    "val_file": "./results/arrhythmia-10000pt-200fs-4s.pkl",
+    "val_patients": 0.10,
+    "val_size": 200000,
+    "batch_size": 256,
     "buffer_size": 100000,
-    "epochs": 100,
-    "steps_per_epoch": 10,
+    "epochs": 200,
+    "steps_per_epoch": 20,
     "val_metric": "loss",
-    "datasets": ["icentia11k"]
+    "datasets": ["icentia11k"],
+    "lr_rate": 5e-3,
+    "augmentations": [
+        {
+            "name": "baseline_wander",
+            "args": {
+                "amplitude": [0.5, 1.0],
+                "frequency": [0.4, 0.5]
+            }
+        },
+        {
+            "name": "motion_noise",
+            "args": {
+                "amplitude": [0.2, 0.4],
+                "frequency": [0.4, 0.6]
+            }
+        },
+        {
+            "name": "burst_noise",
+            "args": {
+                "burst_number": [0, 4],
+                "amplitude": [0.05, 0.5],
+                "frequency": [80, 100]
+            }
+        },
+        {
+            "name": "powerline_noise",
+            "args": {
+                "amplitude": [0.005, 0.01],
+                "frequency": [50, 60]
+            }
+        },
+        {
+            "name": "noise_sources",
+            "args": {
+                "num_sources": [1, 2],
+                "amplitude": [0.04, 0.1],
+                "frequency": [10, 40]
+            }
+        }
+    ]
 }