From bb9d0e04cbc37ed2a0c0f41ecdffa98e0f18200c Mon Sep 17 00:00:00 2001 From: David Chatel Date: Sun, 19 May 2024 07:09:55 +0200 Subject: [PATCH] added bisenetmask and jonathandinumask nodes --- README.md | 1 + nodes.py | 176 ++++++++++++++++++++--------------------- utils.py | 233 +++++++++++++++++++++++++++++------------------------- 3 files changed, 213 insertions(+), 197 deletions(-) diff --git a/README.md b/README.md index 2dfa3fb..4ada328 100644 --- a/README.md +++ b/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-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`. ## Example Workflows diff --git a/nodes.py b/nodes.py index 6718b55..80834e8 100644 --- a/nodes.py +++ b/nodes.py @@ -57,46 +57,6 @@ class OrderedFaceFilter: rest = sorted_faces[:take_start] + sorted_faces[take_start+take_count:] return (filtered, rest) -# class FaceDetailsDEPRECATED: -# @classmethod -# def INPUT_TYPES(cls): -# return { -# 'required': { -# 'faces': ('FACE',), -# 'crop_size': ('INT', {'default': 512, 'min': 512, 'max': 1024, 'step': 128}), -# 'crop_factor': ('FLOAT', {'default': 1.5, 'min': 1.0, 'max': 3, 'step': 0.1}), -# 'mask_type': (mask_types,) -# } -# } - -# RETURN_TYPES = ('IMAGE', 'MASK', 'WARP') -# RETURN_NAMES = ('crops', 'masks', 'warps') -# FUNCTION = 'run' -# CATEGORY = 'facetools' - -# def run(self, faces, crop_size, crop_factor, mask_type): -# if len(faces) == 0: -# empty_crop = torch.zeros((1,512,512,3)) -# empty_mask = torch.zeros((1,512,512)) -# empty_warp = np.array([ -# [1,0,-512], -# [0,1,-512], -# ], dtype=np.float32) -# return (empty_crop, empty_mask, [empty_warp]) - -# crops = [] -# masks = [] -# warps = [] -# for face in faces: -# M, crop = crop_faces(face.image, face, crop_size, crop_factor) -# mask = mask_crop(face, M, crop, mask_type) -# crops.append(np.array(crop) / 255) -# masks.append(np.array(mask)) -# warps.append(M) -# crops = torch.from_numpy(np.array(crops)).type(torch.float32) -# masks = torch.from_numpy(np.array(masks)).type(torch.float32) -# return (crops, masks, warps) - class DetectFaces: @classmethod def INPUT_TYPES(cls): @@ -131,7 +91,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 = image[i] + face.img = image[i] faces.append(face) return (faces,) @@ -166,47 +126,14 @@ class CropFaces: masks = [] warps = [] for face in faces: - 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)) + M, crop = face.crop(crop_size, crop_factor) + mask = mask_crop(face, M, crop, mask_type) + crops.append(np.array(crop[0])) + masks.append(np.array(mask[0])) warps.append(M) crops = torch.from_numpy(np.array(crops)).type(torch.float32) masks = torch.from_numpy(np.array(masks)).type(torch.float32) return (crops, masks, warps) - -# class AlignFacesDEPRECATED: -# @classmethod -# def INPUT_TYPES(cls): -# return { -# 'required': { -# 'insightface': ('INSIGHTFACE',), -# 'image': ('IMAGE',), -# 'threshold': ('FLOAT', {'default': 0.5, 'min': 0.5, 'max': 1.0, 'step': 0.01}), -# 'min_size': ('INT', {'default': 64, 'max': 512, 'step': 8}), -# 'max_size': ('INT', {'default': 512, 'min': 512, 'step': 8}), -# } -# } - -# RETURN_TYPES = ('FACE',) -# RETURN_NAMES = ('faces',) -# FUNCTION = 'run' -# CATEGORY = 'facetools' - -# def run(self, insightface, image, threshold, min_size, max_size): -# faces = [] -# images = (image * 255).type(torch.uint8).numpy() -# for i, img in enumerate(images): -# unfiltered_faces = get_faces(img, insightface) -# for face in unfiltered_faces: -# a, b, c, d = face.bbox -# h = abs(d-b) -# w = abs(c-a) -# if face.det_score >= threshold and (h <= max_size or w <= max_size) and (min_size <= h or min_size <= w): -# face.image_idx = i -# face.image = img -# faces.append(face) -# return (faces,) class WarpFaceBack: RETURN_TYPES = ('IMAGE',) @@ -228,14 +155,7 @@ 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)) + groups[f.image_idx].append((f.img,c,m,w)) results = [] for i, image in enumerate(images): @@ -287,13 +207,89 @@ class MergeWarps: masks = torch.vstack((mask0, mask1)) warps = warp0 + warp1 return (crops, masks, warps) + +class BiSeNetMask: + RETURN_TYPES = ('MASK',) + FUNCTION = 'run' + CATEGORY = 'facetools' + + @classmethod + def INPUT_TYPES(cls): + return { + 'required': { + 'crop': ('IMAGE',), + 'skin': ('BOOLEAN', {'default': True}), + 'left_brow': ('BOOLEAN', {'default': True}), + 'right_brow': ('BOOLEAN', {'default': True}), + 'left_eye': ('BOOLEAN', {'default': True}), + 'right_eye': ('BOOLEAN', {'default': True}), + 'eyeglasses': ('BOOLEAN', {'default': True}), + 'left_ear': ('BOOLEAN', {'default': True}), + 'right_ear': ('BOOLEAN', {'default': True}), + 'earring': ('BOOLEAN', {'default': True}), + 'nose': ('BOOLEAN', {'default': True}), + 'mouth': ('BOOLEAN', {'default': True}), + 'upper_lip': ('BOOLEAN', {'default': True}), + 'lower_lip': ('BOOLEAN', {'default': True}), + 'neck': ('BOOLEAN', {'default': False}), + 'necklace': ('BOOLEAN', {'default': False}), + 'cloth': ('BOOLEAN', {'default': False}), + 'hair': ('BOOLEAN', {'default': False}), + 'hat': ('BOOLEAN', {'default': False}), + } + } + def run(self, crop, skin, left_brow, right_brow, left_eye, right_eye, eyeglasses, + left_ear, right_ear, earring, nose, mouth, upper_lip, lower_lip, + neck, necklace, cloth, hair, hat): + masks = mask_BiSeNet(crop, skin, left_brow, right_brow, left_eye, right_eye, eyeglasses, + left_ear, right_ear, earring, nose, mouth, upper_lip, lower_lip, + neck, necklace, cloth, hair, hat) + return (masks, ) + +class JonathandinuMask: + RETURN_TYPES = ('MASK',) + FUNCTION = 'run' + CATEGORY = 'facetools' + + @classmethod + def INPUT_TYPES(cls): + return { + 'required': { + 'crop': ('IMAGE',), + 'skin': ('BOOLEAN', {'default': True}), + 'nose': ('BOOLEAN', {'default': True}), + 'eyeglasses': ('BOOLEAN', {'default': False}), + 'left_eye': ('BOOLEAN', {'default': True}), + 'right_eye': ('BOOLEAN', {'default': True}), + 'left_brow': ('BOOLEAN', {'default': True}), + 'right_brow': ('BOOLEAN', {'default': True}), + 'left_ear': ('BOOLEAN', {'default': True}), + 'right_ear': ('BOOLEAN', {'default': True}), + 'mouth': ('BOOLEAN', {'default': True}), + 'upper_lip': ('BOOLEAN', {'default': True}), + 'lower_lip': ('BOOLEAN', {'default': True}), + 'hair': ('BOOLEAN', {'default': False}), + 'hat': ('BOOLEAN', {'default': False}), + 'earring': ('BOOLEAN', {'default': False}), + 'necklace': ('BOOLEAN', {'default': False}), + 'neck': ('BOOLEAN', {'default': False}), + 'cloth': ('BOOLEAN', {'default': False}), + } + } + + def run(self, crop, skin, nose, eyeglasses, left_eye, right_eye, left_brow, right_brow, left_ear, right_ear, + mouth, upper_lip, lower_lip, hair, hat, earring, necklace, neck, cloth): + masks = mask_jonathandinu(crop, skin, nose, eyeglasses, left_eye, right_eye, left_brow, right_brow, left_ear, right_ear, + mouth, upper_lip, lower_lip, hair, hat, earring, necklace, neck, cloth) + return (masks, ) + NODE_CLASS_MAPPINGS = { 'DetectFaces': DetectFaces, 'CropFaces': CropFaces, - # 'AlignFaces': AlignFacesDEPRECATED, 'WarpFacesBack': WarpFaceBack, - # 'FaceDetails': FaceDetailsDEPRECATED, + 'BiSeNetMask': BiSeNetMask, + 'JonathandinuMask': JonathandinuMask, 'MergeWarps': MergeWarps, 'GenderFaceFilter': GenderFaceFilter, 'OrderedFaceFilter': OrderedFaceFilter, @@ -302,9 +298,9 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { 'DetectFaces': 'DetectFaces', 'CropFaces': 'CropFaces', - # 'AlignFaces': 'Align Faces (DEPRECATED)', 'WarpFacesBack': 'Warp Faces Back', - # 'FaceDetails': 'Face Details (DEPRECATED)', + 'BiSeNetMask': 'BiSeNet Mask', + 'JonathandinuMask': 'Jonathandinu Mask', 'MergeWarps': 'Merge Warps', 'GenderFaceFilter': 'Gender Face Filter', 'OrderedFaceFilter': 'Ordered Face Filter', diff --git a/utils.py b/utils.py index bc55e8e..37f3595 100644 --- a/utils.py +++ b/utils.py @@ -10,6 +10,7 @@ from folder_paths import models_dir from .BiSeNet import BiSeNet from ultralytics import YOLO from onnxruntime import InferenceSession +from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation def pad_to_stride(image, stride=32): h, w, _ = image.shape @@ -93,10 +94,11 @@ 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.image.numpy(), N, (size, size)) - maskedcrop = cv2.warpAffine(self.img.numpy(), N, (size, size)) + 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#, maskedcrop def detect_faces(img, threshold): img = pad_to_stride(img, stride=32) @@ -113,85 +115,10 @@ def detect_faces(img, threshold): faces.append(face) return faces -### DEPRECATED ### -def squareimg(img): - h, w, _ = img.shape - s = max(h, w) - ph, pw = (s - h) // 2, (s - w) // 2 - img = cv2.copyMakeBorder(img, ph, ph, pw, pw, cv2.BORDER_CONSTANT) - return (img, ph, pw) - -def transform_back(points, rot, pw, ph): - points = np.hstack((points, np.ones((len(points),1)))) @ rot.T - points -= np.array((pw, ph)) - return points - -def get_candidates(img, face_analysis): - simg, ph, pw = squareimg(img) - h, w, _ = img.shape - s = max(h, w) - candidates = [] - for i in range(4): - im = np.rot90(simg, i) - rot = cv2.getRotationMatrix2D((s/2, s/2), -90*i, 1) - dets = face_analysis.get(im) - for face in dets: - face.kps = transform_back(face.kps, rot, pw, ph) - face.landmark_2d_106 = transform_back(face.landmark_2d_106, rot, pw, ph) - - # original bbox and size - a,b,c,d = transform_back(face.bbox.reshape((2,2)), rot, pw, ph).flatten() - a,c = min(a,c),max(a,c) - b,d = min(b,d),max(b,d) - face.w, face.h = abs(c-a), abs(d-b) - - face.centroid = np.array(((a+c)/2, (b+d)/2)) - face.s = max(abs(a-c), abs(b-d)) - - face.angle = i*90 - candidates += dets - return candidates - -def group_faces(candidates): - groups = [] - for face in candidates: - group = None - for g in groups: - if any([np.linalg.norm(face.centroid - other.centroid) < np.sqrt(face.s)*2 for other in g]): - group = g - break - if group is None: - groups.append([face]) - else: - group.append(face) - return groups - -def get_best_faces(groups): - faces = [min(group, key=lambda face:abs(face.pose[2])) for group in groups] - return faces - -def get_faces(img, face_analysis): - candidates = get_candidates(img, face_analysis) - groups = group_faces(candidates) - faces = get_best_faces(groups) - return faces - -def crop_faces(img, face, crop_size, crop_factor): - M = face_align.estimate_norm(face.kps, crop_size, mode=None) - cx, cy = face.kps.mean(axis=0) - T1 = np.array([[1, 0, -cx], [0, 1, -cy], [0, 0, 1]]) - T2 = np.array([[1, 0, cx], [0, 1, cy], [0, 0, 1]]) - S = np.array([[1/crop_factor, 0, 0], [0, 1/crop_factor, 0], [0, 0, 1]]) - M = M @ T2 @ S @ T1 - crop = cv2.warpAffine(img, M, (crop_size,crop_size)) - - return M, crop -### END_DEPRECATED ### - -def get_face_mesh(crop): +def get_face_mesh(crop: torch.Tensor): with mp.solutions.face_mesh.FaceMesh(max_num_faces=10) as face_mesh: - mesh = face_mesh.process(crop) - h, w, _ = crop.shape + mesh = face_mesh.process(crop.mul(255).type(torch.uint8)[0].numpy()) + _, h, w, _ = crop.shape if mesh.multi_face_landmarks is not None: all_pts = np.array([np.array([(w*l.x, h*l.y) for l in lmks.landmark]) for lmks in mesh.multi_face_landmarks], dtype=np.int32) idx = np.argmin(np.abs(all_pts - np.array([w/2,h/2])).sum(axis=(1,2))) @@ -202,7 +129,7 @@ def get_face_mesh(crop): def mask_simple_square(face, M, crop): # rotated bbox and size - h,w = crop.shape[:2] + h,w = crop.shape[1:3] a,b,c,d = face.bbox rect = np.array([ [a,b,1], @@ -214,61 +141,153 @@ def mask_simple_square(face, M, crop): hx, hy = [int(x) for x in np.max(rect, axis=0)] mask = np.zeros((h,w), dtype=np.float32) mask = cv2.rectangle(mask, (lx,ly), (hx,hy), 1, -1) + mask = torch.from_numpy(mask)[None] return mask def mask_convex_hull(face, M, crop): - h,w = crop.shape[:2] + h,w = crop.shape[1:3] points = get_face_mesh(crop) if points is None: return mask_simple_square(face, M, crop) hull = ConvexHull(points) mask = np.zeros((h,w), dtype=np.int32) cv2.fillPoly(mask, [points[hull.vertices,:]], color=1) mask = mask.astype(np.float32) + mask = torch.from_numpy(mask[None]) return mask -def mask_BiSeNet(face, M, crop): +def mask_BiSeNet(crop, + skin=True, + l_brow=True, + r_brow=True, + l_eye=True, + r_eye=True, + eye_g=True, + l_ear=True, + r_ear=True, + ear_r=True, + nose=True, + mouth=True, + u_lip=True, + l_lip=True, + neck=False, + neck_l=False, + cloth=False, + hair=False, + hat=False, + ): with torch.no_grad(): bisenet = BiSeNet(n_classes=19) bisenet.cuda() model_path = os.path.join(models_dir, 'bisenet', '79999_iter.pth') bisenet.load_state_dict(torch.load(model_path)) bisenet.eval() - crop_t = torch.from_numpy(crop)[None].permute(0,3,1,2).contiguous().float().div(255).cuda() + crop_t = crop.permute(0,3,1,2).cuda() segms_t = bisenet(crop_t)[0].argmax(1).float() - # convert segmentation to mask - # 1: skin - # 2: l_brow - # 3: r_brow - # 4: l_eye - # 5: r_eye - # 6: eye_g - # 7: l_ear - # 8: r_ear - # 9: ear_r - # 10: nose - # 11: mouth - # 12: u_lip - # 13: l_lip - # 14: neck - # 15: neck_l - # 16: cloth - # 17: hair - # 18: hat - face_part_ids = torch.tensor([1,2,3,4,5,6,7,8,9,10,11,12,13]).cuda() + + dic = { + 'skin': 1, + 'l_brow': 2, + 'r_brow': 3, + 'l_eye': 4, + 'r_eye': 5, + 'eye_g': 6, + 'l_ear': 7, + 'r_ear': 8, + 'ear_r': 9, + 'nose': 10, + 'mouth': 11, + 'u_lip': 12, + 'l_lip': 13, + 'neck': 14, + 'neck_l': 15, + 'cloth': 16, + 'hair': 17, + 'hat': 18, + } + keep = [] + for k, v in locals().items(): + if k in dic and v: + keep.append(dic[k]) + + face_part_ids = torch.tensor(keep).cuda() segms_t = torch.sum(segms_t.repeat(len(face_part_ids), 1,1,1) == face_part_ids[...,None,None,None], axis=0).float() - mask = segms_t[0].cpu() + mask = segms_t.cpu() + return mask + +def mask_jonathandinu(crop, skin=True, nose=True, eye_g=True, l_eye=True, r_eye=True, l_brow=True, r_brow=True, + l_ear=True, r_ear=True, mouth=True, u_lip=True, l_lip=True, + hair=False, hat=False, ear_r=False, neck_l=False, neck=False, cloth=False): + global jonathandinu_image_processor, jonathandinu_model + + device = ( + "cuda" + # Device for NVIDIA or AMD GPUs + if torch.cuda.is_available() + else "mps" + # Device for Apple Silicon (Metal Performance Shaders) + if torch.backends.mps.is_available() + else "cpu" + ) + + if 'jonathandinu_image_processor' not in globals(): + jonathandinu_image_processor = SegformerImageProcessor.from_pretrained("jonathandinu/face-parsing") + jonathandinu_model = SegformerForSemanticSegmentation.from_pretrained("jonathandinu/face-parsing") + jonathandinu_model.to(device) + + inputs = jonathandinu_image_processor(images=crop.mul(255).type(torch.uint8), return_tensors="pt").to(device) + with torch.no_grad(): + outputs = jonathandinu_model(**inputs) + logits = outputs.logits # shape (batch_size, num_labels, ~height/4, ~width/4) + + # resize output to match input image dimensions + upsampled_logits = tv.transforms.functional.resize(logits, crop.shape[1:3], antialias=True) + + labels = upsampled_logits.argmax(dim=1) + + ids = { + 'skin': 1, + 'nose': 2, + 'eye_g': 3, + 'l_eye': 4, + 'r_eye': 5, + 'l_brow': 6, + 'r_brow': 7, + 'l_ear': 8, + 'r_ear': 9, + 'mouth': 10, + 'u_lip': 11, + 'l_lip': 12, + 'hair': 13, + 'hat': 14, + 'ear_r': 15, + 'neck_l': 16, + 'neck': 17, + 'cloth': 18, + } + keep = [] + for k, v in locals().items(): + if k in ids and v: + keep.append(ids[k]) + face_part_ids = torch.tensor(keep).cuda() + + mask = torch.sum(labels.repeat(len(face_part_ids), 1,1,1) == face_part_ids[...,None,None,None], axis=0).float().cpu() + return mask mask_types = [ 'simple_square', 'convex_hull', 'BiSeNet', + 'jonathandinu', + # 'clean BiSeNet', ] mask_funs = { 'simple_square': mask_simple_square, 'convex_hull': mask_convex_hull, - 'BiSeNet': mask_BiSeNet, + 'BiSeNet': lambda face, M, crop: mask_BiSeNet(crop), + 'jonathandinu': lambda face, M, crop: mask_jonathandinu(crop), + # 'clean BiSeNet': mask_clean_BiSeNet, } def mask_crop(face, M, crop, mask_type):