From 605b429760a883bc387bcd6c94e2cf028caed22e Mon Sep 17 00:00:00 2001
From: David Chatel <dchatel@live.fr>
Date: Wed, 13 Mar 2024 12:20:51 +0100
Subject: [PATCH] fix bug with new mask

---
 nodes.py | 18 +++++++++---------
 utils.py |  3 ++-
 2 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/nodes.py b/nodes.py
index a598dde..001d468 100644
--- a/nodes.py
+++ b/nodes.py
@@ -166,8 +166,8 @@ class CropFaces:
         masks = []
         warps = []
         for face in faces:
-            M, crop = face.crop(crop_size, crop_factor)
-            mask = mask_crop(face, M, crop*255, mask_type)
+            M, crop, maskedcrop = face.crop(crop_size, crop_factor)
+            mask = mask_crop(face, M, maskedcrop, mask_type)
             crops.append(np.array(crop))
             masks.append(np.array(mask))
             warps.append(M)
@@ -228,13 +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)
+            # 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 758f2bc..adb186e 100644
--- a/utils.py
+++ b/utils.py
@@ -94,8 +94,9 @@ class Face:
         T4 = np.array([[1, 0, cx], [0, 1, cy], [0, 0, 1]])
         N = N @ T4 @ S @ T3
         crop = cv2.warpAffine(self.image.numpy(), N, (size, size))
+        maskedcrop = cv2.warpAffine(self.img.numpy(), N, (size, size))
         
-        return N, crop
+        return N, crop, maskedcrop
 
 def detect_faces(img, threshold):
     img = pad_to_stride(img, stride=32)