-
Notifications
You must be signed in to change notification settings - Fork 300
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add python script for face detection (Photomaker V2 needs)
- Loading branch information
bssrdf
committed
Aug 23, 2024
1 parent
89d71d7
commit 161afa8
Showing
1 changed file
with
78 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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') | ||
|
||
|