MorphVideo/eye_shape_classifier.py

"""
Eye Shape Classification for Face Organization

Classifies faces by eye shape characteristics to organize morphing sequences.
"""
import os
import cv2
import numpy as np
from typing import List, Dict, Tuple, Optional
from PIL import Image
import logging
import dlib
from sklearn.cluster import KMeans

logger = logging.getLogger(__name__)

class EyeShapeClassifier:
    """Classifies eye shapes from facial landmarks for video organization."""
    
    def __init__(self, predictor_path: Optional[str] = None):
        """
        Initialize the eye shape classifier.
        
        Args:
            predictor_path: Path to dlib face landmark predictor
        """
        # Use the existing shape predictor from pSp
        if predictor_path is None:
            predictor_path = os.path.join(
                os.path.dirname(os.path.dirname(__file__)),
                'pSp', 'shape_predictor_68_face_landmarks.dat'
            )
        
        if not os.path.exists(predictor_path):
            raise FileNotFoundError(f"Shape predictor not found: {predictor_path}")
        
        self.detector = dlib.get_frontal_face_detector()
        self.predictor = dlib.shape_predictor(predictor_path)
        
        # Eye landmark indices (68-point model)
        self.left_eye_indices = list(range(36, 42))
        self.right_eye_indices = list(range(42, 48))
        
        # Eye shape categories
        self.eye_shapes = {
            'almond': 0,
            'round': 1, 
            'hooded': 2,
            'upturned': 3,
            'downturned': 4,
            'deep_set': 5
        }
    
    def extract_eye_landmarks(self, image_path: str) -> Optional[Dict[str, np.ndarray]]:
        """
        Extract eye landmarks from an image.
        
        Args:
            image_path: Path to the image file
            
        Returns:
            Dictionary with left and right eye landmarks, or None if face not detected
        """
        try:
            # Load image
            img = cv2.imread(image_path)
            if img is None:
                logger.warning(f"Could not load image: {image_path}")
                return None
                
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            
            # Detect faces
            faces = self.detector(gray)
            if len(faces) == 0:
                logger.warning(f"No faces detected in: {image_path}")
                return None
                
            # Use the largest face
            face = max(faces, key=lambda f: f.width() * f.height())
            
            # Get landmarks
            landmarks = self.predictor(gray, face)
            
            # Extract eye coordinates
            left_eye = np.array([[landmarks.part(i).x, landmarks.part(i).y] 
                               for i in self.left_eye_indices])
            right_eye = np.array([[landmarks.part(i).x, landmarks.part(i).y] 
                                for i in self.right_eye_indices])
            
            return {
                'left_eye': left_eye,
                'right_eye': right_eye,
                'face_box': (face.left(), face.top(), face.width(), face.height())
            }
            
        except Exception as e:
            logger.error(f"Error extracting landmarks from {image_path}: {e}")
            return None
    
    def calculate_eye_features(self, landmarks: Dict[str, np.ndarray]) -> np.ndarray:
        """
        Calculate eye shape features from landmarks.
        
        Args:
            landmarks: Dictionary with eye landmarks
            
        Returns:
            Feature vector for eye shape classification
        """
        features = []
        
        for eye_key in ['left_eye', 'right_eye']:
            eye_points = landmarks[eye_key]
            
            # Calculate eye dimensions
            eye_width = np.max(eye_points[:, 0]) - np.min(eye_points[:, 0])
            eye_height = np.max(eye_points[:, 1]) - np.min(eye_points[:, 1])
            
            # Aspect ratio
            aspect_ratio = eye_width / (eye_height + 1e-6)
            
            # Eye curvature (using top and bottom points)
            top_point = eye_points[np.argmin(eye_points[:, 1])]
            bottom_point = eye_points[np.argmax(eye_points[:, 1])]
            left_point = eye_points[np.argmin(eye_points[:, 0])]
            right_point = eye_points[np.argmax(eye_points[:, 0])]
            
            # Calculate angles
            outer_angle = self._calculate_angle(left_point, top_point, right_point)
            inner_curve = self._calculate_curvature(eye_points)
            
            # Relative position features
            centroid = np.mean(eye_points, axis=0)
            relative_height = (centroid[1] - np.min(eye_points[:, 1])) / eye_height
            
            features.extend([
                aspect_ratio,
                outer_angle,
                inner_curve,
                relative_height,
                eye_width,
                eye_height
            ])
        
        return np.array(features)
    
    def _calculate_angle(self, p1: np.ndarray, p2: np.ndarray, p3: np.ndarray) -> float:
        """Calculate angle between three points."""
        v1 = p1 - p2
        v2 = p3 - p2
        
        cos_angle = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-6)
        angle = np.arccos(np.clip(cos_angle, -1.0, 1.0))
        return np.degrees(angle)
    
    def _calculate_curvature(self, points: np.ndarray) -> float:
        """Calculate average curvature of eye shape."""
        # Sort points by x-coordinate for consistent ordering
        sorted_points = points[np.argsort(points[:, 0])]
        
        curvatures = []
        for i in range(1, len(sorted_points) - 1):
            p1, p2, p3 = sorted_points[i-1], sorted_points[i], sorted_points[i+1]
            
            # Calculate curvature using circumcircle radius
            a = np.linalg.norm(p2 - p1)
            b = np.linalg.norm(p3 - p2)
            c = np.linalg.norm(p1 - p3)
            
            # Area using cross product
            area = 0.5 * abs((p2[0] - p1[0]) * (p3[1] - p1[1]) - (p3[0] - p1[0]) * (p2[1] - p1[1]))
            
            if area > 1e-6:
                curvature = (a * b * c) / (4 * area + 1e-6)
                curvatures.append(1.0 / (curvature + 1e-6))
        
        return np.mean(curvatures) if curvatures else 0.0
    
    def classify_faces_by_eye_shape(self, image_paths: List[str], n_clusters: int = 6) -> Dict[str, List[str]]:
        """
        Classify faces by eye shape using clustering.
        
        Args:
            image_paths: List of image file paths
            n_clusters: Number of eye shape clusters
            
        Returns:
            Dictionary mapping eye shape categories to image paths
        """
        features_list = []
        valid_paths = []
        
        logger.info(f"Extracting eye features from {len(image_paths)} images...")
        
        # Extract features from all images
        for img_path in image_paths:
            landmarks = self.extract_eye_landmarks(img_path)
            if landmarks is not None:
                features = self.calculate_eye_features(landmarks)
                features_list.append(features)
                valid_paths.append(img_path)
        
        if len(features_list) < n_clusters:
            logger.warning(f"Not enough valid faces ({len(features_list)}) for {n_clusters} clusters")
            n_clusters = max(1, len(features_list))
        
        logger.info(f"Clustering {len(features_list)} faces into {n_clusters} eye shape groups...")
        
        # Perform clustering
        features_array = np.array(features_list)
        kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
        cluster_labels = kmeans.fit_predict(features_array)
        
        # Group images by cluster
        clustered_faces = {}
        shape_names = list(self.eye_shapes.keys())[:n_clusters]
        
        for i, shape_name in enumerate(shape_names):
            clustered_faces[shape_name] = [
                valid_paths[j] for j, label in enumerate(cluster_labels) if label == i
            ]
        
        # Log results
        for shape, paths in clustered_faces.items():
            logger.info(f"{shape}: {len(paths)} faces")
        
        return clustered_faces
    
    def get_morphing_sequence(self, clustered_faces: Dict[str, List[str]], 
                            faces_per_group: int = 3) -> List[str]:
        """
        Create an optimal morphing sequence across eye shape groups.
        
        Args:
            clustered_faces: Dictionary of eye shape groups
            faces_per_group: Number of faces to select from each group
            
        Returns:
            Ordered list of image paths for smooth morphing
        """
        sequence = []
        
        # Select representative faces from each group
        for shape_name, face_paths in clustered_faces.items():
            if face_paths:
                # Select up to faces_per_group images from this cluster
                selected = face_paths[:min(faces_per_group, len(face_paths))]
                sequence.extend(selected)
        
        logger.info(f"Created morphing sequence with {len(sequence)} faces")
        return sequence