face_recognition.py

"""
ConvRec Face Recognition Inference Module
High-performance face recognition using ConvRec proprietary model
"""

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms, models
from PIL import Image
import numpy as np
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class FaceModel(nn.Module):
    """Face recognition model with ResNet50 backbone"""
    def __init__(self, embedding_size=512):
        super().__init__()
        resnet = models.resnet50(weights=None)
        self.features = nn.Sequential(*list(resnet.children())[:-1])
        self.bn1 = nn.BatchNorm1d(2048)
        self.dropout = nn.Dropout(0.4)
        self.fc = nn.Linear(2048, embedding_size)
        self.bn2 = nn.BatchNorm1d(embedding_size)

    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        x = self.bn1(x)
        x = self.dropout(x)
        x = self.fc(x)
        x = self.bn2(x)
        return F.normalize(x, p=2, dim=1)

class FaceRecognition:
    """Main class for face recognition tasks"""

    def __init__(self, model_path='best_model.pth', device=None):
        """
        Initialize face recognition model

        Args:
            model_path: Path to model checkpoint
            device: Device to use ('cuda', 'cpu', or None for auto)
        """
        if device is None:
            self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        else:
            self.device = torch.device(device)

        self.model = self._load_model(model_path)
        self.transform = self._get_transform()

        logger.info(f"Model loaded on {self.device}")

    def _load_model(self, model_path):
        """Load the trained model"""
        model = FaceModel(embedding_size=512).to(self.device)

        if os.path.exists(model_path):
            checkpoint = torch.load(model_path, map_location=self.device)

            if 'model_state_dict' in checkpoint:
                model.load_state_dict(checkpoint['model_state_dict'])
            elif 'model' in checkpoint:
                model.load_state_dict(checkpoint['model'])
            else:
                model.load_state_dict(checkpoint)

            logger.info(f"Model loaded from {model_path}")
        else:
            logger.warning(f"Model file not found: {model_path}")

        model.eval()
        return model

    def _get_transform(self):
        """Get image preprocessing transform"""
        return transforms.Compose([
            transforms.Resize((120, 120)),
            transforms.CenterCrop((112, 112)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
        ])

    def extract_embedding(self, image_path):
        """
        Extract face embedding from an image

        Args:
            image_path: Path to image file or PIL Image

        Returns:
            512-dimensional embedding vector
        """
        if isinstance(image_path, str):
            image = Image.open(image_path).convert('RGB')
        else:
            image = image_path

        image = self.transform(image).unsqueeze(0).to(self.device)

        with torch.no_grad():
            embedding = self.model(image)

        return embedding.cpu().numpy().squeeze()

    def compute_similarity(self, embedding1, embedding2):
        """
        Compute cosine similarity between two embeddings

        Args:
            embedding1: First embedding vector
            embedding2: Second embedding vector

        Returns:
            Similarity score between -1 and 1
        """
        # Embeddings are already normalized
        return np.dot(embedding1, embedding2)

    def verify_faces(self, image1_path, image2_path, threshold=0.5):
        """
        Verify if two faces belong to the same person

        Args:
            image1_path: Path to first image
            image2_path: Path to second image
            threshold: Similarity threshold for verification

        Returns:
            Dictionary with verification results
        """
        emb1 = self.extract_embedding(image1_path)
        emb2 = self.extract_embedding(image2_path)

        similarity = self.compute_similarity(emb1, emb2)
        is_same = similarity > threshold

        return {
            'is_same': bool(is_same),
            'similarity': float(similarity),
            'threshold': threshold
        }

    def build_gallery(self, folder_path, max_per_person=5):
        """
        Build a gallery of known faces from a folder

        Args:
            folder_path: Path to folder with subfolders for each person
            max_per_person: Maximum images per person to use

        Returns:
            Gallery dictionary with embeddings
        """
        gallery = {}
        folder_path = Path(folder_path)

        for person_dir in folder_path.iterdir():
            if person_dir.is_dir():
                person_name = person_dir.name
                embeddings = []

                image_files = list(person_dir.glob('*.jpg')) + list(person_dir.glob('*.png'))

                for img_path in image_files[:max_per_person]:
                    emb = self.extract_embedding(str(img_path))
                    if emb is not None:
                        embeddings.append(emb)

                if embeddings:
                    # Average embeddings for robustness
                    avg_emb = np.mean(embeddings, axis=0)
                    avg_emb = avg_emb / np.linalg.norm(avg_emb)
                    gallery[person_name] = avg_emb
                    logger.info(f"Added {person_name} to gallery ({len(embeddings)} images)")

        logger.info(f"Gallery built with {len(gallery)} persons")
        return gallery

    def search_in_gallery(self, query_image_path, gallery, top_k=5):
        """
        Search for a face in the gallery

        Args:
            query_image_path: Path to query image
            gallery: Gallery dictionary from build_gallery
            top_k: Number of top matches to return

        Returns:
            List of (person_name, similarity) tuples
        """
        query_emb = self.extract_embedding(query_image_path)

        similarities = {}
        for person_name, gallery_emb in gallery.items():
            sim = self.compute_similarity(query_emb, gallery_emb)
            similarities[person_name] = sim

        # Sort by similarity
        sorted_matches = sorted(similarities.items(), key=lambda x: x[1], reverse=True)

        return sorted_matches[:top_k]

    def are_same_person(self, image1_path, image2_path, threshold=0.5):
        """
        Simple boolean check if two images are the same person

        Args:
            image1_path: Path to first image
            image2_path: Path to second image
            threshold: Similarity threshold

        Returns:
            Boolean indicating if same person
        """
        result = self.verify_faces(image1_path, image2_path, threshold)
        return result['is_same']

# Example usage
if __name__ == '__main__':
    # Initialize model
    fr = FaceRecognition('best_model.pth')

    # Example: Verify two faces
    if os.path.exists('data'):
        # Find some test images
        test_images = list(Path('data').glob('**/*.jpg'))
        if len(test_images) >= 2:
            result = fr.verify_faces(str(test_images[0]), str(test_images[1]))
            print(f"Verification result: {result}")

    print("Face recognition model ready for inference!")