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

added bisenetmask and jonathandinumask nodes

Browse Source
This commit is contained in:
David Chatel
2024-05-19 07:09:55 +02:00
parent 24e02fd1a8
commit bb9d0e04cb
3 changed files with 213 additions and 197 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1
View File
@@ -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
nodes.py
+86 -90
View File
@@ -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',
utils.py
+126 -107
View File
@@ -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):