src/morph_generation/morph_generator.py

"""
MorphGuard Face Morph Generator
===============================
Implements multi-tiered facial landmark extraction, Delaunay triangulation,
piecewise affine triangle warping, and seamless Poisson gradient blending.
"""

import os
import cv2
import math
import uuid
import logging
import numpy as np
from typing import Tuple, List, Dict, Optional

try:
    from scipy.spatial import Delaunay
    from scipy.spatial.distance import euclidean
    SCIPY_AVAILABLE = True
except ImportError:
    Delaunay = None
    euclidean = None
    SCIPY_AVAILABLE = False

logger = logging.getLogger("MorphGenerator")

class FaceMorphGenerator:
    """
    Robust Face Morph Generator supporting:
    - 68-point Dlib landmark extraction (if available)
    - Resilient OpenCV Haar Cascade face + eye fallback
    - Eye alignment & horizon rotation normalization
    - Delaunay piecewise affine triangle warping
    - Seamless Poisson cloning (cv2.seamlessClone)
    """

    def __init__(self, predictor_path: Optional[str] = None):
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        self.eye_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_eye.xml')
        
        self.dlib_available = False
        self.dlib_detector = None
        self.dlib_predictor = None

        # Check for dlib
        try:
            import dlib
            if predictor_path is None:
                candidates = [
                    "models/shape_predictor_68_face_landmarks.dat",
                    "pSp/shape_predictor_68_face_landmarks.dat",
                    os.path.join(os.getcwd(), "models/shape_predictor_68_face_landmarks.dat")
                ]
                for c in candidates:
                    if os.path.exists(c):
                        predictor_path = c
                        break
            
            if predictor_path and os.path.exists(predictor_path):
                self.dlib_detector = dlib.get_frontal_face_detector()
                self.dlib_predictor = dlib.shape_predictor(predictor_path)
                self.dlib_available = True
                logger.info("FaceMorphGenerator: Dlib 68-point landmark predictor active.")
        except Exception as e:
            logger.info(f"FaceMorphGenerator: Dlib not available ({e}), using OpenCV cascade fallback.")

    def detect_face_box(self, img: np.ndarray) -> Optional[Tuple[int, int, int, int]]:
        """Detect the most prominent frontal face bounding box (x, y, w, h)."""
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if len(img.shape) == 3 else img
        faces = self.face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=4, minSize=(80, 80))
        if len(faces) == 0:
            return None
        # Return largest face
        return max(faces, key=lambda f: f[2] * f[3])

    def get_landmarks(self, img: np.ndarray, face_box: Optional[Tuple[int, int, int, int]] = None) -> Optional[np.ndarray]:
        """
        Extract facial landmarks.
        Returns array of shape (N, 2) with points (x, y), including boundary/border anchors.
        """
        h, w = img.shape[:2]
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if len(img.shape) == 3 else img

        if face_box is None:
            face_box = self.detect_face_box(img)

        points = []

        # 1. Try Dlib if active
        if self.dlib_available and self.dlib_detector and self.dlib_predictor:
            try:
                import dlib
                if face_box is not None:
                    x, y, fw, fh = face_box
                    rect = dlib.rectangle(int(x), int(y), int(x + fw), int(y + fh))
                else:
                    dets = self.dlib_detector(gray, 1)
                    rect = dets[0] if len(dets) > 0 else None
                
                if rect is not None:
                    shape = self.dlib_predictor(gray, rect)
                    for i in range(68):
                        points.append([shape.part(i).x, shape.part(i).y])
            except Exception as e:
                logger.warning(f"Dlib landmark detection failed: {e}")

        # 2. Fallback to OpenCV cascade geometric points
        if len(points) == 0:
            if face_box is None:
                # If no face detected, generate centered oval synthetic anchors
                face_box = (int(w * 0.2), int(h * 0.2), int(w * 0.6), int(h * 0.6))
            
            x, y, fw, fh = face_box
            face_roi_gray = gray[y:y+fh, x:x+fw]
            
            # Detect eyes within face ROI
            eyes = self.eye_cascade.detectMultiScale(face_roi_gray, scaleFactor=1.1, minNeighbors=3, minSize=(20, 20))
            if len(eyes) >= 2:
                eyes = sorted(eyes, key=lambda e: e[2] * e[3], reverse=True)[:2]
                if eyes[0][0] > eyes[1][0]:
                    eyes[0], eyes[1] = eyes[1], eyes[0]
                left_eye = (x + eyes[0][0] + eyes[0][2]//2, y + eyes[0][1] + eyes[0][3]//2)
                right_eye = (x + eyes[1][0] + eyes[1][2]//2, y + eyes[1][1] + eyes[1][3]//2)
            else:
                left_eye = (int(x + fw * 0.35), int(y + fh * 0.40))
                right_eye = (int(x + fw * 0.65), int(y + fh * 0.40))

            nose_tip = (int(x + fw * 0.50), int(y + fh * 0.60))
            mouth_left = (int(x + fw * 0.35), int(y + fh * 0.75))
            mouth_right = (int(x + fw * 0.65), int(y + fh * 0.75))
            mouth_center = (int(x + fw * 0.50), int(y + fh * 0.78))
            chin = (int(x + fw * 0.50), int(y + fh * 0.96))
            forehead = (int(x + fw * 0.50), int(y + fh * 0.10))

            # Jawline / perimeter contour points
            contour_pts = [
                (x, int(y + fh * 0.3)), (x, int(y + fh * 0.6)), (int(x + fw * 0.2), int(y + fh * 0.85)),
                chin,
                (int(x + fw * 0.8), int(y + fh * 0.85)), (x + fw, int(y + fh * 0.6)), (x + fw, int(y + fh * 0.3))
            ]

            feature_points = [
                left_eye, right_eye, nose_tip, mouth_left, mouth_right, mouth_center, forehead
            ] + contour_pts
            
            for pt in feature_points:
                points.append([float(pt[0]), float(pt[1])])

        # 3. Add canvas frame boundary anchors to avoid mesh collapse at borders
        border_anchors = [
            [0, 0], [w // 2, 0], [w - 1, 0],
            [0, h // 2], [w - 1, h // 2],
            [0, h - 1], [w // 2, h - 1], [w - 1, h - 1]
        ]
        points.extend(border_anchors)

        return np.array(points, dtype=np.float32)

    @staticmethod
    def align_and_scale_images(img1: np.ndarray, img2: np.ndarray, target_size: Tuple[int, int] = (600, 600)) -> Tuple[np.ndarray, np.ndarray]:
        """Normalize both donor images to the same standard resolution."""
        tw, th = target_size
        scaled1 = cv2.resize(img1, (tw, th), interpolation=cv2.INTER_AREA)
        scaled2 = cv2.resize(img2, (tw, th), interpolation=cv2.INTER_AREA)
        return scaled1, scaled2

    @staticmethod
    def _apply_affine_transform(src: np.ndarray, src_tri: np.ndarray, dst_tri: np.ndarray, size: Tuple[int, int]) -> np.ndarray:
        """Calculate and apply affine transform to warp a triangular patch."""
        warp_mat = cv2.getAffineTransform(np.float32(src_tri), np.float32(dst_tri))
        dst = cv2.warpAffine(src, warp_mat, (size[0], size[1]), None, flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REFLECT_101)
        return dst

    def _morph_triangle(self, img1: np.ndarray, img2: np.ndarray, img_morph: np.ndarray,
                        t1: np.ndarray, t2: np.ndarray, t_morph: np.ndarray, alpha: float):
        """Piecewise warp a single Delaunay triangle from img1 and img2 into img_morph."""
        # Find bounding rectangle for each triangle
        r1 = cv2.boundingRect(np.float32([t1]))
        r2 = cv2.boundingRect(np.float32([t2]))
        r_morph = cv2.boundingRect(np.float32([t_morph]))

        # Offset points by top left of bounding box
        t1_rect = []
        t2_rect = []
        t_morph_rect = []

        for i in range(3):
            t_morph_rect.append(((t_morph[i][0] - r_morph[0]), (t_morph[i][1] - r_morph[1])))
            t1_rect.append(((t1[i][0] - r1[0]), (t1[i][1] - r1[1])))
            t2_rect.append(((t2[i][0] - r2[0]), (t2[i][1] - r2[1])))

        # Create mask for morphed triangle inside bounding box
        mask = np.zeros((r_morph[3], r_morph[2], 3), dtype=np.float32)
        cv2.fillConvexPoly(mask, np.int32(t_morph_rect), (1.0, 1.0, 1.0), 16, 0)

        # Crop input images to bounding rects
        img1_rect = img1[r1[1]:r1[1] + r1[3], r1[0]:r1[0] + r1[2]]
        img2_rect = img2[r2[1]:r2[1] + r2[3], r2[0]:r2[0] + r2[2]]

        size = (r_morph[2], r_morph[3])
        warp_img1 = self._apply_affine_transform(img1_rect, t1_rect, t_morph_rect, size)
        warp_img2 = self._apply_affine_transform(img2_rect, t2_rect, t_morph_rect, size)

        # Alpha blend warped patches
        img_rect = (1.0 - alpha) * warp_img1 + alpha * warp_img2

        # Copy triangular region to morphed output
        img_morph[r_morph[1]:r_morph[1]+r_morph[3], r_morph[0]:r_morph[0]+r_morph[2]] = \
            img_morph[r_morph[1]:r_morph[1]+r_morph[3], r_morph[0]:r_morph[0]+r_morph[2]] * (1 - mask) + img_rect * mask

    @staticmethod
    def _compute_subdiv_triangles(rect: Tuple[int, int, int, int], points: np.ndarray) -> List[List[int]]:
        """Compute Delaunay triangles using cv2.Subdiv2D when scipy is absent."""
        rx, ry, rw, rh = rect
        subdiv = cv2.Subdiv2D((rx, ry, rw, rh))
        
        # Insert points bounded slightly inside rect to avoid out-of-range bounds errors
        valid_indices = []
        for idx, p in enumerate(points):
            px = float(np.clip(p[0], rx + 1, rx + rw - 2))
            py = float(np.clip(p[1], ry + 1, ry + rh - 2))
            try:
                subdiv.insert((px, py))
                valid_indices.append(idx)
            except Exception:
                continue

        triangle_list = subdiv.getTriangleList()
        simplices = []
        for t in triangle_list:
            pts = [(t[0], t[1]), (t[2], t[3]), (t[4], t[5])]
            # Check if all triangle points fall inside the image canvas
            if all(rx <= pt[0] < rx + rw and ry <= pt[1] < ry + rh for pt in pts):
                # Map back to nearest point index
                try:
                    idx1 = int(np.argmin(np.sum((points - pts[0])**2, axis=1)))
                    idx2 = int(np.argmin(np.sum((points - pts[1])**2, axis=1)))
                    idx3 = int(np.argmin(np.sum((points - pts[2])**2, axis=1)))
                    if idx1 != idx2 and idx2 != idx3 and idx1 != idx3:
                        simplices.append([idx1, idx2, idx3])
                except Exception:
                    continue
        return simplices

    def generate_morph(self, img1: np.ndarray, img2: np.ndarray, alpha: float = 0.5, mode: str = "seamless") -> np.ndarray:
        """
        Generate a morphed composite face from img1 and img2.
        
        Args:
            img1: Donor image 1 (BGR uint8)
            img2: Donor image 2 (BGR uint8)
            alpha: Blend ratio in [0.0, 1.0]. 0.0 = pure img1, 1.0 = pure img2.
            mode: "delaunay" (piecewise affine only) or "seamless" (Poisson gradient fusion).

        Returns:
            morphed: Resulting composite image (BGR uint8).
        """
        alpha = float(np.clip(alpha, 0.01, 0.99))
        h, w = 600, 600
        img1, img2 = self.align_and_scale_images(img1, img2, (w, h))

        # Extract landmarks
        points1 = self.get_landmarks(img1)
        points2 = self.get_landmarks(img2)

        if points1 is None or points2 is None:
            raise ValueError("Failed to extract facial landmarks from one or both donor images.")

        # Ensure equal count
        min_pts = min(len(points1), len(points2))
        points1 = points1[:min_pts]
        points2 = points2[:min_pts]

        # Compute intermediate landmark positions
        points_morph = (1.0 - alpha) * points1 + alpha * points2

        # Delaunay triangulation on intermediate points
        if SCIPY_AVAILABLE and Delaunay is not None:
            tri = Delaunay(points_morph)
            simplices = tri.simplices
        else:
            # Fallback if scipy is absent: compute triangles using Subdiv2D
            simplices = self._compute_subdiv_triangles((0, 0, w, h), points_morph)

        # Prepare float images for precision blending
        img1_float = np.float32(img1)
        img2_float = np.float32(img2)
        img_morph_float = np.zeros((h, w, 3), dtype=np.float32)

        if len(simplices) > 0:
            # Warp each triangle
            for simplex in simplices:
                x, y, z = simplex
                t1 = [points1[x], points1[y], points1[z]]
                t2 = [points2[x], points2[y], points2[z]]
                t_morph = [points_morph[x], points_morph[y], points_morph[z]]
                self._morph_triangle(img1_float, img2_float, img_morph_float, t1, t2, t_morph, alpha)
            delaunay_result = np.uint8(np.clip(img_morph_float, 0, 255))
        else:
            # Simple weighted cross-dissolve fallback
            delaunay_result = np.uint8(np.clip((1.0 - alpha) * img1_float + alpha * img2_float, 0, 255))

        if mode == "seamless":
            # Apply Poisson seamless cloning into donor canvas to remove seam artifacts
            try:
                # Compute convex hull of facial points (excluding outer boundary corners)
                facial_points = points_morph[:-8] if len(points_morph) > 8 else points_morph
                hull = cv2.convexHull(np.int32(facial_points))
                
                mask = np.zeros((h, w), dtype=np.uint8)
                cv2.fillConvexPoly(mask, hull, 255)
                
                # Erode mask slightly to avoid perimeter edge bleed
                kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (7, 7))
                mask = cv2.erode(mask, kernel, iterations=2)
                
                # Center for seamless clone
                M = cv2.moments(mask)
                if M["m00"] > 0:
                    center = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
                    # Clone onto primary donor background
                    target_canvas = img1 if alpha <= 0.5 else img2
                    seamless_result = cv2.seamlessClone(delaunay_result, target_canvas, mask, center, cv2.NORMAL_CLONE)
                    return seamless_result
            except Exception as e:
                logger.warning(f"Seamless cloning encountered issue: {e}. Falling back to Delaunay composite.")

        return delaunay_result