#!/usr/bin/env python3 """ ================================================================================ SLEEP APNEA / HYPOPNEA DETECTOR - SHHS DATASET ================================================================================ Entrenamiento completo con arquitectura optimizada 2024-2026: - TCN (dilated convolutions) para patrones locales multiescala - BiGRU para dependencias temporales largas - Temporal Self-Attention para enfocar en eventos relevantes - CWT multi-view branch (opcional) para features espectrales - Focal Loss para desbalance masivo - Augmentación temporal: scaling, jittering, masking - Calibración de threshold en validación - Split por paciente para evitar leakage - Métricas clínicas: sensibilidad, especificidad, AUC-ROC, AUC-PR, AHI Diseñado para Hugging Face Jobs con GPU (T4/A10G). Integra Trackio para monitoreo y push_to_hub para guardar modelo. Uso en Hugging Face Jobs: python train_sleep_apnea.py --data_dir /app/data --output_dir /app/output ================================================================================ """ import os import sys import glob import json import warnings import random import math import argparse from pathlib import Path from collections import defaultdict from typing import Dict, List, Tuple, Optional import numpy as np import scipy.io as sio from scipy import signal as scipy_signal import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts from sklearn.model_selection import GroupShuffleSplit from sklearn.metrics import ( confusion_matrix, roc_auc_score, average_precision_score, accuracy_score, precision_score, recall_score, f1_score, classification_report, roc_curve, precision_recall_curve, cohen_kappa_score ) import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt # Trackio para monitoreo en HF (instalar vía requirements) try: import trackio TRACKIO_AVAILABLE = True except ImportError: TRACKIO_AVAILABLE = False trackio = None # ============================================================================ # CONFIGURACIÓN POR DEFECTO # ============================================================================ class Config: # Rutas data_dir: str = "/app/data" output_dir: str = "/app/output" hub_model_id: Optional[str] = None # Ej: "usuario/sleep-apnea-shhs" # Señales fs: int = 1 channels: List[str] = ["HR", "SaO2"] target_var: str = "Target03" quality_var: str = "OXStat" # Ventaneo temporal (15 min ventanas, 1 min stride) window_size: int = 900 stride: int = 60 label_shift: int = 60 # Shift etiquetas hacia atrás (delay desaturación) min_valid_ratio: float = 0.6 # Preprocesamiento normalize_per_subject: bool = True cwt_scales: List[int] = [4, 8, 16, 32, 64] # Modelo model_name: str = "SleepApneaNetV2" input_channels: int = 2 tcn_channels: List[int] = [32, 64, 128, 256] tcn_kernel: int = 7 gru_hidden: int = 128 gru_layers: int = 2 attention_heads: int = 4 dropout: float = 0.4 use_cwt: bool = False # CWT requiere más VRAM; desactivar en T4 por defecto # Entrenamiento batch_size: int = 32 num_workers: int = 2 epochs: int = 100 lr: float = 1e-3 weight_decay: float = 1e-4 focal_gamma: float = 2.0 focal_alpha: float = 0.25 patience: int = 15 grad_clip: float = 1.0 # Augmentación aug_prob: float = 0.5 aug_scale_range: float = 0.05 aug_jitter_std: float = 0.02 aug_mask_prob: float = 0.1 # Calibración threshold calibrate_threshold: bool = True # Hardware / Reproducibilidad device: str = "cuda" if torch.cuda.is_available() else "cpu" seed: int = 42 # Logging / Trackio trackio_project: str = "sleep-apnea-shhs" trackio_space_id: Optional[str] = None # "usuario/ml-intern-XXXX" run_name: str = "shhs-tcn-bigru" # ============================================================================ # SEMILLAS # ============================================================================ def set_seed(seed: int = 42): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False # ============================================================================ # CARGA .MAT (SHHS) # ============================================================================ def load_mat_file(filepath: str) -> Optional[Dict]: """Carga archivo .mat de SHHS extrayendo HR, SaO2, Target03, OXStat.""" try: mat = sio.loadmat(filepath, squeeze_me=True, struct_as_record=False) data = {} for key in ["HR", "SaO2", "Target03", "OXStat"]: if key in mat: data[key] = mat[key] elif key.lower() in mat: data[key] = mat[key.lower()] # Buscar en estructuras anidadas si no están en raíz if not all(k in data for k in ["HR", "SaO2", "Target03", "OXStat"]): for var in mat.keys(): if not var.startswith('__'): val = mat[var] if hasattr(val, 'dtype') and val.dtype.names is not None: for key in ["HR", "SaO2", "Target03", "OXStat"]: if key in val.dtype.names: item = val[key] data[key] = item.item() if np.isscalar(item) or item.size == 1 else item if not all(k in data for k in ["HR", "SaO2", "Target03", "OXStat"]): return None for key in data: arr = np.asarray(data[key]) if arr.ndim > 1: arr = arr.ravel() data[key] = arr.astype(np.float32) # Alinear longitudes min_len = min(len(data[k]) for k in ["HR", "SaO2", "Target03", "OXStat"]) for key in data: data[key] = data[key][:min_len] return data except Exception as e: warnings.warn(f"Error cargando {filepath}: {e}") return None def parse_patient_id(filename: str) -> str: """Extrae ID paciente de nombre archivo SHHS.""" base = os.path.splitext(os.path.basename(filename))[0] if base.startswith("shhs"): parts = base.split("-") if len(parts) >= 2: return parts[1] return base # ============================================================================ # PREPROCESAMIENTO # ============================================================================ def zscore_normalize(signals: np.ndarray) -> np.ndarray: """Z-score por canal. signals: (T, C)""" out = signals.copy() for c in range(signals.shape[1]): m = out[:, c].mean() s = out[:, c].std() if s > 1e-8: out[:, c] = (out[:, c] - m) / s else: out[:, c] = out[:, c] - m return out def compute_cwt_features(sig: np.ndarray, scales: List[int]) -> np.ndarray: """CWT magnitude para un canal 1D. Retorna (n_scales, T).""" try: # Usar ricker (Mexican hat wavelet) por compatibilidad universal coeffs, _ = scipy_signal.cwt(sig, scipy_signal.ricker, scales) return np.abs(coeffs).astype(np.float32) except Exception: return np.zeros((len(scales), len(sig)), dtype=np.float32) # ============================================================================ # EXTRACCIÓN DE VENTANAS # ============================================================================ def create_windows(data: Dict, config: Config) -> List[Dict]: """Extrae ventanas solapadas con shift de etiquetas para delay desaturación.""" hr = data["HR"] spo2 = data["SaO2"] target = data["Target03"] oxstat = data["OXStat"] n = len(hr) ws = config.window_size if n < ws: return [] # Shift etiquetas hacia atrás (apnea en t se refleja en SpO2 en t+shift) shift = config.label_shift shifted = np.zeros_like(target, dtype=np.int32) if shift > 0 and shift < n: shifted[:n-shift] = (target[shift:] != 0).astype(np.int32) else: shifted = (target != 0).astype(np.int32) raw = np.stack([hr, spo2], axis=1).astype(np.float32) if config.normalize_per_subject: raw = zscore_normalize(raw) windows = [] for start in range(0, n - ws + 1, config.stride): end = start + ws sig_win = raw[start:end] qual_win = oxstat[start:end] tgt_win = shifted[start:end] valid_ratio = float((qual_win == 0).mean()) if valid_ratio < config.min_valid_ratio: continue label = int(tgt_win.any()) win = { "signals": sig_win, "label": label, "valid_ratio": valid_ratio, "start_idx": start, } if config.use_cwt: cwt_feats = [] for c in range(2): cwt_mag = compute_cwt_features(sig_win[:, c], config.cwt_scales) # Downsample CWT a ~300 puntos temporales para reducir VRAM if cwt_mag.shape[1] > 300: pool = cwt_mag.shape[1] // 300 cwt_mag = cwt_mag[:, ::pool] cwt_feats.append(cwt_mag) win["cwt"] = np.concatenate(cwt_feats, axis=0) windows.append(win) return windows # ============================================================================ # DATASET & AUGMENTACIÓN # ============================================================================ class SleepDataset(Dataset): def __init__(self, windows: List[Dict], config: Config, augment: bool = False): self.windows = windows self.config = config self.augment = augment self.labels = [w["label"] for w in windows] def __len__(self): return len(self.windows) def _augment(self, x: np.ndarray) -> np.ndarray: if np.random.rand() > self.config.aug_prob: return x out = x.copy() # Scaling global ±5% if np.random.rand() < 0.5: scale = 1.0 + np.random.uniform(-self.config.aug_scale_range, self.config.aug_scale_range) out *= scale # Jittering gaussiano if np.random.rand() < 0.5: out += np.random.normal(0, self.config.aug_jitter_std, out.shape) # Random temporal masking (simula artefactos oxímetro) if np.random.rand() < 0.3: mask_len = np.random.randint(10, 60) mask_start = np.random.randint(0, max(1, out.shape[0] - mask_len)) out[mask_start:mask_start+mask_len] = 0 return out def __getitem__(self, idx): w = self.windows[idx] sig = w["signals"].copy() if self.augment: sig = self._augment(sig) x = torch.from_numpy(sig).permute(1, 0) # (C, T) y = torch.tensor(w["label"], dtype=torch.long) if self.config.use_cwt and "cwt" in w: cwt = torch.from_numpy(w["cwt"]).unsqueeze(0) # (1, n_scales*2, ~300) return x, cwt, y return x, y def collate_fn(batch): if len(batch[0]) == 3: xs, cwts, ys = [], [], [] for item in batch: xs.append(item[0]); cwts.append(item[1]); ys.append(item[2]) max_cwt = max(c.shape[-1] for c in cwts) cwts_pad = [] for c in cwts: if c.shape[-1] < max_cwt: c = F.pad(c, (0, max_cwt - c.shape[-1])) cwts_pad.append(c) return torch.stack(xs), torch.stack(cwts_pad), torch.stack(ys) else: return torch.stack([item[0] for item in batch]), torch.stack([item[1] for item in batch]) # ============================================================================ # ARQUITECTURA: TCN + BiGRU + Attention + CWT (opcional) # ============================================================================ class TemporalBlock(nn.Module): def __init__(self, in_ch, out_ch, kernel=7, dilation=1, dropout=0.3): super().__init__() pad = (kernel - 1) * dilation // 2 self.conv1 = nn.Conv1d(in_ch, out_ch, kernel, padding=pad, dilation=dilation, bias=False) self.bn1 = nn.BatchNorm1d(out_ch) self.conv2 = nn.Conv1d(out_ch, out_ch, kernel, padding=pad, dilation=dilation, bias=False) self.bn2 = nn.BatchNorm1d(out_ch) self.drop = nn.Dropout(dropout) self.relu = nn.ReLU(inplace=True) self.shortcut = nn.Sequential() if in_ch != out_ch: self.shortcut = nn.Sequential(nn.Conv1d(in_ch, out_ch, 1, bias=False), nn.BatchNorm1d(out_ch)) def forward(self, x): out = self.relu(self.bn1(self.conv1(x))) out = self.bn2(self.conv2(out)) out = self.drop(out) return self.relu(out + self.shortcut(x)) class TCNEncoder(nn.Module): def __init__(self, in_ch, channels, kernel=7, dropout=0.3): super().__init__() layers = [] curr = in_ch for i, out_ch in enumerate(channels): dilation = 2 ** i layers.append(TemporalBlock(curr, out_ch, kernel, dilation, dropout)) layers.append(TemporalBlock(out_ch, out_ch, kernel, dilation, dropout)) curr = out_ch self.net = nn.Sequential(*layers) def forward(self, x): return self.net(x) class TemporalAttention(nn.Module): def __init__(self, d_model, heads=4, dropout=0.1): super().__init__() self.heads = heads self.d_head = d_model // heads self.scale = self.d_head ** -0.5 self.qkv = nn.Linear(d_model, d_model * 3, bias=False) self.proj = nn.Linear(d_model, d_model, bias=False) self.drop = nn.Dropout(dropout) self.norm = nn.LayerNorm(d_model) def forward(self, x): B, L, D = x.shape qkv = self.qkv(x).reshape(B, L, 3, self.heads, self.d_head).permute(2, 0, 3, 1, 4) q, k, v = qkv[0], qkv[1], qkv[2] attn = (q @ k.transpose(-2, -1)) * self.scale attn = attn.softmax(dim=-1) attn = self.drop(attn) out = (attn @ v).transpose(1, 2).reshape(B, L, D) out = self.proj(out) return self.norm(x + out) class CWTBranch(nn.Module): def __init__(self, n_scales_total, out_dim=64): super().__init__() self.conv1 = nn.Conv2d(1, 16, kernel_size=(3, 5), padding=(1, 2)) self.bn1 = nn.BatchNorm2d(16) self.conv2 = nn.Conv2d(16, 32, kernel_size=(3, 5), padding=(1, 2)) self.bn2 = nn.BatchNorm2d(32) # Reduce scales dimension pool_h = max(1, n_scales_total // 4) self.pool = nn.AdaptiveAvgPool2d((pool_h, 1)) self.proj = nn.Linear(32 * pool_h, out_dim) self.relu = nn.ReLU(inplace=True) def forward(self, x): # x: (B, 1, n_scales*2, T_cwt) out = self.relu(self.bn1(self.conv1(x))) out = self.relu(self.bn2(self.conv2(out))) out = self.pool(out) out = out.view(out.size(0), -1) return self.relu(self.proj(out)) class SleepApneaNetV2(nn.Module): def __init__(self, config: Config): super().__init__() self.config = config self.use_cwt = config.use_cwt self.tcn = TCNEncoder(config.input_channels, config.tcn_channels, config.tcn_kernel, config.dropout) tcn_out = config.tcn_channels[-1] self.gru = nn.GRU(tcn_out, config.gru_hidden, config.gru_layers, batch_first=True, bidirectional=True, dropout=config.dropout if config.gru_layers > 1 else 0) gru_out = config.gru_hidden * 2 self.attn = TemporalAttention(gru_out, config.attention_heads, config.dropout) if self.use_cwt: n_scales_total = len(config.cwt_scales) * 2 self.cwt_branch = CWTBranch(n_scales_total, out_dim=64) classifier_in = gru_out * 2 + 64 else: classifier_in = gru_out * 2 self.classifier = nn.Sequential( nn.LayerNorm(classifier_in), nn.Dropout(config.dropout), nn.Linear(classifier_in, 256), nn.ReLU(inplace=True), nn.Dropout(config.dropout), nn.Linear(256, 2), ) def forward(self, x_raw, x_cwt=None): c = self.tcn(x_raw) c = c.permute(0, 2, 1) g, _ = self.gru(c) g = self.attn(g) g_t = g.permute(0, 2, 1) pooled_avg = F.adaptive_avg_pool1d(g_t, 1).squeeze(-1) pooled_max = F.adaptive_max_pool1d(g_t, 1).squeeze(-1) pooled = torch.cat([pooled_avg, pooled_max], dim=-1) if self.use_cwt and x_cwt is not None: cwt_feat = self.cwt_branch(x_cwt) pooled = torch.cat([pooled, cwt_feat], dim=-1) return self.classifier(pooled) # ============================================================================ # FOCAL LOSS # ============================================================================ class FocalLoss(nn.Module): def __init__(self, alpha=0.25, gamma=2.0, weight=None, reduction='mean'): super().__init__() self.alpha = alpha self.gamma = gamma self.weight = weight self.reduction = reduction self.ce = nn.CrossEntropyLoss(weight=weight, reduction='none') def forward(self, inputs, targets): ce_loss = self.ce(inputs, targets) pt = torch.exp(-ce_loss) focal_term = (1 - pt) ** self.gamma loss = self.alpha * focal_term * ce_loss return loss.mean() if self.reduction == 'mean' else loss.sum() # ============================================================================ # MÉTRICAS # ============================================================================ def compute_metrics(y_true, y_pred, y_prob): metrics = { "accuracy": accuracy_score(y_true, y_pred), "precision": precision_score(y_true, y_pred, zero_division=0), "sensitivity": recall_score(y_true, y_pred, zero_division=0), "specificity": recall_score(y_true, y_pred, pos_label=0, zero_division=0), "f1": f1_score(y_true, y_pred, zero_division=0), "cohen_kappa": cohen_kappa_score(y_true, y_pred), } try: metrics["auc_roc"] = roc_auc_score(y_true, y_prob) except ValueError: metrics["auc_roc"] = float('nan') try: metrics["auc_pr"] = average_precision_score(y_true, y_prob) except ValueError: metrics["auc_pr"] = float('nan') metrics["confusion_matrix"] = confusion_matrix(y_true, y_pred).tolist() return metrics # ============================================================================ # EVALUACIÓN # ============================================================================ def evaluate(model, loader, config, criterion, threshold=0.5): model.eval() all_probs, all_preds, all_labels = [], [], [] total_loss = 0.0 with torch.no_grad(): for batch in loader: if len(batch) == 3: x, cwt, y = batch x, cwt, y = x.to(config.device), cwt.to(config.device), y.to(config.device) logits = model(x, cwt) else: x, y = batch x, y = x.to(config.device), y.to(config.device) logits = model(x) loss = criterion(logits, y) total_loss += loss.item() probs = F.softmax(logits, dim=-1)[:, 1].cpu().numpy() preds = (probs > threshold).astype(int) all_probs.append(probs) all_preds.append(preds) all_labels.append(y.cpu().numpy()) all_probs = np.concatenate(all_probs) all_preds = np.concatenate(all_preds) all_labels = np.concatenate(all_labels) avg_loss = total_loss / max(len(loader), 1) metrics = compute_metrics(all_labels, all_preds, all_probs) return avg_loss, metrics, all_probs, all_labels # ============================================================================ # CALIBRACIÓN DE THRESHOLD # ============================================================================ def calibrate_threshold(model, val_loader, config): _, _, probs, labels = evaluate(model, val_loader, config, nn.CrossEntropyLoss(), 0.5) best_thresh = 0.5 best_score = -1 for thresh in np.arange(0.05, 0.95, 0.01): preds = (probs > thresh).astype(int) sens = recall_score(labels, preds, zero_division=0) spec = recall_score(labels, preds, pos_label=0, zero_division=0) if spec >= 0.80 and sens > best_score: best_score = sens best_thresh = thresh if best_score < 0: # Fallback: maximizar F1 best_f1 = -1 for thresh in np.arange(0.05, 0.95, 0.01): preds = (probs > thresh).astype(int) f1 = f1_score(labels, preds, zero_division=0) if f1 > best_f1: best_f1 = f1 best_thresh = thresh print(f"Threshold calibrado: {best_thresh:.2f}") return best_thresh # ============================================================================ # VISUALIZACIÓN # ============================================================================ def plot_curves(history, save_path): epochs = range(1, len(history["train_loss"]) + 1) fig, axes = plt.subplots(2, 2, figsize=(14, 10)) axes[0,0].plot(epochs, history["train_loss"], label="Train", lw=2) axes[0,0].plot(epochs, history["val_loss"], label="Val", lw=2) axes[0,0].set_xlabel("Época"); axes[0,0].set_ylabel("Pérdida") axes[0,0].set_title("Evolución Focal Loss"); axes[0,0].legend(); axes[0,0].grid(True, alpha=0.3) axes[0,1].plot(epochs, history["train_acc"], label="Train", lw=2) axes[0,1].plot(epochs, history["val_acc"], label="Val", lw=2) axes[0,1].set_xlabel("Época"); axes[0,1].set_ylabel("Accuracy (%)") axes[0,1].set_title("Accuracy"); axes[0,1].legend(); axes[0,1].grid(True, alpha=0.3) axes[1,0].plot(epochs, history["train_sens"], label="Train", lw=2) axes[1,0].plot(epochs, history["val_sens"], label="Val", lw=2) axes[1,0].set_xlabel("Época"); axes[1,0].set_ylabel("Sensibilidad (%)") axes[1,0].set_title("Sensibilidad (Recall)"); axes[1,0].legend(); axes[1,0].grid(True, alpha=0.3) axes[1,1].plot(epochs, history["train_spec"], label="Train", lw=2) axes[1,1].plot(epochs, history["val_spec"], label="Val", lw=2) axes[1,1].set_xlabel("Época"); axes[1,1].set_ylabel("Especificidad (%)") axes[1,1].set_title("Especificidad"); axes[1,1].legend(); axes[1,1].grid(True, alpha=0.3) plt.tight_layout(); plt.savefig(save_path, dpi=150); plt.close() print(f"Curvas guardadas: {save_path}") def plot_confusion(cm, save_path): fig, ax = plt.subplots(figsize=(6,5)) im = ax.imshow(cm, cmap='Blues') ax.set_xticks([0,1]); ax.set_yticks([0,1]) ax.set_xticklabels(["Normal","Apnea/Hipopnea"]) ax.set_yticklabels(["Normal","Apnea/Hipopnea"]) ax.set_xlabel("Predicho"); ax.set_ylabel("Real") ax.set_title("Matriz de Confusión (Test)") thresh = cm.max() / 2. for i in range(2): for j in range(2): ax.text(j, i, f"{cm[i,j]}", ha="center", va="center", color="white" if cm[i,j] > thresh else "black", fontsize=14, fontweight='bold') fig.colorbar(im, ax=ax) plt.tight_layout(); plt.savefig(save_path, dpi=150); plt.close() print(f"Matriz guardada: {save_path}") def plot_roc_pr(y_true, y_prob, save_path): fig, axes = plt.subplots(1,2,figsize=(12,5)) fpr, tpr, _ = roc_curve(y_true, y_prob) auc_roc = roc_auc_score(y_true, y_prob) axes[0].plot(fpr, tpr, lw=2, label=f"AUC={auc_roc:.3f}") axes[0].plot([0,1],[0,1],'k--',lw=1) axes[0].set_xlabel("FPR"); axes[0].set_ylabel("TPR") axes[0].set_title("Curva ROC"); axes[0].legend(); axes[0].grid(True, alpha=0.3) prec, rec, _ = precision_recall_curve(y_true, y_prob) auc_pr = average_precision_score(y_true, y_prob) axes[1].plot(rec, prec, lw=2, label=f"AP={auc_pr:.3f}") axes[1].set_xlabel("Recall"); axes[1].set_ylabel("Precision") axes[1].set_title("Curva Precision-Recall"); axes[1].legend(); axes[1].grid(True, alpha=0.3) plt.tight_layout(); plt.savefig(save_path, dpi=150); plt.close() print(f"Curvas ROC/PR guardadas: {save_path}") # ============================================================================ # ENTRENAMIENTO CON TRACKIO # ============================================================================ def train_model(model, train_loader, val_loader, config, class_weights): model.to(config.device) optimizer = AdamW(model.parameters(), lr=config.lr, weight_decay=config.weight_decay) scheduler = CosineAnnealingWarmRestarts(optimizer, T_0=10, T_mult=2) w_tensor = torch.tensor(class_weights, dtype=torch.float32).to(config.device) criterion = FocalLoss(alpha=config.focal_alpha, gamma=config.focal_gamma, weight=w_tensor, reduction='mean') history = {k: [] for k in ["train_loss","val_loss","train_acc","val_acc", "train_sens","val_sens","train_spec","val_spec"]} best_val_f1 = -1.0 best_state = None patience_counter = 0 for epoch in range(1, config.epochs + 1): model.train() train_loss = 0.0 for batch in train_loader: if len(batch) == 3: x, cwt, y = batch x, cwt, y = x.to(config.device), cwt.to(config.device), y.to(config.device) logits = model(x, cwt) else: x, y = batch x, y = x.to(config.device), y.to(config.device) logits = model(x) loss = criterion(logits, y) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), config.grad_clip) optimizer.step() train_loss += loss.item() scheduler.step() t_loss, t_met, _, _ = evaluate(model, train_loader, config, criterion) v_loss, v_met, _, _ = evaluate(model, val_loader, config, criterion) history["train_loss"].append(t_loss) history["val_loss"].append(v_loss) history["train_acc"].append(t_met["accuracy"] * 100) history["val_acc"].append(v_met["accuracy"] * 100) history["train_sens"].append(t_met["sensitivity"] * 100) history["val_sens"].append(v_met["sensitivity"] * 100) history["train_spec"].append(t_met["specificity"] * 100) history["val_spec"].append(v_met["specificity"] * 100) print(f"\nEpoch [{epoch}/{config.epochs}]") print(f" Train Loss: {t_loss:.4f} | Acc: {t_met['accuracy']*100:.2f}% | " f"Sens: {t_met['sensitivity']*100:.2f}% | Spec: {t_met['specificity']*100:.2f}%") print(f" Val Loss: {v_loss:.4f} | Acc: {v_met['accuracy']*100:.2f}% | " f"Sens: {v_met['sensitivity']*100:.2f}% | Spec: {v_met['specificity']*100:.2f}% | " f"F1: {v_met['f1']:.3f} | AUC-ROC: {v_met.get('auc_roc', float('nan')):.3f}") # Trackio logging if TRACKIO_AVAILABLE: try: trackio.log_metrics({ "train_loss": t_loss, "val_loss": v_loss, "train_acc": t_met["accuracy"], "val_acc": v_met["accuracy"], "train_sensitivity": t_met["sensitivity"], "val_sensitivity": v_met["sensitivity"], "train_specificity": t_met["specificity"], "val_specificity": v_met["specificity"], "val_f1": v_met["f1"], "val_auc_roc": v_met.get("auc_roc", 0.0), "epoch": epoch, }) except Exception: pass try: if v_met["f1"] < 0.1 and epoch > 10: trackio.alert("Low F1 Warning", f"val_f1={v_met['f1']:.3f} at epoch {epoch}. " "Model no está detectando eventos. " "Considerar aumentar focal_alpha o oversampling.", level="WARN") if np.isnan(t_loss) or np.isnan(v_loss): trackio.alert("NaN Loss", f"NaN detected at epoch {epoch}. lr={config.lr}. " "Reducir learning rate o revisar datos.", level="ERROR") except Exception: pass if v_met["f1"] > best_val_f1: best_val_f1 = v_met["f1"] patience_counter = 0 best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()} os.makedirs(config.output_dir, exist_ok=True) torch.save({ "epoch": epoch, "model_state_dict": model.state_dict(), "val_f1": v_met["f1"], "config": vars(config), }, os.path.join(config.output_dir, "best_model.pt")) else: patience_counter += 1 if patience_counter >= config.patience: print(f"\nEarly stopping en epoch {epoch}. Mejor F1={best_val_f1:.4f}") if TRACKIO_AVAILABLE: try: trackio.alert("Early Stopping", f"Stopped at epoch {epoch}. Best val_f1={best_val_f1:.4f}", level="INFO") except Exception: pass break if best_state is not None: model.load_state_dict(best_state) model.to(config.device) return model, history # ============================================================================ # AHI COMPUTATION # ============================================================================ def compute_ahi(windows, predictions, stride_min=1.0): n_win = len(windows) n_pos = int(predictions.sum()) total_hours = (n_win * stride_min) / 60.0 return n_pos / total_hours if total_hours > 0 else 0.0 # ============================================================================ # PUSH TO HUB # ============================================================================ def push_model_to_hub(model, config, metrics): """Guarda modelo en Hugging Face Hub si hub_model_id está configurado.""" if config.hub_model_id is None: print("hub_model_id no configurado. Omitiendo push_to_hub.") return try: from huggingface_hub import HfApi, create_repo api = HfApi() # Crear repo si no existe try: create_repo(config.hub_model_id, repo_type="model", exist_ok=True) except Exception: pass # Guardar modelo + config save_path = os.path.join(config.output_dir, "model_for_hub") os.makedirs(save_path, exist_ok=True) torch.save(model.state_dict(), os.path.join(save_path, "pytorch_model.bin")) config_dict = {k: str(v) if isinstance(v, (list, type)) else v for k, v in vars(config).items() if not k.startswith('_')} with open(os.path.join(save_path, "config.json"), "w") as f: json.dump(config_dict, f, indent=2) with open(os.path.join(save_path, "metrics.json"), "w") as f: json.dump(metrics, f, indent=2) api.upload_folder(folder_path=save_path, repo_id=config.hub_model_id, repo_type="model") print(f"Modelo subido a https://huggingface.co/{config.hub_model_id}") except Exception as e: print(f"Error en push_to_hub: {e}") # ============================================================================ # PIPELINE PRINCIPAL # ============================================================================ def run_pipeline(config: Config): set_seed(config.seed) os.makedirs(config.output_dir, exist_ok=True) # Inicializar Trackio if TRACKIO_AVAILABLE: try: trackio.init(project=config.trackio_project, run_name=config.run_name) print(f"Trackio inicializado: proyecto={config.trackio_project}, run={config.run_name}") except Exception as e: print(f"Trackio init error: {e}") else: print("Trackio no instalado (instalar con: pip install trackio)") print("=" * 70) print("DETECTOR DE APNEA DEL SUEÑO - SHHS v2") print("=" * 70) print(f"Dispositivo: {config.device}") print(f"Datos: {config.data_dir}") print(f"Salida: {config.output_dir}") # 1) Descubrir archivos mat_files = sorted(glob.glob(os.path.join(config.data_dir, "*.mat"))) if not mat_files: print(f"ERROR: No .mat en {config.data_dir}") return print(f"\nArchivos .mat: {len(mat_files)}") patient_to_files = defaultdict(list) for f in mat_files: patient_to_files[parse_patient_id(f)].append(f) patients = sorted(patient_to_files.keys()) print(f"Pacientes únicos: {len(patients)}") # 2) Split por paciente n = len(patients) if n < 3: train_patients = patients[:max(1, n-1)] val_patients = patients[max(0, n-2):max(1, n-1)] test_patients = patients[-1:] else: gss = GroupShuffleSplit(n_splits=1, test_size=0.30, random_state=config.seed) train_idx, temp_idx = next(gss.split(patients, groups=patients)) temp_patients = [patients[i] for i in temp_idx] if len(temp_patients) >= 2: gss2 = GroupShuffleSplit(n_splits=1, test_size=0.5, random_state=config.seed) val_idx, test_idx = next(gss2.split(temp_patients, groups=temp_patients)) val_patients = [temp_patients[i] for i in val_idx] test_patients = [temp_patients[i] for i in test_idx] else: val_patients = temp_patients test_patients = temp_patients train_patients = [patients[i] for i in train_idx] print(f"Split: Train={len(train_patients)} pcts, Val={len(val_patients)} pcts, Test={len(test_patients)} pcts") # 3) Procesar grabaciones def process(patient_list, name): wins = [] skipped = 0 for pid in patient_list: for fp in patient_to_files[pid]: data = load_mat_file(fp) if data is None: skipped += 1 continue ww = create_windows(data, config) for w in ww: w["patient_id"] = pid wins.extend(ww) labels = [w["label"] for w in wins] pos = sum(labels) neg = len(labels) - pos print(f" {name}: {len(wins)} ventanas | Pos={pos} ({pos/max(len(labels),1)*100:.2f}%) | " f"Neg={neg} | Skip={skipped}") return wins print("\n--- Extrayendo ventanas ---") train_wins = process(train_patients, "Train") val_wins = process(val_patients, "Val") test_wins = process(test_patients, "Test") if not train_wins or not val_wins or not test_wins: print("ERROR: Dataset vacío.") return # Pesos de clase y_train = np.array([w["label"] for w in train_wins]) counts = np.bincount(y_train, minlength=2) total = counts.sum() weights = total / (counts + 1e-6) weights = weights / weights.sum() * 2 print(f"\nPesos clase: {weights}") pos_ratio = counts[1] / total config.focal_alpha = max(0.1, min(0.9, 1 - pos_ratio)) print(f"Focal alpha: {config.focal_alpha:.3f}") # 4) Dataloaders train_ds = SleepDataset(train_wins, config, augment=True) val_ds = SleepDataset(val_wins, config, augment=False) test_ds = SleepDataset(test_wins, config, augment=False) pin_mem = config.device.startswith('cuda') if len(counts) >= 2: cw = weights[y_train] sampler = WeightedRandomSampler(cw, len(cw), replacement=True) train_loader = DataLoader(train_ds, batch_size=config.batch_size, sampler=sampler, num_workers=config.num_workers, collate_fn=collate_fn, pin_memory=pin_mem) else: train_loader = DataLoader(train_ds, batch_size=config.batch_size, shuffle=True, num_workers=config.num_workers, collate_fn=collate_fn, pin_memory=pin_mem) val_loader = DataLoader(val_ds, batch_size=config.batch_size, shuffle=False, num_workers=config.num_workers, collate_fn=collate_fn, pin_memory=pin_mem) test_loader = DataLoader(test_ds, batch_size=config.batch_size, shuffle=False, num_workers=config.num_workers, collate_fn=collate_fn, pin_memory=pin_mem) # 5) Modelo print("\n--- Arquitectura ---") model = SleepApneaNetV2(config) n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"Modelo: {config.model_name}") print(f"Parámetros: {n_params:,}") dummy_raw = torch.randn(2, 2, config.window_size).to(config.device) dummy_cwt = None if config.use_cwt: n_scales = len(config.cwt_scales) * 2 dummy_cwt = torch.randn(2, 1, n_scales, 300).to(config.device) model.to(config.device) with torch.no_grad(): out = model(dummy_raw, dummy_cwt) print(f"Forward OK: {dummy_raw.shape} → {out.shape}") assert out.shape == (2, 2) # 6) Entrenar print("\n--- Entrenamiento ---") model, history = train_model(model, train_loader, val_loader, config, weights) plot_curves(history, os.path.join(config.output_dir, "training_curves.png")) # 7) Calibrar threshold best_thresh = 0.5 if config.calibrate_threshold and len(val_wins) > 0: print("\n--- Calibración Threshold ---") best_thresh = calibrate_threshold(model, val_loader, config) # 8) Evaluación TEST print("\n" + "=" * 70) print("EVALUACIÓN FINAL - TEST") print("=" * 70) w_tensor = torch.tensor(weights, dtype=torch.float32).to(config.device) criterion = FocalLoss(alpha=config.focal_alpha, gamma=config.focal_gamma, weight=w_tensor, reduction='mean') test_loss, test_met, test_probs, test_labels = evaluate(model, test_loader, config, criterion, threshold=best_thresh) print(f"\nTest Loss: {test_loss:.4f}") print(f"Test Accuracy: {test_met['accuracy']*100:.2f}%") print(f"Test Sensitivity: {test_met['sensitivity']*100:.2f}%") print(f"Test Specificity: {test_met['specificity']*100:.2f}%") print(f"Test Precision: {test_met['precision']*100:.2f}%") print(f"Test F1-Score: {test_met['f1']:.4f}") print(f"Test Cohen's Kappa: {test_met['cohen_kappa']:.4f}") print(f"Test AUC-ROC: {test_met.get('auc_roc', float('nan')):.4f}") print(f"Test AUC-PR: {test_met.get('auc_pr', float('nan')):.4f}") print(f"Threshold: {best_thresh:.2f}") cm = np.array(test_met["confusion_matrix"]) print(f"\nMatriz Confusión:") print(f" Predicho") print(f" Normal Apnea") print(f"Real Normal {cm[0,0]:6d} {cm[0,1]:6d}") print(f" Apnea {cm[1,0]:6d} {cm[1,1]:6d}") preds_thresh = (test_probs > best_thresh).astype(int) print("\nClassification Report:") print(classification_report(test_labels, preds_thresh, target_names=["Normal", "Apnea/Hipopnea"], digits=4)) plot_confusion(cm, os.path.join(config.output_dir, "confusion_matrix.png")) plot_roc_pr(test_labels, test_probs, os.path.join(config.output_dir, "roc_pr_curves.png")) # 9) AHI por paciente print("\n" + "=" * 70) print("AHI (ÍNDICE APNEA-HIPOPNEA)") print("=" * 70) test_by_patient = defaultdict(list) for w in test_wins: test_by_patient[w["patient_id"]].append(w) pred_by_patient = defaultdict(list) idx = 0 model.eval() with torch.no_grad(): for batch in test_loader: if len(batch) == 3: x, cwt, y = batch x = x.to(config.device); cwt = cwt.to(config.device) probs = F.softmax(model(x, cwt), dim=-1)[:, 1].cpu().numpy() else: x, y = batch; x = x.to(config.device) probs = F.softmax(model(x), dim=-1)[:, 1].cpu().numpy() preds = (probs > best_thresh).astype(int) for i in range(len(preds)): if idx < len(test_wins): pred_by_patient[test_wins[idx]["patient_id"]].append(preds[i]) idx += 1 ahi_errors = [] print(f"\n{'Paciente':<18} {'Pred_AHI':>10} {'True_AHI':>10} {'Error':>10}") print("-" * 55) for pid in sorted(test_by_patient.keys())[:30]: wins = test_by_patient[pid] preds = np.array(pred_by_patient.get(pid, [])) if len(preds) == 0: continue pred_ahi = compute_ahi(wins, preds, stride_min=config.stride/config.fs/60) true_ahi = compute_ahi(wins, np.array([w["label"] for w in wins]), stride_min=config.stride/config.fs/60) err = pred_ahi - true_ahi ahi_errors.append(err) print(f"{pid:<18} {pred_ahi:10.2f} {true_ahi:10.2f} {err:10.2f}") if ahi_errors: ahi_errors = np.array(ahi_errors) mae = float(np.abs(ahi_errors).mean()) rmse = float(np.sqrt((ahi_errors**2).mean())) print(f"\nAHI Stats (n={len(ahi_errors)}):") print(f" MAE: {mae:.2f}") print(f" RMSE: {rmse:.2f}") print(f" Bias: {ahi_errors.mean():.2f}") else: mae = rmse = None # 10) Guardar reporte y subir a Hub report = { "model_parameters": n_params, "test_metrics": {k: float(v) if isinstance(v, (float, np.floating)) else v for k, v in test_met.items()}, "ahi_mae": mae, "ahi_rmse": rmse, "threshold": float(best_thresh), "config": {k: str(v) if isinstance(v, (list, type)) else v for k, v in vars(config).items() if not k.startswith('_')}, } report_path = os.path.join(config.output_dir, "final_report.json") with open(report_path, "w") as f: json.dump(report, f, indent=2) print(f"\nReporte: {report_path}") # Push to Hub push_model_to_hub(model, config, report["test_metrics"]) # Trackio final log if TRACKIO_AVAILABLE: try: trackio.log_metrics({ "test_accuracy": test_met["accuracy"], "test_sensitivity": test_met["sensitivity"], "test_specificity": test_met["specificity"], "test_f1": test_met["f1"], "test_auc_roc": test_met.get("auc_roc", 0.0), "test_auc_pr": test_met.get("auc_pr", 0.0), "test_cohen_kappa": test_met["cohen_kappa"], "ahi_mae": mae if mae else 0.0, "ahi_rmse": rmse if rmse else 0.0, "threshold": float(best_thresh), }) trackio.alert("Training Complete", f"Test F1={test_met['f1']:.3f}, Sens={test_met['sensitivity']*100:.1f}%, " f"Spec={test_met['specificity']*100:.1f}%, AUC-ROC={test_met.get('auc_roc', 0):.3f}", level="INFO") except Exception: pass print("\n" + "=" * 70) print("COMPLETADO") print("=" * 70) print(f"Modelo: {os.path.join(config.output_dir, 'best_model.pt')}") print(f"Reporte: {report_path}") if config.hub_model_id: print(f"Hub URL: https://huggingface.co/{config.hub_model_id}") return model, test_met # ============================================================================ # ENTRY POINT # ============================================================================ if __name__ == "__main__": parser = argparse.ArgumentParser(description="Detector Apnea del Sueño SHHS") parser.add_argument("--data_dir", type=str, default="/app/data") parser.add_argument("--output_dir", type=str, default="/app/output") parser.add_argument("--hub_model_id", type=str, default=None) parser.add_argument("--trackio_space_id", type=str, default=None) parser.add_argument("--trackio_project", type=str, default="sleep-apnea-shhs") parser.add_argument("--run_name", type=str, default="shhs-tcn-bigru") parser.add_argument("--epochs", type=int, default=100) parser.add_argument("--batch_size", type=int, default=32) parser.add_argument("--lr", type=float, default=1e-3) parser.add_argument("--window_size", type=int, default=900) parser.add_argument("--stride", type=int, default=60) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--use_cwt", action="store_true", help="Activar branch CWT") parser.add_argument("--patience", type=int, default=15) args = parser.parse_args() cfg = Config() cfg.data_dir = args.data_dir cfg.output_dir = args.output_dir cfg.hub_model_id = args.hub_model_id cfg.trackio_space_id = args.trackio_space_id cfg.trackio_project = args.trackio_project cfg.run_name = args.run_name cfg.epochs = args.epochs cfg.batch_size = args.batch_size cfg.lr = args.lr cfg.window_size = args.window_size cfg.stride = args.stride cfg.seed = args.seed cfg.use_cwt = args.use_cwt cfg.patience = args.patience run_pipeline(cfg)