train_expression_model.py

"""
Train the custom landmark-based expression classifier.

Expected dataset layout (same as FER2013 / RAF-DB folder structure):
  data/
    train/
      angry/   img1.jpg img2.jpg ...
      disgust/ ...
      fear/    ...
      happy/   ...
      neutral/ ...
      sad/     ...
      surprise/...
    val/
      angry/   ...
      ...

Recommended free datasets:
  - FER2013     : https://www.kaggle.com/datasets/msambare/fer2013
  - RAF-DB      : http://www.whdeng.cn/RAF/model1.html
  - AffectNet   : http://mohammadmahoor.com/affectnet/

Usage:
  pip install torch torchvision tqdm scikit-learn mediapipe opencv-python
  python train_expression_model.py --data_dir data --epochs 60

The trained model is saved to models/expression_mlp.pth and is automatically
loaded by the analyzer at startup.
"""

import argparse
import os
import sys

import cv2
import numpy as np

try:
    import mediapipe as mp
    import torch
    import torch.nn as nn
    from sklearn.metrics import classification_report
    from sklearn.preprocessing import LabelEncoder
    from torch.utils.data import DataLoader, TensorDataset
    from tqdm import tqdm
except ImportError as e:
    sys.exit(f"Missing dependency: {e}\nRun: pip install torch mediapipe tqdm scikit-learn")

from expression_model import INPUT_DIM, MODEL_PATH, extract_landmark_features


# ── Model definition (mirrors expression_model.py) ───────────────────────────

class _MLP(nn.Module):
    def __init__(self, n_cls: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(INPUT_DIM, 512), nn.BatchNorm1d(512), nn.GELU(), nn.Dropout(0.30),
            nn.Linear(512, 256),       nn.BatchNorm1d(256), nn.GELU(), nn.Dropout(0.20),
            nn.Linear(256, 128),       nn.GELU(),
            nn.Linear(128, n_cls),
        )
    def forward(self, x): return self.net(x)


# ── Landmark extraction ───────────────────────────────────────────────────────

def extract_dataset(root: str, split: str) -> tuple[np.ndarray, list[str], list[str]]:
    split_dir = os.path.join(root, split)
    if not os.path.isdir(split_dir):
        sys.exit(f"Directory not found: {split_dir}")

    labels = sorted(
        d for d in os.listdir(split_dir) if os.path.isdir(os.path.join(split_dir, d))
    )
    if not labels:
        sys.exit(f"No class subdirectories found in {split_dir}")

    mesh = mp.solutions.face_mesh.FaceMesh(
        static_image_mode=True, max_num_faces=1,
        refine_landmarks=True, min_detection_confidence=0.4,
    )

    X, y = [], []
    skipped = 0

    for label in labels:
        label_dir = os.path.join(split_dir, label)
        files = [f for f in os.listdir(label_dir) if f.lower().endswith((".jpg", ".jpeg", ".png"))]
        print(f"  [{split}/{label}] {len(files)} images")

        for fn in tqdm(files, desc=f"{split}/{label}", leave=False):
            img = cv2.imread(os.path.join(label_dir, fn))
            if img is None:
                skipped += 1
                continue
            rgb    = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            result = mesh.process(rgb)
            if not result.multi_face_landmarks:
                skipped += 1
                continue
            X.append(extract_landmark_features(result.multi_face_landmarks[0]))
            y.append(label)

    mesh.close()
    print(f"  Extracted {len(X)} samples ({skipped} skipped — no face detected)")
    return np.array(X, dtype=np.float32), y, labels


# ── Training ──────────────────────────────────────────────────────────────────

def train(data_dir: str, epochs: int, lr: float, batch_size: int, device_str: str):
    os.makedirs("models", exist_ok=True)
    device = torch.device(device_str if torch.cuda.is_available() else "cpu")
    print(f"\nDevice: {device}")

    print("\nExtracting TRAIN landmarks...")
    X_tr, y_tr_raw, labels = extract_dataset(data_dir, "train")
    print("\nExtracting VAL landmarks...")
    X_vl, y_vl_raw, _      = extract_dataset(data_dir, "val")

    le    = LabelEncoder().fit(labels)
    y_tr  = le.transform(y_tr_raw)
    y_vl  = le.transform(y_vl_raw)

    tr_ds = TensorDataset(torch.tensor(X_tr), torch.tensor(y_tr, dtype=torch.long))
    vl_ds = TensorDataset(torch.tensor(X_vl), torch.tensor(y_vl, dtype=torch.long))
    tr_dl = DataLoader(tr_ds, batch_size=batch_size, shuffle=True,  num_workers=0)
    vl_dl = DataLoader(vl_ds, batch_size=batch_size, shuffle=False, num_workers=0)

    model   = _MLP(n_cls=len(labels)).to(device)
    opt     = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-4)
    sched   = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs)
    loss_fn = nn.CrossEntropyLoss()

    print(f"\nTraining  classes={labels}  samples={len(X_tr)}  val={len(X_vl)}\n")
    best_acc = 0.0

    for epoch in range(1, epochs + 1):
        # ── Train ──
        model.train()
        tr_loss = 0.0
        for Xb, yb in tr_dl:
            Xb, yb = Xb.to(device), yb.to(device)
            opt.zero_grad()
            loss = loss_fn(model(Xb), yb)
            loss.backward()
            opt.step()
            tr_loss += loss.item()
        sched.step()

        # ── Validate ──
        model.eval()
        correct = total = 0
        all_preds, all_true = [], []
        with torch.no_grad():
            for Xb, yb in vl_dl:
                Xb = Xb.to(device)
                preds = model(Xb).argmax(dim=1).cpu()
                correct += (preds == yb).sum().item()
                total   += len(yb)
                all_preds.extend(preds.tolist())
                all_true.extend(yb.tolist())

        val_acc = correct / total
        print(f"Epoch {epoch:03d}/{epochs}  loss={tr_loss/len(tr_dl):.4f}  val_acc={val_acc:.3f}", end="")

        if val_acc > best_acc:
            best_acc = val_acc
            torch.save({"model": model.state_dict(), "labels": list(le.classes_)}, MODEL_PATH)
            print("  ✓ saved", end="")
        print()

    print(f"\nBest val accuracy: {best_acc:.3f}")
    print(f"Model saved to:    {MODEL_PATH}\n")

    # Final classification report
    print(classification_report(all_true, all_preds, target_names=list(le.classes_)))


# ── CLI ───────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    p = argparse.ArgumentParser(description="Train landmark-based expression MLP")
    p.add_argument("--data_dir",   default="data",  help="Root dir with train/ and val/")
    p.add_argument("--epochs",     default=60,  type=int)
    p.add_argument("--lr",         default=1e-3, type=float)
    p.add_argument("--batch_size", default=64,  type=int)
    p.add_argument("--device",     default="cuda", help="cuda or cpu")
    args = p.parse_args()
    train(args.data_dir, args.epochs, args.lr, args.batch_size, args.device)