build_database.py

"""
Build the encrypted face database.

Steps:
  1. Download LFW dataset (if not present)
  2. Generate masked variant
  3. Extract FaceNet features
  4. Fit K-Means cluster model
  5. Encrypt and enroll all templates
"""

import os
import numpy as np
from pathlib import Path
import config
from lfw_dataset import LFWDataset
from facenet_encoder import MaskedFaceNetEncoder
from kmeans_clustering import KMeansClustering
from database import FaceDatabase
from encryption import EncryptionManager


def download_dataset():
    """Download and prepare the LFW dataset."""
    print("="*60)
    print("Step 1: Download LFW Dataset")
    print("="*60)

    lfw = LFWDataset()
    lfw.download()
    lfw.generate_masked_dataset()

    print("Dataset ready.")
    print("="*60)


def extract_features():
    """Extract FaceNet features from enrolled images."""
    print("\n" + "="*60)
    print("Step 2: Extract Features")
    print("="*60)

    output_dir = Path(config.DATABASE_PATH)
    output_dir.mkdir(exist_ok=True)

    features_path = output_dir / 'database_features.npy'
    persons_path = output_dir / 'database_persons.npy'

    if os.path.exists(features_path) and os.path.exists(persons_path):
        print(f"Feature cache found, loading...")
        features = np.load(features_path)
        persons = np.load(persons_path)
        print(f"Loaded {len(features)} features.")
        return features, list(persons)

    lfw = LFWDataset()
    lfw.load(use_masked=False)

    encoder = MaskedFaceNetEncoder(
        use_mask=True,
        model_path=config.MASKED_MODEL_PATH
    )

    print(f"Sampling {config.DATABASE_SIZE} identities...")
    samples = lfw.sample_database_persons(n=config.DATABASE_SIZE)

    print(f"Extracting features...")
    features = []
    persons = []

    for i, (person_name, img_path) in enumerate(samples):
        try:
            img = lfw.load_image(img_path)
            feature = encoder.extract_feature(img)

            if feature is not None:
                features.append(feature)
                persons.append(person_name)

            if (i + 1) % 100 == 0:
                print(f"  Progress: {i+1}/{len(samples)} ({len(features)} successful)")

        except Exception as e:
            print(f"Failed on {img_path}: {e}")

    features = np.array(features)
    persons = np.array(persons)

    print(f"Extracted {len(features)} features.")

    np.save(features_path, features)
    np.save(persons_path, persons)

    print(f"Saved to {features_path}")
    print("="*60)

    return features, list(persons)


def perform_clustering(features):
    """Fit K-Means cluster model on extracted features."""
    print("\n" + "="*60)
    print("Step 3: K-Means Clustering")
    print("="*60)

    model_path = Path(config.DATABASE_PATH) / 'kmeans_model.pkl'

    if os.path.exists(model_path):
        print(f"Cluster model found, loading...")
        kmeans = KMeansClustering.load(model_path)
        return kmeans

    kmeans = KMeansClustering(n_clusters=config.N_CLUSTERS)
    kmeans.fit(features)
    kmeans.save(model_path)

    print("="*60)

    return kmeans


def build_encrypted_database(features, persons, kmeans):
    """Encrypt all templates and build the searchable database."""
    print("\n" + "="*60)
    print("Step 4: Build Encrypted Database")
    print("="*60)

    db_path = Path(config.DATABASE_PATH) / 'encrypted_database.pkl'

    if os.path.exists(db_path):
        print(f"Encrypted database found, loading...")
        return FaceDatabase.load(db_path)

    enc_mgr = EncryptionManager()
    db = FaceDatabase(use_kmeans=True)

    print(f"Encrypting and enrolling {len(features)} templates...")

    for i, (feature, person) in enumerate(zip(features, persons)):
        try:
            normalized_feature = MaskedFaceNetEncoder.normalize(feature)
            enc_template = enc_mgr.encrypt_template_for_enroll(normalized_feature)
            internal_id = enc_mgr.enroll_template(person, enc_template)

            cluster_id = kmeans.predict(feature)

            db.add_entry(
                cluster_id,
                person,
                enc_template,
                feature=feature
            )

            if (i + 1) % 100 == 0:
                print(f"  Progress: {i+1}/{len(features)}")

        except Exception as e:
            print(f"Failed on {person}: {e}")

    db.save(db_path)

    import pickle
    enc_state_path = os.path.join(config.DATABASE_PATH, 'enc_mgr_state.pkl')
    with open(enc_state_path, 'wb') as f:
        pickle.dump({
            'id_counter': enc_mgr.id_counter,
            'id_to_internal_id': enc_mgr.id_to_internal_id
        }, f)
    print(f"Encryption state saved: {enc_state_path}")

    print(f"Encrypted database complete. {db.total_entries} entries.")
    db.print_statistics()

    return db


def main():
    print("="*60)
    print("BioShield: Build Encrypted Database")
    print("="*60)

    download_dataset()
    features, persons = extract_features()
    kmeans = perform_clustering(features)
    db = build_encrypted_database(features, persons, kmeans)

    print("\n" + "="*60)
    print("Build complete.")
    print("="*60)
    db.print_statistics()


if __name__ == "__main__":
    main()