""" 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 # -------------------------------------------------------------------------- # Constants # -------------------------------------------------------------------------- # Default source of the dataset: a Google Drive folder shared as # "Anyone with the link" -> "Viewer". Overridable via --gdrive-url or env var. DEFAULT_GDRIVE_URL = ( "https://drive.google.com/drive/u/1/folders/1CZK9OQzsIM0cBwWeZ9RtIdG7LmebpBr9" ) # Rhythm columns present in the source dataset (used to compute ground truth). GROUND_TRUTH_COLUMNS = ["VT", "SVT", "AFIB", "AFLT"] # Labels the clinician can assign. Order here defines the "1 / 2 / 3" shortcuts. DOCTOR_LABELS = ["VT", "SVT", "Others"] CSV_COLUMNS = [ "sample_index", "dataset", "record", "ground_truth", "doctor_VT", "doctor_SVT", "doctor_Others", "comments", "reviewed", "annotated_at", ] # Injected once into . Handles keyboard shortcuts without interfering # with typing in the comments textbox. KEYBOARD_JS = """ """ 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 acquisition (download from Google Drive, then cache locally) # -------------------------------------------------------------------------- # Files that mark the root of an HF `save_to_disk()` dataset directory. _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 `/test_ds/dataset_info.json` instead of `/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 # ambiguous or already at the right level; let load_from_disk report errors 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) # already downloaded in a prior run 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 # -------------------------------------------------------------------------- # Data / persistence helpers # -------------------------------------------------------------------------- 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 # Dataset size changed since the CSV was created -- rebuild to stay safe. 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 # -------------------------------------------------------------------------- # Rendering helpers # -------------------------------------------------------------------------- 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
amp=%{y:.3f}", ) ) 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"Sample {idx + 1} / {total}", f"Dataset: {row.get('dataset', '')}", f"Record: {row.get('record', '')}", f"Heart Rate: {hr} bpm", f"Ground Truth: {gt_text}", ] items = "".join(f'{f}' for f in fields) return f'

{items}

' 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 ( '

' '

' f"Reviewed{reviewed} / {total} ({pct:.1f}%)" "

" '

' f'

' "

" ) # -------------------------------------------------------------------------- # App # -------------------------------------------------------------------------- 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") # -- Event handlers ------------------------------------------------- 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) # save current sample first 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 # -------------------------------------------------------------------------- # Entry point # -------------------------------------------------------------------------- # Load dataset 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)) # Create annotation directory 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" # Build demo demo = build_app(dataset, csv_path, state_path) # Launch demo.launch( server_name="0.0.0.0", server_port=7860, css=CUSTOM_CSS, head=KEYBOARD_JS, )