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)