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) }