diff --git a/__init__.py b/__init__.py index eac9eff..b1060bc 100644 --- a/__init__.py +++ b/__init__.py @@ -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) diff --git a/nodes.py b/nodes.py index ce3046e..0904da2 100644 --- a/nodes.py +++ b/nodes.py @@ -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", diff --git a/r_facelib/detection/retinaface/retinaface.py b/r_facelib/detection/retinaface/retinaface.py index 35829e0..5b9b9b0 100644 --- a/r_facelib/detection/retinaface/retinaface.py +++ b/r_facelib/detection/retinaface/retinaface.py @@ -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 diff --git a/modules/__init__.py b/r_modules/__init__.py similarity index 100% rename from modules/__init__.py rename to r_modules/__init__.py diff --git a/modules/images.py b/r_modules/images.py similarity index 100% rename from modules/images.py rename to r_modules/images.py diff --git a/modules/processing.py b/r_modules/processing.py similarity index 100% rename from modules/processing.py rename to r_modules/processing.py diff --git a/modules/scripts.py b/r_modules/scripts.py similarity index 100% rename from modules/scripts.py rename to r_modules/scripts.py diff --git a/modules/scripts_postprocessing.py b/r_modules/scripts_postprocessing.py similarity index 100% rename from modules/scripts_postprocessing.py rename to r_modules/scripts_postprocessing.py diff --git a/modules/shared.py b/r_modules/shared.py similarity index 100% rename from modules/shared.py rename to r_modules/shared.py diff --git a/reactor_core/analyzer.py b/reactor_core/analyzer.py new file mode 100644 index 0000000..399f9b1 --- /dev/null +++ b/reactor_core/analyzer.py @@ -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 diff --git a/reactor_core/face_objects.py b/reactor_core/face_objects.py new file mode 100644 index 0000000..82e25d7 --- /dev/null +++ b/reactor_core/face_objects.py @@ -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.") diff --git a/reactor_core/hyperswap.py b/reactor_core/hyperswap.py new file mode 100644 index 0000000..c39db5f --- /dev/null +++ b/reactor_core/hyperswap.py @@ -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) diff --git a/reactor_core/inswap.py b/reactor_core/inswap.py new file mode 100644 index 0000000..7f22095 --- /dev/null +++ b/reactor_core/inswap.py @@ -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 diff --git a/reactor_core/meanshape_68.py b/reactor_core/meanshape_68.py new file mode 100644 index 0000000..7c83041 --- /dev/null +++ b/reactor_core/meanshape_68.py @@ -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) diff --git a/reactor_patcher.py b/reactor_patcher.py deleted file mode 100644 index efa3c9e..0000000 --- a/reactor_patcher.py +++ /dev/null @@ -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) diff --git a/reactor_utils.py b/reactor_utils.py index 61eb283..cdd3eb8 100644 --- a/reactor_utils.py +++ b/reactor_utils.py @@ -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 diff --git a/requirements.txt b/requirements.txt index a58cc12..d01894a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 diff --git a/scripts/reactor_faceswap.py b/scripts/reactor_faceswap.py index cf27d4c..2d9d06e 100644 --- a/scripts/reactor_faceswap.py +++ b/scripts/reactor_faceswap.py @@ -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, diff --git a/scripts/reactor_logger.py b/scripts/reactor_logger.py index b687648..1748bd2 100644 --- a/scripts/reactor_logger.py +++ b/scripts/reactor_logger.py @@ -2,7 +2,7 @@ import logging import copy import sys -from modules import shared +from r_modules import shared from reactor_utils import addLoggingLevel diff --git a/scripts/reactor_sfw.py b/scripts/reactor_sfw.py index 3011bf2..3027642 100644 --- a/scripts/reactor_sfw.py +++ b/scripts/reactor_sfw.py @@ -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) diff --git a/scripts/reactor_swapper.py b/scripts/reactor_swapper.py index 2eda032..a427713 100644 --- a/scripts/reactor_swapper.py +++ b/scripts/reactor_swapper.py @@ -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)