From 49341156b6c503b817653b78c01cff33df02e483 Mon Sep 17 00:00:00 2001
From: David Chatel <dchatel@live.fr>
Date: Wed, 13 Mar 2024 12:02:48 +0100
Subject: [PATCH] added optional mask to detectfaces experimental:
 automatically adds to the mask black pixels adjacent to the edges of crops

---
 nodes.py | 25 +++++++++++++++++++------
 utils.py |  2 +-
 2 files changed, 20 insertions(+), 7 deletions(-)

diff --git a/nodes.py b/nodes.py
index 472a46a..a598dde 100644
--- a/nodes.py
+++ b/nodes.py
@@ -106,6 +106,9 @@ class DetectFaces:
                 'threshold': ('FLOAT', {'default': 0.5, 'min': 0.0, 'max': 1.0, 'step': 0.01}),
                 'min_size': ('INT', {'default': 64, 'max': 512, 'step': 8}),
                 'max_size': ('INT', {'default': 512, 'min': 512, 'step': 8}),
+            },
+            'optional': {
+                'mask': ('MASK',),
             }
         }
     
@@ -114,10 +117,13 @@ class DetectFaces:
     FUNCTION = 'run'
     CATEGORY = 'facetools'
 
-    def run(self, image, threshold, min_size, max_size):
+    def run(self, image, threshold, min_size, max_size, mask=None):
         faces = []
-        images = (image * 255).type(torch.uint8)
-        for i, img in enumerate(images):
+        masked = image
+        if mask is not None:
+            masked = image * tv.transforms.functional.resize(1-mask, image.shape[1:3])[..., None]
+        masked = (masked * 255).type(torch.uint8)
+        for i, img in enumerate(masked):
             unfiltered_faces = detect_faces(img, threshold)
             for face in unfiltered_faces:
                 a, b, c, d = face.bbox
@@ -125,7 +131,7 @@ class DetectFaces:
                 w = abs(c-a)
                 if (h <= max_size or w <= max_size) and (min_size <= h or min_size <= w):
                     face.image_idx = i
-                    face.image = img
+                    face.image = image[i]
                     faces.append(face)
         return (faces,)
 
@@ -161,8 +167,8 @@ class CropFaces:
         warps = []
         for face in faces:
             M, crop = face.crop(crop_size, crop_factor)
-            mask = mask_crop(face, M, crop, mask_type)
-            crops.append(np.array(crop) / 255)
+            mask = mask_crop(face, M, crop*255, mask_type)
+            crops.append(np.array(crop))
             masks.append(np.array(mask))
             warps.append(M)
         crops = torch.from_numpy(np.array(crops)).type(torch.float32)
@@ -222,6 +228,13 @@ class WarpFaceBack:
     def run(self, images, face, crop, mask, warp):
         groups = defaultdict(list)
         for f,c,m,w in zip(face, crop, mask, warp):
+            gray = cv2.cvtColor(c.numpy(), cv2.COLOR_RGB2GRAY)
+            _, gray = cv2.threshold(gray, 0.01, 1, cv2.THRESH_BINARY)
+            gray = gray.astype(np.uint8)
+            cnts, _ = cv2.findContours(gray, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+            largeCnts = list(filter(lambda x: cv2.contourArea(x) > 10000, cnts))
+            gray = cv2.drawContours(gray, largeCnts, -1, 1, -1)
+            m *= torch.from_numpy(gray)
             groups[f.image_idx].append((f.image,c,m,w))
 
         results = []
diff --git a/utils.py b/utils.py
index 79f1ed8..758f2bc 100644
--- a/utils.py
+++ b/utils.py
@@ -93,7 +93,7 @@ class Face:
         T3 = np.array([[1, 0, -cx], [0, 1, -cy], [0, 0, 1]])
         T4 = np.array([[1, 0, cx], [0, 1, cy], [0, 0, 1]])
         N = N @ T4 @ S @ T3
-        crop = cv2.warpAffine(self.img.numpy(), N, (size, size))
+        crop = cv2.warpAffine(self.image.numpy(), N, (size, size))
         
         return N, crop