face_recognition.py

import os
import glob
import numpy as np
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from PIL import Image as PILImage


from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.decomposition import PCA
from sklearn.svm import SVC
from sklearn.metrics import classification_report, confusion_matrix
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
import joblib
import time
import json
#constants
MODEL_PATH = os.path.join(os.path.dirname(__file__), "model")
STATIC_PATH = os.path.join(os.path.dirname(__file__), "static")
N_COMPONENTS = 100
IMG_SIZE = (62, 47)
#logic ensure the directory 
def ensure_dirs():
    os.makedirs(MODEL_PATH, exist_ok=True)
    os.makedirs(STATIC_PATH, exist_ok=True)
#kaggle
def load_kaggle_faces(max_per_class=800):
    try:
        import kagglehub
        ds_path = kagglehub.dataset_download("ashwingupta3012/human-faces")
    except Exception:
        return None, None, None


    exts = ("*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp")
    all_imgs = []
    for ext in exts:
        all_imgs.extend(glob.glob(os.path.join(ds_path, "**", ext), recursive=True))
    if not all_imgs:
        return None, None, None
    sub_dirs = [d for d in os.listdir(ds_path)
                if os.path.isdir(os.path.join(ds_path, d)) and not d.startswith('.')]


    person_dirs = {}
    for sd in sub_dirs:
        sd_path = os.path.join(ds_path, sd)
        imgs_in_dir = []
        for ext in exts:
            imgs_in_dir.extend(glob.glob(os.path.join(sd_path, ext)))
            imgs_in_dir.extend(glob.glob(os.path.join(sd_path, "**", ext), recursive=True))
        if len(imgs_in_dir) >= 5:
            person_dirs[sd] = imgs_in_dir[:max_per_class]


    if len(person_dirs) >= 2:
        return _load_labeled(person_dirs)
    else:
        return _load_flat(all_imgs, max_per_class * 2)
def _load_labeled(person_dirs):
    images, labels, names = [], [], sorted(person_dirs.keys())
    h, w = IMG_SIZE
    for idx, name in enumerate(names):
        for p in person_dirs[name]:
            try:
                img = PILImage.open(p).convert("L").resize((w, h))
                images.append(np.array(img, dtype=np.float64))
                labels.append(idx)
            except Exception:
                continue
    X = np.array(images)
    y = np.array(labels)
    target_names = np.array(names)
    return X, y, target_names




def _load_flat(all_paths, limit):
    h, w = IMG_SIZE
    images = []
    for p in all_paths[:limit]:
        try:
            img = PILImage.open(p).convert("L").resize((w, h))
            images.append(np.array(img, dtype=np.float64))
        except Exception:
            continue
    X = np.array(images)
    return X, None, None




def load_lfw(min_faces=40):
    from sklearn.datasets import fetch_lfw_people
    lfw = fetch_lfw_people(min_faces_per_person=min_faces, resize=0.5)
    X = lfw.images
    y = lfw.target
    target_names = lfw.target_names
    return X, y, target_names




def build_pipeline(n_components=N_COMPONENTS):
    return Pipeline([
        ("scaler", StandardScaler()),
        ("pca", PCA(n_components=n_components, whiten=True, svd_solver="randomized")),
        ("svm", SVC(kernel="rbf", class_weight="balanced", probability=True)),
    ])




def train(X_train, y_train, n_components=N_COMPONENTS):
    n_comp = min(n_components, X_train.shape[0], X_train.shape[1])
    pipe = build_pipeline(n_comp)
    param_grid = {
        "svm__C": [1, 5, 10],
        "svm__gamma": [0.001, 0.005, 0.01],
    }
    t0 = time.time()
    clf = GridSearchCV(pipe, param_grid, cv=3, n_jobs=-1, verbose=0)
    clf.fit(X_train, y_train)
    elapsed = time.time() - t0
    return clf




def evaluate(clf, X_test, y_test, target_names):
    y_pred = clf.predict(X_test)
    report_dict = classification_report(y_test, y_pred,
                                        target_names=target_names,
                                        output_dict=True)
    cm = confusion_matrix(y_test, y_pred)
    return y_pred, report_dict, cm




def plot_eigenfaces(pca, h, w, n=12):
    fig, axes = plt.subplots(2, 6, figsize=(14, 5),
                             subplot_kw={"xticks": [], "yticks": []})
    fig.suptitle("Top Eigenfaces (Principal Components)", fontsize=14, y=1.02)
    for i, ax in enumerate(axes.flat):
        if i < min(n, pca.n_components_):
            ax.imshow(pca.components_[i].reshape(h, w), cmap="bone")
            ax.set_title(f"PC {i+1}", fontsize=10)
        else:
            ax.axis("off")
    fig.tight_layout()
    path = os.path.join(STATIC_PATH, "eigenfaces.png")
    fig.savefig(path, dpi=120, bbox_inches="tight")
    plt.close(fig)




def plot_predictions(images_test, y_test, y_pred, target_names, n=15):
    fig, axes = plt.subplots(3, 5, figsize=(14, 9),
                             subplot_kw={"xticks": [], "yticks": []})
    fig.suptitle("Predictions (green = correct, red = wrong)", fontsize=14, y=1.01)
    for i, ax in enumerate(axes.flat):
        if i < min(n, len(y_test)):
            ax.imshow(images_test[i], cmap="gray")
            pred_name = target_names[y_pred[i]].split()[-1] if ' ' in target_names[y_pred[i]] else target_names[y_pred[i]]
            true_name = target_names[y_test[i]].split()[-1] if ' ' in target_names[y_test[i]] else target_names[y_test[i]]
            color = "#2ecc71" if y_pred[i] == y_test[i] else "#e74c3c"
            ax.set_title(f"pred: {pred_name}\ntrue: {true_name}",
                         fontsize=9, color=color, fontweight="bold")
    fig.tight_layout()
    path = os.path.join(STATIC_PATH, "predictions.png")
    fig.savefig(path, dpi=120, bbox_inches="tight")
    plt.close(fig)




def plot_confusion_matrix(cm, target_names):
    fig, ax = plt.subplots(figsize=(max(8, len(target_names)),
                                     max(6, len(target_names) * 0.8)))
    short = [n.split()[-1] if ' ' in n else n for n in target_names]
    im = ax.imshow(cm, interpolation="nearest", cmap="Blues")
    fig.colorbar(im, ax=ax, shrink=0.8)
    ax.set(xticks=range(len(short)), yticks=range(len(short)),
           xticklabels=short, yticklabels=short,
           xlabel="Predicted", ylabel="True",
           title="Confusion Matrix")
    plt.setp(ax.get_xticklabels(), rotation=45, ha="right")
    thresh = cm.max() / 2
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            ax.text(j, i, str(cm[i, j]),
                    ha="center", va="center",
                    color="white" if cm[i, j] > thresh else "black",
                    fontsize=9)
    fig.tight_layout()
    path = os.path.join(STATIC_PATH, "confusion_matrix.png")
    fig.savefig(path, dpi=120, bbox_inches="tight")
    plt.close(fig)




def plot_sample_faces(images, title="Sample Faces from Dataset", n=15):
    fig, axes = plt.subplots(3, 5, figsize=(14, 9),
                             subplot_kw={"xticks": [], "yticks": []})
    fig.suptitle(title, fontsize=14, y=1.01)
    indices = np.random.choice(len(images), min(n, len(images)), replace=False)
    for i, ax in enumerate(axes.flat):
        if i < len(indices):
            ax.imshow(images[indices[i]], cmap="gray")
        else:
            ax.axis("off")
    fig.tight_layout()
    path = os.path.join(STATIC_PATH, "sample_faces.png")
    fig.savefig(path, dpi=120, bbox_inches="tight")
    plt.close(fig)




def save_model(clf, target_names, h, w, data_source="kaggle"):
    ensure_dirs()
    joblib.dump(clf, os.path.join(MODEL_PATH, "face_clf.pkl"))
    meta = {
        "target_names": list(target_names),
        "h": h, "w": w,
        "data_source": data_source,
    }
    with open(os.path.join(MODEL_PATH, "meta.json"), "w") as f:
        json.dump(meta, f)




def load_model():
    clf = joblib.load(os.path.join(MODEL_PATH, "face_clf.pkl"))
    with open(os.path.join(MODEL_PATH, "meta.json")) as f:
        meta = json.load(f)
    return clf, meta




def run_full_pipeline():
    ensure_dirs()
    h, w = IMG_SIZE
    data_source = "kaggle"


    X_images, y, target_names = load_kaggle_faces()


    if X_images is None or y is None:
        X_images, y, target_names = load_lfw(min_faces=40)
        data_source = "lfw"


    plot_sample_faces(X_images, title=f"Sample Faces ({data_source.upper()} Dataset)")


    X = X_images.reshape(len(X_images), -1)


    X_train, X_test, y_train, y_test, img_train, img_test = train_test_split(
        X, y, X_images, test_size=0.25, random_state=42, stratify=y
    )


    clf = train(X_train, y_train)


    y_pred, report_dict, cm = evaluate(clf, X_test, y_test, target_names)


    pca = clf.best_estimator_.named_steps["pca"]
    plot_eigenfaces(pca, h, w)
    plot_predictions(img_test, y_test, y_pred, target_names)
    plot_confusion_matrix(cm, target_names)


    save_model(clf, target_names, h, w, data_source)


    summary = {
        "accuracy": round(report_dict["accuracy"], 4),
        "n_classes": len(target_names),
        "n_train": len(y_train),
        "n_test": len(y_test),
        "data_source": data_source,
        "best_params": clf.best_params_,
        "per_class": {
            name: {
                "precision": round(report_dict[name]["precision"], 3),
                "recall": round(report_dict[name]["recall"], 3),
                "f1": round(report_dict[name]["f1-score"], 3),
                "support": int(report_dict[name]["support"]),
            }
            for name in target_names
        },
    }
    with open(os.path.join(STATIC_PATH, "summary.json"), "w") as f:
        json.dump(summary, f, indent=2)


    return clf, summary




if __name__ == "__main__":
    run_full_pipeline()