src/train_model.py

import argparse
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import classification_report, confusion_matrix

import tensorflow as tf
from tensorflow.keras.applications import MobileNetV2
from tensorflow.keras.applications.mobilenet_v2 import preprocess_input
from tensorflow.keras.layers import AveragePooling2D, Dense, Dropout, Flatten, Input
from tensorflow.keras.models import Model
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.callbacks import (EarlyStopping, ModelCheckpoint,
                                        ReduceLROnPlateau, TensorBoard)

from preprocess import (
    split_dataset, build_generators,
    IMG_SIZE, BATCH_SIZE, CLASSES,
)

BASE_DIR    = Path(__file__).resolve().parent.parent
DATA_DIR    = BASE_DIR / "dataset" / "data"
SPLIT_DIR   = BASE_DIR / "dataset" / "split"
MODELS_DIR  = BASE_DIR / "models"


# ── Model ─────────────────────────────────────────────────────────────────────

def build_model():
    base = MobileNetV2(
        weights="imagenet",
        include_top=False,
        input_tensor=Input(shape=(*IMG_SIZE, 3)),
    )
    for layer in base.layers:
        layer.trainable = False

    head = AveragePooling2D(pool_size=(7, 7))(base.output)
    head = Flatten()(head)
    head = Dense(128, activation="relu")(head)
    head = Dropout(0.5)(head)
    head = Dense(2, activation="softmax")(head)

    return Model(inputs=base.input, outputs=head), base


# ── Plotting ──────────────────────────────────────────────────────────────────

def plot_history(history, title: str, save_path: Path) -> None:
    fig, axes = plt.subplots(1, 2, figsize=(13, 4))

    axes[0].plot(history.history["loss"],     label="train")
    axes[0].plot(history.history["val_loss"], label="val")
    axes[0].set_title(f"{title} — Loss")
    axes[0].set_xlabel("Epoch"); axes[0].set_ylabel("Loss")
    axes[0].legend()

    axes[1].plot(history.history["accuracy"],     label="train")
    axes[1].plot(history.history["val_accuracy"], label="val")
    axes[1].set_title(f"{title} — Accuracy")
    axes[1].set_xlabel("Epoch"); axes[1].set_ylabel("Accuracy")
    axes[1].legend()

    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close()
    print(f"[INFO] Plot saved → {save_path}")


def plot_confusion_matrix(y_true, y_pred, class_names, save_path: Path) -> None:
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(6, 5))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues",
                xticklabels=class_names, yticklabels=class_names)
    plt.title("Confusion Matrix")
    plt.ylabel("True Label"); plt.xlabel("Predicted Label")
    plt.tight_layout()
    plt.savefig(save_path, dpi=150, bbox_inches="tight")
    plt.close()
    print(f"[INFO] Confusion matrix saved → {save_path}")


# ── Main ──────────────────────────────────────────────────────────────────────

def main(args):
    # 1. Split dataset into train / val / test folders
    print("[INFO] Preparing dataset split ...")
    split_dataset(
        src_dir=Path(args.dataset),
        dest_dir=SPLIT_DIR,
        val_split=0.15,
        test_split=0.15,
        seed=42,
        overwrite=args.overwrite_split,
    )

    # 2. Build generators
    train_gen, val_gen, test_gen, class_indices = build_generators(
        split_dir=SPLIT_DIR,
        batch_size=BATCH_SIZE,
        img_size=IMG_SIZE,
    )
    class_names = [k for k, _ in sorted(class_indices.items(), key=lambda x: x[1])]
    print(f"[INFO] Class indices: {class_indices}")
    print(f"[INFO] Steps — train={len(train_gen)} val={len(val_gen)} test={len(test_gen)}")

    # 3. Build model
    model, base = build_model()
    MODELS_DIR.mkdir(parents=True, exist_ok=True)
    ckpt_path = str(MODELS_DIR / "mask_detector_best.keras")

    # ── Phase 1: Warm-up (head only) ────────────────────────────────────────
    print("\n[INFO] Phase 1 — training head (base frozen) ...")
    model.compile(
        optimizer=Adam(learning_rate=1e-3),
        loss="categorical_crossentropy",
        metrics=["accuracy"],
    )
    H1 = model.fit(
        train_gen,
        validation_data=val_gen,
        epochs=args.warmup_epochs,
        callbacks=[
            EarlyStopping(patience=5, restore_best_weights=True, verbose=1),
            ModelCheckpoint(ckpt_path, save_best_only=True, verbose=1),
            ReduceLROnPlateau(factor=0.5, patience=3, verbose=1),
            TensorBoard(log_dir=str(BASE_DIR / "logs" / "warmup")),
        ],
    )
    plot_history(H1, "Warm-up", BASE_DIR / "warmup_plot.png")

    # ── Phase 2: Fine-tune (unfreeze top 20 base layers) ────────────────────
    print("\n[INFO] Phase 2 — fine-tuning top layers ...")
    for layer in base.layers[-20:]:
        layer.trainable = True

    model.compile(
        optimizer=Adam(learning_rate=1e-4),
        loss="categorical_crossentropy",
        metrics=["accuracy"],
    )
    H2 = model.fit(
        train_gen,
        validation_data=val_gen,
        epochs=args.finetune_epochs,
        callbacks=[
            EarlyStopping(patience=7, restore_best_weights=True, verbose=1),
            ModelCheckpoint(ckpt_path, save_best_only=True, verbose=1),
            ReduceLROnPlateau(factor=0.3, patience=4, verbose=1),
            TensorBoard(log_dir=str(BASE_DIR / "logs" / "finetune")),
        ],
    )
    plot_history(H2, "Fine-tune", BASE_DIR / "finetune_plot.png")

    # ── Evaluation ───────────────────────────────────────────────────────────
    print("\n[INFO] Evaluating on test set ...")
    test_gen.reset()
    pred_probs  = model.predict(test_gen, verbose=1)
    pred_labels = np.argmax(pred_probs, axis=1)
    true_labels = test_gen.classes

    print("\n" + classification_report(true_labels, pred_labels,
                                       target_names=class_names))
    plot_confusion_matrix(true_labels, pred_labels, class_names,
                          BASE_DIR / "confusion_matrix.png")

    # ── Save final model ─────────────────────────────────────────────────────
    out_path = Path(args.model)
    out_path.parent.mkdir(parents=True, exist_ok=True)
    model.save(str(out_path))
    print(f"[INFO] Final model saved → {out_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset",         default=str(DATA_DIR),
                        help="Path to raw dataset (with_mask / without_mask folders)")
    parser.add_argument("--model",           default=str(MODELS_DIR / "mask_detector.keras"),
                        help="Output path for the saved model")
    parser.add_argument("--warmup-epochs",   type=int, default=10)
    parser.add_argument("--finetune-epochs", type=int, default=20)
    parser.add_argument("--overwrite-split", action="store_true",
                        help="Re-split dataset even if split folder already exists")
    args = parser.parse_args()
    main(args)