Skip to content

Commit

Permalink
add python script for face detection (Photomaker V2 needs)
Browse files Browse the repository at this point in the history
  • Loading branch information
bssrdf committed Aug 23, 2024
1 parent 89d71d7 commit 161afa8
Showing 1 changed file with 78 additions and 0 deletions.
78 changes: 78 additions & 0 deletions face_detect.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
# -*- coding: utf-8 -*-
"""
Created on Sun Aug 18 15:52:46 2024
@author: merli
"""
import os
import sys

import numpy as np
import torch
from diffusers.utils import load_image
# pip install insightface==0.7.3
from insightface.app import FaceAnalysis
from insightface.data import get_image as ins_get_image
from safetensors.torch import save_file

###
# https://github.com/cubiq/ComfyUI_IPAdapter_plus/issues/165#issue-2055829543
###
class FaceAnalysis2(FaceAnalysis):
# NOTE: allows setting det_size for each detection call.
# the model allows it but the wrapping code from insightface
# doesn't show it, and people end up loading duplicate models
# for different sizes where there is absolutely no need to
def get(self, img, max_num=0, det_size=(640, 640)):
if det_size is not None:
self.det_model.input_size = det_size

return super().get(img, max_num)

def analyze_faces(face_analysis: FaceAnalysis, img_data: np.ndarray, det_size=(640, 640)):
# NOTE: try detect faces, if no faces detected, lower det_size until it does
detection_sizes = [None] + [(size, size) for size in range(640, 256, -64)] + [(256, 256)]

for size in detection_sizes:
faces = face_analysis.get(img_data, det_size=size)
if len(faces) > 0:
return faces

return []

if __name__ == "__main__":
#face_detector = FaceAnalysis2(providers=['CUDAExecutionProvider'], allowed_modules=['detection', 'recognition'])
face_detector = FaceAnalysis2(providers=['CPUExecutionProvider'], allowed_modules=['detection', 'recognition'])
face_detector.prepare(ctx_id=0, det_size=(640, 640))
#input_folder_name = './scarletthead_woman'
#input_folder_name = './Elizabeth'
input_folder_name = sys.argv[1]
image_basename_list = os.listdir(input_folder_name)
image_path_list = sorted([os.path.join(input_folder_name, basename) for basename in image_basename_list])

input_id_images = []
for image_path in image_path_list:
input_id_images.append(load_image(image_path))

id_embed_list = []

for img in input_id_images:
img = np.array(img)
img = img[:, :, ::-1]
faces = analyze_faces(face_detector, img)
if len(faces) > 0:
id_embed_list.append(torch.from_numpy((faces[0]['embedding'])))

if len(id_embed_list) == 0:
raise ValueError(f"No face detected in input image pool")

id_embeds = torch.stack(id_embed_list)

for r in id_embeds:
print(r)
#torch.save(id_embeds, input_folder_name+'/id_embeds.pt');
weights = dict()
weights["id_embeds"] = id_embeds
save_file(weights, input_folder_name+'/id_embeds.safetensors')


0 comments on commit 161afa8

Please sign in to comment.