Updated GenderFaceFilter node

Author: dchatel

README.md

@@ -5,6 +5,7 @@ These custom nodes provide a rotation aware face extraction, paste back, and var
 ![Comparison](examples/comparison.jpg)
 
 ## Patch notes
+- 2024-05-22 - Updated GenderFaceFilter node.
 - 2024-05-19 - Added BiSeNetMask and JonathandinuMask nodes. Careful about JonathandinuMask, it's more accurate than BiSeNet, but it takes more memory; you can get out of memory more easily with it.
 - 2024-03-10 - Added nodes to detect faces using `face_yolov8m` instead of `insightface`.
 

nodes.py

@@ -1,14 +1,15 @@
 import torch
 from collections import defaultdict
 from .utils import *
+from transformers import pipeline
 
 class GenderFaceFilter:
     @classmethod
     def INPUT_TYPES(cls):
         return {
             'required': {
                 'faces': ('FACE',),
-                'gender': (['male', 'female'],)
+                'gender': (['man', 'woman'],)
             }
         }
 
@@ -18,11 +19,16 @@ def INPUT_TYPES(cls):
     CATEGORY = 'facetools'
 
     def run(self, faces, gender):
-        gid = 0 if gender == 'female' else 1
         filtered = []
         rest = []
+        pipe = pipeline('image-classification', model='dima806/man_woman_face_image_detection', device=0)
         for face in faces:
-            if face.gender == gid:
+            _, im = face.crop(224, 1.2)
+            im = im.permute(0,3,1,2)[0]
+            im = tv.transforms.functional.resize(im, (224,224))
+            r = pipe(tv.transforms.functional.to_pil_image(im))
+            idx = np.argmax([i['score'] for i in r])
+            if r[idx]['label'] == gender:
                 filtered.append(face)
             else:
                 rest.append(face)

utils.py

@@ -96,9 +96,8 @@ def crop(self, size, crop_factor):
         N = N @ T4 @ S @ T3
         crop = cv2.warpAffine(self.img.numpy(), N, (size, size))
         crop = torch.from_numpy(crop)[None]
-        # maskedcrop = cv2.warpAffine(self.img.numpy(), N, (size, size))
 
-        return N, crop#, maskedcrop
+        return N, crop
 
 def detect_faces(img, threshold):
     img = pad_to_stride(img, stride=32)