wink_enhancer.py

import cv2
import numpy as np

class WinkQualityEnhancer:
    """
    High-speed OpenCV/NumPy post-processor for Wink-level visual enhancement:
    1. Real Skin Grain & Texture (Frequency Separation)
    2. Localized Eye & Lip Sharpening / Sparkle Boost
    3. LAB CLAHE Lighting & Micro-Contrast Tone Balance
    """
    def __init__(self):
        self.clahe_eye = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(4, 4))
        self.clahe_lab = cv2.createCLAHE(clipLimit=1.5, tileGridSize=(8, 8))

    def apply_skin_grain(self, restored_face: np.ndarray, cropped_original: np.ndarray, skin_mask: np.ndarray = None, grain_amount: float = 0.15) -> np.ndarray:
        """
        Extract high-frequency texture from cropped_original and inject into restored_face
        to eliminate plastic/soapy skin look while preserving AI face restoration.
        """
        if grain_amount <= 0.0 or cropped_original is None:
            return restored_face
            
        try:
            # Ensure same dimensions
            if restored_face.shape[:2] != cropped_original.shape[:2]:
                cropped_orig_resized = cv2.resize(cropped_original, (restored_face.shape[1], restored_face.shape[0]), interpolation=cv2.INTER_LANCZOS4)
            else:
                cropped_orig_resized = cropped_original

            # Frequency Separation: Extract high-frequency details from original
            orig_blur = cv2.GaussianBlur(cropped_orig_resized, (5, 5), 0)
            high_freq = cv2.subtract(cropped_orig_resized.astype(np.int16), orig_blur.astype(np.int16))
            
            # Scale high frequency grain
            grain_layer = (high_freq * grain_amount).clip(-128, 127)

            if skin_mask is not None:
                skin_mask_2d = np.squeeze(skin_mask)
                if skin_mask_2d.ndim == 2:
                    if skin_mask_2d.shape != restored_face.shape[:2]:
                        skin_mask_resized = cv2.resize(skin_mask_2d.astype(np.uint8), (restored_face.shape[1], restored_face.shape[0]), interpolation=cv2.INTER_NEAREST)
                    else:
                        skin_mask_resized = skin_mask_2d
                    
                    # Skin category in facexlib parse mask is index 1
                    skin_binary = (skin_mask_resized == 1).astype(np.float32)
                    # Smooth mask edge
                    skin_binary = cv2.GaussianBlur(skin_binary, (5, 5), 0)[:, :, np.newaxis]
                    blended = restored_face.astype(np.float32) + grain_layer * skin_binary
                else:
                    blended = restored_face.astype(np.float32) + grain_layer
            else:
                blended = restored_face.astype(np.float32) + grain_layer

            return np.clip(blended, 0, 255).astype(np.uint8)
        except Exception as e:
            print(f"[WinkEnhancer] Skin grain warning: {e}")
            return restored_face

    def enhance_eyes_and_lips(self, face_img: np.ndarray, parse_mask: np.ndarray = None, enable_eyes: bool = True, enable_lips: bool = True) -> np.ndarray:
        """
        Enhance eyes (catchlight, contrast, sharpness) and lips using facial parsing mask.
        """
        if parse_mask is None:
            # Fallback: General soft unsharp mask on entire face
            blur = cv2.GaussianBlur(face_img, (0, 0), 2.0)
            return cv2.addWeighted(face_img, 1.15, blur, -0.15, 0)

        try:
            parse_mask_2d = np.squeeze(parse_mask)
            if parse_mask_2d.ndim != 2:
                blur = cv2.GaussianBlur(face_img, (0, 0), 2.0)
                return cv2.addWeighted(face_img, 1.15, blur, -0.15, 0)

            h, w = face_img.shape[:2]
            if parse_mask_2d.shape[:2] != (h, w):
                parse_mask_res = cv2.resize(parse_mask_2d.astype(np.uint8), (w, h), interpolation=cv2.INTER_NEAREST)
            else:
                parse_mask_res = parse_mask_2d


            # Facial feature mask IDs in facexlib:
            # 4: Left Eye, 5: Right Eye, 6: Glasses, 11: Upper Lip, 12: Lower Lip, 13: Inner Mouth
            eye_mask = ((parse_mask_res == 4) | (parse_mask_res == 5) | (parse_mask_res == 6)).astype(np.uint8)
            lip_mask = ((parse_mask_res == 11) | (parse_mask_res == 12) | (parse_mask_res == 13)).astype(np.uint8)

            result = face_img.copy()

            # 1. Enhance Eyes: CLAHE on L channel + Unsharp Masking
            if enable_eyes and np.any(eye_mask):
                # Expand eye mask slightly for seamless blending
                kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
                eye_mask_dilated = cv2.dilate(eye_mask, kernel, iterations=1)
                
                # Convert to LAB for luminance contrast
                lab = cv2.cvtColor(result, cv2.COLOR_BGR2LAB)
                l, a, b = cv2.split(lab)
                
                l_eye_clahe = self.clahe_eye.apply(l)
                l_blended = np.where(eye_mask_dilated == 1, l_eye_clahe, l)
                lab_enhanced = cv2.merge([l_blended, a, b])
                result = cv2.cvtColor(lab_enhanced, cv2.COLOR_LAB2BGR)
                
                # Unsharp Mask on Eyes
                eye_blur = cv2.GaussianBlur(result, (3, 3), 0)
                eye_sharp = cv2.addWeighted(result, 1.3, eye_blur, -0.3, 0)
                
                eye_mask_float = cv2.GaussianBlur(eye_mask_dilated.astype(np.float32), (3, 3), 0)[:, :, np.newaxis]
                result = (result * (1.0 - eye_mask_float) + eye_sharp * eye_mask_float).astype(np.uint8)

            # 2. Enhance Lips: Subtle contrast and saturation boost
            if enable_lips and np.any(lip_mask):
                lip_mask_float = cv2.GaussianBlur(lip_mask.astype(np.float32), (3, 3), 0)[:, :, np.newaxis]
                hsv = cv2.cvtColor(result, cv2.COLOR_BGR2HSV).astype(np.float32)
                hsv[:, :, 1] = np.where(lip_mask == 1, np.clip(hsv[:, :, 1] * 1.1, 0, 255), hsv[:, :, 1]) # Boost saturation slightly
                lip_enhanced = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
                result = (result * (1.0 - lip_mask_float) + lip_enhanced * lip_mask_float).astype(np.uint8)

            return result
        except Exception as e:
            print(f"[WinkEnhancer] Eye/Lip enhancement warning: {e}")
            return face_img

    def balance_skin_tone_lab(self, face_img: np.ndarray) -> np.ndarray:
        """
        Apply CLAHE on the L channel of LAB space to balance skin lighting, micro-contrast, and dynamic range.
        """
        try:
            lab = cv2.cvtColor(face_img, cv2.COLOR_BGR2LAB)
            l, a, b = cv2.split(lab)
            l_clahe = self.clahe_lab.apply(l)
            # Soft blend to avoid over-exposure
            l_final = cv2.addWeighted(l, 0.6, l_clahe, 0.4, 0)
            lab_balanced = cv2.merge([l_final, a, b])
            return cv2.cvtColor(lab_balanced, cv2.COLOR_LAB2BGR)
        except Exception as e:
            print(f"[WinkEnhancer] LAB tone balance warning: {e}")
            return face_img

    def match_color_reinhard(self, target_img: np.ndarray, source_img: np.ndarray, blend: float = 0.5) -> np.ndarray:
        """
        Reinhard Color Transfer: Match color statistics (mean and std dev in LAB space)
        of target_img (restored AI face) to source_img (original cropped face/neck).
        """
        if source_img is None or blend <= 0.0:
            return target_img

        try:
            if target_img.shape[:2] != source_img.shape[:2]:
                source_res = cv2.resize(source_img, (target_img.shape[1], target_img.shape[0]), interpolation=cv2.INTER_LANCZOS4)
            else:
                source_res = source_img

            target_lab = cv2.cvtColor(target_img, cv2.COLOR_BGR2LAB).astype(np.float32)
            source_lab = cv2.cvtColor(source_res, cv2.COLOR_BGR2LAB).astype(np.float32)

            t_mean, t_std = cv2.meanStdDev(target_lab)
            s_mean, s_std = cv2.meanStdDev(source_lab)

            t_mean = t_mean.flatten()
            t_std = np.maximum(t_std.flatten(), 1e-5)
            s_mean = s_mean.flatten()
            s_std = s_std.flatten()

            res_lab = np.zeros_like(target_lab)
            for i in range(3):
                res_lab[:, :, i] = ((target_lab[:, :, i] - t_mean[i]) * (s_std[i] / t_std[i])) + s_mean[i]

            res_lab = np.clip(res_lab, 0, 255).astype(np.uint8)
            matched_bgr = cv2.cvtColor(res_lab, cv2.COLOR_LAB2BGR)

            return cv2.addWeighted(target_img, 1.0 - blend, matched_bgr, blend, 0)
        except Exception as e:
            print(f"[WinkEnhancer] Color match warning: {e}")
            return target_img

    def apply_adaptive_sharpening(self, img: np.ndarray, sharpen_amount: float = 0.2) -> np.ndarray:
        """
        Multi-Scale Edge-Aware Sharpening:
        Extracts structural edge mask using Sobel magnitude and applies dual-scale
        Unsharp Masking (fine micro-details + coarse structural edges) without halos.
        """
        if sharpen_amount <= 0.0 or img is None:
            return img

        try:
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            
            # Sobel edge magnitude
            grad_x = cv2.Sobel(gray, cv2.CV_32F, 1, 0, ksize=3)
            grad_y = cv2.Sobel(gray, cv2.CV_32F, 0, 1, ksize=3)
            edge_mag = cv2.magnitude(grad_x, grad_y)
            edge_norm = cv2.normalize(edge_mag, None, 0.0, 1.0, cv2.NORM_MINMAX)[:, :, np.newaxis]
            
            # Dual-scale Unsharp Masking
            blur_fine = cv2.GaussianBlur(img, (3, 3), 1.0)
            blur_coarse = cv2.GaussianBlur(img, (7, 7), 3.0)
            
            sharp_fine = cv2.addWeighted(img, 1.0 + sharpen_amount, blur_fine, -sharpen_amount, 0)
            sharp_coarse = cv2.addWeighted(img, 1.0 + (sharpen_amount * 0.5), blur_coarse, -(sharpen_amount * 0.5), 0)
            
            # Blend sharp layers weighted by edge mask
            out = img.astype(np.float32) * (1.0 - edge_norm) + (sharp_fine.astype(np.float32) * 0.7 + sharp_coarse.astype(np.float32) * 0.3) * edge_norm
            return np.clip(out, 0, 255).astype(np.uint8)
        except Exception as e:
            print(f"[WinkEnhancer] Adaptive sharpening warning: {e}")
            return img

    def enhance_face(self, restored_face: np.ndarray, cropped_original: np.ndarray = None, parse_mask: np.ndarray = None, wink_mode: bool = True, eye_enhancement: bool = True, skin_grain: float = 0.15, color_match: bool = True, enable_eyes: bool = True, enable_lips: bool = True, enable_skin: bool = True, sharpen_amount: float = 0.2) -> np.ndarray:
        """
        Master method to execute Wink-level enhancement pipeline on a restored face crop.
        """
        if not wink_mode:
            return restored_face

        out_face = restored_face.copy()

        # Step A: Reinhard Color Transfer (Auto Skin Tone Alignment to original face/neck)
        if color_match and cropped_original is not None:
            out_face = self.match_color_reinhard(out_face, cropped_original, blend=0.4)

        # Step B: Skin tone & micro-contrast balance
        out_face = self.balance_skin_tone_lab(out_face)

        # Step C: Eye & Lip local enhancement
        if eye_enhancement and (enable_eyes or enable_lips):
            out_face = self.enhance_eyes_and_lips(out_face, parse_mask=parse_mask, enable_eyes=enable_eyes, enable_lips=enable_lips)

        # Step D: Multi-Scale Edge-Aware Adaptive Sharpening
        if sharpen_amount > 0.0:
            out_face = self.apply_adaptive_sharpening(out_face, sharpen_amount=sharpen_amount)

        # Step E: Real Skin Grain Injection (Frequency Separation)
        if enable_skin and skin_grain > 0.0 and cropped_original is not None:
            out_face = self.apply_skin_grain(out_face, cropped_original, skin_mask=parse_mask, grain_amount=skin_grain)

        return out_face


    def calculate_sharpness(self, img: np.ndarray) -> float:
        """Calculate image sharpness using Variance of Laplacian."""
        if img is None:
            return 0.0
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if len(img.shape) == 3 else img
        return float(cv2.Laplacian(gray, cv2.CV_64F).var())

    def calculate_quality_report(self, orig_img: np.ndarray, enhanced_img: np.ndarray, face_count: int = 0) -> dict:
        """
        Generate AI Quality Score & Comparison metrics report.
        """
        orig_sharpness = self.calculate_sharpness(orig_img)
        enh_sharpness = self.calculate_sharpness(enhanced_img)
        
        sharpness_gain_pct = ((enh_sharpness - orig_sharpness) / max(orig_sharpness, 1e-5)) * 100.0
        sharpness_gain_pct = float(np.clip(sharpness_gain_pct, 0.0, 1000.0))

        # Skin tone fidelity score (using LAB luminance correlation)
        try:
            o_res = cv2.resize(orig_img, (enhanced_img.shape[1], enhanced_img.shape[0]))
            o_lab = cv2.cvtColor(o_res, cv2.COLOR_BGR2LAB).astype(np.float32)
            e_lab = cv2.cvtColor(enhanced_img, cv2.COLOR_BGR2LAB).astype(np.float32)
            diff = np.mean(np.abs(o_lab[:, :, 1:] - e_lab[:, :, 1:]))
            tone_fidelity_pct = float(np.clip(100.0 - (diff * 1.5), 70.0, 99.9))
        except Exception:
            tone_fidelity_pct = 95.0

        return {
            'orig_sharpness': round(orig_sharpness, 1),
            'enh_sharpness': round(enh_sharpness, 1),
            'sharpness_gain_pct': round(sharpness_gain_pct, 1),
            'face_count': face_count,
            'tone_fidelity_pct': round(tone_fidelity_pct, 1)
        }