add insightface analysis
This commit is contained in:
matt3o
+121
-923
File diff suppressed because it is too large
Load Diff
+46
-26
@@ -1,27 +1,45 @@
|
||||
import dlib
|
||||
from insightface.app import FaceAnalysis
|
||||
import torch
|
||||
import torchvision.transforms.v2 as T
|
||||
import os
|
||||
import folder_paths
|
||||
import numpy as np
|
||||
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")
|
||||
|
||||
class FaceAnalysisModels:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {}}
|
||||
return {"required": {
|
||||
"library": (["dlib", "insightface"], ),
|
||||
"provider": (["CPU", "CUDA", "DirectML", "OpenVINO", "ROCM", "CoreML"], ),
|
||||
}}
|
||||
|
||||
RETURN_TYPES = ("ANALYSIS_MODELS", )
|
||||
FUNCTION = "load_models"
|
||||
CATEGORY = "FaceAnalysis"
|
||||
|
||||
def load_models(self):
|
||||
return ({
|
||||
"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")),
|
||||
}, )
|
||||
def load_models(self, library, provider):
|
||||
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))
|
||||
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")),
|
||||
}
|
||||
|
||||
return (out, )
|
||||
|
||||
class FaceEmbedDistance:
|
||||
@classmethod
|
||||
@@ -31,8 +49,8 @@ class FaceEmbedDistance:
|
||||
"analysis_models": ("ANALYSIS_MODELS", ),
|
||||
"reference": ("IMAGE", ),
|
||||
"image": ("IMAGE", ),
|
||||
"filter_thresh_eucl": ("FLOAT", { "default": 1.0, "min": 0.001, "max": 1.0, "step": 0.001 }),
|
||||
"filter_thresh_cos": ("FLOAT", { "default": 1.0, "min": 0.001, "max": 1.0, "step": 0.001 }),
|
||||
"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 }),
|
||||
"generate_image_overlay": ("BOOLEAN", { "default": True })
|
||||
},
|
||||
}
|
||||
@@ -49,9 +67,7 @@ class FaceEmbedDistance:
|
||||
background_color = ImageColor.getrgb("#000000AA")
|
||||
txt_height = font.getmask("Q").getbbox()[3] + font.getmetrics()[1]
|
||||
|
||||
self.detector = analysis_models.get("detector")
|
||||
self.shape_predict = analysis_models.get("shape_predict")
|
||||
self.face_recog = analysis_models.get("face_recog")
|
||||
self.analysis_models = analysis_models
|
||||
|
||||
ref = np.array(T.ToPILImage()(reference[0].permute(2, 0, 1)).convert('RGB'))
|
||||
ref = self.get_descriptor(ref)
|
||||
@@ -71,27 +87,21 @@ class FaceEmbedDistance:
|
||||
eucl_dist = 1.0
|
||||
cos_dist = 1.0
|
||||
else:
|
||||
if ref == img: # Same face
|
||||
if np.array_equal(ref, img): # Same face
|
||||
eucl_dist = 0.0
|
||||
cos_dist = 0.0
|
||||
else:
|
||||
eucl_dist = np.linalg.norm(np.array(ref) - np.array(img))
|
||||
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:
|
||||
out_eucl.append(eucl_dist)
|
||||
out_cos.append(cos_dist)
|
||||
|
||||
print(f"\033[96mFace Analysis: Euclidean: {eucl_dist}, Cosine: {cos_dist}\033[0m")
|
||||
|
||||
eucl_dist = round(eucl_dist, 3)
|
||||
cos_dist = round(cos_dist, 3)
|
||||
|
||||
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: {eucl_dist} | COS-1: {cos_dist}", font=font, fill=(255, 255, 255, 255))
|
||||
draw.text((0, 0), f"EUC: {round(eucl_dist, 3)} | COS: {round(cos_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)
|
||||
@@ -99,6 +109,9 @@ class FaceEmbedDistance:
|
||||
else:
|
||||
out.append(i)
|
||||
|
||||
out_eucl.append(eucl_dist)
|
||||
out_cos.append(cos_dist)
|
||||
|
||||
if not out:
|
||||
raise Exception('No image matches the filter criteria.')
|
||||
|
||||
@@ -115,12 +128,19 @@ class FaceEmbedDistance:
|
||||
return(img, out_eucl, out_cos, csv,)
|
||||
|
||||
def get_descriptor(self, image):
|
||||
faces = self.detector(image)
|
||||
if len(faces) > 0:
|
||||
shape = self.shape_predict(image, faces[0])
|
||||
return self.face_recog.compute_face_descriptor(image, shape)
|
||||
embeds = None
|
||||
|
||||
return None
|
||||
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))
|
||||
|
||||
return embeds
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"FaceEmbedDistance": FaceEmbedDistance,
|
||||
|
||||
@@ -1 +1,3 @@
|
||||
dlib
|
||||
onnxruntime
|
||||
insightface
|
||||
Reference in New Issue
Block a user