Spaces:
Runtime error
Runtime error
File size: 6,061 Bytes
832948a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """Phase 4: Stage 1 binary classifier -- Rest vs Active."""
import os
import numpy as np
import joblib
from pathlib import Path
from sklearn.svm import SVC
from sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score
from sklearn.model_selection import cross_val_score, GroupKFold
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import seaborn as sns
PROJECT_ROOT = Path(__file__).resolve().parent.parent
MODELS_DIR = PROJECT_ROOT / "models"
RESULTS_DIR = PROJECT_ROOT / "results"
def to_binary(label_str):
"""Map label string to binary: 0=Rest, 1=Active."""
return 0 if label_str == "Relax" else 1
def train_stage1():
"""Train and evaluate Stage 1 binary classifier."""
MODELS_DIR.mkdir(parents=True, exist_ok=True)
RESULTS_DIR.mkdir(parents=True, exist_ok=True)
# Load features
data = np.load(str(PROJECT_ROOT / "features.npz"), allow_pickle=True)
X = data["X"]
y_str = data["y"]
subjects = data["subjects"]
# Binary labels
y_binary = np.array([to_binary(s) for s in y_str])
print(f"Total samples: {len(y_binary)}")
print(f" Rest (0): {np.sum(y_binary == 0)}")
print(f" Active (1): {np.sum(y_binary == 1)}")
# Cross-subject split: hold out ~20% of subjects
unique_subjects = sorted(set(subjects))
n_test = max(1, len(unique_subjects) // 5)
test_subjects = set(unique_subjects[-n_test:])
train_subjects = set(unique_subjects[:-n_test])
print(f"\nSubject split:")
print(f" Train subjects ({len(train_subjects)}): {sorted(train_subjects)}")
print(f" Test subjects ({len(test_subjects)}): {sorted(test_subjects)}")
train_mask = np.array([s in train_subjects for s in subjects])
test_mask = np.array([s in test_subjects for s in subjects])
X_train, y_train = X[train_mask], y_binary[train_mask]
X_test, y_test = X[test_mask], y_binary[test_mask]
subjects_train = subjects[train_mask]
print(f" Train samples: {len(y_train)} (Rest: {np.sum(y_train==0)}, Active: {np.sum(y_train==1)})")
print(f" Test samples: {len(y_test)} (Rest: {np.sum(y_test==0)}, Active: {np.sum(y_test==1)})")
# Verify no subject leakage
assert len(set(subjects[train_mask]) & set(subjects[test_mask])) == 0, "Subject leakage!"
# Models to try
models = {
"SVM_RBF": Pipeline([
("scaler", StandardScaler()),
("clf", SVC(kernel="rbf", probability=True, class_weight="balanced", random_state=42))
]),
"RandomForest": Pipeline([
("scaler", StandardScaler()),
("clf", RandomForestClassifier(n_estimators=200, class_weight="balanced", random_state=42))
]),
"GradientBoosting": Pipeline([
("scaler", StandardScaler()),
("clf", GradientBoostingClassifier(n_estimators=200, random_state=42))
]),
}
# 5-fold GroupKFold cross-validation on training set
print("\n--- Cross-Validation (5-fold GroupKFold) ---")
gkf = GroupKFold(n_splits=min(5, len(train_subjects)))
cv_results = {}
for name, model in models.items():
scores = cross_val_score(
model, X_train, y_train, cv=gkf, groups=subjects_train, scoring="accuracy"
)
cv_results[name] = scores
print(f" {name}: {scores.mean():.3f} +/- {scores.std():.3f}")
# Train on full training set and evaluate on test set
print("\n--- Test Set Evaluation ---")
test_results = {}
for name, model in models.items():
model.fit(X_train, y_train)
y_pred = model.predict(X_test)
acc = accuracy_score(y_test, y_pred)
test_results[name] = (acc, model, y_pred)
print(f" {name}: {acc:.3f}")
# Pick best model
best_name = max(test_results, key=lambda k: test_results[k][0])
best_acc, best_model, best_pred = test_results[best_name]
print(f"\nBest model: {best_name} (accuracy: {best_acc:.3f})")
# Save best model
model_path = MODELS_DIR / "stage1_binary.pkl"
joblib.dump(best_model, str(model_path))
print(f"Saved to {model_path}")
# Classification report
report = classification_report(y_test, best_pred, target_names=["Rest", "Active"])
print(f"\nClassification Report:\n{report}")
with open(str(RESULTS_DIR / "stage1_report.txt"), "w") as f:
f.write(f"Stage 1 Binary Classifier: {best_name}\n")
f.write(f"Test Accuracy: {best_acc:.3f}\n\n")
f.write(f"Train subjects: {sorted(train_subjects)}\n")
f.write(f"Test subjects: {sorted(test_subjects)}\n\n")
f.write(f"Cross-validation results:\n")
for name, scores in cv_results.items():
f.write(f" {name}: {scores.mean():.3f} +/- {scores.std():.3f}\n")
f.write(f"\nTest set results:\n")
for name, (acc, _, _) in test_results.items():
f.write(f" {name}: {acc:.3f}\n")
f.write(f"\nBest: {best_name}\n\n")
f.write(f"Classification Report:\n{report}\n")
# Confusion matrix plot
cm = confusion_matrix(y_test, best_pred)
fig, ax = plt.subplots(figsize=(6, 5))
sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
xticklabels=["Rest", "Active"],
yticklabels=["Rest", "Active"], ax=ax)
ax.set_title(f"Stage 1: {best_name} (Acc: {best_acc:.3f})")
ax.set_xlabel("Predicted")
ax.set_ylabel("True")
plt.tight_layout()
plt.savefig(str(RESULTS_DIR / "stage1_confusion.png"), dpi=150)
plt.close()
print(f"Confusion matrix saved to results/stage1_confusion.png")
# Save the train/test subject split for consistency with Stage 2
np.savez(
str(PROJECT_ROOT / "subject_split.npz"),
train_subjects=np.array(sorted(train_subjects)),
test_subjects=np.array(sorted(test_subjects)),
)
return best_model, best_acc
if __name__ == "__main__":
train_stage1()
|