From dcb46ec908f346b611887f17d2a843fb7a180f6d Mon Sep 17 00:00:00 2001 From: Asger Gitz-Johansen Date: Wed, 29 Jul 2026 22:25:15 +0200 Subject: [PATCH] feat: first training attempt It still sucks. I'm sure the the feature embedding is terribly wrong. --- .gitignore | 4 + pyproject.toml | 5 +- scripts/train_from_data.sh | 88 +++++ src/sleep_detection/main.py | 342 ++++++++++++++++++- src/sleep_detection/sleep_epoch_classifer.py | 40 +++ 5 files changed, 476 insertions(+), 3 deletions(-) create mode 100755 scripts/train_from_data.sh create mode 100644 src/sleep_detection/sleep_epoch_classifer.py diff --git a/.gitignore b/.gitignore index 505a3b1..94ce3d9 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,7 @@ wheels/ # Virtual environments .venv + +# Models +**/*.pt +*.pt diff --git a/pyproject.toml b/pyproject.toml index 320eff7..8ddfdac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,10 @@ authors = [ ] requires-python = ">=3.12" dependencies = [ - "torch" + "torch", + "numpy", + "pandas", + "loguru" ] [project.scripts] diff --git a/scripts/train_from_data.sh b/scripts/train_from_data.sh new file mode 100755 index 0000000..9fc5da3 --- /dev/null +++ b/scripts/train_from_data.sh @@ -0,0 +1,88 @@ +#!/usr/bin/env bash + +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +DATA_DIR="${ROOT_DIR}/data" +OUTPUT_FILE="${ROOT_DIR}/sleep-model.pt" +MODEL_FILE="${OUTPUT_FILE}" + +if [[ ${1:-} == "-h" || ${1:-} == "--help" ]]; then + cat <<'EOF' +Usage: scripts/train_from_data.sh [-o OUTPUT_FILE] + +Trains the model on complete session triplets found under data/{motion,heart_rate,labels}/. +EOF + exit 0 +fi + +if [[ ${1:-} == "-o" || ${1:-} == "--output" ]]; then + OUTPUT_FILE="${2:?Missing output file after -o/--output}" + MODEL_FILE="${OUTPUT_FILE}" +fi + +shopt -s nullglob + +checkpoint_class_count() { + python - "$1" <<'PY' +from pathlib import Path +import sys +import torch + +payload = torch.load(Path(sys.argv[1]), weights_only=True) +state_dict = payload.get("model_state_dict", payload) +print(state_dict["net.2.bias"].shape[0]) +PY +} + +motions=("${DATA_DIR}/motion"/*_acceleration.txt) + +subject_ids=() + +for motion_file in "${motions[@]}"; do + sample_id="$(basename "${motion_file}" _acceleration.txt)" + heartrate_file="${DATA_DIR}/heart_rate/${sample_id}_heartrate.txt" + label_file="${DATA_DIR}/labels/${sample_id}_labeled_sleep.txt" + + if [[ -f "${heartrate_file}" && -f "${label_file}" ]]; then + subject_ids+=("${sample_id}") + fi +done + +if [[ ${#subject_ids[@]} -eq 0 ]]; then + printf 'No complete data triplets found under %s\n' "${DATA_DIR}" >&2 + exit 1 +fi + +mapfile -t subject_ids < <(printf '%s\n' "${subject_ids[@]}" | sort -u) + +printf 'Training on %d subject(s)\n' "${#subject_ids[@]}" + +export PYTHONPATH="${ROOT_DIR}/src${PYTHONPATH:+:${PYTHONPATH}}" + +for subject_id in "${subject_ids[@]}"; do + acceleration_file="${DATA_DIR}/motion/${subject_id}_acceleration.txt" + heartrate_file="${DATA_DIR}/heart_rate/${subject_id}_heartrate.txt" + label_file="${DATA_DIR}/labels/${subject_id}_labeled_sleep.txt" + resume_args=() + + printf 'Training subject %s\n' "${subject_id}" + + if [[ -f "${MODEL_FILE}" ]]; then + checkpoint_classes="$(checkpoint_class_count "${MODEL_FILE}")" + if [[ "${checkpoint_classes}" == "3" ]]; then + resume_args=(-i "${MODEL_FILE}") + else + printf 'Skipping checkpoint for subject %s: model classes=%s, expected=3\n' "${subject_id}" "${checkpoint_classes}" + fi + fi + + args=(python -m sleep_detection train -a "${acceleration_file}" -H "${heartrate_file}" -l "${label_file}" -o "${MODEL_FILE}") + if [[ ${#resume_args[@]} -gt 0 ]]; then + args+=("${resume_args[@]}") + fi + + "${args[@]}" +done + +printf 'Saved trained model to %s\n' "${MODEL_FILE}" diff --git a/src/sleep_detection/main.py b/src/sleep_detection/main.py index 1138c3e..b14e0d5 100644 --- a/src/sleep_detection/main.py +++ b/src/sleep_detection/main.py @@ -1,6 +1,344 @@ """Main Entrypoint codebase.""" +import sys +from argparse import ArgumentParser +from argparse import Namespace +from collections.abc import Sequence +from pathlib import Path -def main(): +import numpy +import pandas +import torch +from loguru import logger +from torch import Tensor +from torch.nn import CrossEntropyLoss +from torch.nn import Module +from torch.optim import SGD + +from sleep_detection.sleep_epoch_classifer import SleepEpochClassifier + +SLEEP_STAGE_LABELS = [0, 1, 2] + + +def normalize_sleep_stage(stage) -> float: + """Map raw sleep stages to the 3-class target space. + + Raw labels: + 0 -> 0 + 1, 2, 3, 4 -> 1 + 5, 6 -> 2 + Other values are treated as unknown. + """ + if pandas.isna(stage): + return numpy.nan + + value = int(stage) + if value == 0: + return 0.0 + if 1 <= value <= 4: + return 1.0 + if 5 <= value <= 6: + return 2.0 + + return numpy.nan + + +@logger.catch +def main() -> None: """The Main entrypoint function.""" - print("Hello from sleep-detection!") + cli = parse_args() + logger.info("welcome") + if cli.command == "train": + sessions = load_sessions(cli.a, cli.H, cli.l) + epochs, features, target, label_mapping = featureize_sessions(sessions) + + c = SleepEpochClassifier( + n_features=features.shape[1], + n_classes=len(SLEEP_STAGE_LABELS), + hidden_size=8, + ) + if cli.input is not None: + c.load(cli.input) + + logger.info(f"Using a {c.__class__.__name__}") + + train(features, target, c) + save_network(cli.output, c, label_mapping) + return + + if cli.command == "execute": + adf = load_acceleration_data(cli.a) + hdf = load_heartrate_data(cli.H) + ldf = load_labels_data(cli.l) + + df = adf.join(hdf, how="outer").join(ldf, how="outer").sort_index() + epochs, features, target, label_mapping = featureize(df) + + c = SleepEpochClassifier( + n_features=features.shape[1], + n_classes=len(SLEEP_STAGE_LABELS), + hidden_size=8, + ) + loaded_label_mapping = c.load(cli.model) + logger.info(f"Using a {c.__class__.__name__}") + execute(epochs, features, target, c, loaded_label_mapping or label_mapping) + return + + raise ValueError(f"Unknown command: {cli.command}") + + +def parse_args() -> Namespace: + parser = ArgumentParser(description="Training program.") + parser.add_argument("-v", "--verbose", help="Increase verbosity.", action="store_true") + parser.add_argument("-V", "--version", action="store_true", help="Show version") + command = parser.add_subparsers(dest="command") + + train = command.add_parser("train", help="Train the model.") + train.add_argument("-a", type=Path, nargs="+", required=True, help="Acceleration data file(s).") + train.add_argument("-H", type=Path, nargs="+", required=True, help="Heartrate data file(s).") + train.add_argument("-l", type=Path, nargs="+", required=True, help="Labels data file(s).") + train.add_argument("-o", "--output", type=Path, default=None, required=True, help="Output .pt file.") + train.add_argument("-i", "--input", type=Path, default=None, required=False, help="Input .pt file.") + + execute = command.add_parser("execute", help="Execute the model.") + execute.add_argument("-m", "--model", type=Path, default=None, required=True, help="Input .pr file to use.") + execute.add_argument("-a", type=Path, default=None, required=True, help="Acceleration data file.") + execute.add_argument("-H", type=Path, default=None, required=True, help="Heartrate data file.") + execute.add_argument("-l", type=Path, default=None, required=True, help="Labels data file.") + + return parser.parse_args() + + +def featureize(data: pandas.DataFrame) -> tuple[pandas.DataFrame, Tensor, Tensor, list]: + """Convert loaded data into input and target tensors. + + Args: + data: Loaded data, should contain the keys: [timestamp, heartrate, motion, x, y, z, stage] + + """ + logger.info(f"Featurizing the data...") + epochs = build_epochs(data) + + epochs["hr_std"] = epochs["hr_std"].fillna(0.0) + epochs = epochs.dropna(subset=["hr_mean", "motion_mean", "motion_max", "stage"]) + + epochs["seconds_since_start"] = epochs.index.total_seconds() + + CONTEXT = 2 # epochs each side -> 5 epochs total = 25 min window + FEATURE_COLS = ["hr_mean", "hr_std", "motion_mean", "motion_max", "seconds_since_start"] + + features = epochs[FEATURE_COLS].to_numpy(dtype=numpy.float32) + labels = epochs["stage"].astype(numpy.int64).to_numpy() + + X, y = [], [] + for i in range(CONTEXT, len(epochs) - CONTEXT): + window = features[i - CONTEXT : i + CONTEXT + 1] # shape: (2*CONTEXT+1, n_feats) + X.append(window.flatten()) # -> 1D vector + y.append(labels[i]) + + if not X: + raise ValueError("Not enough complete epochs to build a training set.") + + X = torch.tensor(numpy.stack(X), dtype=torch.float32) # shape: (n_epochs, (2*CONTEXT+1)*n_feats) + y = torch.tensor(y, dtype=torch.long) # shape: (n_epochs,) + return epochs, X, y, SLEEP_STAGE_LABELS + + +def featureize_sessions(dataframes: Sequence[pandas.DataFrame]) -> tuple[pandas.DataFrame, Tensor, Tensor, list]: + """Convert multiple loaded sessions into input and target tensors.""" + epochs_per_session = [] + for data in dataframes: + epochs = build_epochs(data) + + epochs["hr_std"] = epochs["hr_std"].fillna(0.0) + epochs = epochs.dropna(subset=["hr_mean", "motion_mean", "motion_max", "stage"]) + epochs["seconds_since_start"] = epochs.index.total_seconds() + epochs_per_session.append(epochs) + + if not epochs_per_session: + raise ValueError("No sessions provided.") + + combined_epochs = pandas.concat(epochs_per_session) + + CONTEXT = 2 # epochs each side -> 5 epochs total = 25 min window + FEATURE_COLS = ["hr_mean", "hr_std", "motion_mean", "motion_max", "seconds_since_start"] + + X, y = [], [] + for epochs in epochs_per_session: + features = epochs[FEATURE_COLS].to_numpy(dtype=numpy.float32) + labels = epochs["stage"].astype(numpy.int64).to_numpy() + + for i in range(CONTEXT, len(epochs) - CONTEXT): + window = features[i - CONTEXT : i + CONTEXT + 1] # shape: (2*CONTEXT+1, n_feats) + X.append(window.flatten()) # -> 1D vector + y.append(labels[i]) + + if not X: + raise ValueError("Not enough complete epochs to build a training set.") + + X = torch.tensor(numpy.stack(X), dtype=torch.float32) # shape: (n_epochs, (2*CONTEXT+1)*n_feats) + y = torch.tensor(y, dtype=torch.long) # shape: (n_epochs,) + return combined_epochs, X, y, SLEEP_STAGE_LABELS + + +def load_sessions( + acceleration_files: Sequence[Path], heartrate_files: Sequence[Path], label_files: Sequence[Path] +) -> list[pandas.DataFrame]: + """Load multiple session triplets.""" + if not (len(acceleration_files) == len(heartrate_files) == len(label_files)): + raise ValueError("Acceleration, heartrate, and label file counts must match.") + + sessions = [] + for acceleration_file, heartrate_file, label_file in zip(acceleration_files, heartrate_files, label_files): + adf = load_acceleration_data(acceleration_file) + hdf = load_heartrate_data(heartrate_file) + ldf = load_labels_data(label_file) + sessions.append(adf.join(hdf, how="outer").join(ldf, how="outer").sort_index()) + + return sessions + + +def build_epochs(data: pandas.DataFrame) -> pandas.DataFrame: + """Aggregate raw signals into 5-minute epochs.""" + EPOCH_LEN = pandas.Timedelta(minutes=5) + + return data.resample(EPOCH_LEN).agg( + hr_mean=("hr", "mean"), + hr_std=("hr", "std"), + motion_mean=("motion", "mean"), + motion_max=("motion", "max"), + stage=("stage", lambda s: s.dropna().mode().iloc[0] if len(s.dropna()) else numpy.nan), + ) + + +def execute(epochs: pandas.DataFrame, input: Tensor, target: Tensor, classifier: Module, label_mapping: list) -> None: + """Run the model and emit per-epoch predictions.""" + classifier.eval() + with torch.no_grad(): + output = classifier(input) + predicted = output.argmax(dim=1) + + decoded_predicted = [label_mapping[index] for index in predicted.tolist()] + decoded_target = [label_mapping[index] for index in target.tolist()] + + report = pandas.DataFrame( + { + "timestamp": epochs.index[2 : len(epochs) - 2].astype("timedelta64[s]") / numpy.timedelta64(1, "s"), + "predicted": decoded_predicted, + "actual": decoded_target, + } + ) + report.to_csv(sys.stdout, index=False) + + accuracy = (predicted == target).float().mean().item() + logger.info(f"accuracy={accuracy:.3f}") + + +def load_acceleration_data(data_file: Path) -> pandas.DataFrame: + """Load acceleration data file. + + Expected input format: + + Args: + data_file: The filepath for the file containing acceleration data. + + Returns: + A pandas data frame with the loaded acceleration data. + + """ + logger.info(f"Loading acceleration data file {data_file}") + df = pandas.read_csv(data_file, sep=r"\s+", names=["timestamp", "x", "y", "z"], header=None) + df["timestamp"] = pandas.to_numeric(df["timestamp"], errors="coerce") + df = df.dropna(subset=["timestamp"]) + df = df.sort_values("timestamp").set_index(pandas.to_timedelta(df["timestamp"], unit="s")) + df = df.drop(columns=["timestamp"]) + df.index.name = "timestamp" + # df["motion"] = df[["x", "y", "z"]].pow(2).sum(axis=1).pow(0.5) + df["motion"] = df[["x", "y", "z"]].diff().pow(2).sum(axis=1).pow(0.5) + return df + + +def load_heartrate_data(data_file: Path) -> pandas.DataFrame: + """Load heart rate data file. + + Expected input format: , + + Args: + data_file: The filepath for the file containing heartrate data. + + Returns: + A pandas data frame with the loaded heartrate data. + + """ + logger.info(f"Loading heartrate date file {data_file}") + df = pandas.read_csv(data_file, sep=",", names=["timestamp", "hr"], header=None) + df["timestamp"] = pandas.to_numeric(df["timestamp"], errors="coerce") + df = df.dropna(subset=["timestamp"]) + df = df.sort_values("timestamp").set_index(pandas.to_timedelta(df["timestamp"], unit="s")) + df = df.drop(columns=["timestamp"]) + df.index.name = "timestamp" + return df + + +def load_labels_data(data_file: Path) -> pandas.DataFrame: + """Load label data file. + + Expected input format: ,