diff --git a/README.md b/README.md index 948c91d..37b7d9c 100644 --- a/README.md +++ b/README.md @@ -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). diff --git a/faceanalysis.py b/faceanalysis.py index 2f8a350..a9c97c6 100644 --- a/faceanalysis.py +++ b/faceanalysis.py @@ -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", }