nodes/dchatel-comfyui_facetools
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

added yolov face detection

Browse Source
This commit is contained in:
David Chatel
2024-03-10 11:54:14 +01:00
parent f68841dde9
commit 93d05edcc5
3 changed files with 194 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+8 -2
View File
@@ -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.
- 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.
nodes.py
+85 -9
View File
@@ -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',
utils.py
+101 -2
View File
@@ -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):