Merge branch 'cleanup_and_refactor' into evolve
This commit is contained in:
+3
-31
@@ -1,39 +1,11 @@
|
||||
import sys
|
||||
import os
|
||||
|
||||
# Добавляем путь расширения, чтобы Питон видел наши папки (r_facelib, scripts и т.д.)
|
||||
repo_dir = os.path.dirname(os.path.realpath(__file__))
|
||||
sys.path.insert(0, repo_dir)
|
||||
original_modules = sys.modules.copy()
|
||||
if repo_dir not in sys.path:
|
||||
sys.path.insert(0, repo_dir)
|
||||
|
||||
# Place aside existing modules if using a1111 web ui
|
||||
modules_used = [
|
||||
"modules",
|
||||
"modules.images",
|
||||
"modules.processing",
|
||||
"modules.scripts_postprocessing",
|
||||
"modules.scripts",
|
||||
"modules.shared",
|
||||
]
|
||||
original_webui_modules = {}
|
||||
for module in modules_used:
|
||||
if module in sys.modules:
|
||||
original_webui_modules[module] = sys.modules.pop(module)
|
||||
|
||||
# Proceed with node setup
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
|
||||
# Clean up imports
|
||||
# Remove repo directory from path
|
||||
sys.path.remove(repo_dir)
|
||||
# Remove any new modules
|
||||
modules_to_remove = []
|
||||
for module in sys.modules:
|
||||
if module not in original_modules and not module.startswith("google.protobuf") and not module.startswith("onnx") and not module.startswith("cv2"):
|
||||
modules_to_remove.append(module)
|
||||
for module in modules_to_remove:
|
||||
del sys.modules[module]
|
||||
|
||||
# Restore original modules
|
||||
sys.modules.update(original_webui_modules)
|
||||
|
||||
@@ -7,6 +7,7 @@ import torchvision.transforms as T
|
||||
from torchvision.transforms.functional import normalize
|
||||
from torchvision.ops import masks_to_boxes
|
||||
|
||||
import onnxruntime
|
||||
import numpy as np
|
||||
import cv2
|
||||
import math
|
||||
@@ -14,12 +15,12 @@ from typing import List
|
||||
from PIL import Image
|
||||
import io
|
||||
from scipy import stats
|
||||
from insightface.app.common import Face
|
||||
from reactor_core.face_objects import Face
|
||||
from segment_anything import sam_model_registry
|
||||
|
||||
from modules.processing import ProcessingImg2Img
|
||||
from modules.shared import state
|
||||
# from comfy_extras.chainner_models import model_loading
|
||||
from r_modules.processing import ProcessingImg2Img
|
||||
from r_modules.shared import state
|
||||
|
||||
import comfy.model_management as model_management
|
||||
import comfy.utils
|
||||
import folder_paths
|
||||
@@ -55,7 +56,7 @@ from reactor_utils import (
|
||||
progress_bar,
|
||||
progress_bar_reset
|
||||
)
|
||||
from reactor_patcher import apply_patch
|
||||
|
||||
from r_facelib.utils.face_restoration_helper import FaceRestoreHelper
|
||||
from r_basicsr.utils.registry import ARCH_REGISTRY
|
||||
import scripts.r_archs.codeformer_arch
|
||||
@@ -91,6 +92,17 @@ if "ultralytics" not in folder_paths.folder_names_and_paths:
|
||||
if "sams" not in folder_paths.folder_names_and_paths:
|
||||
add_folder_path_and_extensions("sams", [os.path.join(models_dir, "sams")], folder_paths.supported_pt_extensions)
|
||||
|
||||
def apply_log_level(console_log_level):
|
||||
if console_log_level == 0:
|
||||
logger.setLevel(logging.WARNING)
|
||||
onnxruntime.set_default_logger_severity(3) # Убивает ворнинги ORT
|
||||
elif console_log_level == 1:
|
||||
logger.setLevel(logging.STATUS)
|
||||
onnxruntime.set_default_logger_severity(3)
|
||||
elif console_log_level == 2:
|
||||
logger.setLevel(logging.INFO)
|
||||
onnxruntime.set_default_logger_severity(0)
|
||||
|
||||
def get_facemodels():
|
||||
models_path = os.path.join(FACE_MODELS_PATH, "*")
|
||||
models = glob.glob(models_path)
|
||||
@@ -423,7 +435,7 @@ class reactor:
|
||||
if faces_order is None:
|
||||
faces_order = self.faces_order
|
||||
|
||||
apply_patch(console_log_level)
|
||||
apply_log_level(console_log_level)
|
||||
|
||||
if not enabled:
|
||||
return (input_image,face_model)
|
||||
@@ -666,7 +678,7 @@ class ReActorWeight:
|
||||
|
||||
images_list: List[Image.Image] = []
|
||||
|
||||
apply_patch(1)
|
||||
apply_log_level(0)
|
||||
|
||||
if len(images) > 0:
|
||||
|
||||
@@ -762,7 +774,7 @@ class BuildFaceModel:
|
||||
faces = []
|
||||
embeddings = []
|
||||
|
||||
apply_patch(1)
|
||||
apply_log_level(0)
|
||||
|
||||
if images is not None:
|
||||
images_list: List[Image.Image] = batch_tensor_to_pil(images)
|
||||
@@ -861,7 +873,7 @@ class SaveFaceModel:
|
||||
if save_mode and image is not None:
|
||||
source = tensor_to_pil(image)
|
||||
source = cv2.cvtColor(np.array(source), cv2.COLOR_RGB2BGR)
|
||||
apply_patch(1)
|
||||
apply_log_level(0)
|
||||
logger.status("Building Face Model...")
|
||||
face_model_raw = analyze_faces(source, det_size)
|
||||
if len(face_model_raw) == 0:
|
||||
@@ -1640,6 +1652,7 @@ class ReActorFaceBoost:
|
||||
}
|
||||
return (face_boost, )
|
||||
|
||||
|
||||
class ReActorUnload:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
@@ -1658,6 +1671,59 @@ class ReActorUnload:
|
||||
return (trigger,)
|
||||
|
||||
|
||||
class ReActorFaceSimilarity:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"image1": ("IMAGE",),
|
||||
"image2": ("IMAGE",),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("FLOAT", "STRING")
|
||||
RETURN_NAMES = ("similarity_float", "similarity_text")
|
||||
FUNCTION = "compare_faces"
|
||||
CATEGORY = "🌌 ReActor"
|
||||
|
||||
def compare_faces(self, image1, image2):
|
||||
|
||||
apply_log_level(0)
|
||||
|
||||
# 1. Конвертируем тензоры ComfyUI в формат OpenCV (BGR)
|
||||
img1_cv = 255. * image1[0].cpu().numpy()
|
||||
img1_cv = cv2.cvtColor(img1_cv.astype(np.uint8), cv2.COLOR_RGB2BGR)
|
||||
|
||||
img2_cv = 255. * image2[0].cpu().numpy()
|
||||
img2_cv = cv2.cvtColor(img2_cv.astype(np.uint8), cv2.COLOR_RGB2BGR)
|
||||
|
||||
# 2. Ищем лица через
|
||||
faces1 = analyze_faces(img1_cv, det_size=(640, 640))
|
||||
faces2 = analyze_faces(img2_cv, det_size=(640, 640))
|
||||
|
||||
# 3. Защита от отсутствия лиц
|
||||
if not faces1 or not faces2:
|
||||
return (0.0, "Face not found in one or both images")
|
||||
|
||||
# Берем первые найденные лица
|
||||
face1 = faces1[0]
|
||||
face2 = faces2[0]
|
||||
|
||||
# 4. Вычисляем косинусное сходство (Cosine Similarity)
|
||||
emb1 = face1.normed_embedding
|
||||
emb2 = face2.normed_embedding
|
||||
|
||||
# Скалярное произведение нормализованных векторов
|
||||
similarity = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
|
||||
|
||||
# 5. Форматируем результат
|
||||
sim_float = float(similarity)
|
||||
sim_float = max(0.0, min(1.0, sim_float))
|
||||
sim_text = f"{sim_float * 100:.2f}%"
|
||||
|
||||
return (sim_float, sim_text)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
# --- MAIN NODES ---
|
||||
"ReActorFaceSwap": reactor,
|
||||
@@ -1674,6 +1740,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
# --- Additional Nodes ---
|
||||
"ReActorRestoreFace": RestoreFace,
|
||||
"ReActorRestoreFaceAdvanced": RestoreFaceAdvanced,
|
||||
"ReActorFaceSimilarity": ReActorFaceSimilarity,
|
||||
"ReActorImageDublicator": ImageDublicator,
|
||||
"ImageRGBA2RGB": ImageRGBA2RGB,
|
||||
"ReActorUnload": ReActorUnload,
|
||||
@@ -1695,6 +1762,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
# --- Additional Nodes ---
|
||||
"ReActorRestoreFace": "Restore Face 🌌 ReActor",
|
||||
"ReActorRestoreFaceAdvanced": "Restore Face Advanced 🌌 ReActor",
|
||||
"ReActorFaceSimilarity": "Face Similarity 🌌 ReActor",
|
||||
"ReActorImageDublicator": "Image Dublicator (List) 🌌 ReActor",
|
||||
"ImageRGBA2RGB": "Convert RGBA to RGB 🌌 ReActor",
|
||||
"ReActorUnload": "Unload ReActor Models 🌌 ReActor",
|
||||
|
||||
@@ -6,7 +6,7 @@ import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
from torchvision.models._utils import IntermediateLayerGetter as IntermediateLayerGetter
|
||||
|
||||
from modules import shared
|
||||
from r_modules import shared
|
||||
|
||||
from r_facelib.detection.align_trans import get_reference_facial_points, warp_and_crop_face
|
||||
from r_facelib.detection.retinaface.retinaface_net import FPN, SSH, MobileNetV1, make_bbox_head, make_class_head, make_landmark_head
|
||||
|
||||
@@ -0,0 +1,103 @@
|
||||
import os
|
||||
import zipfile
|
||||
from reactor_utils import download # Твоя функция скачивания из utils
|
||||
from .face_objects import Face
|
||||
from .inswap import SCRFD, ArcFaceONNX, Attribute, Landmark
|
||||
from scripts.reactor_logger import logger
|
||||
|
||||
class ReActorFaceAnalysis:
|
||||
"""
|
||||
Главный класс-оркестратор.
|
||||
Берет картинку, находит лица, определяет пол/возраст и вычисляет эмбеддинги.
|
||||
"""
|
||||
def __init__(self, name="buffalo_l", root="./models/insightface", providers=None):
|
||||
self.name = name
|
||||
self.root = root
|
||||
self.providers = providers or ["CPUExecutionProvider"]
|
||||
self.models = {}
|
||||
|
||||
model_dir = os.path.join(root, "models", name)
|
||||
os.makedirs(model_dir, exist_ok=True)
|
||||
|
||||
det_file = os.path.join(model_dir, "det_10g.onnx")
|
||||
rec_file = os.path.join(model_dir, "w600k_r50.onnx")
|
||||
attr_file = os.path.join(model_dir, "genderage.onnx")
|
||||
lmk2d_file = os.path.join(model_dir, "2d106det.onnx")
|
||||
lmk3d_file = os.path.join(model_dir, "1k3d68.onnx")
|
||||
|
||||
# Если файлов нет - качаем архив
|
||||
if not (os.path.exists(det_file) and os.path.exists(rec_file) and os.path.exists(attr_file)):
|
||||
zip_url = "https://huggingface.co/datasets/Gourieff/ReActor/resolve/main/models/buffalo_l.zip"
|
||||
zip_path = os.path.join(model_dir, f"{name}.zip")
|
||||
|
||||
logger.status(f"Downloading {name} models archive...")
|
||||
download(zip_url, zip_path, f"{name}.zip")
|
||||
|
||||
logger.status(f"Extracting {name} models...")
|
||||
try:
|
||||
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
|
||||
zip_ref.extractall(model_dir)
|
||||
logger.status("Extraction completed!")
|
||||
except zipfile.BadZipFile:
|
||||
logger.error("Downloaded zip file is corrupted. Please try again.")
|
||||
finally:
|
||||
# В любом случае пытаемся удалить архив, чтобы не занимать место
|
||||
if os.path.exists(zip_path):
|
||||
os.remove(zip_path)
|
||||
|
||||
# Инициализируем только те модели, которые физически есть в папке
|
||||
if os.path.exists(det_file):
|
||||
self.models["detection"] = SCRFD(det_file, providers=self.providers)
|
||||
if os.path.exists(rec_file):
|
||||
self.models["recognition"] = ArcFaceONNX(rec_file, providers=self.providers)
|
||||
if os.path.exists(attr_file):
|
||||
self.models["attribute"] = Attribute(attr_file, providers=self.providers)
|
||||
if os.path.exists(lmk2d_file):
|
||||
self.models["landmark_2d"] = Landmark(lmk2d_file, providers=self.providers)
|
||||
if os.path.exists(lmk3d_file):
|
||||
self.models["landmark_3d"] = Landmark(lmk3d_file, providers=self.providers)
|
||||
|
||||
if "detection" not in self.models:
|
||||
raise FileNotFoundError(
|
||||
f"Detection model (det_10g.onnx) not found at {det_file}. "
|
||||
"Please ensure the buffalo_l models are downloaded and extracted properly."
|
||||
)
|
||||
|
||||
def prepare(self, ctx_id=0, det_size=(640, 640), det_thresh=0.5):
|
||||
self.det_size = det_size
|
||||
self.det_thresh = det_thresh
|
||||
|
||||
def get(self, img, max_num=0):
|
||||
bboxes, kpss = self.models["detection"].detect(
|
||||
img,
|
||||
det_thresh=self.det_thresh,
|
||||
input_size=self.det_size,
|
||||
max_num=max_num
|
||||
)
|
||||
|
||||
if bboxes.shape[0] == 0:
|
||||
return []
|
||||
|
||||
ret = []
|
||||
for i in range(bboxes.shape[0]):
|
||||
bbox = bboxes[i, 0:4]
|
||||
det_score = bboxes[i, 4]
|
||||
kps = kpss[i] if kpss is not None else None
|
||||
|
||||
face = Face(bbox=bbox, kps=kps, det_score=det_score)
|
||||
|
||||
if "attribute" in self.models:
|
||||
self.models["attribute"].get(img, face)
|
||||
|
||||
if "recognition" in self.models:
|
||||
self.models["recognition"].get(img, face)
|
||||
|
||||
if "landmark_2d" in self.models:
|
||||
self.models["landmark_2d"].get(img, face)
|
||||
|
||||
if "landmark_3d" in self.models:
|
||||
self.models["landmark_3d"].get(img, face)
|
||||
|
||||
ret.append(face)
|
||||
|
||||
return ret
|
||||
@@ -0,0 +1,82 @@
|
||||
import numpy as np
|
||||
import onnxruntime as ort
|
||||
import os
|
||||
|
||||
class Face(dict):
|
||||
"""
|
||||
Класс-хранилище данных о лице.
|
||||
Наследуется от dict для полной обратной совместимости с кодом,
|
||||
который ожидает доступ по ключу (например, face['bbox']).
|
||||
"""
|
||||
def __init__(self, d=None, **kwargs):
|
||||
if d is None:
|
||||
d = {}
|
||||
if kwargs:
|
||||
d.update(**kwargs)
|
||||
for k, v in d.items():
|
||||
setattr(self, k, v)
|
||||
# Инициализируем родительский словарь
|
||||
super().__init__(d)
|
||||
|
||||
def __setattr__(self, name, value):
|
||||
# Если это массив, делаем копию, чтобы избежать багов с мутацией по ссылке
|
||||
if isinstance(value, (list, tuple)):
|
||||
value = [x for x in value]
|
||||
elif isinstance(value, np.ndarray):
|
||||
value = value.copy()
|
||||
|
||||
super().__setattr__(name, value)
|
||||
super().__setitem__(name, value)
|
||||
|
||||
def __setitem__(self, key, value):
|
||||
super().__setitem__(key, value)
|
||||
super().__setattr__(key, value)
|
||||
|
||||
@property
|
||||
def sex(self):
|
||||
"""Возвращает 'M' или 'F' на основе числового значения gender"""
|
||||
gender = self.get('gender', None)
|
||||
if gender is None:
|
||||
return None
|
||||
return 'M' if gender == 1 else 'F'
|
||||
|
||||
@property
|
||||
def normed_embedding(self):
|
||||
"""Автоматически нормализует эмбеддинг для ArcFace / INSwapper"""
|
||||
embedding = self.get('embedding', None)
|
||||
if embedding is None:
|
||||
return None
|
||||
norm = np.linalg.norm(embedding)
|
||||
if norm == 0:
|
||||
return embedding
|
||||
return embedding / norm
|
||||
|
||||
|
||||
class BaseONNXModel:
|
||||
"""
|
||||
Базовый класс для всех моделей (Детектор, ArcFace, INSwapper).
|
||||
Берет на себя рутину по открытию сессий и чтению входов/выходов.
|
||||
"""
|
||||
def __init__(self, model_file, providers=None):
|
||||
self.model_file = model_file
|
||||
self.providers = providers or ["CPUExecutionProvider"]
|
||||
|
||||
if not os.path.exists(self.model_file):
|
||||
raise FileNotFoundError(f"Model file not found: {self.model_file}")
|
||||
|
||||
self.session = ort.InferenceSession(self.model_file, providers=self.providers)
|
||||
|
||||
# Получаем параметры входов
|
||||
self.inputs = self.session.get_inputs()
|
||||
self.input_names = [inp.name for inp in self.inputs]
|
||||
|
||||
# Обычно нас интересует шейп первого входа (например, батч, каналы, высота, ширина)
|
||||
self.input_shape = self.inputs[0].shape
|
||||
|
||||
# Получаем параметры выходов
|
||||
self.outputs = self.session.get_outputs()
|
||||
self.output_names = [out.name for out in self.outputs]
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
"""Этот метод будет переопределен в классах-наследниках"""
|
||||
raise NotImplementedError("Forward method must be implemented by subclasses.")
|
||||
@@ -0,0 +1,181 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
from .face_objects import BaseONNXModel
|
||||
|
||||
class HyperSwapper(BaseONNXModel):
|
||||
"""Класс для работы с моделями семейства Hyperswap"""
|
||||
|
||||
def __init__(self, model_file, providers=None):
|
||||
super().__init__(model_file, providers)
|
||||
|
||||
# Функция для получения 5 ключевых точек из объекта Face
|
||||
def get_landmarks_5(self, face):
|
||||
if hasattr(face, 'landmark_5') and face.landmark_5 is not None:
|
||||
return face.landmark_5
|
||||
elif hasattr(face, 'kps') and face.kps is not None:
|
||||
return face.kps
|
||||
elif hasattr(face, 'landmark') and face.landmark is not None:
|
||||
if face.landmark.shape[0] >= 68:
|
||||
idxs = [36, 45, 30, 48, 54]
|
||||
return face.landmark[idxs]
|
||||
return None
|
||||
|
||||
# Функция для вычисления аффинного преобразования
|
||||
def get_affine_transform(self, src_pts, dst_pts):
|
||||
M, _ = cv2.estimateAffinePartial2D(src_pts, dst_pts)
|
||||
return M
|
||||
|
||||
# Создаём градиентную маску овальной формы без обрезки
|
||||
def create_gradient_mask(self, crop_size=256):
|
||||
# 1. Создаём пустую маску (все пиксели = 0)
|
||||
mask = np.zeros((crop_size, crop_size), dtype=np.float32)
|
||||
|
||||
# 2. Определяем центр и размеры эллипса
|
||||
center = (crop_size // 2, crop_size // 2)
|
||||
axes = (int(crop_size * 0.35), int(crop_size * 0.4))
|
||||
|
||||
# 3. Рисуем эллипс (заполняем белым цветом, значение=1.0)
|
||||
cv2.ellipse(
|
||||
mask, # Массив для рисования
|
||||
center, # Центр эллипса
|
||||
axes, # Полуоси (ширина, высота)
|
||||
angle=0, # Угол поворота
|
||||
startAngle=0, # Начальный угол дуги
|
||||
endAngle=360, # Конечный угол дуги (360 = полный эллипс)
|
||||
color=1.0, # Значение для заполнения (белый = 1.0)
|
||||
thickness=-1 # -1 = заполнить всю область эллипса
|
||||
)
|
||||
|
||||
# 4. Применяем размытие для плавных краёв
|
||||
blur_ksize = 15 # Нечётное число, чтобы ядро было симметричным
|
||||
mask = cv2.GaussianBlur(mask, (blur_ksize, blur_ksize), 0)
|
||||
|
||||
# 5. Ограничим значения в диапазоне [0, 1]
|
||||
mask = np.clip(mask, 0, 1)
|
||||
|
||||
return mask
|
||||
|
||||
def paste_back(self, target_img, swapped_face, M, crop_size=256):
|
||||
|
||||
# 1. Создание мягкой маски (Эрозия + Размытие)
|
||||
mask = self.create_gradient_mask(crop_size)
|
||||
|
||||
# Преобразуем в трехканальную маску
|
||||
mask_3c = np.stack([mask] * 3, axis=2)
|
||||
|
||||
# 2. Получаем размеры целевого изображения
|
||||
h, w = target_img.shape[:2]
|
||||
|
||||
# 3. Нормализация swapped_face к float32 [0,1] для warp
|
||||
swapped_face_norm = swapped_face.astype(np.float32) / 255.0
|
||||
mask_norm = mask_3c.astype(np.float32) # Маска уже [0,1]
|
||||
|
||||
# 4. Обратное преобразование (WARP_INVERSE_MAP) для лица И маски
|
||||
# Используем BORDER_CONSTANT с borderValue=0.5 (серый, чтобы избежать синих/зеленых артефактов)
|
||||
warped_face = cv2.warpAffine(
|
||||
swapped_face_norm,
|
||||
M,
|
||||
(w, h),
|
||||
flags=cv2.INTER_LANCZOS4 | cv2.WARP_INVERSE_MAP,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0.5
|
||||
)
|
||||
|
||||
# Для маски (INTER_CUBIC — плавные границы)
|
||||
warped_mask = cv2.warpAffine(
|
||||
mask_norm,
|
||||
M,
|
||||
(w, h),
|
||||
flags=cv2.INTER_CUBIC | cv2.WARP_INVERSE_MAP,
|
||||
borderMode=cv2.BORDER_CONSTANT,
|
||||
borderValue=0.0 # Маска: 0 за пределами
|
||||
)
|
||||
|
||||
# 5. Обработка после warp: Clip, NaN fix
|
||||
warped_face = np.clip(warped_face, 0, 1) # Убираем отрицательные
|
||||
warped_face = np.nan_to_num(warped_face, nan=0.5) # NaN -> серый
|
||||
|
||||
warped_mask = np.clip(warped_mask, 0, 1)
|
||||
warped_mask = np.nan_to_num(warped_mask, nan=0.0)
|
||||
|
||||
# 6. Дополнительное размытие для устранения артефактов
|
||||
warped_mask = cv2.GaussianBlur(warped_mask, (3, 3), 0)
|
||||
|
||||
# 7. Плавное наложение в float32
|
||||
target_float = target_img.astype(np.float32) / 255.0
|
||||
result_float = target_float * (1.0 - warped_mask) + warped_face * warped_mask
|
||||
|
||||
# 8. Обратная нормализация к uint8
|
||||
result = (result_float * 255).clip(0, 255).astype(np.uint8)
|
||||
|
||||
return result
|
||||
|
||||
def visualize_points(self, img, points, color=(0, 255, 0)):
|
||||
img = img.copy()
|
||||
for p in points:
|
||||
cv2.circle(img, tuple(p.astype(int)), 3, color, -1)
|
||||
|
||||
# Итоговая функция run_hyperswap (get) с аффинным преобразованием
|
||||
def get(self, img, target_face, source_face, paste_back=True):
|
||||
# 1. Подготовка эмбеддинга
|
||||
source_embedding = source_face.normed_embedding.reshape(1, -1).astype(np.float32)
|
||||
|
||||
# 2. Получаем 5 точек target
|
||||
target_landmarks_5 = self.get_landmarks_5(target_face)
|
||||
# self.visualize_points(img, target_landmarks_5, (0, 255, 0)) # не для продакшена
|
||||
|
||||
if target_landmarks_5 is None:
|
||||
return img if paste_back else (None, None)
|
||||
|
||||
# 3. Определение эталонных точек для выравнивания 256x256 (FFHQ Alignment)
|
||||
std_landmarks_256 = np.array([
|
||||
[ 84.87, 105.94], # Левый глаз
|
||||
[171.13, 105.94], # Правый глаз
|
||||
[128.00, 146.66], # Кончик носа
|
||||
[ 96.95, 188.64], # Левый уголок рта
|
||||
[159.05, 188.64] # Правый уголок рта
|
||||
], dtype=np.float32)
|
||||
|
||||
# Вычисляем аффинную матрицу
|
||||
M = self.get_affine_transform(target_landmarks_5.astype(np.float32), std_landmarks_256)
|
||||
|
||||
# Применяем аффинное преобразование с новой матрицей M
|
||||
crop = cv2.warpAffine(img, M, (256, 256), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REFLECT)
|
||||
|
||||
# 4. Преобразуем crop для модели
|
||||
crop_input = crop[:, :, ::-1].astype(np.float32) / 255.0 # RGB -> [0,1]
|
||||
crop_input = (crop_input - 0.5) / 0.5 # Нормализация
|
||||
crop_input = crop_input.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32)
|
||||
|
||||
# 5. Инференс
|
||||
try:
|
||||
output = self.session.run(None, {'source': source_embedding, 'target': crop_input})[0][0]
|
||||
except:
|
||||
return img if paste_back else (None, None)
|
||||
|
||||
if isinstance(output, np.ndarray):
|
||||
# устранение NaN и бесконечностей
|
||||
output = np.nan_to_num(output, nan=0.0, posinf=255.0, neginf=0.0)
|
||||
|
||||
# если диапазон похож на [-1,1] → нормализуем в [0,255]
|
||||
if output.min() < 0.0 or output.max() <= 1.5:
|
||||
output = ((output + 1.0) / 2.0 * 255.0)
|
||||
# жёсткое ограничение диапазона и тип для OpenCV
|
||||
output = np.clip(output, 0, 255).astype(np.uint8).copy()
|
||||
|
||||
# защита от повторного использования буфера (inplace CPU bug)
|
||||
try:
|
||||
output.setflags(write=True)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 6. Обратная нормализация
|
||||
output = output.transpose(1, 2, 0) # CHW -> HWC
|
||||
output = output[:, :, ::-1] # BGR -> RGB
|
||||
|
||||
# 7. Возвращаем результат в зависимости от флага paste_back
|
||||
if not paste_back:
|
||||
return output, M # Возвращаем только кроп лица (256x256) и матрицу M
|
||||
|
||||
# Если нужна полная вклейка в исходное изображение:
|
||||
return self.paste_back(img, output, M, crop_size=256)
|
||||
@@ -0,0 +1,440 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
from .face_objects import BaseONNXModel
|
||||
from .meanshape_68 import MEANSHAPE_68
|
||||
|
||||
# --- Математика для 3D позы ---
|
||||
|
||||
def estimate_affine_matrix_3d23d(X, Y):
|
||||
''' Вычисляет аффинную матрицу трансформации 3D -> 3D '''
|
||||
X_homo = np.hstack((X, np.ones([X.shape[0], 1])))
|
||||
P = np.linalg.lstsq(X_homo, Y, rcond=None)[0].T
|
||||
return P
|
||||
|
||||
def P2sRt(P):
|
||||
''' Разбивает матрицу проекции '''
|
||||
t1 = np.linalg.norm(P[:,0])
|
||||
t2 = np.linalg.norm(P[:,1])
|
||||
t3 = np.linalg.norm(P[:,2])
|
||||
s = (t1 + t2 + t3) / 3.0
|
||||
P1 = P / s
|
||||
R = P1[:, 0:3]
|
||||
t = P1[:, 3]
|
||||
return s, R, t
|
||||
|
||||
def matrix2angle(R):
|
||||
''' Превращает матрицу поворота в углы Эйлера (pitch, yaw, roll) '''
|
||||
if R[2,0] != 1 and R[2,0] != -1:
|
||||
pitch = -np.arcsin(R[2,0])
|
||||
yaw = np.arctan2(R[2,1]/np.cos(pitch), R[2,2]/np.cos(pitch))
|
||||
roll = np.arctan2(R[1,0]/np.cos(pitch), R[0,0]/np.cos(pitch))
|
||||
else:
|
||||
yaw = 0
|
||||
if R[2,0] == -1:
|
||||
pitch = np.pi/2
|
||||
roll = yaw + np.arctan2(R[0,1], R[0,2])
|
||||
else:
|
||||
pitch = -np.pi/2
|
||||
roll = -yaw + np.arctan2(-R[0,1], -R[0,2])
|
||||
return pitch, yaw, roll
|
||||
|
||||
# --- Вспомогательные функции ---
|
||||
|
||||
def distance2bbox(points, distance, max_shape=None):
|
||||
x1 = points[:, 0] - distance[:, 0]
|
||||
y1 = points[:, 1] - distance[:, 1]
|
||||
x2 = points[:, 0] + distance[:, 2]
|
||||
y2 = points[:, 1] + distance[:, 3]
|
||||
if max_shape is not None:
|
||||
x1 = np.clip(x1, 0, max_shape[1])
|
||||
y1 = np.clip(y1, 0, max_shape[0])
|
||||
x2 = np.clip(x2, 0, max_shape[1])
|
||||
y2 = np.clip(y2, 0, max_shape[0])
|
||||
return np.stack([x1, y1, x2, y2], axis=-1)
|
||||
|
||||
def distance2kps(points, distance, max_shape=None):
|
||||
preds = []
|
||||
for i in range(0, distance.shape[1], 2):
|
||||
px = points[:, i%2] + distance[:, i]
|
||||
py = points[:, i%2+1] + distance[:, i+1]
|
||||
if max_shape is not None:
|
||||
px = np.clip(px, 0, max_shape[1])
|
||||
py = np.clip(py, 0, max_shape[0])
|
||||
preds.append(px)
|
||||
preds.append(py)
|
||||
return np.stack(preds, axis=-1)
|
||||
|
||||
# Стандартные точки FFHQ/ArcFace для матрицы трансформации 112x112
|
||||
ARCFACE_STD_POINTS = np.array([
|
||||
[38.2946, 51.6963],
|
||||
[73.5318, 51.5014],
|
||||
[56.0252, 71.7366],
|
||||
[41.5493, 92.3655],
|
||||
[70.7299, 92.2041]
|
||||
], dtype=np.float32)
|
||||
|
||||
def norm_crop(img, landmark, image_size=112):
|
||||
"""Выравнивает и обрезает лицо (для ArcFace)"""
|
||||
M, _ = cv2.estimateAffinePartial2D(landmark, ARCFACE_STD_POINTS)
|
||||
warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)
|
||||
return warped
|
||||
|
||||
# --- Модели ---
|
||||
|
||||
class SCRFD(BaseONNXModel):
|
||||
"""Детектор лиц (находит bbox и 5 ключевых точек)"""
|
||||
def __init__(self, model_file, providers=None):
|
||||
super().__init__(model_file, providers)
|
||||
self.batched = len(self.outputs[0].shape) == 3
|
||||
self.input_mean = 127.5
|
||||
self.input_std = 128.0
|
||||
self.use_kps = len(self.outputs) in [9, 15]
|
||||
self.fmc = 5 if len(self.outputs) in [10, 15] else 3
|
||||
self._feat_stride_fpn = [8, 16, 32, 64, 128][:self.fmc]
|
||||
self._num_anchors = 2 if self.fmc == 3 else 1
|
||||
self.center_cache = {}
|
||||
|
||||
def forward(self, img, threshold):
|
||||
scores_list, bboxes_list, kpss_list = [], [], []
|
||||
input_size = tuple(img.shape[0:2][::-1])
|
||||
blob = cv2.dnn.blobFromImage(img, 1.0/self.input_std, input_size,
|
||||
(self.input_mean, self.input_mean, self.input_mean), swapRB=True)
|
||||
|
||||
net_outs = self.session.run(self.output_names, {self.input_names[0]: blob})
|
||||
input_height, input_width = blob.shape[2], blob.shape[3]
|
||||
|
||||
for idx, stride in enumerate(self._feat_stride_fpn):
|
||||
if self.batched:
|
||||
scores = net_outs[idx][0]
|
||||
bbox_preds = net_outs[idx + self.fmc][0] * stride
|
||||
if self.use_kps:
|
||||
kps_preds = net_outs[idx + self.fmc * 2][0] * stride
|
||||
else:
|
||||
scores = net_outs[idx]
|
||||
bbox_preds = net_outs[idx + self.fmc] * stride
|
||||
if self.use_kps:
|
||||
kps_preds = net_outs[idx + self.fmc * 2] * stride
|
||||
|
||||
height, width = input_height // stride, input_width // stride
|
||||
key = (height, width, stride)
|
||||
|
||||
if key in self.center_cache:
|
||||
anchor_centers = self.center_cache[key]
|
||||
else:
|
||||
anchor_centers = np.stack(np.mgrid[:height, :width][::-1], axis=-1).astype(np.float32)
|
||||
anchor_centers = (anchor_centers * stride).reshape((-1, 2))
|
||||
if self._num_anchors > 1:
|
||||
anchor_centers = np.stack([anchor_centers] * self._num_anchors, axis=1).reshape((-1, 2))
|
||||
if len(self.center_cache) < 100:
|
||||
self.center_cache[key] = anchor_centers
|
||||
|
||||
pos_inds = np.where(scores >= threshold)[0]
|
||||
bboxes = distance2bbox(anchor_centers, bbox_preds)
|
||||
|
||||
scores_list.append(scores[pos_inds])
|
||||
bboxes_list.append(bboxes[pos_inds])
|
||||
|
||||
if self.use_kps:
|
||||
kpss = distance2kps(anchor_centers, kps_preds)
|
||||
kpss = kpss.reshape((kpss.shape[0], -1, 2))
|
||||
kpss_list.append(kpss[pos_inds])
|
||||
|
||||
return scores_list, bboxes_list, kpss_list
|
||||
|
||||
def detect(self, img, det_thresh=0.5, input_size=(640, 640), max_num=0):
|
||||
im_ratio = float(img.shape[0]) / img.shape[1]
|
||||
model_ratio = float(input_size[1]) / input_size[0]
|
||||
|
||||
if im_ratio > model_ratio:
|
||||
new_height = input_size[1]
|
||||
new_width = int(new_height / im_ratio)
|
||||
else:
|
||||
new_width = input_size[0]
|
||||
new_height = int(new_width * im_ratio)
|
||||
|
||||
det_scale = float(new_height) / img.shape[0]
|
||||
resized_img = cv2.resize(img, (new_width, new_height))
|
||||
det_img = np.zeros((input_size[1], input_size[0], 3), dtype=np.uint8)
|
||||
det_img[:new_height, :new_width, :] = resized_img
|
||||
|
||||
scores_list, bboxes_list, kpss_list = self.forward(det_img, det_thresh)
|
||||
|
||||
scores = np.vstack(scores_list).ravel()
|
||||
order = scores.argsort()[::-1]
|
||||
bboxes = np.vstack(bboxes_list) / det_scale
|
||||
|
||||
if self.use_kps:
|
||||
kpss = np.vstack(kpss_list) / det_scale
|
||||
|
||||
pre_det = np.hstack((bboxes, scores[:, None])).astype(np.float32, copy=False)
|
||||
pre_det = pre_det[order, :]
|
||||
keep = self.nms(pre_det)
|
||||
|
||||
det = pre_det[keep, :]
|
||||
kpss = kpss[order, :, :][keep, :, :] if self.use_kps else None
|
||||
|
||||
if max_num > 0 and det.shape[0] > max_num:
|
||||
area = (det[:, 2] - det[:, 0]) * (det[:, 3] - det[:, 1])
|
||||
img_center = img.shape[0] // 2, img.shape[1] // 2
|
||||
offsets = np.vstack([
|
||||
(det[:, 0] + det[:, 2]) / 2 - img_center[1],
|
||||
(det[:, 1] + det[:, 3]) / 2 - img_center[0]
|
||||
])
|
||||
offset_dist_squared = np.sum(np.power(offsets, 2.0), 0)
|
||||
values = area - offset_dist_squared * 2.0
|
||||
bindex = np.argsort(values)[::-1][:max_num]
|
||||
det = det[bindex, :]
|
||||
if kpss is not None:
|
||||
kpss = kpss[bindex, :]
|
||||
|
||||
return det, kpss
|
||||
|
||||
def nms(self, dets, nms_thresh=0.4):
|
||||
x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4]
|
||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
order = scores.argsort()[::-1]
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
||||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
||||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
||||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
||||
w = np.maximum(0.0, xx2 - xx1 + 1)
|
||||
h = np.maximum(0.0, yy2 - yy1 + 1)
|
||||
inter = w * h
|
||||
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
||||
inds = np.where(ovr <= nms_thresh)[0]
|
||||
order = order[inds + 1]
|
||||
return keep
|
||||
|
||||
|
||||
class ArcFaceONNX(BaseONNXModel):
|
||||
"""Распознаватель лиц (выдает вектор/эмбеддинг)"""
|
||||
def __init__(self, model_file, providers=None):
|
||||
super().__init__(model_file, providers)
|
||||
self.input_mean = 127.5
|
||||
self.input_std = 127.5
|
||||
self.input_size = tuple(self.input_shape[2:4][::-1])
|
||||
|
||||
def get(self, img, face):
|
||||
aimg = norm_crop(img, landmark=face.kps, image_size=self.input_size[0])
|
||||
blob = cv2.dnn.blobFromImage(aimg, 1.0 / self.input_std, self.input_size,
|
||||
(self.input_mean, self.input_mean, self.input_mean), swapRB=True)
|
||||
net_out = self.session.run(self.output_names, {self.input_names[0]: blob})[0]
|
||||
face.embedding = net_out.flatten()
|
||||
return face.embedding
|
||||
|
||||
|
||||
class Attribute(BaseONNXModel):
|
||||
"""Анализатор атрибутов (выдает пол и возраст)"""
|
||||
def __init__(self, model_file, providers=None):
|
||||
super().__init__(model_file, providers)
|
||||
self.input_mean = 127.5
|
||||
self.input_std = 128.0
|
||||
self.input_size = tuple(self.input_shape[2:4][::-1])
|
||||
|
||||
def get(self, img, face):
|
||||
bbox = face.bbox
|
||||
w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1])
|
||||
center = ((bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2)
|
||||
_scale = self.input_size[0] / (max(w, h) * 1.5)
|
||||
|
||||
# Простая трансформация для Attribute (не требует 5 точек, только центр и масштаб)
|
||||
M = np.array([
|
||||
[_scale, 0, self.input_size[0] * 0.5 - center[0] * _scale],
|
||||
[0, _scale, self.input_size[1] * 0.5 - center[1] * _scale]
|
||||
], dtype=np.float32)
|
||||
|
||||
aimg = cv2.warpAffine(img, M, self.input_size, borderValue=0.0)
|
||||
blob = cv2.dnn.blobFromImage(aimg, 1.0 / self.input_std, self.input_size,
|
||||
(self.input_mean, self.input_mean, self.input_mean), swapRB=True)
|
||||
|
||||
pred = self.session.run(self.output_names, {self.input_names[0]: blob})[0][0]
|
||||
|
||||
# Получаем гендер и возраст
|
||||
gender = int(np.argmax(pred[:2]))
|
||||
age = int(np.round(pred[2] * 100))
|
||||
face.gender = gender
|
||||
face.age = age
|
||||
return gender, age
|
||||
|
||||
class INSwapper(BaseONNXModel):
|
||||
"""Свопер лиц (модели inswapper_128, reswapper)"""
|
||||
def __init__(self, model_file, providers=None):
|
||||
super().__init__(model_file, providers)
|
||||
self.input_mean = 0.0
|
||||
self.input_std = 255.0
|
||||
self.input_size = tuple(self.input_shape[2:4][::-1])
|
||||
|
||||
# Хак для экономии памяти: импортируем onnx только здесь,
|
||||
# читаем нужную матрицу emap и сразу выгружаем тяжелую модель из RAM.
|
||||
import onnx
|
||||
from onnx import numpy_helper
|
||||
model = onnx.load(self.model_file, load_external_data=False)
|
||||
self.emap = numpy_helper.to_array(model.graph.initializer[-1])
|
||||
del model
|
||||
|
||||
def get(self, img, target_face, source_face, paste_back=True):
|
||||
# 1. Идеальное позиционирование (1 в 1 как в оригинальном C++ Insightface)
|
||||
# ВАЖНО: Insightface центрирует лицо для INSwapper ТОЛЬКО по оси X!
|
||||
# По оси Y оно остается прижатым выше, сохраняя оригинальные пропорции.
|
||||
ratio = float(self.input_size[0]) / 128.0
|
||||
diff_x = 8.0 * ratio
|
||||
|
||||
src_pts = ARCFACE_STD_POINTS.copy() * ratio
|
||||
src_pts[:, 0] += diff_x # Смещаем ТОЛЬКО координаты X!
|
||||
|
||||
# 2. Вычисляем аффинную матрицу родным методом OpenCV
|
||||
M, _ = cv2.estimateAffinePartial2D(target_face.kps, src_pts)
|
||||
|
||||
# 3. Кропаем и выравниваем лицо
|
||||
aimg = cv2.warpAffine(img, M, self.input_size, borderValue=0.0)
|
||||
blob = cv2.dnn.blobFromImage(aimg, 1.0 / self.input_std, self.input_size,
|
||||
(self.input_mean, self.input_mean, self.input_mean), swapRB=True)
|
||||
|
||||
# 4. Подготавливаем эмбеддинг донора
|
||||
latent = source_face.normed_embedding.reshape((1, -1))
|
||||
latent = np.dot(latent, self.emap)
|
||||
latent /= np.linalg.norm(latent)
|
||||
|
||||
# 5. Инференс
|
||||
pred = self.session.run(self.output_names, {
|
||||
self.input_names[0]: blob,
|
||||
self.input_names[1]: latent.astype(np.float32)
|
||||
})[0]
|
||||
|
||||
img_fake = pred.transpose((0, 2, 3, 1))[0]
|
||||
bgr_fake = np.clip(255 * img_fake, 0, 255).astype(np.uint8)[:, :, ::-1]
|
||||
|
||||
if not paste_back:
|
||||
return bgr_fake, M
|
||||
|
||||
# 6. Обратная вклейка (Paste Back)
|
||||
target_img = img
|
||||
fake_diff = bgr_fake.astype(np.float32) - aimg.astype(np.float32)
|
||||
fake_diff = np.abs(fake_diff).mean(axis=2)
|
||||
|
||||
# Обрезаем края
|
||||
fake_diff[:2, :] = 0
|
||||
fake_diff[-2:, :] = 0
|
||||
fake_diff[:, :2] = 0
|
||||
fake_diff[:, -2:] = 0
|
||||
|
||||
IM = cv2.invertAffineTransform(M)
|
||||
img_white = np.full((aimg.shape[0], aimg.shape[1]), 255, dtype=np.float32)
|
||||
|
||||
# Возвращаем в исходную перспективу
|
||||
bgr_fake_warped = cv2.warpAffine(bgr_fake, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0)
|
||||
img_white_warped = cv2.warpAffine(img_white, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0)
|
||||
fake_diff_warped = cv2.warpAffine(fake_diff, IM, (target_img.shape[1], target_img.shape[0]), borderValue=0.0)
|
||||
|
||||
img_white_warped[img_white_warped > 20] = 255
|
||||
fthresh = 10
|
||||
fake_diff_warped[fake_diff_warped < fthresh] = 0
|
||||
fake_diff_warped[fake_diff_warped >= fthresh] = 255
|
||||
|
||||
img_mask = img_white_warped
|
||||
mask_h_inds, mask_w_inds = np.where(img_mask == 255)
|
||||
|
||||
# Защита от пустой маски
|
||||
if len(mask_h_inds) > 0 and len(mask_w_inds) > 0:
|
||||
mask_h = np.max(mask_h_inds) - np.min(mask_h_inds)
|
||||
mask_w = np.max(mask_w_inds) - np.min(mask_w_inds)
|
||||
mask_size = int(np.sqrt(mask_h * mask_w))
|
||||
|
||||
k = max(mask_size // 10, 10)
|
||||
kernel = np.ones((k, k), np.uint8)
|
||||
img_mask = cv2.erode(img_mask, kernel, iterations=1)
|
||||
|
||||
kernel = np.ones((2, 2), np.uint8)
|
||||
fake_diff_warped = cv2.dilate(fake_diff_warped, kernel, iterations=1)
|
||||
|
||||
k = max(mask_size // 20, 5)
|
||||
blur_size = (k * 2 + 1, k * 2 + 1)
|
||||
img_mask = cv2.GaussianBlur(img_mask, blur_size, 0)
|
||||
|
||||
k = 5
|
||||
blur_size = (k * 2 + 1, k * 2 + 1)
|
||||
fake_diff_warped = cv2.GaussianBlur(fake_diff_warped, blur_size, 0)
|
||||
|
||||
img_mask /= 255.0
|
||||
img_mask = np.reshape(img_mask, [img_mask.shape[0], img_mask.shape[1], 1])
|
||||
|
||||
fake_merged = img_mask * bgr_fake_warped + (1.0 - img_mask) * target_img.astype(np.float32)
|
||||
return fake_merged.astype(np.uint8)
|
||||
|
||||
class Landmark(BaseONNXModel):
|
||||
"""Извлекает 106 (2D) или 68 (3D) точек лица"""
|
||||
def __init__(self, model_file, providers=None):
|
||||
super().__init__(model_file, providers)
|
||||
self.input_mean = 127.5
|
||||
self.input_std = 128.0
|
||||
self.input_size = tuple(self.input_shape[2:4][::-1])
|
||||
|
||||
output_shape = self.outputs[0].shape
|
||||
# Определяем, какая это модель (3D или 2D) по размеру выхода
|
||||
if output_shape[1] == 3309:
|
||||
self.lmk_dim = 3
|
||||
self.lmk_num = 68
|
||||
self.taskname = 'landmark_3d_68'
|
||||
else:
|
||||
self.lmk_dim = 2
|
||||
self.lmk_num = output_shape[1] // self.lmk_dim
|
||||
self.taskname = f'landmark_2d_{self.lmk_num}'
|
||||
|
||||
def get(self, img, face):
|
||||
bbox = face.bbox
|
||||
w, h = (bbox[2] - bbox[0]), (bbox[3] - bbox[1])
|
||||
center = ((bbox[2] + bbox[0]) / 2, (bbox[3] + bbox[1]) / 2)
|
||||
_scale = self.input_size[0] / (max(w, h) * 1.5)
|
||||
|
||||
# Матрица трансформации (выравнивание по центру bbox)
|
||||
M = np.array([
|
||||
[_scale, 0, self.input_size[0] * 0.5 - center[0] * _scale],
|
||||
[0, _scale, self.input_size[1] * 0.5 - center[1] * _scale]
|
||||
], dtype=np.float32)
|
||||
|
||||
aimg = cv2.warpAffine(img, M, self.input_size, borderValue=0.0)
|
||||
blob = cv2.dnn.blobFromImage(aimg, 1.0 / self.input_std, self.input_size,
|
||||
(self.input_mean, self.input_mean, self.input_mean), swapRB=True)
|
||||
|
||||
pred = self.session.run(self.output_names, {self.input_names[0]: blob})[0][0]
|
||||
|
||||
if pred.shape[0] >= 3000:
|
||||
pred = pred.reshape((-1, 3))
|
||||
else:
|
||||
pred = pred.reshape((-1, 2))
|
||||
|
||||
if self.lmk_num < pred.shape[0]:
|
||||
pred = pred[-self.lmk_num:, :]
|
||||
|
||||
# Денормализация точек в размер модели
|
||||
pred[:, 0:2] += 1
|
||||
pred[:, 0:2] *= (self.input_size[0] // 2)
|
||||
if pred.shape[1] == 3:
|
||||
pred[:, 2] *= (self.input_size[0] // 2)
|
||||
|
||||
# Обратная трансформация точек на оригинальное изображение
|
||||
IM = cv2.invertAffineTransform(M)
|
||||
pred_xy = pred[:, 0:2]
|
||||
pred_xy = np.hstack((pred_xy, np.ones((pred_xy.shape[0], 1)))) # Добавляем гомогенную координату
|
||||
pred_xy = np.dot(IM, pred_xy.T).T
|
||||
|
||||
if pred.shape[1] == 3:
|
||||
pred = np.hstack((pred_xy, pred[:, 2:3])) # Возвращаем Z
|
||||
else:
|
||||
pred = pred_xy
|
||||
|
||||
# Сохраняем в объект Face под правильным именем
|
||||
setattr(face, self.taskname, pred)
|
||||
|
||||
# Честный расчет 3D позы
|
||||
if self.taskname == 'landmark_3d_68':
|
||||
P = estimate_affine_matrix_3d23d(MEANSHAPE_68, pred)
|
||||
_, R, _ = P2sRt(P)
|
||||
rx, ry, rz = matrix2angle(R)
|
||||
face.pose = np.array([rx, ry, rz], dtype=np.float32)
|
||||
|
||||
return pred
|
||||
@@ -0,0 +1,348 @@
|
||||
import numpy as np
|
||||
|
||||
# эталонная 3D-модель "усредненного" человеческого лица (68 точек в 3D-пространстве)
|
||||
|
||||
matrix = [
|
||||
[
|
||||
-0.6266950368881226,
|
||||
-0.2926996946334839,
|
||||
-0.3140018582344055
|
||||
],
|
||||
[
|
||||
-0.5996649265289307,
|
||||
-0.12250272184610367,
|
||||
-0.2924409508705139
|
||||
],
|
||||
[
|
||||
-0.571026086807251,
|
||||
0.05118739977478981,
|
||||
-0.2549465000629425
|
||||
],
|
||||
[
|
||||
-0.5338566899299622,
|
||||
0.21284838020801544,
|
||||
-0.1866333782672882
|
||||
],
|
||||
[
|
||||
-0.4797332286834717,
|
||||
0.3506459593772888,
|
||||
-0.047699932008981705
|
||||
],
|
||||
[
|
||||
-0.3957556486129761,
|
||||
0.4465116560459137,
|
||||
0.07328663021326065
|
||||
],
|
||||
[
|
||||
-0.29880842566490173,
|
||||
0.5106605887413025,
|
||||
0.17977793514728546
|
||||
],
|
||||
[
|
||||
-0.18838681280612946,
|
||||
0.554440438747406,
|
||||
0.3166600465774536
|
||||
],
|
||||
[
|
||||
0.0014708322705700994,
|
||||
0.584439218044281,
|
||||
0.38841578364372253
|
||||
],
|
||||
[
|
||||
0.19099053740501404,
|
||||
0.5517070889472961,
|
||||
0.31433114409446716
|
||||
],
|
||||
[
|
||||
0.3269285559654236,
|
||||
0.48957568407058716,
|
||||
0.16839952766895294
|
||||
],
|
||||
[
|
||||
0.4400261342525482,
|
||||
0.4023579955101013,
|
||||
0.03596242889761925
|
||||
],
|
||||
[
|
||||
0.5068787336349487,
|
||||
0.3116249442100525,
|
||||
-0.09476063400506973
|
||||
],
|
||||
[
|
||||
0.540894627571106,
|
||||
0.20452618598937988,
|
||||
-0.20267142355442047
|
||||
],
|
||||
[
|
||||
0.574118435382843,
|
||||
0.04570329561829567,
|
||||
-0.2841764986515045
|
||||
],
|
||||
[
|
||||
0.5991416573524475,
|
||||
-0.14585931599140167,
|
||||
-0.296495646238327
|
||||
],
|
||||
[
|
||||
0.6275436282157898,
|
||||
-0.30774804949760437,
|
||||
-0.30199483036994934
|
||||
],
|
||||
[
|
||||
-0.47466811537742615,
|
||||
-0.4376046359539032,
|
||||
0.23648710548877716
|
||||
],
|
||||
[
|
||||
-0.4166599214076996,
|
||||
-0.4717560112476349,
|
||||
0.3159925639629364
|
||||
],
|
||||
[
|
||||
-0.3475387990474701,
|
||||
-0.4840780198574066,
|
||||
0.366113543510437
|
||||
],
|
||||
[
|
||||
-0.26063987612724304,
|
||||
-0.4763909876346588,
|
||||
0.3992317318916321
|
||||
],
|
||||
[
|
||||
-0.16712301969528198,
|
||||
-0.45777544379234314,
|
||||
0.4166091978549957
|
||||
],
|
||||
[
|
||||
0.1231740415096283,
|
||||
-0.45874011516571045,
|
||||
0.4250619411468506
|
||||
],
|
||||
[
|
||||
0.2063615769147873,
|
||||
-0.4804193377494812,
|
||||
0.41578391194343567
|
||||
],
|
||||
[
|
||||
0.28667303919792175,
|
||||
-0.4901494085788727,
|
||||
0.3919537663459778
|
||||
],
|
||||
[
|
||||
0.3623969256877899,
|
||||
-0.4768601655960083,
|
||||
0.3527863025665283
|
||||
],
|
||||
[
|
||||
0.425568163394928,
|
||||
-0.4500580430030823,
|
||||
0.2953187823295593
|
||||
],
|
||||
[
|
||||
-0.007627937477082014,
|
||||
-0.3230886459350586,
|
||||
0.46194377541542053
|
||||
],
|
||||
[
|
||||
-0.007876846939325333,
|
||||
-0.255738765001297,
|
||||
0.5104694366455078
|
||||
],
|
||||
[
|
||||
-0.007687545381486416,
|
||||
-0.19917990267276764,
|
||||
0.5525456666946411
|
||||
],
|
||||
[
|
||||
-0.007345120422542095,
|
||||
-0.14261691272258759,
|
||||
0.5986667275428772
|
||||
],
|
||||
[
|
||||
-0.144961878657341,
|
||||
0.033126939088106155,
|
||||
0.4196909964084625
|
||||
],
|
||||
[
|
||||
-0.08434253931045532,
|
||||
0.03127360716462135,
|
||||
0.4731740951538086
|
||||
],
|
||||
[
|
||||
-0.005499685648828745,
|
||||
0.03975145146250725,
|
||||
0.5146695971488953
|
||||
],
|
||||
[
|
||||
0.06346292793750763,
|
||||
0.04613516479730606,
|
||||
0.4792240560054779
|
||||
],
|
||||
[
|
||||
0.13398391008377075,
|
||||
0.02204025723040104,
|
||||
0.41907867789268494
|
||||
],
|
||||
[
|
||||
-0.38675081729888916,
|
||||
-0.3133975565433502,
|
||||
0.259631872177124
|
||||
],
|
||||
[
|
||||
-0.3166607916355133,
|
||||
-0.35007160902023315,
|
||||
0.3285270035266876
|
||||
],
|
||||
[
|
||||
-0.2341379076242447,
|
||||
-0.35491427779197693,
|
||||
0.3334933817386627
|
||||
],
|
||||
[
|
||||
-0.15516234934329987,
|
||||
-0.31524932384490967,
|
||||
0.3143281042575836
|
||||
],
|
||||
[
|
||||
-0.23092176020145416,
|
||||
-0.28427040576934814,
|
||||
0.3255828619003296
|
||||
],
|
||||
[
|
||||
-0.3175090253353119,
|
||||
-0.28516536951065063,
|
||||
0.3098825216293335
|
||||
],
|
||||
[
|
||||
0.13895708322525024,
|
||||
-0.30982404947280884,
|
||||
0.3182835876941681
|
||||
],
|
||||
[
|
||||
0.21945592761039734,
|
||||
-0.35319215059280396,
|
||||
0.33802759647369385
|
||||
],
|
||||
[
|
||||
0.30174651741981506,
|
||||
-0.349665105342865,
|
||||
0.333102285861969
|
||||
],
|
||||
[
|
||||
0.37665316462516785,
|
||||
-0.31351813673973083,
|
||||
0.26322856545448303
|
||||
],
|
||||
[
|
||||
0.2966947853565216,
|
||||
-0.2871439754962921,
|
||||
0.3220130503177643
|
||||
],
|
||||
[
|
||||
0.21462441980838776,
|
||||
-0.2905276417732239,
|
||||
0.33124178647994995
|
||||
],
|
||||
[
|
||||
-0.20143844187259674,
|
||||
0.23736143112182617,
|
||||
0.37953662872314453
|
||||
],
|
||||
[
|
||||
-0.13732077181339264,
|
||||
0.18578563630580902,
|
||||
0.46525245904922485
|
||||
],
|
||||
[
|
||||
-0.07648587226867676,
|
||||
0.15119342505931854,
|
||||
0.5035716891288757
|
||||
],
|
||||
[
|
||||
-0.002535885199904442,
|
||||
0.16872699558734894,
|
||||
0.5164376497268677
|
||||
],
|
||||
[
|
||||
0.06442102789878845,
|
||||
0.15088020265102386,
|
||||
0.5045241117477417
|
||||
],
|
||||
[
|
||||
0.126465305685997,
|
||||
0.17947596311569214,
|
||||
0.4685916602611542
|
||||
],
|
||||
[
|
||||
0.21824784576892853,
|
||||
0.23899227380752563,
|
||||
0.37567368149757385
|
||||
],
|
||||
[
|
||||
0.13288259506225586,
|
||||
0.28392839431762695,
|
||||
0.4400508999824524
|
||||
],
|
||||
[
|
||||
0.06802233308553696,
|
||||
0.2973543107509613,
|
||||
0.4774041771888733
|
||||
],
|
||||
[
|
||||
-0.0004690653004217893,
|
||||
0.30004069209098816,
|
||||
0.4871094226837158
|
||||
],
|
||||
[
|
||||
-0.06934267282485962,
|
||||
0.29696860909461975,
|
||||
0.4805404841899872
|
||||
],
|
||||
[
|
||||
-0.14252015948295593,
|
||||
0.2743033170700073,
|
||||
0.43808019161224365
|
||||
],
|
||||
[
|
||||
-0.1781347393989563,
|
||||
0.23059049248695374,
|
||||
0.3963589668273926
|
||||
],
|
||||
[
|
||||
-0.07403063774108887,
|
||||
0.2147187888622284,
|
||||
0.4653262794017792
|
||||
],
|
||||
[
|
||||
-0.002636224264279008,
|
||||
0.21414154767990112,
|
||||
0.4832296073436737
|
||||
],
|
||||
[
|
||||
0.05981616675853729,
|
||||
0.21076396107673645,
|
||||
0.47224411368370056
|
||||
],
|
||||
[
|
||||
0.16690002381801605,
|
||||
0.23127099871635437,
|
||||
0.396789014339447
|
||||
],
|
||||
[
|
||||
0.059809282422065735,
|
||||
0.22376428544521332,
|
||||
0.46641337871551514
|
||||
],
|
||||
[
|
||||
-0.0014343614457175136,
|
||||
0.2257590889930725,
|
||||
0.4752439558506012
|
||||
],
|
||||
[
|
||||
-0.07522077113389969,
|
||||
0.23065608739852905,
|
||||
0.4671475291252136
|
||||
]
|
||||
]
|
||||
|
||||
MEANSHAPE_68 = np.array(matrix, dtype=np.float32)
|
||||
@@ -1,161 +0,0 @@
|
||||
import os.path as osp
|
||||
import glob
|
||||
import logging
|
||||
import insightface
|
||||
from insightface.model_zoo.model_zoo import ModelRouter, PickableInferenceSession
|
||||
from insightface.model_zoo.retinaface import RetinaFace
|
||||
from insightface.model_zoo.landmark import Landmark
|
||||
from insightface.model_zoo.attribute import Attribute
|
||||
from insightface.model_zoo.inswapper import INSwapper
|
||||
from insightface.model_zoo.arcface_onnx import ArcFaceONNX
|
||||
from insightface.app import FaceAnalysis
|
||||
from insightface.utils import DEFAULT_MP_NAME, ensure_available
|
||||
from insightface.model_zoo import model_zoo
|
||||
import onnxruntime
|
||||
import onnx
|
||||
from onnx import numpy_helper
|
||||
from scripts.reactor_logger import logger
|
||||
|
||||
|
||||
def patched_get_model_log(self, **kwargs):
|
||||
session = PickableInferenceSession(self.onnx_file, **kwargs)
|
||||
print(f'Applied providers: {session._providers}, with options: {session._provider_options}')
|
||||
inputs = session.get_inputs()
|
||||
input_cfg = inputs[0]
|
||||
input_shape = input_cfg.shape
|
||||
outputs = session.get_outputs()
|
||||
|
||||
if len(outputs) >= 5:
|
||||
return RetinaFace(model_file=self.onnx_file, session=session)
|
||||
elif input_shape[2] == 192 and input_shape[3] == 192:
|
||||
return Landmark(model_file=self.onnx_file, session=session)
|
||||
elif input_shape[2] == 96 and input_shape[3] == 96:
|
||||
return Attribute(model_file=self.onnx_file, session=session)
|
||||
elif len(inputs) == 2 and input_shape[2] == 128 and input_shape[3] == 128:
|
||||
return INSwapper(model_file=self.onnx_file, session=session)
|
||||
elif len(inputs) == 2 and input_shape[2] == 256 and input_shape[3] == 256:
|
||||
return INSwapper(model_file=self.onnx_file, session=session)
|
||||
elif input_shape[2] == input_shape[3] and input_shape[2] >= 112 and input_shape[2] % 16 == 0:
|
||||
return ArcFaceONNX(model_file=self.onnx_file, session=session)
|
||||
else:
|
||||
return None
|
||||
|
||||
def patched_get_model(self, **kwargs):
|
||||
session = PickableInferenceSession(self.onnx_file, **kwargs)
|
||||
inputs = session.get_inputs()
|
||||
input_cfg = inputs[0]
|
||||
input_shape = input_cfg.shape
|
||||
outputs = session.get_outputs()
|
||||
|
||||
if len(outputs) >= 5:
|
||||
return RetinaFace(model_file=self.onnx_file, session=session)
|
||||
elif input_shape[2] == 192 and input_shape[3] == 192:
|
||||
return Landmark(model_file=self.onnx_file, session=session)
|
||||
elif input_shape[2] == 96 and input_shape[3] == 96:
|
||||
return Attribute(model_file=self.onnx_file, session=session)
|
||||
elif len(inputs) == 2 and input_shape[2] == 128 and input_shape[3] == 128:
|
||||
return INSwapper(model_file=self.onnx_file, session=session)
|
||||
elif len(inputs) == 2 and input_shape[2] == 256 and input_shape[3] == 256:
|
||||
return INSwapper(model_file=self.onnx_file, session=session)
|
||||
elif input_shape[2] == input_shape[3] and input_shape[2] >= 112 and input_shape[2] % 16 == 0:
|
||||
return ArcFaceONNX(model_file=self.onnx_file, session=session)
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def patched_faceanalysis_init(self, name=DEFAULT_MP_NAME, root='~/.insightface', allowed_modules=None, **kwargs):
|
||||
onnxruntime.set_default_logger_severity(3)
|
||||
self.models = {}
|
||||
self.model_dir = ensure_available('models', name, root=root)
|
||||
onnx_files = glob.glob(osp.join(self.model_dir, '*.onnx'))
|
||||
onnx_files = sorted(onnx_files)
|
||||
for onnx_file in onnx_files:
|
||||
model = model_zoo.get_model(onnx_file, **kwargs)
|
||||
if model is None:
|
||||
print('model not recognized:', onnx_file)
|
||||
elif allowed_modules is not None and model.taskname not in allowed_modules:
|
||||
print('model ignore:', onnx_file, model.taskname)
|
||||
del model
|
||||
elif model.taskname not in self.models and (allowed_modules is None or model.taskname in allowed_modules):
|
||||
self.models[model.taskname] = model
|
||||
else:
|
||||
print('duplicated model task type, ignore:', onnx_file, model.taskname)
|
||||
del model
|
||||
assert 'detection' in self.models
|
||||
self.det_model = self.models['detection']
|
||||
|
||||
|
||||
def patched_faceanalysis_prepare(self, ctx_id, det_thresh=0.5, det_size=(640, 640)):
|
||||
self.det_thresh = det_thresh
|
||||
assert det_size is not None
|
||||
self.det_size = det_size
|
||||
for taskname, model in self.models.items():
|
||||
if taskname == 'detection':
|
||||
model.prepare(ctx_id, input_size=det_size, det_thresh=det_thresh)
|
||||
else:
|
||||
model.prepare(ctx_id)
|
||||
|
||||
|
||||
def patched_inswapper_init(self, model_file=None, session=None):
|
||||
self.model_file = model_file
|
||||
self.session = session
|
||||
model = onnx.load(self.model_file)
|
||||
graph = model.graph
|
||||
self.emap = numpy_helper.to_array(graph.initializer[-1])
|
||||
self.input_mean = 0.0
|
||||
self.input_std = 255.0
|
||||
if self.session is None:
|
||||
self.session = onnxruntime.InferenceSession(self.model_file, None)
|
||||
inputs = self.session.get_inputs()
|
||||
self.input_names = []
|
||||
for inp in inputs:
|
||||
self.input_names.append(inp.name)
|
||||
outputs = self.session.get_outputs()
|
||||
output_names = []
|
||||
for out in outputs:
|
||||
output_names.append(out.name)
|
||||
self.output_names = output_names
|
||||
assert len(self.output_names) == 1
|
||||
input_cfg = inputs[0]
|
||||
input_shape = input_cfg.shape
|
||||
self.input_shape = input_shape
|
||||
self.input_size = tuple(input_shape[2:4][::-1])
|
||||
|
||||
|
||||
def pathced_retinaface_prepare(self, ctx_id, **kwargs):
|
||||
if ctx_id<0:
|
||||
self.session.set_providers(['CPUExecutionProvider'])
|
||||
nms_thresh = kwargs.get('nms_thresh', None)
|
||||
if nms_thresh is not None:
|
||||
self.nms_thresh = nms_thresh
|
||||
det_thresh = kwargs.get('det_thresh', None)
|
||||
if det_thresh is not None:
|
||||
self.det_thresh = det_thresh
|
||||
input_size = kwargs.get('input_size', None)
|
||||
if input_size is not None and self.input_size is None:
|
||||
self.input_size = input_size
|
||||
|
||||
|
||||
def patch_insightface(get_model, faceanalysis_init, faceanalysis_prepare, inswapper_init, retinaface_prepare):
|
||||
insightface.model_zoo.model_zoo.ModelRouter.get_model = get_model
|
||||
insightface.app.FaceAnalysis.__init__ = faceanalysis_init
|
||||
insightface.app.FaceAnalysis.prepare = faceanalysis_prepare
|
||||
insightface.model_zoo.inswapper.INSwapper.__init__ = inswapper_init
|
||||
insightface.model_zoo.retinaface.RetinaFace.prepare = retinaface_prepare
|
||||
|
||||
|
||||
# original_functions = [ModelRouter.get_model, FaceAnalysis.__init__, FaceAnalysis.prepare, INSwapper.__init__, RetinaFace.prepare]
|
||||
original_functions = [patched_get_model_log, FaceAnalysis.__init__, FaceAnalysis.prepare, INSwapper.__init__, RetinaFace.prepare]
|
||||
patched_functions = [patched_get_model, patched_faceanalysis_init, patched_faceanalysis_prepare, patched_inswapper_init, pathced_retinaface_prepare]
|
||||
|
||||
|
||||
def apply_patch(console_log_level):
|
||||
if console_log_level == 0:
|
||||
patch_insightface(*patched_functions)
|
||||
logger.setLevel(logging.WARNING)
|
||||
elif console_log_level == 1:
|
||||
patch_insightface(*patched_functions)
|
||||
logger.setLevel(logging.STATUS)
|
||||
elif console_log_level == 2:
|
||||
patch_insightface(*original_functions)
|
||||
logger.setLevel(logging.INFO)
|
||||
+1
-1
@@ -7,7 +7,7 @@ import cv2
|
||||
import math
|
||||
import logging
|
||||
import hashlib
|
||||
from insightface.app.common import Face
|
||||
from reactor_core.face_objects import Face
|
||||
from safetensors.torch import save_file, safe_open
|
||||
from tqdm import tqdm
|
||||
import urllib.request
|
||||
|
||||
+1
-2
@@ -1,7 +1,6 @@
|
||||
albumentations>=1.4.16
|
||||
insightface==0.7.3
|
||||
onnx>=1.14.0
|
||||
opencv-python>=4.7.0.72
|
||||
numpy==1.26.4
|
||||
numpy
|
||||
segment_anything
|
||||
ultralytics
|
||||
|
||||
@@ -2,14 +2,14 @@ import os, glob
|
||||
|
||||
from PIL import Image
|
||||
|
||||
import modules.scripts as scripts
|
||||
import r_modules.scripts as scripts
|
||||
# from modules.upscaler import Upscaler, UpscalerData
|
||||
from modules import scripts, scripts_postprocessing
|
||||
from modules.processing import (
|
||||
from r_modules import scripts, scripts_postprocessing
|
||||
from r_modules.processing import (
|
||||
Processing,
|
||||
ProcessingImg2Img,
|
||||
)
|
||||
from modules.shared import state
|
||||
from r_modules.shared import state
|
||||
from scripts.reactor_logger import logger
|
||||
from scripts.reactor_swapper import (
|
||||
swap_face,
|
||||
|
||||
@@ -2,7 +2,7 @@ import logging
|
||||
import copy
|
||||
import sys
|
||||
|
||||
from modules import shared
|
||||
from r_modules import shared
|
||||
from reactor_utils import addLoggingLevel
|
||||
|
||||
|
||||
|
||||
@@ -30,7 +30,7 @@ def ensure_nsfw_model(nsfwdet_model_path):
|
||||
MODEL_EXISTS = True if downloaded == 3 else False
|
||||
return MODEL_EXISTS
|
||||
|
||||
SCORE = 0.969
|
||||
SCORE = 0.979
|
||||
|
||||
logging.getLogger("transformers").setLevel(logging.ERROR)
|
||||
|
||||
|
||||
+20
-183
@@ -6,15 +6,15 @@ import cv2
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
import onnxruntime as ort
|
||||
|
||||
import insightface
|
||||
from insightface.app.common import Face
|
||||
from reactor_core.analyzer import ReActorFaceAnalysis
|
||||
from reactor_core.face_objects import Face
|
||||
from reactor_core.inswap import INSwapper
|
||||
from reactor_core.hyperswap import HyperSwapper
|
||||
import torch
|
||||
|
||||
import folder_paths
|
||||
import comfy.model_management as model_management
|
||||
from modules.shared import state
|
||||
from r_modules.shared import state
|
||||
|
||||
from scripts.reactor_logger import logger
|
||||
from reactor_utils import (
|
||||
@@ -80,11 +80,6 @@ TARGET_IMAGE_LIST_HASH = []
|
||||
|
||||
def unload_model(model):
|
||||
if model is not None:
|
||||
# check if model has unload method
|
||||
# if "unload" in model:
|
||||
# model.unload()
|
||||
# if "model_unload" in model:
|
||||
# model.model_unload()
|
||||
del model
|
||||
return None
|
||||
|
||||
@@ -102,7 +97,7 @@ def getAnalysisModel(det_size = (640, 640)):
|
||||
global ANALYSIS_MODELS
|
||||
ANALYSIS_MODEL = ANALYSIS_MODELS[str(det_size[0])]
|
||||
if ANALYSIS_MODEL is None:
|
||||
ANALYSIS_MODEL = insightface.app.FaceAnalysis(
|
||||
ANALYSIS_MODEL = ReActorFaceAnalysis(
|
||||
name="buffalo_l", providers=providers, root=insightface_path
|
||||
)
|
||||
ANALYSIS_MODEL.prepare(ctx_id=0, det_size=det_size)
|
||||
@@ -116,163 +111,16 @@ def getFaceSwapModel(model_path: str):
|
||||
FS_MODEL = unload_model(FS_MODEL)
|
||||
|
||||
model_filename = os.path.basename(model_path)
|
||||
if "hyperswap" in model_filename.lower():
|
||||
if "hyperswap" in model_filename.lower(): # Если это Hyperswap
|
||||
model_path = os.path.join(folder_paths.models_dir, "hyperswap", model_filename)
|
||||
FS_MODEL = ort.InferenceSession(model_path, providers=providers)
|
||||
elif "reswapper" in model_filename.lower():
|
||||
model_path = os.path.join(folder_paths.models_dir, "reswapper", model_filename)
|
||||
FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers)
|
||||
else:
|
||||
FS_MODEL = insightface.model_zoo.get_model(model_path, providers=providers)
|
||||
FS_MODEL = HyperSwapper(model_path, providers=providers)
|
||||
else: # Если это INSwapper / Reswapper
|
||||
if "reswapper" in model_filename.lower():
|
||||
model_path = os.path.join(folder_paths.models_dir, "reswapper", model_filename)
|
||||
FS_MODEL = INSwapper(model_path, providers=providers)
|
||||
|
||||
return FS_MODEL
|
||||
|
||||
|
||||
# Функция для получения 5 ключевых точек из объекта Face
|
||||
def get_landmarks_5(face):
|
||||
if hasattr(face, 'landmark_5') and face.landmark_5 is not None:
|
||||
return face.landmark_5
|
||||
elif hasattr(face, 'kps') and face.kps is not None:
|
||||
return face.kps
|
||||
elif hasattr(face, 'landmark') and face.landmark is not None:
|
||||
if face.landmark.shape[0] >= 68:
|
||||
idxs = [36, 45, 30, 48, 54]
|
||||
return face.landmark[idxs]
|
||||
return None
|
||||
|
||||
# Функция для вычисления аффинного преобразования
|
||||
def get_affine_transform(src_pts, dst_pts):
|
||||
M, _ = cv2.estimateAffinePartial2D(src_pts, dst_pts)
|
||||
return M
|
||||
|
||||
# Создаём градиентную маску овальной формы без обрезки
|
||||
def create_gradient_mask(crop_size=256):
|
||||
# 1. Создаём пустую маску (все пиксели = 0)
|
||||
mask = np.zeros((crop_size, crop_size), dtype=np.float32)
|
||||
|
||||
# 2. Определяем центр и размеры эллипса
|
||||
center = (crop_size // 2, crop_size // 2)
|
||||
axes = (int(crop_size * 0.35), int(crop_size * 0.4))
|
||||
|
||||
# 3. Рисуем эллипс (заполняем белым цветом, значение=1.0)
|
||||
cv2.ellipse(
|
||||
mask, # Массив для рисования
|
||||
center, # Центр эллипса
|
||||
axes, # Полуоси (ширина, высота)
|
||||
angle=0, # Угол поворота
|
||||
startAngle=0, # Начальный угол дуги
|
||||
endAngle=360, # Конечный угол дуги (360 = полный эллипс)
|
||||
color=1.0, # Значение для заполнения (белый = 1.0)
|
||||
thickness=-1 # -1 = заполнить всю область эллипса
|
||||
)
|
||||
|
||||
# 4. Применяем размытие для плавных краёв
|
||||
blur_ksize = 15 # Нечётное число, чтобы ядро было симметричным
|
||||
mask = cv2.GaussianBlur(mask, (blur_ksize, blur_ksize), 0)
|
||||
|
||||
# 5. Ограничим значения в диапазоне [0, 1]
|
||||
mask = np.clip(mask, 0, 1)
|
||||
|
||||
return mask
|
||||
|
||||
def paste_back(target_img, swapped_face, M, crop_size=256):
|
||||
|
||||
# 1. Создание мягкой маски (Эрозия + Размытие)
|
||||
mask = create_gradient_mask(crop_size)
|
||||
|
||||
# Преобразуем в трехканальную маску
|
||||
mask_3c = np.stack([mask] * 3, axis=2)
|
||||
|
||||
# 2. Получаем размеры целевого изображения
|
||||
h, w = target_img.shape[:2]
|
||||
|
||||
# 3. Обратное преобразование (WARP_INVERSE_MAP) для лица И маски
|
||||
# Для лица (INTER_LANCZOS4 — высококачественная интерполяция)
|
||||
inv_face = cv2.warpAffine(
|
||||
swapped_face.astype(np.float32),
|
||||
M,
|
||||
(w, h),
|
||||
flags=cv2.INTER_LANCZOS4 | cv2.WARP_INVERSE_MAP,
|
||||
borderMode=cv2.BORDER_TRANSPARENT
|
||||
)
|
||||
|
||||
# Для маски (INTER_CUBIC — плавные границы)
|
||||
inv_mask = cv2.warpAffine(
|
||||
mask_3c,
|
||||
M,
|
||||
(w, h),
|
||||
flags=cv2.INTER_CUBIC | cv2.WARP_INVERSE_MAP,
|
||||
borderMode=cv2.BORDER_TRANSPARENT
|
||||
)
|
||||
|
||||
# 4. Ограничение значений маски [0, 1]
|
||||
inv_mask = np.clip(inv_mask, 0, 1)
|
||||
|
||||
# 5. Дополнительное размытие для устранения артефактов
|
||||
inv_mask = cv2.GaussianBlur(inv_mask, (3, 3), 0)
|
||||
|
||||
# 6. Плавное наложение
|
||||
target_img_float = target_img.astype(np.float32)
|
||||
inv_face_float = inv_face.astype(np.float32)
|
||||
result = target_img_float * (1.0 - inv_mask) + inv_face_float * inv_mask
|
||||
|
||||
# 7. Ограничение результата [0, 255]
|
||||
result = np.clip(result, 0, 255).astype(np.uint8)
|
||||
|
||||
return result
|
||||
|
||||
def visualize_points(img, points, color=(0, 255, 0)):
|
||||
img = img.copy()
|
||||
for p in points:
|
||||
cv2.circle(img, tuple(p.astype(int)), 3, color, -1)
|
||||
|
||||
# Итоговая функция run_hyperswap с аффинным преобразованием
|
||||
def run_hyperswap(session, source_face, target_face, target_img):
|
||||
# 1. Подготовка эмбеддинга
|
||||
source_embedding = source_face.normed_embedding.reshape(1, -1).astype(np.float32)
|
||||
|
||||
# 2. Получаем 5 точек target
|
||||
target_landmarks_5 = get_landmarks_5(target_face)
|
||||
visualize_points(target_img, target_landmarks_5, (0, 255, 0)) # Зеленые точки
|
||||
|
||||
if target_landmarks_5 is None:
|
||||
return None, None
|
||||
|
||||
# 3. Определение эталонных точек для выравнивания 256x256 (FFHQ Alignment)
|
||||
std_landmarks_256 = np.array([
|
||||
[ 84.87, 105.94], # Левый глаз
|
||||
[171.13, 105.94], # Правый глаз
|
||||
[128.00, 146.66], # Кончик носа
|
||||
[ 96.95, 188.64], # Левый уголок рта
|
||||
[159.05, 188.64] # Правый уголок рта
|
||||
], dtype=np.float32)
|
||||
|
||||
# Вычисляем аффинную матрицу
|
||||
M = get_affine_transform(target_landmarks_5.astype(np.float32), std_landmarks_256)
|
||||
|
||||
# Применяем аффинное преобразование с новой матрицей M
|
||||
crop = cv2.warpAffine(target_img, M, (256, 256), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REFLECT)
|
||||
|
||||
# 4. Преобразуем crop для модели
|
||||
crop_input = crop[:, :, ::-1].astype(np.float32) / 255.0 # RGB -> [0,1]
|
||||
crop_input = (crop_input - 0.5) / 0.5 # Нормализация
|
||||
crop_input = crop_input.transpose(2, 0, 1)[np.newaxis, ...].astype(np.float32)
|
||||
|
||||
# 5. Инференс
|
||||
try:
|
||||
output = session.run(None, {'source': source_embedding, 'target': crop_input})[0][0]
|
||||
except:
|
||||
return target_img
|
||||
|
||||
# 6. Обратная нормализация
|
||||
output = (output * 0.5 + 0.5) * 255.0 # [-1..1] -> [0..255]
|
||||
output = np.clip(output, 0, 255).astype(np.uint8)
|
||||
output = output.transpose(1, 2, 0) # CHW -> HWC
|
||||
output = output[:, :, ::-1] # BGR -> RGB
|
||||
|
||||
return output, M # Возвращаем лицо (256x256) и матрицу M
|
||||
|
||||
|
||||
def sort_by_order(face, order: str):
|
||||
if order == "left-right":
|
||||
return sorted(face, key=lambda x: x.bbox[0])
|
||||
@@ -370,7 +218,6 @@ def get_face_single(img_data: np.ndarray, face, face_index=0, det_size=(640, 640
|
||||
try:
|
||||
faces_sorted = sort_by_order(face, order)
|
||||
return faces_sorted[face_index], 0, face_index
|
||||
# return sorted(face, key=lambda x: x.bbox[0])[face_index], 0
|
||||
except IndexError:
|
||||
return None, 0, None
|
||||
|
||||
@@ -471,13 +318,13 @@ def swap_face(
|
||||
logger.status("Using Hashed Target Face(s) Model...")
|
||||
target_faces = TARGET_FACES
|
||||
|
||||
# No use in trying to swap faces if no faces are found, enhancement
|
||||
# No use in trying to swap faces if no faces are found
|
||||
if len(target_faces) == 0:
|
||||
logger.status("Cannot detect any Target, skipping swapping...")
|
||||
return result_image, bbox, swapped_indexes
|
||||
|
||||
if source_img is not None:
|
||||
# separated management of wrong_gender between source and target, enhancement
|
||||
# separated management of wrong_gender between source and target
|
||||
source_face, src_wrong_gender, source_face_index = get_face_single(source_img, source_faces, face_index=source_faces_index[0], gender_source=gender_source, order=faces_order[1])
|
||||
else:
|
||||
# source_face = sorted(source_faces, key=lambda x: x.bbox[0])[source_faces_index[0]]
|
||||
@@ -500,7 +347,7 @@ def swap_face(
|
||||
source_face_idx = 0
|
||||
|
||||
for face_num in faces_index:
|
||||
# No use in trying to swap faces if no further faces are found, enhancement
|
||||
# No use in trying to swap faces if no further faces are found
|
||||
if face_num >= len(target_faces):
|
||||
logger.status("Checked all existing target faces, skipping swapping...")
|
||||
break
|
||||
@@ -513,11 +360,7 @@ def swap_face(
|
||||
target_face, wrong_gender, target_face_index = get_face_single(target_img, target_faces, face_index=face_num, gender_target=gender_target, order=faces_order[0])
|
||||
if target_face is not None and wrong_gender == 0:
|
||||
logger.status(f"Swapping...")
|
||||
if "hyperswap" in model:
|
||||
swapped_face_256, M = run_hyperswap(face_swapper, source_face, target_face, result)
|
||||
if swapped_face_256 is not None:
|
||||
result = paste_back(result, swapped_face_256, M, crop_size=256)
|
||||
elif face_boost_enabled:
|
||||
if face_boost_enabled and "hyperswap" not in model:
|
||||
logger.status(f"Face Boost is enabled (inswapper/reswapper only)")
|
||||
bgr_fake, M = face_swapper.get(result, target_face, source_face, paste_back=False)
|
||||
bgr_fake, scale = restorer.get_restored_face(bgr_fake, face_restore_model, face_restore_visibility, codeformer_weight, interpolation)
|
||||
@@ -669,7 +512,6 @@ def swap_face_many(
|
||||
|
||||
|
||||
# logger.status(f"Analyzing Target Image {i}...")
|
||||
# target_face = analyze_faces(target_img)
|
||||
if target_face is not None:
|
||||
target_faces.append(target_face)
|
||||
|
||||
@@ -677,13 +519,13 @@ def swap_face_many(
|
||||
|
||||
progress_bar_reset(pbar)
|
||||
|
||||
# No use in trying to swap faces if no faces are found, enhancement
|
||||
# No use in trying to swap faces if no faces are found
|
||||
if len(target_faces) == 0:
|
||||
logger.status("Cannot detect any Target, skipping swapping...")
|
||||
return result_images, bbox, swapped_indexes
|
||||
|
||||
if source_img is not None:
|
||||
# separated management of wrong_gender between source and target, enhancement
|
||||
# separated management of wrong_gender between source and target
|
||||
source_face, src_wrong_gender, source_face_index = get_face_single(source_img, source_faces, face_index=source_faces_index[0], gender_source=gender_source, order=faces_order[1])
|
||||
else:
|
||||
# source_face = sorted(source_faces, key=lambda x: x.bbox[0])[source_faces_index[0]]
|
||||
@@ -709,7 +551,7 @@ def swap_face_many(
|
||||
|
||||
logger.status(f"Swapping...")
|
||||
for face_num in faces_index:
|
||||
# No use in trying to swap faces if no further faces are found, enhancement
|
||||
# No use in trying to swap faces if no further faces are found
|
||||
if face_num >= len(target_faces):
|
||||
logger.status("Checked all existing target faces, skipping swapping...")
|
||||
break
|
||||
@@ -720,16 +562,11 @@ def swap_face_many(
|
||||
|
||||
if source_face is not None and src_wrong_gender == 0:
|
||||
# Reading results to make current face swap on a previous face result
|
||||
# logger.status(f"Swapping...")
|
||||
for i, (target_img, target_face) in enumerate(zip(results, target_faces)):
|
||||
target_face_single, wrong_gender, target_face_index = get_face_single(target_img, target_face, face_index=face_num, gender_target=gender_target, order=faces_order[0])
|
||||
if target_face_single is not None and wrong_gender == 0:
|
||||
result = target_img
|
||||
if "hyperswap" in model:
|
||||
swapped_face_256, M = run_hyperswap(face_swapper, source_face, target_face_single, result)
|
||||
if swapped_face_256 is not None:
|
||||
result = paste_back(result, swapped_face_256, M, crop_size=256)
|
||||
elif face_boost_enabled:
|
||||
if face_boost_enabled and "hyperswap" not in model:
|
||||
logger.status(f"Face Boost is enabled (inswapper/reswapper only)")
|
||||
bgr_fake, M = face_swapper.get(target_img, target_face_single, source_face, paste_back=False)
|
||||
bgr_fake, scale = restorer.get_restored_face(bgr_fake, face_restore_model, face_restore_visibility, codeformer_weight, interpolation)
|
||||
|
||||
Reference in New Issue
Block a user