nodes/cubiq-ComfyUI_FaceAnalysis
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

nodes refactoring

Browse Source
This commit is contained in:
matt3o
2024-05-13 10:03:38 +02:00
parent 2cc9a76db2
commit 3dc3877dda
2 changed files with 249 additions and 107 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+1 -1
View File
@@ -8,7 +8,7 @@ The best way to evaluate generated faces is to first send a batch of 3 reference
You need to install either InsightFace or Dlib (or both).
For DLIB download [Shape Predictor](https://huggingface.co/matt3ounstable/dlib_predictor_recognition/resolve/main/shape_predictor_68_face_landmarks.dat?download=true) and the [Face Recognition](https://huggingface.co/matt3ounstable/dlib_predictor_recognition/resolve/main/dlib_face_recognition_resnet_model_v1.dat?download=true) models and place them into the `dlib` directory.
For DLIB download [Shape Predictor](https://huggingface.co/matt3ounstable/dlib_predictor_recognition/resolve/main/shape_predictor_68_face_landmarks.dat?download=true), [Face Predictor 5 landmarks](https://huggingface.co/matt3ounstable/dlib_predictor_recognition/resolve/main/shape_predictor_5_face_landmarks.dat?download=true) and the [Face Recognition](https://huggingface.co/matt3ounstable/dlib_predictor_recognition/resolve/main/dlib_face_recognition_resnet_model_v1.dat?download=true) models and place them into the `dlib` directory.
Precompiled Dlib for windows can be found [here](https://github.com/z-mahmud22/Dlib_Windows_Python3.x).
faceanalysis.py
+248 -106
View File
@@ -15,16 +15,10 @@ except ImportError:
if not IS_DLIB_INSTALLED and not IS_INSIGHTFACE_INSTALLED:
raise Exception("Please install either dlib or insightface to use this node.")
INSTALLED_LIBRARIES = []
if IS_DLIB_INSTALLED:
INSTALLED_LIBRARIES.append("dlib")
if IS_INSIGHTFACE_INSTALLED:
INSTALLED_LIBRARIES.append("insightface")
import torch
import torch.nn.functional as F
#import torch.nn.functional as F
import torchvision.transforms.v2 as T
import comfy.utils
#import comfy.utils
import os
import folder_paths
import numpy as np
@@ -33,11 +27,133 @@ from PIL import Image, ImageDraw, ImageFont, ImageColor
DLIB_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "dlib")
INSIGHTFACE_DIR = os.path.join(folder_paths.models_dir, "insightface")
THRESHOLDS = { # from DeepFace
"VGG-Face": {"cosine": 0.68, "euclidean": 1.17, "L2_norm": 1.17},
"Facenet": {"cosine": 0.40, "euclidean": 10, "L2_norm": 0.80},
"Facenet512": {"cosine": 0.30, "euclidean": 23.56, "L2_norm": 1.04},
"ArcFace": {"cosine": 0.68, "euclidean": 4.15, "L2_norm": 1.13},
"Dlib": {"cosine": 0.07, "euclidean": 0.6, "L2_norm": 0.4},
"SFace": {"cosine": 0.593, "euclidean": 10.734, "L2_norm": 1.055},
"OpenFace": {"cosine": 0.10, "euclidean": 0.55, "L2_norm": 0.55},
"DeepFace": {"cosine": 0.23, "euclidean": 64, "L2_norm": 0.64},
"DeepID": {"cosine": 0.015, "euclidean": 45, "L2_norm": 0.17},
"GhostFaceNet": {"cosine": 0.65, "euclidean": 35.71, "L2_norm": 1.10},
}
def tensor_to_image(image):
return np.array(T.ToPILImage()(image.permute(2, 0, 1)).convert('RGB'))
def image_to_tensor(image):
return T.ToTensor()(image).permute(1, 2, 0)
#return T.ToTensor()(Image.fromarray(image)).permute(1, 2, 0)
class InsightFace:
def __init__(self, provider="CPU", name="buffalo_l"):
self.face_analysis = FaceAnalysis(name=name, root=INSIGHTFACE_DIR, providers=[provider + 'ExecutionProvider',])
self.face_analysis.prepare(ctx_id=0, det_size=(640, 640))
self.thresholds = THRESHOLDS["ArcFace"]
def get_face(self, image):
for size in [(size, size) for size in range(640, 256, -64)]:
self.face_analysis.det_model.input_size = size
faces = self.face_analysis.get(image)
if len(faces) > 0:
return sorted(faces, key=lambda x:(x['bbox'][2]-x['bbox'][0])*(x['bbox'][3]-x['bbox'][1]), reverse=True)
return None
def get_embeds(self, image):
face = self.get_face(image)
if face is not None:
face = face[0].normed_embedding
return face
def get_bbox(self, image, padding=0, padding_percent=0):
faces = self.get_face(np.array(image))
img = []
x = []
y = []
w = []
h = []
for face in faces:
x1, y1, x2, y2 = face['bbox']
width = x2 - x1
height = y2 - y1
x1 = int(max(0, x1 - int(width * padding_percent) - padding))
y1 = int(max(0, y1 - int(height * padding_percent) - padding))
x2 = int(min(image.width, x2 + int(width * padding_percent) + padding))
y2 = int(min(image.height, y2 + int(height * padding_percent) + padding))
crop = image.crop((x1, y1, x2, y2))
img.append(T.ToTensor()(crop).permute(1, 2, 0).unsqueeze(0))
x.append(x1)
y.append(y1)
w.append(x2 - x1)
h.append(y2 - y1)
return (img, x, y, w, h)
class DLib:
def __init__(self):
self.face_detector = dlib.get_frontal_face_detector()
# check if the models are available
if not os.path.exists(os.path.join(DLIB_DIR, "shape_predictor_5_face_landmarks.dat")):
raise Exception("The 5 point landmark model is not available. Please download it from https://huggingface.co/matt3ounstable/dlib_predictor_recognition/blob/main/shape_predictor_5_face_landmarks.dat")
if not os.path.exists(os.path.join(DLIB_DIR, "dlib_face_recognition_resnet_model_v1.dat")):
raise Exception("The face recognition model is not available. Please download it from https://huggingface.co/matt3ounstable/dlib_predictor_recognition/blob/main/dlib_face_recognition_resnet_model_v1.dat")
self.shape_predictor = dlib.shape_predictor(os.path.join(DLIB_DIR, "shape_predictor_5_face_landmarks.dat"))
self.face_recognition = dlib.face_recognition_model_v1(os.path.join(DLIB_DIR, "dlib_face_recognition_resnet_model_v1.dat"))
self.thresholds = THRESHOLDS["Dlib"]
def get_face(self, image):
faces = self.face_detector(np.array(image), 1)
if len(faces) > 0:
return sorted(faces, key=lambda x: x.area(), reverse=True)
return None
def get_embeds(self, image):
faces = self.get_face(image)
if faces is not None:
shape = self.shape_predictor(image, faces[0])
faces = np.array(self.face_recognition.compute_face_descriptor(image, shape))
return faces
def get_bbox(self, image, padding=0, padding_percent=0):
faces = self.get_face(image)
img = []
x = []
y = []
w = []
h = []
for face in faces:
x1 = max(0, face.left() - int(face.width() * padding_percent) - padding)
y1 = max(0, face.top() - int(face.height() * padding_percent) - padding)
x2 = min(image.width, face.right() + int(face.width() * padding_percent) + padding)
y2 = min(image.height, face.bottom() + int(face.height() * padding_percent) + padding)
crop = image.crop((x1, y1, x2, y2))
img.append(T.ToTensor()(crop).permute(1, 2, 0).unsqueeze(0))
x.append(x1)
y.append(y1)
w.append(x2 - x1)
h.append(y2 - y1)
return (img, x, y, w, h)
def get_landmarks(self, image):
faces = self.get_face(image)
if faces is not None:
shape = self.shape_predictor(image, faces[0])
return shape
return None
class FaceAnalysisModels:
@classmethod
def INPUT_TYPES(s):
libraries = []
if IS_INSIGHTFACE_INSTALLED:
libraries.append("insightface")
if IS_DLIB_INSTALLED:
libraries.append("dlib")
return {"required": {
"library": (INSTALLED_LIBRARIES, ),
"library": (libraries, ),
"provider": (["CPU", "CUDA", "DirectML", "OpenVINO", "ROCM", "CoreML"], ),
}}
@@ -49,30 +165,12 @@ class FaceAnalysisModels:
out = {}
if library == "insightface":
out = {
"library": library,
"detector": FaceAnalysis(name="buffalo_l", root=INSIGHTFACE_DIR, providers=[provider + 'ExecutionProvider',])
}
out["detector"].prepare(ctx_id=0, det_size=(640, 640))
out = InsightFace(provider)
else:
out = {
"library": library,
"detector": dlib.get_frontal_face_detector(),
"shape_predict": dlib.shape_predictor(os.path.join(DLIB_DIR, "shape_predictor_68_face_landmarks.dat")),
"face_recog": dlib.face_recognition_model_v1(os.path.join(DLIB_DIR, "dlib_face_recognition_resnet_model_v1.dat")),
}
out = DLib()
return (out, )
def crop_face(image, x, y, w, h, padding=0):
x = max(0, x - padding)
y = max(0, y - padding)
w = min(image.width, w + 2 * padding)
h = min(image.height, h + 2 * padding)
return image.crop((x, y, x + w, y + h))
class FaceBoundingBox:
@classmethod
def INPUT_TYPES(s):
@@ -81,6 +179,7 @@ class FaceBoundingBox:
"analysis_models": ("ANALYSIS_MODELS", ),
"image": ("IMAGE", ),
"padding": ("INT", { "default": 0, "min": 0, "max": 4096, "step": 1 }),
"padding_percent": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 2.0, "step": 0.05 }),
"index": ("INT", { "default": -1, "min": -1, "max": 4096, "step": 1 }),
},
}
@@ -91,7 +190,7 @@ class FaceBoundingBox:
CATEGORY = "FaceAnalysis"
OUTPUT_IS_LIST = (True, True, True, True, True,)
def bbox(self, analysis_models, image, padding, index=-1):
def bbox(self, analysis_models, image, padding, padding_percent, index=-1):
out_img = []
out_x = []
out_y = []
@@ -99,39 +198,13 @@ class FaceBoundingBox:
out_h = []
for i in image:
img = T.ToPILImage()(i.permute(2, 0, 1)).convert('RGB')
if analysis_models["library"] == "insightface":
faces = analysis_models["detector"].get(np.array(img))
for face in faces:
x, y, w, h = face.bbox.astype(int)
w = w - x
h = h - y
x = max(0, x - padding)
y = max(0, y - padding)
w = min(img.width, w + 2 * padding)
h = min(img.height, h + 2 * padding)
crop = img.crop((x, y, x + w, y + h))
out_img.append(T.ToTensor()(crop).permute(1, 2, 0).unsqueeze(0))
out_x.append(x)
out_y.append(y)
out_w.append(w)
out_h.append(h)
else:
faces = analysis_models["detector"](np.array(img), 1)
for face in faces:
x, y, w, h = face.left(), face.top(), face.width(), face.height()
x = max(0, x - padding)
y = max(0, y - padding)
w = min(img.width, w + 2 * padding)
h = min(img.height, h + 2 * padding)
crop = img.crop((x, y, x + w, y + h))
out_img.append(T.ToTensor()(crop).permute(1, 2, 0).unsqueeze(0))
out_x.append(x)
out_y.append(y)
out_w.append(w)
out_h.append(h)
i = T.ToPILImage()(i.permute(2, 0, 1)).convert('RGB')
img, x, y, w, h = analysis_models.get_bbox(i, padding, padding_percent)
out_img.extend(img)
out_x.extend(x)
out_y.extend(y)
out_w.extend(w)
out_h.extend(h)
if not out_img:
raise Exception('No face detected in image.')
@@ -165,30 +238,31 @@ class FaceEmbedDistance:
"analysis_models": ("ANALYSIS_MODELS", ),
"reference": ("IMAGE", ),
"image": ("IMAGE", ),
"filter_thresh_eucl": ("FLOAT", { "default": 1.0, "min": 0.001, "max": 2.0, "step": 0.001 }),
"filter_thresh_cos": ("FLOAT", { "default": 1.0, "min": 0.001, "max": 2.0, "step": 0.001 }),
"similarity_metric": (["L2_norm", "cosine", "euclidean"], ),
"filter_thresh": ("FLOAT", { "default": 100.0, "min": 0.001, "max": 100.0, "step": 0.001 }),
"filter_best": ("INT", { "default": 0, "min": 0, "max": 4096, "step": 1 }),
"generate_image_overlay": ("BOOLEAN", { "default": True }),
},
}
RETURN_TYPES = ("IMAGE", "FLOAT", "FLOAT")
RETURN_NAMES = ("IMAGE", "euclidean", "cosine")
OUTPUT_IS_LIST = (False, True, True)
RETURN_TYPES = ("IMAGE", "FLOAT")
RETURN_NAMES = ("IMAGE", "distance")
FUNCTION = "analize"
CATEGORY = "FaceAnalysis"
def analize(self, analysis_models, reference, image, filter_thresh_eucl=1.0, filter_thresh_cos=1.0, filter_best=0, generate_image_overlay=True):
def analize(self, analysis_models, reference, image, similarity_metric, filter_thresh, filter_best, generate_image_overlay=True):
if generate_image_overlay:
font = ImageFont.truetype(os.path.join(os.path.dirname(os.path.realpath(__file__)), "Inconsolata.otf"), 32)
background_color = ImageColor.getrgb("#000000AA")
txt_height = font.getmask("Q").getbbox()[3] + font.getmetrics()[1]
self.analysis_models = analysis_models
if filter_thresh == 0.0:
filter_thresh = analysis_models.thresholds[similarity_metric]
# you can send multiple reference images in which case the embeddings are averaged
ref = []
for i in reference:
ref_emb = self.get_descriptor(np.array(T.ToPILImage()(i.permute(2, 0, 1).cpu()).convert('RGB')))
ref_emb = analysis_models.get_embeds(np.array(T.ToPILImage()(i.permute(2, 0, 1)).convert('RGB')))
if ref_emb is not None:
ref.append(torch.from_numpy(ref_emb))
@@ -199,33 +273,43 @@ class FaceEmbedDistance:
ref = np.array(torch.mean(ref, dim=0))
out = []
out_eucl = []
out_cos = []
out_dist = []
for i in image:
img = np.array(T.ToPILImage()(i.permute(2, 0, 1).cpu()).convert('RGB'))
img = np.array(T.ToPILImage()(i.permute(2, 0, 1)).convert('RGB'))
img = self.get_descriptor(img)
img = analysis_models.get_embeds(img)
if img is None: # No face detected
eucl_dist = 1.0
cos_dist = 1.0
dist = 100.0
norm_dist = 0
else:
if np.array_equal(ref, img): # Same face
eucl_dist = 0.0
cos_dist = 0.0
dist = 0.0
norm_dist = 1
else:
eucl_dist = np.float64(np.linalg.norm(ref - img))
cos_dist = 1 - np.dot(ref, img) / (np.linalg.norm(ref) * np.linalg.norm(img))
if eucl_dist <= filter_thresh_eucl and cos_dist <= filter_thresh_cos:
print(f"\033[96mFace Analysis: Euclidean: {eucl_dist}, Cosine: {cos_dist}\033[0m")
if similarity_metric == "L2_norm":
#dist = euclidean_distance(ref, img, True)
ref = ref / np.linalg.norm(ref)
img = img / np.linalg.norm(img)
dist = np.float64(np.linalg.norm(ref - img))
elif similarity_metric == "cosine":
dist = np.float64(1 - np.dot(ref, img) / (np.linalg.norm(ref) * np.linalg.norm(img)))
#dist = cos_distance(ref, img)
else:
#dist = euclidean_distance(ref, img)
dist = np.float64(np.linalg.norm(ref - img))
norm_dist = min(1.0, 1 / analysis_models.thresholds[similarity_metric] * dist)
if dist <= filter_thresh:
print(f"\033[96mFace Analysis: value: {dist}, normalized: {norm_dist}\033[0m")
if generate_image_overlay:
tmp = T.ToPILImage()(i.permute(2, 0, 1)).convert('RGBA')
txt = Image.new('RGBA', (image.shape[2], txt_height), color=background_color)
draw = ImageDraw.Draw(txt)
draw.text((0, 0), f"EUC: {round(eucl_dist, 3)} | COS: {round(cos_dist, 3)}", font=font, fill=(255, 255, 255, 255))
draw.text((0, 0), f"VALUE: {round(dist, 3)} | DIST: {round(norm_dist, 3)}", font=font, fill=(255, 255, 255, 255))
composite = Image.new('RGBA', tmp.size)
composite.paste(txt, (0, tmp.height - txt.height))
composite = Image.alpha_composite(tmp, composite)
@@ -233,8 +317,7 @@ class FaceEmbedDistance:
else:
out.append(i)
out_eucl.append(eucl_dist)
out_cos.append(cos_dist)
out_dist.append(dist)
if not out:
raise Exception('No image matches the filter criteria.')
@@ -242,41 +325,100 @@ class FaceEmbedDistance:
# filter out the best matches
if filter_best > 0:
out = np.array(out)
out_eucl = np.array(out_eucl)
out_cos = np.array(out_cos)
idx = np.argsort((out_eucl + out_cos) / 2)
out_dist = np.array(out_dist)
idx = np.argsort(out_dist)
out = torch.from_numpy(out[idx][:filter_best])
out_eucl = out_eucl[idx][:filter_best].tolist()
out_cos = out_cos[idx][:filter_best].tolist()
out_dist = out_dist[idx][:filter_best].tolist()
if isinstance(out, list):
out = torch.stack(out)
return(out, out_eucl, out_cos,)
def get_descriptor(self, image):
embeds = None
return(out, out_dist,)
if self.analysis_models["library"] == "insightface":
faces = self.analysis_models["detector"].get(image)
if len(faces) > 0:
embeds = faces[0].normed_embedding
else:
faces = self.analysis_models["detector"](image)
if len(faces) > 0:
shape = self.analysis_models["shape_predict"](image, faces[0])
embeds = np.array(self.analysis_models["face_recog"].compute_face_descriptor(image, shape))
class FaceAlign:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"analysis_models": ("ANALYSIS_MODELS", ),
"image_from": ("IMAGE", ),
}, "optional": {
"image_to": ("IMAGE", ),
}
}
return embeds
RETURN_TYPES = ("IMAGE", )
FUNCTION = "align"
CATEGORY = "FaceAnalysis"
def align(self, analysis_models, image_from, image_to=None):
image_from = tensor_to_image(image_from[0])
shape = analysis_models.get_landmarks(image_from)
r_eye_from = (
int((shape.part(2).x + shape.part(3).x) // 2),
int((shape.part(2).y + shape.part(3).y) // 2)
)
l_eye_from = (
int((shape.part(0).x + shape.part(1).x) // 2),
int((shape.part(0).y + shape.part(1).y) // 2)
)
angle = float(np.degrees(np.arctan2(l_eye_from[1] - r_eye_from[1], l_eye_from[0] - r_eye_from[0])))
if image_to is not None:
image_to = tensor_to_image(image_to[0])
shape = analysis_models.get_landmarks(image_to)
r_eye_to = (
int((shape.part(2).x + shape.part(3).x) // 2),
int((shape.part(2).y + shape.part(3).y) // 2)
)
l_eye_to = (
int((shape.part(0).x + shape.part(1).x) // 2),
int((shape.part(0).y + shape.part(1).y) // 2)
)
angle -= float(np.degrees(np.arctan2(l_eye_to[1] - r_eye_to[1], l_eye_to[0] - r_eye_to[0])))
# rotate the image
image_from = Image.fromarray(image_from).rotate(angle)
image_from = image_to_tensor(image_from).unsqueeze(0)
#img = np.array(Image.fromarray(image_from).rotate(angle))
#img = image_to_tensor(img).unsqueeze(0)
return (image_from, )
"""
def cos_distance(source, test):
a = np.matmul(np.transpose(source), test)
b = np.sum(np.multiply(source, source))
c = np.sum(np.multiply(test, test))
return np.float64(1 - (a / (np.sqrt(b) * np.sqrt(c))))
def euclidean_distance(source, test, norm=False):
if norm:
source = l2_normalize(source)
test = l2_normalize(test)
dist = source - test
dist = np.sum(np.multiply(dist, dist))
dist = np.sqrt(dist)
return np.float64(dist)
def l2_normalize(x):
return x / np.sqrt(np.sum(np.multiply(x, x)))
"""
NODE_CLASS_MAPPINGS = {
"FaceEmbedDistance": FaceEmbedDistance,
"FaceAnalysisModels": FaceAnalysisModels,
"FaceBoundingBox": FaceBoundingBox,
#"FaceAlign": FaceAlign,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"FaceEmbedDistance": "Face Embeds Distance",
"FaceAnalysisModels": "Face Analysis Models",
"FaceBoundingBox": "Face Bounding Box",
#"FaceAlign": "Face Align",
}