multiple faces detection support

DZFaceDetailer.py

@@ -11,7 +11,7 @@ from folder_paths import base_path
 
 face_model_path = os.path.join(base_path, "models/DZ-FaceDetailer/yolo/face_yolov8n.pt")
 MASK_CONTROL = ["dilate", "erode", "disabled"]
-MASK_TYPE = ["box", "face"]
+MASK_TYPE = ["face", "box"]
 
 class FaceDetailer:
     @classmethod
@@ -52,12 +52,12 @@ class FaceDetailer:
         # Process the face mesh or make the face box for masking
         if mask_type == "box":
             try:
-                final_mask = facebox_mask(img, mask_type)
+                final_mask = facebox_mask(img)
             except:
                 return (latent_image, )
         else:
             try:
-                final_mask = facemesh_mask(img, mask_type)
+                final_mask = facemesh_mask(img)
             except:
                 return (latent_image, )
         # Erode/Dilate mask
@@ -86,67 +86,121 @@ class FaceDetailer:
         return (latent[0], final_mask,)
 
 
-def facebox_mask(image, mask_type):
+def facebox_mask(image):
+    # Create an empty image with alpha
+    mask = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
+
     # setup yolov8n face detection model
     face_model = YOLO(face_model_path)
     face_bbox = face_model(image)
     boxes = face_bbox[0].boxes
-    box = boxes[0].xyxy
+    # box = boxes[0].xyxy
-    x_min, y_min, x_max, y_max = box[0].tolist()
+    for box in boxes.xyxy:
+        x_min, y_min, x_max, y_max = box.tolist()
+        # Calculate the center of the bounding box
+        center_x = (x_min + x_max) / 2
+        center_y = (y_min + y_max) / 2
 
-    # Calculate the center of the bounding box
+        # Calcule the maximum width and height
-    center_x = (x_min + x_max) / 2
+        width = x_max - x_min
-    center_y = (y_min + y_max) / 2
+        height = y_max - y_min
+        max_size = max(width, height)
 
-    # Calcule the maximum width and height
+        # Get the new WxH for a ratio of 1:1
-    width = x_max - x_min
+        new_width = max_size
-    height = y_max - y_min
+        new_height = max_size
-    max_size = max(width, height)
 
-    # Get the new WxH for a ratio of 1:1
+        # Calculate the new coordinates
-    new_width = max_size
+        new_x_min = int(center_x - new_width / 2)
-    new_height = max_size
+        new_y_min = int(center_y - new_height / 2)
+        new_x_max = int(center_x + new_width / 2)
+        new_y_max = int(center_y + new_height / 2)
 
-    # Calculate the new coordinates
+        # print((new_x_min, new_y_min), (new_x_max, new_y_max))
-    new_x_min = int(center_x - new_width / 2)
+        # set the square in the face location
-    new_y_min = int(center_y - new_height / 2)
+        cv2.rectangle(mask, (new_x_min, new_y_min), (new_x_max, new_y_max), (0, 0, 0, 255), -1)
-    new_x_max = int(center_x + new_width / 2)
+        testing = image.copy()
-    new_y_max = int(center_y + new_height / 2)
+        testing = testing[new_y_min:new_y_max, new_x_min:new_x_max, :]
+        cv2.imwrite("xd.jpg", testing)
 
-    # Create an empty image with alpha and set the square in the face location
-    mask = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
-    cv2.rectangle(mask, (new_x_min, new_y_min), (new_x_max, new_y_max), (0, 0, 0, 255), -1)
     mask[:, :, 3] = ~mask[:, :, 3]  # invert the mask
 
     return mask
 
 
-def facemesh_mask(image, mask_type):
+def facemesh_mask(image):
-    mp_face_mesh = solutions.face_mesh
-    face_mesh = mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1)
-    results = face_mesh.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
-    if results.multi_face_landmarks:
-        for face_landmarks in results.multi_face_landmarks:
-            # List of detected face points
-            points = []
-            for landmark in face_landmarks.landmark:
-                cx, cy = int(
-                    landmark.x * image.shape[1]), int(landmark.y * image.shape[0])
-                points.append([cx, cy])
 
-            # Empty image with the same shape as input
+    faces_mask = []
-            mask = np.zeros(
-                (image.shape[0], image.shape[1], 4), dtype=np.uint8)
 
-            # Obtain the countour of the face
+    # Empty image with the same shape as input
-            convex_hull = cv2.convexHull(np.array(points))
+    mask = np.zeros(
+        (image.shape[0], image.shape[1], 4), dtype=np.uint8)
+    
+    # setup yolov8n face detection model
+    face_model = YOLO(face_model_path)
+    face_bbox = face_model(image)
+    boxes = face_bbox[0].boxes
+    # box = boxes[0].xyxy
+    for box in boxes.xyxy:
+        x_min, y_min, x_max, y_max = box.tolist()
+        # Calculate the center of the bounding box
+        center_x = (x_min + x_max) / 2
+        center_y = (y_min + y_max) / 2
 
-            # Fill the contour and store it in alpha for the mask
+        # Calcule the maximum width and height
-            cv2.fillConvexPoly(mask, convex_hull, (0, 0, 0, 255))
+        width = x_max - x_min
-            mask[:, :, 3] = ~mask[:, :, 3]
+        height = y_max - y_min
+        max_size = max(width, height)
 
-            return mask
+        # Get the new WxH for a ratio of 1:1
+        new_width = max_size
+        new_height = max_size
 
+        # Calculate the new coordinates
+        new_x_min = int(center_x - new_width / 2)
+        new_y_min = int(center_y - new_height / 2)
+        new_x_max = int(center_x + new_width / 2)
+        new_y_max = int(center_y + new_height / 2)
+
+        # print((new_x_min, new_y_min), (new_x_max, new_y_max))
+        # set the square in the face location
+        face = image[new_y_min:new_y_max, new_x_min:new_x_max, :]
+
+        mp_face_mesh = solutions.face_mesh
+        face_mesh = mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=1, min_detection_confidence=0.5)
+        results = face_mesh.process(cv2.cvtColor(face, cv2.COLOR_BGR2RGB))
+        if results.multi_face_landmarks:
+            for face_landmarks in results.multi_face_landmarks:
+                # List of detected face points
+                points = []
+                for landmark in face_landmarks.landmark:
+                    cx, cy = int(
+                        landmark.x * face.shape[1]), int(landmark.y * face.shape[0])
+                    points.append([cx, cy])
+
+                face_mask = np.zeros((face.shape[0], face.shape[1], 4), dtype=np.uint8)
+
+                # Obtain the countour of the face
+                convex_hull = cv2.convexHull(np.array(points))
+                
+                # Fill the contour and store it in alpha for the mask
+                cv2.fillConvexPoly(face_mask, convex_hull, (0, 0, 0, 255))
+
+                faces_mask.append([face_mask, [new_x_min, new_x_max, new_y_min, new_y_max]])
+            
+    for face_mask in faces_mask:
+        paste_numpy_images(mask, face_mask[0], face_mask[1][0], face_mask[1][1], face_mask[1][2], face_mask[1][3])
+
+    print(f"{len(faces_mask)} faces detected")
+    mask[:, :, 3] = ~mask[:, :, 3]
+    return mask
+
+
+def paste_numpy_images(target_image, source_image, x_min, x_max, y_min, y_max):
+    # Paste the source image into the target image at the specified coordinates
+    target_image[y_min:y_max, x_min:x_max, :] = source_image
+
+    return target_image
 
 def erode_mask(mask, dilate):
     # I use erode function because the mask is inverted
@@ -185,6 +239,5 @@ def image2nparray(image, BGR):
 
 def set_mask(samples, mask):
     s = samples.copy()
-    print(s)
     s["noise_mask"] = mask.reshape((-1, 1, mask.shape[-2], mask.shape[-1]))
     return s

workflow_example.json

@@ -544,32 +544,6 @@
       "color": "#322",
       "bgcolor": "#533"
     },
-    {
-      "id": 23,
-      "type": "PreviewImage",
-      "pos": [
-        2905.4059446271544,
-        391.4753301214381
-      ],
-      "size": {
-        "0": 300.9404296875,
-        "1": 440.5021057128906
-      },
-      "flags": {},
-      "order": 26,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "images",
-          "type": "IMAGE",
-          "link": 43,
-          "slot_index": 0
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "PreviewImage"
-      }
-    },
     {
       "id": 24,
       "type": "Reroute",
@@ -607,31 +581,6 @@
         "horizontal": false
       }
     },
-    {
-      "id": 14,
-      "type": "PreviewImage",
-      "pos": [
-        3271.3467063344615,
-        384.1296736506407
-      ],
-      "size": {
-        "0": 318.26812744140625,
-        "1": 461.947265625
-      },
-      "flags": {},
-      "order": 20,
-      "mode": 0,
-      "inputs": [
-        {
-          "name": "images",
-          "type": "IMAGE",
-          "link": 47
-        }
-      ],
-      "properties": {
-        "Node name for S&R": "PreviewImage"
-      }
-    },
     {
       "id": 17,
       "type": "DZ_Face_Detailer",
@@ -702,7 +651,7 @@
         "Node name for S&R": "DZ_Face_Detailer"
       },
       "widgets_values": [
-        48098681251663,
+        465341715093210,
         "randomize",
         20,
         7,
@@ -957,7 +906,7 @@
         "Node name for S&R": "KSampler"
       },
       "widgets_values": [
-        662123511475493,
+        888117471012813,
         "randomize",
         15,
         7,
@@ -1145,6 +1094,57 @@
       ],
       "color": "#223",
       "bgcolor": "#335"
+    },
+    {
+      "id": 14,
+      "type": "PreviewImage",
+      "pos": [
+        3255,
+        420
+      ],
+      "size": [
+        331.1569279423511,
+        354.9488064968368
+      ],
+      "flags": {},
+      "order": 20,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 47
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      }
+    },
+    {
+      "id": 23,
+      "type": "PreviewImage",
+      "pos": [
+        2858.480640089445,
+        425.4401372391366
+      ],
+      "size": [
+        325.9292706156648,
+        359.02643171887235
+      ],
+      "flags": {},
+      "order": 26,
+      "mode": 0,
+      "inputs": [
+        {
+          "name": "images",
+          "type": "IMAGE",
+          "link": 43,
+          "slot_index": 0
+        }
+      ],
+      "properties": {
+        "Node name for S&R": "PreviewImage"
+      }
     }
   ],
   "links": [
@@ -1468,10 +1468,10 @@
     {
       "title": "Detailer",
       "bounding": [
-        2887,
+        2816,
-        286,
+        278,
-        337,
+        411,
-        592
+        596
       ],
       "color": "#8A8",
       "font_size": 24