| |
| """Evaluate a face-recognition backbone on the LFW 6000 verification pairs.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import csv |
| import json |
| import os |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
| from typing import Iterable |
|
|
| import matplotlib |
|
|
| matplotlib.use("Agg") |
|
|
| import matplotlib.pyplot as plt |
| import numpy as np |
| import torch |
| from facenet_pytorch import MTCNN, fixed_image_standardization |
| from PIL import Image |
| from sklearn.metrics import auc, confusion_matrix, roc_curve |
| from torch import nn |
| from torch.utils.data import DataLoader, Dataset |
| from torchvision import transforms |
| from tqdm import tqdm |
|
|
| try: |
| from face_backbones import build_backbone, canonical_backbone_name |
| except ModuleNotFoundError: |
| from .face_backbones import build_backbone, canonical_backbone_name |
|
|
|
|
| PROJECT_ROOT = Path(__file__).resolve().parents[1] |
| DEFAULT_LFW_ROOT = PROJECT_ROOT / "data" / "raw" / "lfw-deepfunneled" |
| DEFAULT_PAIRS = PROJECT_ROOT / "design" / "lfw_test_pair.txt" |
| DEFAULT_OUTPUT = PROJECT_ROOT / "results" |
| DEFAULT_MODEL_DIR = PROJECT_ROOT / "models" |
|
|
|
|
| @dataclass |
| class Pair: |
| image_a: str |
| image_b: str |
| label: int |
|
|
|
|
| @dataclass |
| class FoldResult: |
| fold: int |
| threshold: float |
| accuracy: float |
| true_positive: int |
| false_positive: int |
| true_negative: int |
| false_negative: int |
|
|
|
|
| class LFWImageDataset(Dataset): |
| def __init__( |
| self, |
| root: Path, |
| rel_paths: list[str], |
| preprocess: str, |
| image_size: int, |
| mtcnn_margin: int, |
| device: torch.device, |
| ): |
| self.root = root |
| self.rel_paths = rel_paths |
| self.preprocess = preprocess |
| self.image_size = image_size |
| self.device = device |
| self.resize = transforms.Resize((image_size, image_size), antialias=True) |
| self.to_tensor = transforms.PILToTensor() |
| self.mtcnn = None |
| if preprocess == "mtcnn": |
| self.mtcnn = MTCNN(image_size=image_size, margin=mtcnn_margin, post_process=True, device=device) |
|
|
| def __len__(self) -> int: |
| return len(self.rel_paths) |
|
|
| def __getitem__(self, index: int) -> tuple[str, torch.Tensor, bool]: |
| rel_path = self.rel_paths[index] |
| image_path = self.root / rel_path |
| with Image.open(image_path) as img: |
| img = img.convert("RGB") |
| if self.preprocess == "resize": |
| tensor = self.to_tensor(self.resize(img)).float() |
| tensor = fixed_image_standardization(tensor) |
| return rel_path, tensor, True |
|
|
| assert self.mtcnn is not None |
| face = self.mtcnn(img) |
| if face is None: |
| tensor = self.to_tensor(self.resize(img)).float() |
| tensor = fixed_image_standardization(tensor) |
| return rel_path, tensor, False |
| return rel_path, face.cpu(), True |
|
|
|
|
| def parse_pairs(path: Path) -> list[Pair]: |
| pairs: list[Pair] = [] |
| with path.open("r", encoding="utf-8") as f: |
| for line_no, line in enumerate(f, start=1): |
| parts = line.strip().split() |
| if not parts: |
| continue |
| if len(parts) != 3: |
| raise ValueError(f"Invalid pair line {line_no}: {line!r}") |
| pairs.append(Pair(parts[0], parts[1], int(parts[2]))) |
| if len(pairs) != 6000: |
| raise ValueError(f"Expected 6000 LFW pairs, got {len(pairs)}") |
| labels = np.array([p.label for p in pairs]) |
| if labels.sum() != 3000 or len(labels) - labels.sum() != 3000: |
| raise ValueError("Expected 3000 positive and 3000 negative pairs") |
| return pairs |
|
|
|
|
| def unique_image_paths(pairs: Iterable[Pair]) -> list[str]: |
| paths = sorted({p.image_a for p in pairs} | {p.image_b for p in pairs}) |
| return paths |
|
|
|
|
| def collate_images(batch: list[tuple[str, torch.Tensor, bool]]) -> tuple[list[str], torch.Tensor, list[bool]]: |
| paths, tensors, detected = zip(*batch) |
| return list(paths), torch.stack(list(tensors), dim=0), list(detected) |
|
|
|
|
| def build_model( |
| model_name: str, |
| backbone_name: str, |
| device: torch.device, |
| checkpoint: Path | None = None, |
| ) -> tuple[nn.Module, str]: |
| if checkpoint: |
| payload = torch.load(checkpoint, map_location="cpu") |
| checkpoint_backbone = canonical_backbone_name(payload.get("backbone", payload.get("model_arch", backbone_name))) |
| model, _, canonical = build_backbone(checkpoint_backbone, pretrained_model=None) |
| state = payload.get("backbone_state_dict", payload.get("model_state_dict", payload)) |
| missing, unexpected = model.load_state_dict(state, strict=False) |
| allowed_unexpected = [k for k in unexpected if k.startswith("logits.")] |
| if missing or len(allowed_unexpected) != len(unexpected): |
| raise RuntimeError(f"Checkpoint load issue. missing={missing}, unexpected={unexpected}") |
| else: |
| model, _, canonical = build_backbone(backbone_name, pretrained_model=model_name) |
| return model.eval().to(device), canonical |
|
|
|
|
| def extract_embeddings( |
| model: nn.Module, |
| dataset: LFWImageDataset, |
| batch_size: int, |
| num_workers: int, |
| device: torch.device, |
| tta_flip: bool, |
| ) -> tuple[dict[str, np.ndarray], dict[str, bool]]: |
| loader = DataLoader( |
| dataset, |
| batch_size=batch_size, |
| shuffle=False, |
| num_workers=num_workers, |
| collate_fn=collate_images, |
| pin_memory=(device.type == "cuda"), |
| ) |
| embeddings: dict[str, np.ndarray] = {} |
| detected: dict[str, bool] = {} |
| with torch.inference_mode(): |
| for rel_paths, batch, batch_detected in tqdm(loader, desc="Extract embeddings"): |
| batch = batch.to(device, non_blocking=True) |
| if tta_flip: |
| flipped = torch.flip(batch, dims=[3]) |
| out_all = model(torch.cat([batch, flipped], dim=0)) |
| out = (out_all[: batch.size(0)] + out_all[batch.size(0) :]) / 2.0 |
| else: |
| out = model(batch) |
| out = torch.nn.functional.normalize(out, p=2, dim=1) |
| out_np = out.detach().cpu().numpy().astype(np.float32) |
| for rel_path, emb, ok in zip(rel_paths, out_np, batch_detected): |
| embeddings[rel_path] = emb |
| detected[rel_path] = bool(ok) |
| return embeddings, detected |
|
|
|
|
| def pair_scores(pairs: list[Pair], embeddings: dict[str, np.ndarray]) -> tuple[np.ndarray, np.ndarray]: |
| scores = np.empty(len(pairs), dtype=np.float32) |
| labels = np.empty(len(pairs), dtype=np.int64) |
| for i, pair in enumerate(pairs): |
| emb_a = embeddings[pair.image_a] |
| emb_b = embeddings[pair.image_b] |
| scores[i] = float(np.dot(emb_a, emb_b)) |
| labels[i] = pair.label |
| return scores, labels |
|
|
|
|
| def best_threshold(scores: np.ndarray, labels: np.ndarray) -> tuple[float, float]: |
| order = np.argsort(scores) |
| sorted_scores = scores[order] |
| candidates = np.empty(len(sorted_scores) + 1, dtype=np.float32) |
| candidates[0] = sorted_scores[0] - 1e-6 |
| candidates[-1] = sorted_scores[-1] + 1e-6 |
| if len(sorted_scores) > 1: |
| candidates[1:-1] = (sorted_scores[:-1] + sorted_scores[1:]) / 2.0 |
|
|
| best_acc = -1.0 |
| best_thr = float(candidates[0]) |
| for threshold in candidates: |
| pred = (scores >= threshold).astype(np.int64) |
| acc = float((pred == labels).mean()) |
| if acc > best_acc: |
| best_acc = acc |
| best_thr = float(threshold) |
| return best_thr, best_acc |
|
|
|
|
| def fold_indices(num_pairs: int = 6000, folds: int = 10) -> list[np.ndarray]: |
| if num_pairs != 6000 or folds != 10: |
| raise ValueError("This helper is for the official 6000-pair, 10-fold LFW protocol") |
| fold_ids = [] |
| for fold in range(folds): |
| pos = np.arange(fold * 300, (fold + 1) * 300) |
| neg = np.arange(3000 + fold * 300, 3000 + (fold + 1) * 300) |
| fold_ids.append(np.concatenate([pos, neg])) |
| return fold_ids |
|
|
|
|
| def evaluate_10_fold(scores: np.ndarray, labels: np.ndarray) -> tuple[list[FoldResult], np.ndarray]: |
| ids = fold_indices(len(scores), 10) |
| all_indices = np.arange(len(scores)) |
| predictions = np.zeros_like(labels) |
| results: list[FoldResult] = [] |
| for fold, test_idx in enumerate(ids, start=1): |
| train_idx = np.setdiff1d(all_indices, test_idx) |
| threshold, _ = best_threshold(scores[train_idx], labels[train_idx]) |
| pred = (scores[test_idx] >= threshold).astype(np.int64) |
| predictions[test_idx] = pred |
| tn, fp, fn, tp = confusion_matrix(labels[test_idx], pred, labels=[0, 1]).ravel() |
| results.append( |
| FoldResult( |
| fold=fold, |
| threshold=threshold, |
| accuracy=float((pred == labels[test_idx]).mean()), |
| true_positive=int(tp), |
| false_positive=int(fp), |
| true_negative=int(tn), |
| false_negative=int(fn), |
| ) |
| ) |
| return results, predictions |
|
|
|
|
| def save_scores_csv(path: Path, pairs: list[Pair], scores: np.ndarray, labels: np.ndarray) -> None: |
| with path.open("w", encoding="utf-8", newline="") as f: |
| writer = csv.writer(f) |
| writer.writerow(["image_a", "image_b", "label", "cosine_score"]) |
| for pair, score, label in zip(pairs, scores, labels): |
| writer.writerow([pair.image_a, pair.image_b, int(label), f"{float(score):.8f}"]) |
|
|
|
|
| def save_fold_csv(path: Path, fold_results: list[FoldResult]) -> None: |
| with path.open("w", encoding="utf-8", newline="") as f: |
| writer = csv.DictWriter(f, fieldnames=list(asdict(fold_results[0]).keys())) |
| writer.writeheader() |
| for result in fold_results: |
| writer.writerow(asdict(result)) |
|
|
|
|
| def plot_roc(path: Path, labels: np.ndarray, scores: np.ndarray) -> float: |
| fpr, tpr, _ = roc_curve(labels, scores) |
| roc_auc = float(auc(fpr, tpr)) |
| fig, ax = plt.subplots(figsize=(6, 5)) |
| ax.plot(fpr, tpr, label=f"AUC = {roc_auc:.4f}", linewidth=2) |
| ax.plot([0, 1], [0, 1], linestyle="--", color="gray", linewidth=1) |
| ax.set_xlabel("False Positive Rate") |
| ax.set_ylabel("True Positive Rate") |
| ax.set_title("LFW ROC Curve") |
| ax.legend(loc="lower right") |
| ax.grid(alpha=0.25) |
| fig.tight_layout() |
| fig.savefig(path, dpi=200) |
| plt.close(fig) |
| return roc_auc |
|
|
|
|
| def plot_confusion(path: Path, labels: np.ndarray, predictions: np.ndarray) -> list[list[int]]: |
| cm = confusion_matrix(labels, predictions, labels=[0, 1]) |
| fig, ax = plt.subplots(figsize=(5, 4)) |
| im = ax.imshow(cm, cmap="Blues") |
| ax.set_xticks([0, 1], labels=["Different", "Same"]) |
| ax.set_yticks([0, 1], labels=["Different", "Same"]) |
| ax.set_xlabel("Predicted") |
| ax.set_ylabel("Ground Truth") |
| ax.set_title("LFW Confusion Matrix") |
| for i in range(2): |
| for j in range(2): |
| ax.text(j, i, str(int(cm[i, j])), ha="center", va="center", color="black") |
| fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) |
| fig.tight_layout() |
| fig.savefig(path, dpi=200) |
| plt.close(fig) |
| return cm.astype(int).tolist() |
|
|
|
|
| def plot_score_histogram(path: Path, labels: np.ndarray, scores: np.ndarray, threshold: float) -> None: |
| fig, ax = plt.subplots(figsize=(7, 4)) |
| ax.hist(scores[labels == 1], bins=60, alpha=0.75, label="Same person") |
| ax.hist(scores[labels == 0], bins=60, alpha=0.75, label="Different people") |
| ax.axvline(threshold, color="black", linestyle="--", linewidth=1.5, label=f"Global threshold {threshold:.4f}") |
| ax.set_xlabel("Cosine Similarity") |
| ax.set_ylabel("Pair Count") |
| ax.set_title("LFW Pair Score Distribution") |
| ax.legend() |
| fig.tight_layout() |
| fig.savefig(path, dpi=200) |
| plt.close(fig) |
|
|
|
|
| def verify_files_exist(lfw_root: Path, pairs: list[Pair]) -> None: |
| missing = [] |
| for rel_path in unique_image_paths(pairs): |
| if not (lfw_root / rel_path).is_file(): |
| missing.append(rel_path) |
| if len(missing) >= 10: |
| break |
| if missing: |
| raise FileNotFoundError(f"Missing LFW images under {lfw_root}: {missing}") |
|
|
|
|
| def save_model_artifact(model_name: str, model_dir: Path, checkpoint: Path | None = None) -> Path: |
| if checkpoint is not None: |
| return checkpoint.resolve() |
| checkpoint_name = { |
| "casia-webface": "20180408-102900-casia-webface.pt", |
| "vggface2": "20180402-114759-vggface2.pt", |
| }[model_name] |
| source = Path(os.environ.get("TORCH_HOME", str(model_dir / "torch"))) / "checkpoints" / checkpoint_name |
| target = model_dir / f"facenet_{model_name.replace('-', '_')}.pth" |
| if source.exists(): |
| import shutil |
|
|
| shutil.copy2(source, target) |
| return target |
|
|
|
|
| def write_metrics( |
| path: Path, |
| args: argparse.Namespace, |
| labels: np.ndarray, |
| scores: np.ndarray, |
| predictions: np.ndarray, |
| fold_results: list[FoldResult], |
| global_threshold: float, |
| global_accuracy: float, |
| roc_auc: float, |
| confusion: list[list[int]], |
| detected: dict[str, bool], |
| model_artifact: Path, |
| backbone_name: str, |
| ) -> None: |
| fold_accuracies = [r.accuracy for r in fold_results] |
| metrics = { |
| "model": "local_checkpoint" if args.checkpoint else args.model, |
| "pretrained_model_arg": None if args.checkpoint else args.model, |
| "backbone": backbone_name, |
| "checkpoint": str(args.checkpoint) if args.checkpoint else None, |
| "preprocess": args.preprocess, |
| "mtcnn_margin": int(args.mtcnn_margin), |
| "tta_flip": bool(args.tta_flip), |
| "image_size": int(args.image_size), |
| "device": args.device, |
| "lfw_root": str(args.lfw_root), |
| "pairs_file": str(args.pairs_file), |
| "num_pairs": int(len(labels)), |
| "num_positive_pairs": int(labels.sum()), |
| "num_negative_pairs": int(len(labels) - labels.sum()), |
| "num_unique_images": int(len(detected)), |
| "mtcnn_or_resize_success_rate": float(np.mean(list(detected.values()))), |
| "ten_fold_accuracy_mean": float(np.mean(fold_accuracies)), |
| "ten_fold_accuracy_std": float(np.std(fold_accuracies)), |
| "ten_fold_accuracies": fold_accuracies, |
| "global_best_threshold": float(global_threshold), |
| "global_best_accuracy": float(global_accuracy), |
| "roc_auc": float(roc_auc), |
| "confusion_matrix_10_fold": confusion, |
| "model_artifact": str(model_artifact), |
| } |
| with path.open("w", encoding="utf-8") as f: |
| json.dump(metrics, f, indent=2, ensure_ascii=False) |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--lfw-root", type=Path, default=DEFAULT_LFW_ROOT) |
| parser.add_argument("--pairs-file", type=Path, default=DEFAULT_PAIRS) |
| parser.add_argument("--output-dir", type=Path, default=DEFAULT_OUTPUT) |
| parser.add_argument("--model-dir", type=Path, default=DEFAULT_MODEL_DIR) |
| parser.add_argument("--model", choices=["casia-webface", "vggface2"], default="casia-webface") |
| parser.add_argument( |
| "--backbone", |
| choices=["inception_resnet_v1", "ir_resnet18", "ir_resnet34"], |
| default="inception_resnet_v1", |
| help="used only without --checkpoint; checkpoints carry their own backbone", |
| ) |
| parser.add_argument("--checkpoint", type=Path, default=None, help="local scratch-trained backbone checkpoint") |
| parser.add_argument("--preprocess", choices=["resize", "mtcnn"], default="resize") |
| parser.add_argument("--image-size", type=int, default=160) |
| parser.add_argument("--mtcnn-margin", type=int, default=0) |
| parser.add_argument("--batch-size", type=int, default=512) |
| parser.add_argument("--num-workers", type=int, default=4) |
| parser.add_argument("--tta-flip", action="store_true", help="average embeddings from original and horizontal flip") |
| parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu") |
| return parser.parse_args() |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| project_torch_home = args.model_dir / "torch" |
| os.environ.setdefault("TORCH_HOME", str(project_torch_home)) |
| args.output_dir.mkdir(parents=True, exist_ok=True) |
| args.model_dir.mkdir(parents=True, exist_ok=True) |
|
|
| pairs = parse_pairs(args.pairs_file) |
| verify_files_exist(args.lfw_root, pairs) |
|
|
| device = torch.device(args.device) |
| model, backbone_name = build_model(args.model, args.backbone, device, args.checkpoint) |
| rel_paths = unique_image_paths(pairs) |
| dataset = LFWImageDataset(args.lfw_root, rel_paths, args.preprocess, args.image_size, args.mtcnn_margin, device) |
| embeddings, detected = extract_embeddings(model, dataset, args.batch_size, args.num_workers, device, args.tta_flip) |
| scores, labels = pair_scores(pairs, embeddings) |
|
|
| fold_results, predictions = evaluate_10_fold(scores, labels) |
| global_threshold, global_accuracy = best_threshold(scores, labels) |
| roc_auc = plot_roc(args.output_dir / "roc_curve.png", labels, scores) |
| confusion = plot_confusion(args.output_dir / "confusion_matrix.png", labels, predictions) |
| plot_score_histogram(args.output_dir / "score_histogram.png", labels, scores, global_threshold) |
|
|
| save_scores_csv(args.output_dir / "pair_scores.csv", pairs, scores, labels) |
| save_fold_csv(args.output_dir / "fold_metrics.csv", fold_results) |
| np.savez_compressed( |
| args.output_dir / "lfw_embeddings.npz", |
| paths=np.array(rel_paths), |
| embeddings=np.stack([embeddings[p] for p in rel_paths]), |
| ) |
| model_artifact = save_model_artifact(args.model, args.model_dir, args.checkpoint) |
| write_metrics( |
| args.output_dir / "metrics.json", |
| args, |
| labels, |
| scores, |
| predictions, |
| fold_results, |
| global_threshold, |
| global_accuracy, |
| roc_auc, |
| confusion, |
| detected, |
| model_artifact, |
| backbone_name, |
| ) |
|
|
| mean_acc = np.mean([r.accuracy for r in fold_results]) |
| std_acc = np.std([r.accuracy for r in fold_results]) |
| print(f"LFW 10-fold accuracy: {mean_acc:.4%} ± {std_acc:.4%}") |
| print(f"ROC AUC: {roc_auc:.6f}") |
| print(f"Global best accuracy: {global_accuracy:.4%} @ threshold {global_threshold:.6f}") |
| print(f"Results written to: {args.output_dir}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
