Update utils.py

utils.py

@@ -9,40 +9,18 @@ from PIL import Image, ImageEnhance, ImageFilter, ImageDraw
 from config import COLOR_MATCH_CONFIG, FACE_MASK_CONFIG, AGE_BRACKETS
 
 
-# ============================================
-# NEW: Type Safety Helpers
-# ============================================
-
 def ensure_int(value):
-    """
-    Convert numpy.int64 or similar to Python int.
-    Prevents tensor construction errors with PIL dimensions.
-    """
+    """Convert numpy.int64 or similar to Python int"""
     if isinstance(value, (int, float)):
         return int(value)
-    if hasattr(value, 'item'):
-        return int(value.item())
-    return int(value)
+    return int(value.item()) if hasattr(value, 'item') else int(value)
 
 
 def safe_image_size(image):
-    """
-    Get image size as pure Python ints (not numpy.int64).
-    Prevents errors when using PIL dimensions in tensor operations.
-    
-    Args:
-        image: PIL Image
-    
-    Returns:
-        Tuple of (width, height) as Python ints
-    """
+    """Get image size as pure Python ints to avoid numpy.int64 issues"""
     return (ensure_int(image.width), ensure_int(image.height))
 
 
-# ============================================
-# Original Utility Functions
-# ============================================
-
 def sanitize_text(text):
     """
     Remove or replace problematic characters (emojis, special unicode) 
@@ -253,29 +231,15 @@ def create_face_mask(image, face_bbox, feather=None):
     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.
-    
-    Args:
-        image_pil: PIL Image
-        kps: Keypoints array from InsightFace
-        color_list: List of colors for different keypoints
-    
-    Returns:
-        PIL Image with keypoints drawn
-    """
-    import cv2
-    import numpy as np
-    from PIL import Image
-    
+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)
 
-    # Convert PIL to OpenCV
-    out_img = np.array(image_pil)
-    
+    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]]
@@ -284,9 +248,10 @@ def draw_kps(image_pil, kps, color_list=[(255,0,0), (0,255,0), (0,0,255), (255,2
         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)
+        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):
@@ -294,56 +259,49 @@ def draw_kps(image_pil, kps, color_list=[(255,0,0), (0,255,0), (0,0,255), (255,2
         x, y = kp
         out_img = cv2.circle(out_img.copy(), (int(x), int(y)), 10, color, -1)
 
-    # Convert back to PIL
-    return Image.fromarray(out_img.astype(np.uint8))
+    out_img_pil = Image.fromarray(out_img.astype(np.uint8))
+    return out_img_pil
 
 
 def get_facial_attributes(face):
     """
-    Extract facial attributes from InsightFace detection.
-    
-    Args:
-        face: InsightFace face detection object
-    
-    Returns:
-        Dictionary of facial attributes
+    Extract comprehensive facial attributes.
+    Returns dict with age, gender, expression, quality metrics.
     """
     attributes = {
-        'description': [],
         'age': None,
         'gender': None,
         'expression': None,
-        'pose_angle': None,
-        'quality': None
+        'quality': 1.0,
+        'pose_angle': 0,
+        'description': []
     }
     
-    # Age
+    # Age extraction
     try:
         if hasattr(face, 'age'):
             age = int(face.age)
             attributes['age'] = age
-            
-            # Age bracket
             for min_age, max_age, label in AGE_BRACKETS:
                 if min_age <= age < max_age:
                     attributes['description'].append(label)
                     break
-    except (ValueError, TypeError, AttributeError):
-        pass
+    except (ValueError, TypeError, AttributeError) as e:
+        print(f"[WARNING] Age extraction failed: {e}")
     
-    # Gender
+    # Gender extraction
     try:
         if hasattr(face, 'gender'):
             gender_code = int(face.gender)
             attributes['gender'] = gender_code
             if gender_code == 1:
-                attributes['description'].append('male')
+                attributes['description'].append("male")
             elif gender_code == 0:
-                attributes['description'].append('female')
-    except (ValueError, TypeError, AttributeError):
-        pass
+                attributes['description'].append("female")
+    except (ValueError, TypeError, AttributeError) as e:
+        print(f"[WARNING] Gender extraction failed: {e}")
     
-    # Expression
+    # Expression/emotion detection (if available)
     try:
         if hasattr(face, 'emotion'):
             # Some InsightFace models provide emotion
@@ -464,7 +422,7 @@ def calculate_optimal_size(original_width, original_height, recommended_sizes=No
     Returns:
         Tuple of (optimal_width, optimal_height) as multiples of 64
     """
-    # Ensure pure Python ints
+    # Ensure pure Python ints to avoid numpy.int64 issues
     original_width = ensure_int(original_width)
     original_height = ensure_int(original_height)
     
@@ -565,4 +523,4 @@ def enhance_face_crop(face_crop):
     return face_crop_final
 
 
-print("[OK] Utilities loaded with type safety helpers")
+print("[OK] Utilities loaded")