From 93d05edcc5e8c3f35bfe6718dfff09eafc6e68c0 Mon Sep 17 00:00:00 2001 From: David Chatel Date: Sun, 10 Mar 2024 11:54:14 +0100 Subject: [PATCH] added yolov face detection --- README.md | 10 ++++-- nodes.py | 94 ++++++++++++++++++++++++++++++++++++++++++++----- utils.py | 103 ++++++++++++++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 194 insertions(+), 13 deletions(-) diff --git a/README.md b/README.md index fe1daf1..9c36d2a 100644 --- a/README.md +++ b/README.md @@ -4,6 +4,9 @@ These custom nodes provide a rotation aware face extraction, paste back, and var ![Comparison](examples/comparison.jpg) +## Patch notes +- 2024-03-10 - Added nodes to detect faces using `face_yolov8m` instead of `insightface`. + ## Example Workflows - Full inpainting workflow with two controlnets which allows to get as high as 1.0 denoise strength without messing things up. ![Full Inpainting](examples/full_inpainting.png) @@ -46,5 +49,8 @@ These custom nodes provide a rotation aware face extraction, paste back, and var - I made some changes in `FaceDetailer` and `WarpFacesBack` since I've created this one, so it probably doesn't work anymore. Needs more testing and more work. ## Installation -You will need [ComfyUI-Impact-Pack](https://github.com/ltdrdata/) for `Load InsightFace` node and [comfyui_controlnet_aux](https://github.com/Fannovel16/comfyui_controlnet_aux) for MediaPipe library (which is required for convex_hull masks) and `MediaPipe Face Mesh` node if you want to use that controlnet. -You will also need to download the [BiSeNet model](https://drive.google.com/open?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812) and save it in `ComfyUI/models/bisenet` to use occlusion aware masks. \ No newline at end of file +- Download [bluefoxcreation/FaceAlignment](https://huggingface.co/bluefoxcreation/FaceAlignment) model into `ComfyUI/models/landmarks` +- (Only if you want to use occlusion aware masks) Download [BiSeNet model](https://drive.google.com/open?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812) into `ComfyUI/models/bisenet` + +###### DEPRECATED: +- You need [ComfyUI-Impact-Pack](https://github.com/ltdrdata/) for `Load InsightFace` node and [comfyui_controlnet_aux](https://github.com/Fannovel16/comfyui_controlnet_aux) for MediaPipe library (which is required for convex_hull masks) and `MediaPipe Face Mesh` node if you want to use that controlnet. \ No newline at end of file diff --git a/nodes.py b/nodes.py index 6ae9f92..46a2730 100644 --- a/nodes.py +++ b/nodes.py @@ -57,7 +57,7 @@ class OrderedFaceFilter: rest = sorted_faces[:take_start] + sorted_faces[take_start+take_count:] return (filtered, rest) -class FaceDetails: +class FaceDetailsDEPRECATED: @classmethod def INPUT_TYPES(cls): return { @@ -93,11 +93,83 @@ class FaceDetails: crops.append(np.array(crop) / 255) masks.append(np.array(mask)) warps.append(M) - crops = torch.from_numpy(np.array(crops)) - masks = torch.from_numpy(np.array(masks)) + 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 AlignFaces: +class DetectFaces: + @classmethod + def INPUT_TYPES(cls): + return { + 'required': { + 'image': ('IMAGE',), + '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}), + } + } + + RETURN_TYPES = ('FACE',) + RETURN_NAMES = ('faces',) + FUNCTION = 'run' + CATEGORY = 'facetools' + + def run(self, image, threshold, min_size, max_size): + faces = [] + images = (image * 255).type(torch.uint8) + for i, img in enumerate(images): + unfiltered_faces = detect_faces(img, threshold) + for face in unfiltered_faces: + a, b, c, d = face.bbox + h = abs(d-b) + 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 + faces.append(face) + return (faces,) + +class CropFaces: + @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 = face.crop(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 AlignFacesDEPRECATED: @classmethod def INPUT_TYPES(cls): return { @@ -173,7 +245,7 @@ class WarpFaceBack: for single_crop, single_mask, single_warp in zip(crop, warped_masks, warp) ], axis=0) / np.maximum(1, full_mask) full_mask = np.minimum(1, full_mask) - result = (swapped + (1 - full_mask) * image / 255) + result = (swapped + (1 - full_mask) * image.numpy() / 255) result = torch.from_numpy(result) results.append(result) @@ -205,18 +277,22 @@ class MergeWarps: return (crops, masks, warps) NODE_CLASS_MAPPINGS = { - 'AlignFaces': AlignFaces, + 'DetectFaces': DetectFaces, + 'CropFaces': CropFaces, + 'AlignFaces': AlignFacesDEPRECATED, 'WarpFacesBack': WarpFaceBack, - 'FaceDetails': FaceDetails, + 'FaceDetails': FaceDetailsDEPRECATED, 'MergeWarps': MergeWarps, 'GenderFaceFilter': GenderFaceFilter, 'OrderedFaceFilter': OrderedFaceFilter, } NODE_DISPLAY_NAME_MAPPINGS = { - 'AlignFaces': 'Align Faces', + 'DetectFaces': 'DetectFaces', + 'CropFaces': 'CropFaces', + 'AlignFaces': 'Align Faces (DEPRECATED)', 'WarpFacesBack': 'Warp Faces Back', - 'FaceDetails': 'Face Details', + 'FaceDetails': 'Face Details (DEPRECATED)', 'MergeWarps': 'Merge Warps', 'GenderFaceFilter': 'Gender Face Filter', 'OrderedFaceFilter': 'Ordered Face Filter', diff --git a/utils.py b/utils.py index b659d1a..a3e4bd9 100644 --- a/utils.py +++ b/utils.py @@ -1,5 +1,6 @@ import os import torch +import torchvision as tv import numpy as np import cv2 from insightface.utils import face_align @@ -7,7 +8,105 @@ import mediapipe as mp from scipy.spatial import ConvexHull from folder_paths import models_dir from .BiSeNet import BiSeNet +from ultralytics import YOLO +from onnxruntime import InferenceSession +from collections import namedtuple +def resize(img, size): + h, w, _ = img.shape + s = max(h, w) + scale_factor = s / size + ph, pw = (s - h) // 2, (s - w) // 2 + pad = tv.transforms.Pad((pw, ph)) + resize = tv.transforms.Resize(size=(size, size), antialias=True) + img = resize(pad(img.permute(2,0,1))).permute(1,2,0) + return img, scale_factor, ph, pw + +class Models: + @classmethod + def yolo(cls, img, threshold): + if '_yolo' not in cls.__dict__: + cls._yolo = YOLO(os.path.join(models_dir,'ultralytics','bbox','face_yolov8m.pt')) + dets = cls._yolo(img, conf=threshold)[0] + return dets + @classmethod + def lmk(cls, crop): + if '_lmk' not in cls.__dict__: + cls._lmk = InferenceSession(os.path.join(models_dir, 'landmarks', 'fan2_68_landmark.onnx')) + lmk = cls._lmk.run(None, {'input': crop})[0] + return lmk + +def get_submatrix_with_padding(img, a, b, c, d): + pl = -min(a, 0) + pt = -min(b, 0) + pr = -min(img.shape[1] - c, 0) + pb = -min(img.shape[0] - d, 0) + a, b, c, d = max(a, 0), max(b, 0), min(c, img.shape[1]), min(d, img.shape[0]) + + submatrix = img[b:d, a:c] + pad = tv.transforms.Pad((pl, pt, pr, pb)) + submatrix = pad(submatrix) + + return submatrix + +class Face: + def __init__(self, img, a, b, c, d) -> None: + self.img = img + lmk = None + best_score = 0 + i = 0 + crop = get_submatrix_with_padding(self.img, a, b, c, d) + for curr_i in range(4): + rcrop, s, ph, pw = resize(crop.rot90(curr_i), 256) + rcrop = (rcrop[None] / 255).permute(0,3,1,2).type(torch.float32).numpy() + curr_lmk = Models.lmk(rcrop) + score = np.mean(curr_lmk[0,:,2]) + if score > best_score: + best_score = score + lmk = curr_lmk + i = curr_i + + self.bbox = (a,b,c,d) + self.confidence = best_score + + self.kps = np.vstack([ + lmk[0,[37,38,40,41],:2].mean(axis=0), + lmk[0,[43,44,46,47],:2].mean(axis=0), + lmk[0,[30,48,54],:2] + ]) * 4 * s + + self.T2 = np.array([[1, 0, -a], [0, 1, -b], [0, 0, 1]]) + rot = cv2.getRotationMatrix2D((128*s,128*s), 90*i, 1) + self.R = np.vstack((rot, np.array((0,0,1)))) + + def crop(self, size, crop_factor): + S = np.array([[1/crop_factor, 0, 0], [0, 1/crop_factor, 0], [0, 0, 1]]) + M = face_align.estimate_norm(self.kps, size) + N = M @ self.R @ self.T2 + cx, cy = np.array((size/2, size/2, 1)) @ cv2.invertAffineTransform(M @ self.R @ self.T2).T + 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)) + + return N, crop + +def detect_faces(img, threshold): + dets = Models.yolo((img[None] / 255).permute(0,3,1,2), threshold) + boxes = (dets.boxes.xyxy.reshape(-1,2,2)).reshape(-1,4) + a,b,c,d = [int(x) for x in boxes[0]] + faces = [] + for (a,b,c,d), box in zip(boxes.type(torch.int).cpu().numpy(), dets.boxes): + cx, cy = (a+c)/2, (b+d)/2 + r = np.sqrt((c-a)**2 + (d-b)**2) / 2 + + a,b,c,d = [int(x) for x in (cx - r, cy - r, cx + r, cy + r)] + face = Face(img, a, b, c, d) + + faces.append(face) + return faces + +### DEPRECATED ### def squareimg(img): h, w, _ = img.shape s = max(h, w) @@ -80,6 +179,7 @@ def crop_faces(img, face, crop_size, crop_factor): crop = cv2.warpAffine(img, M, (crop_size,crop_size)) return M, crop +### END_DEPRECATED ### def get_face_mesh(crop): with mp.solutions.face_mesh.FaceMesh(max_num_faces=10) as face_mesh: @@ -104,9 +204,8 @@ def mask_simple_square(face, M, crop): ]) @ M.T lx, ly = [int(x) for x in np.min(rect, axis=0)] hx, hy = [int(x) for x in np.max(rect, axis=0)] - mask = np.zeros((512,512), dtype=np.int32) + mask = np.zeros((512,512), dtype=np.float32) mask = cv2.rectangle(mask, (lx,ly), (hx,hy), 1, -1) - mask = mask.astype(np.float32) return mask def mask_convex_hull(face, M, crop):