add face shape detection

app.py

@@ -3,12 +3,17 @@ import cvzone
 import numpy as np
 import os
 import gradio as gr
+import dlib
 from datetime import datetime
 
 # Load the YuNet model
 model_path = 'face_detection_yunet_2023mar.onnx'
 face_detector = cv2.FaceDetectorYN.create(model_path, "", (320, 320))
 
+# Load dlib's shape predictor
+shape_predictor_path = 'shape_predictor_68_face_landmarks.dat'
+shape_predictor = dlib.shape_predictor(shape_predictor_path)
+
 # Initialize the glass number
 num = 1
 overlay = cv2.imread(f'glasses/glass{num}.png', cv2.IMREAD_UNCHANGED)
@@ -20,6 +25,19 @@ def count_files_in_directory(directory):
         file_count += len(files)
     return file_count
 
+# Determine face shape
+def determine_face_shape(landmarks):
+    # Example logic to determine face shape based on landmarks
+    # This is a simplified version and may need adjustments
+    jaw_width = np.linalg.norm(landmarks[0] - landmarks[16])
+    face_height = np.linalg.norm(landmarks[8] - landmarks[27])
+    if jaw_width / face_height > 1.5:
+        return "Round"
+    elif jaw_width / face_height < 1.2:
+        return "Oval"
+    else:
+        return "Square"
+
 directory_path = 'glasses'
 total_glass_num = count_files_in_directory(directory_path)
 
@@ -32,62 +50,59 @@ def change_glasses():
     overlay = cv2.imread(f'glasses/glass{num}.png', cv2.IMREAD_UNCHANGED)
     return overlay
 
-# Process frame for overlay
+# Process frame for overlay and face shape detection
 def process_frame(frame):
     global overlay
-    # Ensure the frame is writable
     frame = np.array(frame, copy=True)
-
     height, width = frame.shape[:2]
     face_detector.setInputSize((width, height))
     _, faces = face_detector.detect(frame)
 
+    face_shape = "Unknown"
     if faces is not None:
         for face in faces:
             x, y, w, h = face[:4].astype(int)
-            face_landmarks = face[4:14].reshape(5, 2).astype(int)  # Facial landmarks
+            face_landmarks = face[4:14].reshape(5, 2).astype(int)
+
+            # Convert to dlib rectangle
+            dlib_rect = dlib.rectangle(x, y, x + w, y + h)
+            landmarks = shape_predictor(frame, dlib_rect)
+            landmarks = np.array([(p.x, p.y) for p in landmarks.parts()])
+
+            # Determine face shape
+            face_shape = determine_face_shape(landmarks)
 
             # Get the nose position
             nose_x, nose_y = face_landmarks[2].astype(int)
-            # Left and right eye positions
             left_eye_x, left_eye_y = face_landmarks[0].astype(int)
             right_eye_x, right_eye_y = face_landmarks[1].astype(int)
 
-            # Calculate the midpoint between the eyes
             eye_center_x = (left_eye_x + right_eye_x) // 2
             eye_center_y = (left_eye_y + right_eye_y) // 2
 
-            # Calculate the angle of rotation
             delta_x = right_eye_x - left_eye_x
             delta_y = right_eye_y - left_eye_y
             angle = np.degrees(np.arctan2(delta_y, delta_x))
-
-            # Negate the angle to rotate in the opposite direction
             angle = -angle
 
-            # Resize the overlay
             overlay_resize = cv2.resize(overlay, (int(w * 1.15), int(h * 0.8)))
-
-            # Rotate the overlay
             overlay_center = (overlay_resize.shape[1] // 2, overlay_resize.shape[0] // 2)
             rotation_matrix = cv2.getRotationMatrix2D(overlay_center, angle, 1.0)
             overlay_rotated = cv2.warpAffine(
-                overlay_resize, rotation_matrix, 
+                overlay_resize, rotation_matrix,
                 (overlay_resize.shape[1], overlay_resize.shape[0]),
                 flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0)
             )
 
-            # Calculate the position to center the glasses on the eyes
             overlay_x = eye_center_x - overlay_rotated.shape[1] // 2
             overlay_y = eye_center_y - overlay_rotated.shape[0] // 2
 
-            # Overlay the glasses
             try:
                 frame = cvzone.overlayPNG(frame, overlay_rotated, [overlay_x, overlay_y])
             except Exception as e:
                 print(f"Error overlaying glasses: {e}")
-                
-    return frame
+
+    return frame, face_shape
 
 # Transform function
 def transform_cv2(frame, transform):
@@ -148,9 +163,9 @@ def save_frame(frame):
 
 # Gradio webcam input
 def webcam_input(frame, transform):
-    frame = process_frame(frame)
+    frame, face_shape = process_frame(frame)
     frame = transform_cv2(frame, transform)
-    return frame
+    return frame, face_shape
 
 # Gradio Interface
 with gr.Blocks() as demo:
@@ -159,13 +174,16 @@ with gr.Blocks() as demo:
             transform = gr.Dropdown(choices=["cartoon", "edges", "none"],
                                     value="none", label="Transformation")
             input_img = gr.Image(sources=["webcam"], type="numpy", streaming=True)
+            face_shape_output = gr.Textbox(label="Detected Face Shape")
             next_button = gr.Button("Next Glasses")
-            save_button = gr.Button("Save as a Picture") 
-    input_img.stream(webcam_input, [input_img, transform], [input_img], time_limit=30, stream_every=0.1)
+            save_button = gr.Button("Save as a Picture")
+
+    input_img.stream(webcam_input, [input_img, transform], [input_img, face_shape_output], time_limit=30, stream_every=0.1)
     with gr.Row():
         next_button.click(change_glasses, [], [])
     with gr.Row():
         save_button.click(save_frame, [input_img], [])
+
 if __name__ == "__main__":
     demo.launch(share=True)
 

requirements.txt

@@ -1,4 +1,5 @@
 gradio
 cvzone
 opencv-python
-numpy
+numpy
+dlib