""" Utility functions for Pixagram AI Pixel Art Generator """ import numpy as np import cv2 import math from PIL import Image, ImageEnhance, ImageFilter, ImageDraw from config import COLOR_MATCH_CONFIG, FACE_MASK_CONFIG, AGE_BRACKETS def sanitize_text(text): """ Remove or replace problematic characters (emojis, special unicode) that might cause encoding errors. """ if not text: return text try: # Encode/decode to remove invalid bytes text = text.encode('utf-8', errors='ignore').decode('utf-8') # Keep only characters within safe unicode range text = ''.join(char for char in text if ord(char) < 65536) except Exception as e: print(f"[WARNING] Text sanitization warning: {e}") return text def color_match_lab(target, source, preserve_saturation=True): """ LAB color space matching for better skin tones with saturation preservation. GENTLE version to prevent color fading. Args: target: Target image to adjust source: Source image to match colors from preserve_saturation: If True, preserves original saturation levels """ try: target_lab = cv2.cvtColor(target.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32) source_lab = cv2.cvtColor(source.astype(np.uint8), cv2.COLOR_RGB2LAB).astype(np.float32) result_lab = np.copy(target_lab) # Very gentle L channel matching t_mean, t_std = target_lab[:,:,0].mean(), target_lab[:,:,0].std() s_mean, s_std = source_lab[:,:,0].mean(), source_lab[:,:,0].std() if t_std > 1e-6: matched = (target_lab[:,:,0] - t_mean) * (s_std / t_std) * 0.5 + s_mean result_lab[:,:,0] = target_lab[:,:,0] * (1 - COLOR_MATCH_CONFIG['lab_lightness_blend']) + matched * COLOR_MATCH_CONFIG['lab_lightness_blend'] if preserve_saturation: # Minimal adjustment to A and B channels for i in [1, 2]: t_mean, t_std = target_lab[:,:,i].mean(), target_lab[:,:,i].std() s_mean, s_std = source_lab[:,:,i].mean(), source_lab[:,:,i].std() if t_std > 1e-6: matched = (target_lab[:,:,i] - t_mean) * (s_std / t_std) + s_mean blend_factor = COLOR_MATCH_CONFIG['lab_color_blend_preserved'] result_lab[:,:,i] = target_lab[:,:,i] * (1 - blend_factor) + matched * blend_factor else: # Gentle full matching for i in [1, 2]: t_mean, t_std = target_lab[:,:,i].mean(), target_lab[:,:,i].std() s_mean, s_std = source_lab[:,:,i].mean(), source_lab[:,:,i].std() if t_std > 1e-6: matched = (target_lab[:,:,i] - t_mean) * (s_std / t_std) + s_mean blend_factor = COLOR_MATCH_CONFIG['lab_color_blend_full'] result_lab[:,:,i] = target_lab[:,:,i] * (1 - blend_factor) + matched * blend_factor return cv2.cvtColor(result_lab.astype(np.uint8), cv2.COLOR_LAB2RGB) except Exception as e: print(f"LAB conversion error: {e}") return target.astype(np.uint8) def enhance_saturation(image, boost=1.05): """ Minimal saturation enhancement (disabled by default). Args: image: PIL Image boost: Saturation multiplier (1.0 = no change, >1.0 = more saturated) """ if boost <= 1.0: return image enhancer = ImageEnhance.Color(image) return enhancer.enhance(boost) def enhanced_color_match(target_img, source_img, face_bbox=None, preserve_vibrance=False): """ Enhanced color matching with face-aware processing. Very gentle to prevent color fading. Args: target_img: Generated image to adjust source_img: Original image to match colors from face_bbox: Optional [x1, y1, x2, y2] for face region preserve_vibrance: If True, adds minimal saturation boost (disabled by default) """ try: target = np.array(target_img).astype(np.float32) source = np.array(source_img).astype(np.float32) if face_bbox is not None: # Create face mask x1, y1, x2, y2 = [int(c) for c in face_bbox] x1, y1 = max(0, x1), max(0, y1) x2, y2 = min(target.shape[1], x2), min(target.shape[0], y2) face_mask = np.zeros((target.shape[0], target.shape[1]), dtype=np.float32) face_mask[y1:y2, x1:x2] = 1.0 # Blur mask for smooth transition face_mask = cv2.GaussianBlur( face_mask, COLOR_MATCH_CONFIG['gaussian_blur_kernel'], COLOR_MATCH_CONFIG['gaussian_blur_sigma'] ) face_mask = face_mask[:, :, np.newaxis] # Match colors for face region with saturation preservation if y2 > y1 and x2 > x1: face_result = color_match_lab( target[y1:y2, x1:x2], source[y1:y2, x1:x2], preserve_saturation=True ) target[y1:y2, x1:x2] = face_result # Blend with original using mask result = target * face_mask + target * (1 - face_mask) else: result = color_match_lab(target, source, preserve_saturation=True) else: # Standard LAB color matching with saturation preservation result = color_match_lab(target, source, preserve_saturation=True) result_img = Image.fromarray(result.astype(np.uint8)) # NO saturation boost by default if preserve_vibrance: result_img = enhance_saturation(result_img, boost=COLOR_MATCH_CONFIG['saturation_boost']) return result_img except Exception as e: print(f"Enhanced color matching failed: {e}, returning target image") return target_img def color_match(target_img, source_img, mode='mkl'): """ Legacy color matching function - kept for compatibility. Use enhanced_color_match for better results. """ try: target = np.array(target_img).astype(np.float32) source = np.array(source_img).astype(np.float32) if mode == 'simple': result = np.zeros_like(target) for i in range(3): t_mean, t_std = target[:,:,i].mean(), target[:,:,i].std() s_mean, s_std = source[:,:,i].mean(), source[:,:,i].std() result[:,:,i] = (target[:,:,i] - t_mean) * (s_std / (t_std + 1e-6)) + s_mean result[:,:,i] = np.clip(result[:,:,i], 0, 255) elif mode == 'mkl': result = color_match_lab(target, source) else: # pdf mode result = np.zeros_like(target) for i in range(3): result[:,:,i] = np.interp( target[:,:,i].flatten(), np.linspace(target[:,:,i].min(), target[:,:,i].max(), 256), np.linspace(source[:,:,i].min(), source[:,:,i].max(), 256) ).reshape(target[:,:,i].shape) return Image.fromarray(result.astype(np.uint8)) except Exception as e: print(f"Color matching failed: {e}, returning target image") return target_img def create_face_mask(image, face_bbox, feather=None): """ Create a soft mask around the detected face for better blending. Args: image: PIL Image face_bbox: [x1, y1, x2, y2] feather: blur radius for soft edges (uses config default if None) """ if feather is None: feather = FACE_MASK_CONFIG['feather'] mask = Image.new('L', image.size, 0) draw = ImageDraw.Draw(mask) # Expand bbox slightly x1, y1, x2, y2 = face_bbox padding = int((x2 - x1) * FACE_MASK_CONFIG['padding']) x1 = max(0, x1 - padding) y1 = max(0, y1 - padding) x2 = min(image.width, x2 + padding) y2 = min(image.height, y2 + padding) # Draw ellipse for more natural face shape draw.ellipse([x1, y1, x2, y2], fill=255) # Apply gaussian blur for soft edges mask = mask.filter(ImageFilter.GaussianBlur(feather)) return mask def draw_kps(image_pil, kps, color_list=[(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255)]): """Draw facial keypoints on image for InstantID ControlNet""" stickwidth = 4 limbSeq = np.array([[0, 2], [1, 2], [3, 2], [4, 2]]) kps = np.array(kps) w, h = image_pil.size out_img = np.zeros([h, w, 3]) for i in range(len(limbSeq)): index = limbSeq[i] color = color_list[index[0]] x = kps[index][:, 0] y = kps[index][:, 1] length = ((x[0] - x[1]) ** 2 + (y[0] - y[1]) ** 2) ** 0.5 angle = math.degrees(math.atan2(y[0] - y[1], x[0] - x[1])) polygon = cv2.ellipse2Poly( (int(np.mean(x)), int(np.mean(y))), (int(length / 2), stickwidth), int(angle), 0, 360, 1 ) out_img = cv2.fillConvexPoly(out_img.copy(), polygon, color) out_img = (out_img * 0.6).astype(np.uint8) for idx_kp, kp in enumerate(kps): color = color_list[idx_kp] x, y = kp out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1) out_img_pil = Image.fromarray(out_img.astype(np.uint8)) return out_img_pil def get_facial_attributes(face): """ Extract comprehensive facial attributes. Returns dict with age, gender, expression, quality metrics. """ attributes = { 'age': None, 'gender': None, 'expression': None, 'quality': 1.0, 'pose_angle': 0, 'description': [] } # Age extraction try: if hasattr(face, 'age'): age = int(face.age) attributes['age'] = age for min_age, max_age, label in AGE_BRACKETS: if min_age <= age < max_age: attributes['description'].append(label) break except (ValueError, TypeError, AttributeError) as e: print(f"[WARNING] Age extraction failed: {e}") # Gender extraction try: if hasattr(face, 'gender'): gender_code = int(face.gender) attributes['gender'] = gender_code if gender_code == 1: attributes['description'].append("male") elif gender_code == 0: attributes['description'].append("female") except (ValueError, TypeError, AttributeError) as e: print(f"[WARNING] Gender extraction failed: {e}") # Expression/emotion detection (if available) try: if hasattr(face, 'emotion'): # Some InsightFace models provide emotion emotion = face.emotion if isinstance(emotion, (list, tuple)) and len(emotion) > 0: emotions = ['neutral', 'happiness', 'surprise', 'sadness', 'anger', 'disgust', 'fear'] emotion_idx = int(np.argmax(emotion)) emotion_name = emotions[emotion_idx] if emotion_idx < len(emotions) else 'neutral' confidence = float(emotion[emotion_idx]) if confidence > 0.4: # Only add if confident expression_desc = None if emotion_name == 'happiness': expression_desc = 'smiling' elif emotion_name == 'surprise': expression_desc = 'surprised expression' elif emotion_name == 'sadness': expression_desc = 'sad expression' elif emotion_name == 'anger': expression_desc = 'angry expression' elif emotion_name == 'neutral': expression_desc = 'neutral expression' # Add other emotions like 'disgust' or 'fear' if desired if expression_desc: attributes['expression'] = expression_desc # Only add non-neutral expressions to the prompt description if emotion_name != 'neutral': if expression_desc not in attributes['description']: attributes['description'].append(expression_desc) except (ValueError, TypeError, AttributeError, IndexError) as e: # Expression not available in this model pass # Pose angle (profile detection) try: if hasattr(face, 'pose'): pose = face.pose if len(pose) > 1: yaw = float(pose[1]) attributes['pose_angle'] = abs(yaw) except (ValueError, TypeError, AttributeError, IndexError): pass # Detection quality try: if hasattr(face, 'det_score'): attributes['quality'] = float(face.det_score) except (ValueError, TypeError, AttributeError): pass return attributes def build_enhanced_prompt(base_prompt, facial_attributes, trigger_word): """ Build enhanced prompt with facial attributes intelligently integrated. """ prompt = base_prompt descriptions = facial_attributes['description'] if not descriptions: return base_prompt # Check if demographics already in prompt prompt_lower = prompt.lower() has_demographics = any(desc.lower() in prompt_lower for desc in descriptions) if not has_demographics: # Insert after trigger word for better integration demographic_str = ", ".join(descriptions) + " person" prompt = prompt.replace( trigger_word, f"{trigger_word}, {demographic_str}", 1 ) age = facial_attributes.get('age') quality = facial_attributes.get('quality') expression = facial_attributes.get('expression') print(f"[FACE] Detected: {', '.join(descriptions)}") print(f" Age: {age if age else 'N/A'}, Quality: {quality:.2f}") if expression: print(f" Expression: {expression}") return prompt def get_demographic_description(age, gender_code): """ Legacy function - kept for compatibility. Use get_facial_attributes() for new code. """ demo_desc = [] if age is not None: try: age_int = int(age) for min_age, max_age, label in AGE_BRACKETS: if min_age <= age_int < max_age: demo_desc.append(label) break except (ValueError, TypeError): pass if gender_code is not None: try: if int(gender_code) == 1: demo_desc.append("male") elif int(gender_code) == 0: demo_desc.append("female") except (ValueError, TypeError): pass return demo_desc def calculate_optimal_size(original_width, original_height, recommended_sizes=None, max_dimension=1536): """ Calculate optimal size maintaining aspect ratio with dimensions as multiples of 64. This updated version supports ANY aspect ratio (not just predefined ones), while ensuring dimensions are multiples of 64 and keeping total pixels reasonable. Args: original_width: Original image width original_height: Original image height recommended_sizes: Optional list of (width, height) tuples (legacy support) max_dimension: Maximum allowed dimension (default 1536) Returns: Tuple of (optimal_width, optimal_height) as multiples of 64 """ aspect_ratio = original_width / original_height # Legacy mode: use recommended sizes if provided if recommended_sizes is not None: best_match = None best_diff = float('inf') for width, height in recommended_sizes: rec_aspect = width / height diff = abs(rec_aspect - aspect_ratio) if diff < best_diff: best_diff = diff best_match = (width, height) # Ensure dimensions are multiples of 64 width, height = best_match width = int((width // 64) * 64) height = int((height // 64) * 64) return width, height # NEW: Support any aspect ratio # Strategy: Keep aspect ratio, scale to reasonable total pixels, round to multiples of 64 # Target total pixels (around 1 megapixel for SDXL, adjustable) target_pixels = 1024 * 1024 # ~1MP, good balance for SDXL # Calculate dimensions that maintain aspect ratio and hit target pixels # width * height = target_pixels # width / height = aspect_ratio # => width = aspect_ratio * height # => aspect_ratio * height^2 = target_pixels # => height = sqrt(target_pixels / aspect_ratio) optimal_height = math.sqrt(target_pixels / aspect_ratio) optimal_width = optimal_height * aspect_ratio # Ensure we don't exceed max_dimension if optimal_width > max_dimension: optimal_width = max_dimension optimal_height = optimal_width / aspect_ratio if optimal_height > max_dimension: optimal_height = max_dimension optimal_width = optimal_height * aspect_ratio # Round to nearest multiple of 64 width = int(round(optimal_width / 64) * 64) height = int(round(optimal_height / 64) * 64) # Ensure minimum size (at least 512 on shortest side) min_dimension = 512 if min(width, height) < min_dimension: if width < height: width = min_dimension height = int(round((width / aspect_ratio) / 64) * 64) else: height = min_dimension width = int(round((height * aspect_ratio) / 64) * 64) # Final safety check: ensure multiples of 64 width = max(64, int((width // 64) * 64)) height = max(64, int((height // 64) * 64)) print(f"[SIZING] Aspect ratio: {aspect_ratio:.3f}, Output: {width}x{height} ({width*height/1e6:.2f}MP)") return width, height def enhance_face_crop(face_crop): """ Multi-stage enhancement for better feature preservation. Args: face_crop: PIL Image of face region Returns: Enhanced PIL Image """ # Stage 1: Resize to optimal size for CLIP (224x224) face_crop_resized = face_crop.resize((224, 224), Image.LANCZOS) # Stage 2: Enhance sharpness (helps with facial features) enhancer = ImageEnhance.Sharpness(face_crop_resized) face_crop_sharp = enhancer.enhance(1.5) # Stage 3: Enhance contrast slightly (helps with lighting) enhancer = ImageEnhance.Contrast(face_crop_sharp) face_crop_enhanced = enhancer.enhance(1.1) # Stage 4: Slight brightness adjustment to normalize lighting enhancer = ImageEnhance.Brightness(face_crop_enhanced) face_crop_final = enhancer.enhance(1.05) return face_crop_final print("[OK] Utilities loaded")