| """ |
| ECG Annotation Application |
| =========================== |
| |
| A lightweight Gradio tool that lets a clinician review ECG rhythm segments |
| (one 10-second Lead II strip at a time) and correct mislabeled rhythm |
| annotations (VT / SVT / Others), with comments, autosave, and resume support. |
| |
| Run: |
| python ecg_annotation_app.py --dataset-path /path/to/test_ds |
| |
| See README.md for full instructions. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| from datetime import datetime, timezone |
| from pathlib import Path |
| from typing import Any |
|
|
| import gradio as gr |
| import gdown |
| import numpy as np |
| import pandas as pd |
| import plotly.graph_objects as go |
| from datasets import load_from_disk |
|
|
| |
| |
| |
|
|
| |
| |
| DEFAULT_GDRIVE_URL = ( |
| "https://drive.google.com/drive/u/1/folders/1CZK9OQzsIM0cBwWeZ9RtIdG7LmebpBr9" |
| ) |
|
|
| |
| GROUND_TRUTH_COLUMNS = ["VT", "SVT", "AFIB", "AFLT"] |
|
|
| |
| DOCTOR_LABELS = ["VT", "SVT", "Others"] |
|
|
| CSV_COLUMNS = [ |
| "sample_index", |
| "dataset", |
| "record", |
| "ground_truth", |
| "doctor_VT", |
| "doctor_SVT", |
| "doctor_Others", |
| "comments", |
| "reviewed", |
| "annotated_at", |
| ] |
|
|
| |
| |
| KEYBOARD_JS = """ |
| <script> |
| document.addEventListener("keydown", function (e) { |
| const active = document.activeElement; |
| const tag = active ? active.tagName.toLowerCase() : ""; |
| // Never hijack keys while the user is typing (comments box, etc.) |
| if (tag === "textarea" || tag === "input") return; |
| |
| if (e.key === "ArrowRight") { |
| document.getElementById("next-btn")?.click(); |
| } else if (e.key === "ArrowLeft") { |
| document.getElementById("prev-btn")?.click(); |
| } else if (["1", "2", "3"].includes(e.key)) { |
| const group = document.getElementById("rhythm-checkboxes"); |
| if (!group) return; |
| const inputs = group.querySelectorAll('input[type="checkbox"]'); |
| const idx = parseInt(e.key, 10) - 1; |
| inputs[idx]?.click(); |
| } |
| }); |
| </script> |
| """ |
|
|
| CUSTOM_CSS = """ |
| .gradio-container { max-width: 1500px !important; width: 97% !important; margin: auto !important; |
| padding-top: 6px !important; } |
| #rhythm-checkboxes label { font-size: 1.05em; } |
| #meta-line { font-size: 1.05em; padding: 8px 12px; border-radius: 8px; |
| background: rgba(128,128,128,0.08); margin-bottom: 4px; } |
| .gradio-container .form { gap: 6px !important; } |
| footer { display: none !important; } |
| """ |
|
|
|
|
| |
| |
| |
|
|
| |
| _DATASET_MARKER_FILES = {"dataset_info.json", "dataset_dict.json", "state.json"} |
|
|
|
|
| def _is_dataset_root(path: Path) -> bool: |
| return any((path / marker).exists() for marker in _DATASET_MARKER_FILES) |
|
|
|
|
| def _find_dataset_root(path: Path) -> Path: |
| """ |
| Google Drive folder downloads sometimes wrap the dataset in an extra |
| directory level (e.g. downloading folder "test_ds" produces |
| `<output>/test_ds/dataset_info.json` instead of `<output>/dataset_info.json`). |
| Walk down while there's exactly one subdirectory and no dataset marker |
| file at the current level, so `load_from_disk` gets the right path either way. |
| """ |
| current = path |
| while not _is_dataset_root(current): |
| subdirs = [p for p in current.iterdir() if p.is_dir()] |
| if len(subdirs) == 1: |
| current = subdirs[0] |
| else: |
| break |
| return current |
|
|
|
|
| def ensure_dataset_available(local_path: Path, gdrive_url: str | None) -> Path: |
| """ |
| Return a local path to the dataset, downloading it from Google Drive first |
| if it isn't already cached on disk. This lets the app (and a hosted Space) |
| fetch the dataset at runtime instead of committing large binary files to |
| version control. |
| |
| The Google Drive folder must be shared as "Anyone with the link" -> "Viewer". |
| """ |
| if local_path.exists() and any(local_path.iterdir()): |
| return _find_dataset_root(local_path) |
|
|
| if not gdrive_url: |
| raise FileNotFoundError( |
| f"No dataset found at '{local_path}' and no --gdrive-url was provided." |
| ) |
|
|
| print(f"Dataset not found at '{local_path}' -- downloading from Google Drive...") |
| local_path.mkdir(parents=True, exist_ok=True) |
| gdown.download_folder(url=gdrive_url, output=str(local_path), quiet=False, use_cookies=False) |
|
|
| root = _find_dataset_root(local_path) |
| if not _is_dataset_root(root): |
| raise FileNotFoundError( |
| f"Downloaded files from Google Drive but couldn't find a dataset at " |
| f"'{root}'. Check that the shared folder contains the dataset produced " |
| f"by `save_to_disk()` (it should have a dataset_info.json / state.json)." |
| ) |
| return root |
|
|
|
|
| |
| |
| |
|
|
| def ground_truth_labels(row: dict[str, Any]) -> list[str]: |
| """Return the list of positive rhythm columns for a dataset row.""" |
| return [col for col in GROUND_TRUTH_COLUMNS if int(row.get(col, 0) or 0) == 1] |
|
|
|
|
| def init_annotations(dataset, csv_path: Path) -> pd.DataFrame: |
| """Load an existing annotation CSV, or create a fresh one for the dataset.""" |
| if csv_path.exists(): |
| df = pd.read_csv(csv_path) |
| if len(df) == len(dataset): |
| df["reviewed"] = df["reviewed"].astype(bool) |
| for label in DOCTOR_LABELS: |
| df[f"doctor_{label}"] = df[f"doctor_{label}"].astype(bool) |
| df["comments"] = df["comments"].fillna("").astype(str) |
| return df |
| |
|
|
| rows = [] |
| for i in range(len(dataset)): |
| row = dataset[i] |
| rows.append( |
| { |
| "sample_index": i, |
| "dataset": row.get("dataset", ""), |
| "record": row.get("record", ""), |
| "ground_truth": ",".join(ground_truth_labels(row)), |
| "doctor_VT": False, |
| "doctor_SVT": False, |
| "doctor_Others": False, |
| "comments": "", |
| "reviewed": False, |
| "annotated_at": "", |
| } |
| ) |
| df = pd.DataFrame(rows, columns=CSV_COLUMNS) |
| df.to_csv(csv_path, index=False) |
| return df |
|
|
|
|
| def save_annotations(df: pd.DataFrame, csv_path: Path) -> None: |
| df.to_csv(csv_path, index=False) |
|
|
|
|
| def load_state(state_path: Path) -> dict: |
| if not state_path.exists(): |
| return {} |
| try: |
| return json.loads(state_path.read_text()) |
| except (json.JSONDecodeError, OSError): |
| return {} |
|
|
|
|
| def save_state(state_path: Path, sample_index: int) -> None: |
| state = { |
| "last_opened_sample": sample_index, |
| "last_saved": datetime.now(timezone.utc).isoformat(timespec="seconds"), |
| } |
| state_path.write_text(json.dumps(state, indent=2)) |
|
|
|
|
| def first_unreviewed_index(df: pd.DataFrame) -> int: |
| """Index of the first sample not yet reviewed, or 0 if all are done.""" |
| unreviewed = df.index[~df["reviewed"].astype(bool)] |
| return int(unreviewed[0]) if len(unreviewed) else 0 |
|
|
|
|
| |
| |
| |
|
|
| def build_ecg_figure(signal: np.ndarray, fs: float, dataset_name: str, record: str) -> go.Figure: |
| """Build an interactive, theme-agnostic Plotly strip of a Lead II signal.""" |
| t = np.arange(len(signal)) / fs if fs else np.arange(len(signal)) |
|
|
| fig = go.Figure() |
| fig.add_trace( |
| go.Scatter( |
| x=t, |
| y=signal, |
| mode="lines", |
| line=dict(width=1.4, color="#e63946"), |
| name="Lead II", |
| hovertemplate="t=%{x:.2f}s<br>amp=%{y:.3f}<extra></extra>", |
| ) |
| ) |
| fig.update_layout( |
| title=f"{dataset_name} — Record {record} (Lead II)", |
| xaxis_title="Time (s)", |
| yaxis_title="Amplitude (mV)", |
| margin=dict(l=55, r=20, t=36, b=36), |
| height=360, |
| dragmode="pan", |
| hovermode="x unified", |
| paper_bgcolor="rgba(0,0,0,0)", |
| plot_bgcolor="rgba(0,0,0,0)", |
| font=dict(color="#888888"), |
| showlegend=False, |
| ) |
| fig.update_xaxes(showgrid=True, gridcolor="rgba(128,128,128,0.25)", zeroline=False) |
| fig.update_yaxes(showgrid=True, gridcolor="rgba(128,128,128,0.25)", zeroline=False) |
| return fig |
|
|
|
|
| def render_meta_line(idx: int, total: int, row: dict, gt_labels: list[str]) -> str: |
| """Single-line metadata bar: sample counter, dataset, record, HR, ground truth.""" |
| hr = row.get("HR", "?") |
| gt_text = ", ".join(gt_labels) if gt_labels else "—" |
| fields = [ |
| f"<b>Sample</b> {idx + 1} / {total}", |
| f"<b>Dataset:</b> {row.get('dataset', '')}", |
| f"<b>Record:</b> {row.get('record', '')}", |
| f"<b>Heart Rate:</b> {hr} bpm", |
| f"<b>Ground Truth:</b> {gt_text}", |
| ] |
| items = "".join(f'<span style="white-space:nowrap;">{f}</span>' for f in fields) |
| return f'<div id="meta-line" style="display:flex;flex-wrap:wrap;gap:24px;align-items:center;">{items}</div>' |
|
|
|
|
| def render_progress(df: pd.DataFrame) -> str: |
| total = len(df) |
| reviewed = int(df["reviewed"].astype(bool).sum()) |
| pct = (reviewed / total * 100) if total else 0.0 |
| return ( |
| '<div style="width:100%;">' |
| '<div style="display:flex;justify-content:space-between;font-size:0.9em;' |
| 'margin-bottom:4px;opacity:0.85;">' |
| f"<span>Reviewed</span><span>{reviewed} / {total} ({pct:.1f}%)</span>" |
| "</div>" |
| '<div style="background:rgba(128,128,128,0.25);border-radius:6px;height:10px;' |
| 'overflow:hidden;">' |
| f'<div style="background:#2a9d8f;width:{pct:.1f}%;height:100%;' |
| 'transition:width 0.2s;"></div>' |
| "</div></div>" |
| ) |
|
|
|
|
| |
| |
| |
|
|
| def build_app(dataset, csv_path: Path, state_path: Path) -> gr.Blocks: |
| total = len(dataset) |
| annotations = init_annotations(dataset, csv_path) |
|
|
| state = load_state(state_path) |
| start_idx = state.get("last_opened_sample") |
| if start_idx is None or not (0 <= start_idx < total): |
| start_idx = first_unreviewed_index(annotations) |
|
|
| def checked_labels(df: pd.DataFrame, idx: int) -> list[str]: |
| ann = df.loc[df["sample_index"] == idx].iloc[0] |
| return [label for label in DOCTOR_LABELS if bool(ann[f"doctor_{label}"])] |
|
|
| def comment_for(df: pd.DataFrame, idx: int) -> str: |
| ann = df.loc[df["sample_index"] == idx].iloc[0] |
| value = ann["comments"] |
| return "" if pd.isna(value) else str(value) |
|
|
| def load_sample(idx: int, df: pd.DataFrame): |
| row = dataset[idx] |
| signal = np.asarray(row["II"], dtype=float) |
| fs = float(row.get("fs", 500) or 500) |
| fig = build_ecg_figure(signal, fs, row.get("dataset", ""), row.get("record", "")) |
| meta = render_meta_line(idx, total, row, ground_truth_labels(row)) |
| prog = render_progress(df) |
| return meta, fig, checked_labels(df, idx), comment_for(df, idx), prog |
|
|
| def persist(idx: int, df: pd.DataFrame, checked: list[str], comment: str) -> pd.DataFrame: |
| """Write the current UI state for `idx` into the dataframe and save to disk.""" |
| mask = df["sample_index"] == idx |
| for label in DOCTOR_LABELS: |
| df.loc[mask, f"doctor_{label}"] = label in checked |
| df.loc[mask, "comments"] = comment |
| df.loc[mask, "reviewed"] = True |
| df.loc[mask, "annotated_at"] = datetime.now(timezone.utc).isoformat(timespec="seconds") |
| save_annotations(df, csv_path) |
| return df |
|
|
| with gr.Blocks(title="ECG Annotation Review", fill_width=True) as demo: |
| df_state = gr.State(annotations) |
| idx_state = gr.State(start_idx) |
|
|
| meta_html = gr.HTML() |
| plot = gr.Plot(label=None, show_label=False) |
|
|
| with gr.Row(): |
| with gr.Column(scale=2): |
| checkboxes = gr.CheckboxGroup( |
| choices=DOCTOR_LABELS, |
| label="Doctor Annotation (shortcuts: 1 / 2 / 3)", |
| elem_id="rhythm-checkboxes", |
| ) |
| with gr.Column(scale=3): |
| comments = gr.Textbox( |
| label="Comments", |
| lines=2, |
| max_lines=2, |
| placeholder="Optional notes for this sample...", |
| ) |
|
|
| with gr.Row(): |
| prev_btn = gr.Button("◀ Previous", elem_id="prev-btn") |
| progress_html = gr.HTML() |
| next_btn = gr.Button("Next ▶", elem_id="next-btn", variant="primary") |
|
|
| |
|
|
| def on_load(idx, df): |
| return load_sample(idx, df) |
|
|
| demo.load( |
| on_load, |
| inputs=[idx_state, df_state], |
| outputs=[meta_html, plot, checkboxes, comments, progress_html], |
| ) |
|
|
| def on_annotation_change(idx, df, checked, comment): |
| df = persist(idx, df, checked, comment) |
| save_state(state_path, idx) |
| return df, render_progress(df) |
|
|
| checkboxes.change( |
| on_annotation_change, |
| inputs=[idx_state, df_state, checkboxes, comments], |
| outputs=[df_state, progress_html], |
| ) |
| comments.change( |
| on_annotation_change, |
| inputs=[idx_state, df_state, checkboxes, comments], |
| outputs=[df_state, progress_html], |
| ) |
|
|
| def go_to(offset: int): |
| def _handler(idx, df, checked, comment): |
| df = persist(idx, df, checked, comment) |
| new_idx = min(max(idx + offset, 0), total - 1) |
| save_state(state_path, new_idx) |
| meta, fig, chk, com, prog = load_sample(new_idx, df) |
| return new_idx, df, meta, fig, chk, com, prog |
|
|
| return _handler |
|
|
| shared_outputs = [idx_state, df_state, meta_html, plot, checkboxes, comments, progress_html] |
| shared_inputs = [idx_state, df_state, checkboxes, comments] |
|
|
| next_btn.click(go_to(1), inputs=shared_inputs, outputs=shared_outputs) |
| prev_btn.click(go_to(-1), inputs=shared_inputs, outputs=shared_outputs) |
|
|
| return demo |
|
|
|
|
| |
| |
| |
| |
| dataset_root = ensure_dataset_available( |
| Path(os.environ.get("ECG_DATASET_PATH", "./test_ds")), |
| os.environ.get("ECG_GDRIVE_URL", DEFAULT_GDRIVE_URL), |
| ) |
|
|
| dataset = load_from_disk(str(dataset_root)) |
|
|
| |
| data_dir = Path("./annotation_data") |
| data_dir.mkdir(parents=True, exist_ok=True) |
|
|
| csv_path = data_dir / "annotations.csv" |
| state_path = data_dir / "state.json" |
|
|
| |
| demo = build_app(dataset, csv_path, state_path) |
|
|
| |
| demo.launch( |
| server_name="0.0.0.0", |
| server_port=7860, |
| css=CUSTOM_CSS, |
| head=KEYBOARD_JS, |
| ) |
