374 lines
13 KiB
Python
374 lines
13 KiB
Python
import torch
|
|
|
|
from deepface import DeepFace
|
|
import numpy as np
|
|
import os
|
|
|
|
import comfy.utils
|
|
import folder_paths
|
|
|
|
def comfy_image_from_deepface_image(deepface_image):
|
|
image_data = np.array(deepface_image).astype(np.float32)
|
|
return torch.from_numpy(image_data)[None,]
|
|
|
|
def deepface_image_from_comfy_image(comfy_image):
|
|
image_data = np.clip(255 * comfy_image.cpu().numpy(), 0, 255).astype(np.uint8)
|
|
image_data = image_data[:, :, :3] # Ensure RGBA is reduced to RGB
|
|
return image_data[:, :, ::-1] # Convert RGB to BGR
|
|
|
|
def pad_deepface_image(deepface_image):
|
|
h, w = deepface_image.shape[:2]
|
|
pad = int(max(h, w) * 0.4)
|
|
padded = np.pad(
|
|
deepface_image,
|
|
((pad, pad), (pad, pad), (0, 0)),
|
|
mode="constant",
|
|
constant_values=0,
|
|
)
|
|
return padded, pad
|
|
|
|
def prepare_deepface_home():
|
|
deepface_path = os.path.join(folder_paths.models_dir, "deepface")
|
|
|
|
# Deepface requires a specific structure within the DEEPFACE_HOME directory
|
|
deepface_dot_path = os.path.join(deepface_path, ".deepface")
|
|
deepface_weights_path = os.path.join(deepface_dot_path, "weights")
|
|
if not os.path.exists(deepface_weights_path):
|
|
os.makedirs(deepface_weights_path)
|
|
|
|
os.environ["DEEPFACE_HOME"] = deepface_path
|
|
|
|
def result_from_images_with_measurements(images_with_measurements, sort_by):
|
|
images_with_measurements.sort(key=lambda row: row[1] if sort_by == "distance" else (-row[2], row[1]))
|
|
|
|
images = [row[0] for row in images_with_measurements]
|
|
distances = [row[1] for row in images_with_measurements]
|
|
ratios = [row[2] for row in images_with_measurements]
|
|
|
|
if len(images) > 0:
|
|
return torch.stack(images, dim=0), distances, ratios
|
|
else:
|
|
return (None, None, None,)
|
|
|
|
class AverageListNode:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"numbers": ("NUMBER",)
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("FLOAT",)
|
|
FUNCTION = "average"
|
|
|
|
def average(self, numbers):
|
|
if not numbers:
|
|
return (0.0,)
|
|
return (sum(numbers) / len(numbers),)
|
|
|
|
class DeepfaceExtractFacesNode:
|
|
def __init__(self):
|
|
prepare_deepface_home()
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"target_face_size": ("INT", {
|
|
"default": 224,
|
|
"min": 32,
|
|
"max": 1024,
|
|
"step": 32,
|
|
"display": "number"
|
|
}),
|
|
"detector_backend": ([
|
|
"opencv",
|
|
"ssd",
|
|
"dlib",
|
|
"mtcnn",
|
|
"retinaface",
|
|
"mediapipe",
|
|
"yolov8",
|
|
"yunet",
|
|
"fastmtcnn",
|
|
], {
|
|
"default": "retinaface",
|
|
}),
|
|
"enforce_detection": ("BOOLEAN", {
|
|
"default": False,
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("face_images",)
|
|
|
|
FUNCTION = "run"
|
|
|
|
CATEGORY = "deepface"
|
|
|
|
def run(self, images, target_face_size, detector_backend, enforce_detection):
|
|
target_size = (target_face_size, target_face_size)
|
|
|
|
progress_bar = comfy.utils.ProgressBar(len(images))
|
|
|
|
output_images = []
|
|
for image in images:
|
|
progress_bar.update(1)
|
|
|
|
image = deepface_image_from_comfy_image(image)
|
|
|
|
try:
|
|
detected_faces = DeepFace.extract_faces(
|
|
image,
|
|
detector_backend=detector_backend,
|
|
enforce_detection=enforce_detection,
|
|
target_size=target_size,
|
|
)
|
|
except ValueError:
|
|
detected_faces = []
|
|
|
|
for detected_face in detected_faces:
|
|
# print(detected_face["confidence"])
|
|
face_image = comfy_image_from_deepface_image(detected_face["face"])
|
|
output_images.append(face_image)
|
|
|
|
if len(output_images) > 0:
|
|
return (torch.cat(output_images, dim=0),)
|
|
else:
|
|
return (None,)
|
|
|
|
class DeepfaceVerifyNode:
|
|
def __init__(self):
|
|
prepare_deepface_home()
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"reference_images": ("IMAGE",),
|
|
"distance_threshold": ("FLOAT", {
|
|
"default": 0.6,
|
|
"display": "number",
|
|
"minimum": 0.0,
|
|
"maximum": 1.0,
|
|
"step": 0.01,
|
|
}),
|
|
"ratio_threshold": ("FLOAT", {
|
|
"default": 0.7,
|
|
"display": "number",
|
|
"minimum": 0.0,
|
|
"maximum": 1.0,
|
|
"step": 0.01,
|
|
}),
|
|
"detector_backend": ([
|
|
"opencv",
|
|
"ssd",
|
|
"dlib",
|
|
"mtcnn",
|
|
"retinaface",
|
|
"mediapipe",
|
|
"yolov8",
|
|
"yunet",
|
|
"fastmtcnn",
|
|
], {
|
|
"default": "ssd",
|
|
}),
|
|
"model_name": ([
|
|
"VGG-Face",
|
|
"Facenet",
|
|
"Facenet512",
|
|
"OpenFace",
|
|
"DeepFace",
|
|
"DeepID",
|
|
"ArcFace",
|
|
"Dlib",
|
|
"SFace",
|
|
], {
|
|
"default": "VGG-Face",
|
|
}),
|
|
"sort_by": ([
|
|
"distance",
|
|
"ratio"
|
|
], {
|
|
"default": "distance",
|
|
})
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "NUMBER", "NUMBER", "IMAGE", "NUMBER", "NUMBER", "BOUNDING_BOX",)
|
|
RETURN_NAMES = (
|
|
"verified_images",
|
|
"verified_image_distances",
|
|
"verified_image_verified_ratios",
|
|
"rejected_images",
|
|
"rejected_image_distances",
|
|
"rejected_image_verified_ratios",
|
|
"bboxes",
|
|
)
|
|
|
|
FUNCTION = "run"
|
|
|
|
CATEGORY = "deepface"
|
|
|
|
def run(self, images, reference_images, distance_threshold, ratio_threshold, detector_backend, model_name, sort_by):
|
|
deepface_reference_images = []
|
|
for reference_image in reference_images:
|
|
deepface_reference_images.append(deepface_image_from_comfy_image(reference_image))
|
|
|
|
total_steps = len(deepface_reference_images) * len(images)
|
|
progress_bar = comfy.utils.ProgressBar(total_steps)
|
|
|
|
rejected_image_tuples = []
|
|
verified_image_tuples = []
|
|
bboxes = []
|
|
image_counter = 0
|
|
for image in images:
|
|
print(f"Deepface verify { image_counter + 1 }/{ len(images) }")
|
|
|
|
comparison_image = deepface_image_from_comfy_image(image)
|
|
|
|
image_bbox = None
|
|
padding_attempted = False
|
|
padding_offset = 0
|
|
reference_image_counter = 1
|
|
total_distance = 0
|
|
verified_images_count = 0
|
|
for deepface_reference_image in deepface_reference_images:
|
|
progress_bar.update(1)
|
|
|
|
result = None
|
|
try:
|
|
result = DeepFace.verify(
|
|
deepface_reference_image,
|
|
comparison_image,
|
|
detector_backend=detector_backend,
|
|
enforce_detection=False,
|
|
model_name=model_name
|
|
)
|
|
except ValueError as e:
|
|
if padding_attempted:
|
|
print(f" Reference image { reference_image_counter }/{ len(reference_images) }: ERROR: { e }")
|
|
continue
|
|
|
|
# The input face may be too large to detect; pad it and retry
|
|
padding_attempted = True
|
|
padded_image, pad = pad_deepface_image(comparison_image)
|
|
try:
|
|
result = DeepFace.verify(
|
|
deepface_reference_image,
|
|
padded_image,
|
|
detector_backend=detector_backend,
|
|
enforce_detection=False,
|
|
model_name=model_name
|
|
)
|
|
except ValueError as retry_error:
|
|
print(f" Reference image { reference_image_counter }/{ len(reference_images) }: ERROR (after padding retry): { retry_error }")
|
|
continue
|
|
|
|
# Padding worked: reuse the padded image for the rest of the comparisons
|
|
print(f" Reference image { reference_image_counter }/{ len(reference_images) }: retried with padding")
|
|
comparison_image = padded_image
|
|
padding_offset = pad
|
|
|
|
if image_bbox is None:
|
|
facial_area = result.get("facial_areas", {}).get("img2")
|
|
if facial_area is not None:
|
|
image_bbox = {
|
|
"x": facial_area["x"] - padding_offset,
|
|
"y": facial_area["y"] - padding_offset,
|
|
"width": facial_area["w"],
|
|
"height": facial_area["h"],
|
|
}
|
|
|
|
distance = result["distance"]
|
|
is_verified = result["verified"]
|
|
print(f" Reference image { reference_image_counter }/{ len(reference_images) }: distance={ distance } verified={ is_verified }")
|
|
reference_image_counter += 1
|
|
total_distance += distance
|
|
if is_verified:
|
|
verified_images_count += 1
|
|
|
|
average_distance = total_distance / len(deepface_reference_images)
|
|
verified_ratio = round(verified_images_count / len(deepface_reference_images), 2)
|
|
print(f"Average distance: { average_distance } ({ verified_ratio } verified)")
|
|
|
|
if average_distance < distance_threshold and verified_ratio >= ratio_threshold:
|
|
verified_image_tuples.append((image, average_distance, verified_ratio))
|
|
else:
|
|
rejected_image_tuples.append((image, average_distance, verified_ratio))
|
|
|
|
# Emit one frame of detections per image (in input order) so the
|
|
# bboxes line up with the corresponding `images` when drawn. A frame
|
|
# is an empty list when no face was detected for that image.
|
|
bboxes.append([image_bbox] if image_bbox is not None else [])
|
|
|
|
image_counter += 1
|
|
|
|
return result_from_images_with_measurements(verified_image_tuples, sort_by) + result_from_images_with_measurements(rejected_image_tuples, sort_by) + (bboxes,)
|
|
|
|
class DeepfaceAnalyzeNode:
|
|
def __init__(self):
|
|
prepare_deepface_home()
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"actions": (["age", "gender", "race", "emotion"], {
|
|
"default": ["age", "gender", "race", "emotion"],
|
|
"multiselect": True
|
|
}),
|
|
"detector_backend": ([
|
|
"opencv", "ssd", "dlib", "mtcnn", "retinaface", "mediapipe", "yolov8", "yunet", "fastmtcnn"
|
|
], {
|
|
"default": "retinaface",
|
|
}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("JSON",)
|
|
RETURN_NAMES = ("analysis_results",)
|
|
FUNCTION = "run"
|
|
CATEGORY = "deepface"
|
|
|
|
def run(self, images, actions, detector_backend):
|
|
results = []
|
|
for image in images:
|
|
np_image = deepface_image_from_comfy_image(image)
|
|
try:
|
|
analysis = DeepFace.analyze(
|
|
np_image,
|
|
actions=actions,
|
|
detector_backend=detector_backend,
|
|
enforce_detection=False
|
|
)
|
|
results.append(analysis)
|
|
except Exception as e:
|
|
results.append({"error": str(e)})
|
|
return (results,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"AverageList": AverageListNode,
|
|
"DeepfaceExtractFaces": DeepfaceExtractFacesNode,
|
|
"DeepfaceVerify": DeepfaceVerifyNode,
|
|
"DeepfaceAnalyze": DeepfaceAnalyzeNode,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"AverageList": "Average List",
|
|
"DeepfaceExtractFaces": "Deepface Extract Faces",
|
|
"DeepfaceVerify": "Deepface Verify",
|
|
"DeepfaceAnalyze": "Deepface Analyze",
|
|
}
|