updated to add the race, gender and image segmentation

app.py

@@ -62,7 +62,7 @@ def classify_hair_color(image):
     # Run the pipeline on the uploaded image
     output = pipe(image)
     
-    print("Pipeline output for hir color:", output)
+    print("Pipeline output for hair color:", output)
     # Format the output to be compatible with gr.outputs.Label
     formatted_output = {item['label']: item['score'] for item in output}
     
@@ -96,13 +96,39 @@ def classify_eye_color(image):
     
     return formatted_output
     
-
 def process_gradio_image(pil_image):
     # Convert PIL image to NumPy array
     image = np.array(pil_image)
     image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  # Convert RGB (from PIL) to BGR (OpenCV default)
     return image
+
+def classify_race(image):
+    '''
+    "0": "East Asian",
+    "1": "Indian",
+    "2": "Black",
+    "3": "White",
+    "4": "Middle Eastern",
+    "5": "Latino_Hispanic",
+    "6": "Southeast Asian"
+    '''
+    pipe = pipeline("image-classification", model="crangana/trained-race")
+    # Run the pipeline on the uploaded image
+    output = pipe("face_region.jpg")
+
+    # Format the output to be compatible with gr.outputs.Label
+    formatted_output = {item['label']: item['score'] for item in output}
     
+    return formatted_output
+
+def classify_gender(image):
+    pipe = pipeline("image-classification", model="rizvandwiki/gender-classification")
+    output = pipe("face_region.jpg")
+    # Format the output to be compatible with gr.outputs.Label
+    formatted_output = {item['label']: item['score'] for item in output}
+    return formatted_output
+
+
 def classify_image_with_multiple_models(image):
     create_eye_region(image)
     face_shape_result = classify_face_shape(image)
@@ -112,9 +138,10 @@ def classify_image_with_multiple_models(image):
     hair_color_results = classify_hair_color(image)
     eye_shape = classify_eye_shape(image)
     eye_color = classify_eye_color(image)
+    race = classify_race(image)
+    gender = classify_gender(image)
 
-    return face_shape_result, age_result, skin_type_result, acne_results, hair_color_results, eye_shape, eye_color
-
+    return face_shape_result, age_result, skin_type_result, acne_results, hair_color_results, eye_shape, eye_color, race, gender, Image.open("segmented_face.jpg")
 
 def create_eye_region(image):
     # Load the pre-trained face detector
@@ -132,12 +159,22 @@ def create_eye_region(image):
         # Draw a rectangle around the face
         cv2.rectangle(image, (x, y), (x + w, y + h), (255, 0, 0), 2)
 
+        # Extract the face region
+        face_roi = image[y:y + h, x:x + w]
+        cv2.imwrite('face_region.jpg', face_roi)
+
         # Region of Interest (ROI) for the face
         roi_gray = gray[y:y + h, x:x + w]
         roi_color = image[y:y + h, x:x + w]
 
         # Detect eyes in the face ROI
-        eyes = eye_cascade.detectMultiScale(roi_gray)
+        eyes = eye_cascade.detectMultiScale(roi_gray, scaleFactor=1.1, minNeighbors=10, minSize=(20, 20))
+
+        eye_positions = []
+        for (ex, ey, ew, eh) in eyes:
+            # Ensure eyes are within the upper half of the face region
+            if ey + eh < h // 2:
+                eye_positions.append((ex, ey, ew, eh))
 
         for (ex, ey, ew, eh) in eyes:
             # Draw a rectangle around the eyes
@@ -151,15 +188,15 @@ def create_eye_region(image):
             avg_color = np.mean(eye_roi, axis=(0, 1))
 
             # Classify eye color based on average color
-            if avg_color[0] > avg_color[1] and avg_color[0] > avg_color[2]:
-                color = "Brown"
-            elif avg_color[1] > avg_color[0] and avg_color[1] > avg_color[2]:
-                color = "Green"
-            else:
-                color = "Blue"
+            #if avg_color[0] > avg_color[1] and avg_color[0] > avg_color[2]:
+            #    color = "Brown"
+            #elif avg_color[1] > avg_color[0] and avg_color[1] > avg_color[2]:
+            #    color = "Green"
+            #else:
+            #    color = "Blue"
 
             # Display the eye color
-            cv2.putText(image, color, (ex, ey - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
+            #cv2.putText(image, color, (ex, ey - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
 
             cv2.imwrite('segmented_face.jpg', image)
 
@@ -176,9 +213,12 @@ demo = gr.Interface(
         gr.Label(num_top_classes=5, label="Hair Color"), 
         gr.Label(num_top_classes=4, label="Eye Shape"),
         gr.Label(num_top_classes=5, label="Eye Color"),
+        gr.Label(num_top_classes=7, label="Race"),
+        gr.Label(num_top_classes=2, label="Gender"),
+        gr.Image(type="pil", label="Segmented Face", value="segmented_face.jpg")  # Provide the path to the image
     ],
     title="Multiple Model Classification",
-    description="Upload an image to classify the face using mutiple classification models"
+    description="Upload an image to classify the face using multiple classification models"
 )
     
 #demo.launch(auth=("admin", "pass1234"))