Upload reactor_swapper.py

custom_nodes/comfyui-reactor-node/scripts/reactor_swapper.py

@@ -199,70 +199,60 @@ def create_gradient_mask(crop_size=256):
     
 #### 1. Используйте `cv2.BORDER_TRANSPARENT` (OpenCV ≥ 4.5)  
 # Этот флаг позволяет **не заполнять** области за пределами маски никаким цветом (пиксели остаются `0` или "прозрачные").    
-def paste_back(target_img, swapped_face, M, crop_size=256):
-    # Улучшенная функция paste_back с идеальной овальной маской и исправлениями артефактов
+def paste_back(target_img, swapped_face, M, bbox, crop_size=256):
+    # Улучшенная функция paste_back с овальной маской на основе bbox (без warp маски)
     
     # target_img: Исходное изображение (BGR, numpy, uint8)
     # swapped_face: Результат работы модели (256x256, BGR, uint8)
-    # M: Матрица аффинного преобразования (Target -> Crop), но здесь используется M_inv из run_hyperswap
+    # M: Матрица аффинного преобразования (M_inv из run_hyperswap)
+    # bbox: Bounding box целевого лица [x1, y1, x2, y2]
     # crop_size: Размер кропа (для HyperSwap это 256)
     
-    # 1. Создание мягкой маски (Эрозия + Размытие)
-    mask = create_gradient_mask(crop_size)
+    # 1. Получаем размеры целевого изображения
+    h, w = target_img.shape[:2]
+    
+    # 2. Создаем овальную маску в пространстве оригинального изображения на основе bbox
+    mask = np.zeros((h, w), dtype=np.float32)
+    center = ((bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2)  # Центр bbox
+    axes = ((bbox[2] - bbox[0]) / 2 * 0.8, (bbox[3] - bbox[1]) / 2 * 0.8)  # Оси эллипса (с запасом 0.8)
+    cv2.ellipse(mask, (int(center[0]), int(center[1])), (int(axes[0]), int(axes[1])), 
+                angle=0, startAngle=0, endAngle=360, color=1.0, thickness=-1)
+    
+    # Применяем размытие для плавных краев
+    mask = cv2.GaussianBlur(mask, (15, 15), 0)
+    mask = np.clip(mask, 0, 1)
     
     # Преобразуем в трехканальную маску
     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 (серый, чтобы избежать синих/зеленых артефактов)
+    # 4. Обратное преобразование для лица
     warped_face = cv2.warpAffine(
         swapped_face_norm,
-        M,  # Это M_inv из run_hyperswap
+        M,
         (w, h),
         flags=cv2.INTER_LANCZOS4 | cv2.WARP_INVERSE_MAP,
         borderMode=cv2.BORDER_CONSTANT,
-        borderValue=0.5  # Серый фон вместо черного/белого
-    )
-    
-    warped_mask = cv2.warpAffine(
-        mask_norm,
-        M,  # Это M_inv из run_hyperswap
-        (w, h),
-        flags=cv2.INTER_CUBIC | cv2.WARP_INVERSE_MAP,
-        borderMode=cv2.BORDER_CONSTANT,
-        borderValue=0.0  # Маска: 0 за пределами
+        borderValue=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)
+    warped_face = np.clip(warped_face, 0, 1)
+    warped_face = np.nan_to_num(warped_face, nan=0.0)
     
-    # 6. Дополнительное размытие для устранения артефактов (опционально, но помогает)
-    warped_mask = cv2.GaussianBlur(warped_mask, (3, 3), 0)
-    
-    # Отладочные логи (добавьте после warp)
-    logger.debug("Warped face shape: %s | Min: %s | Max: %s | NaN count: %s", 
-                 warped_face.shape, warped_face.min(), warped_face.max(), np.isnan(warped_face).sum())
-    logger.debug("Warped mask shape: %s | Min: %s | Max: %s | NaN count: %s", 
-                 warped_mask.shape, warped_mask.min(), warped_mask.max(), np.isnan(warped_mask).sum())
-    
-    # 7. Плавное наложение в float32
+    # 6. Плавн��е наложение
     target_float = target_img.astype(np.float32) / 255.0
-    result_float = target_float * (1.0 - warped_mask) + warped_face * warped_mask
+    result_float = target_float * (1.0 - mask_3c) + warped_face * mask_3c
     
-    # 8. Обратная нормализация к uint8
+    # 7. Обратная нормализация к uint8
     result = (result_float * 255).clip(0, 255).astype(np.uint8)
     
+    logger.debug("Warped face shape: %s | Min: %s | Max: %s | NaN count: %s", 
+                 warped_face.shape, warped_face.min(), warped_face.max(), np.isnan(warped_face).sum())
+    logger.debug("Mask shape: %s | Min: %s | Max: %s", 
+                 mask_3c.shape, mask_3c.min(), mask_3c.max())
     logger.debug("Final result: shape %s | Min: %s | Max: %s", result.shape, result.min(), result.max())
     
     return result
@@ -304,12 +294,10 @@ def run_hyperswap(session, source_face, target_face, target_img):
 
     # Вычисляем аффинную матрицу
     M = get_affine_transform(target_landmarks_5.astype(np.float32), std_landmarks_256)
+    M_inv = cv2.invertAffineTransform(M)  # Обратная для paste
     logger.debug("Affine Matrix M (used for cropping):\n%s", M)
+    logger.debug("Inverse Affine Matrix M_inv (used for paste):\n%s", M_inv)
 
-#### Что проверить:
-# Матрица `M` не должна содержать `NaN` или бесконечности.
-# Если матрица нулевая или искаженная — проблема в точках `target_landmarks_5`.
-    
     # Применяем аффинное преобразование с новой матрицей M
     crop = cv2.warpAffine(target_img, M, (256, 256), flags=cv2.INTER_CUBIC, borderMode=cv2.BORDER_REFLECT)
 
@@ -340,7 +328,6 @@ def run_hyperswap(session, source_face, target_face, target_img):
 
     # --- CPU FLOAT NORMALIZATION FIX ---
     # предотвращает появление "синей кожи" и "шума" при работе на CPU
-    # (адаптировано из патча patch_cpu_fix.diff)
     if isinstance(output, np.ndarray):
         # устранение NaN и бесконечностей
         output = np.nan_to_num(output, nan=0.0, posinf=255.0, neginf=0.0)
@@ -362,7 +349,7 @@ def run_hyperswap(session, source_face, target_face, target_img):
     # (ваш код без изменений, но без старой денормализации)
     output = output.transpose(1, 2, 0)  # CHW -> HWC
     output = output[:, :, ::-1]  # RGB -> BGR (Убедитесь, что это BGR, если вход был BGR)
-    logger.debug("Output after denormalization: Min: %s | Max: %s", output.min(), output.max())
+    logger.debug("Output after denormalization: Min: %s | Max: %s | NaN count: %s", output.min(), output.max(), np.isnan(output).sum())
     
     # Визуализация после денормализации
     #### Что проверить:
@@ -371,8 +358,9 @@ def run_hyperswap(session, source_face, target_face, target_img):
     logger.debug("Output after denormalization: Min: %s | Max: %s", output.min(), output.max())
   # cv2.imshow("Output After Denormalization", output)
   # cv2.waitKey(1)
-       
-    return output, M # Возвращаем лицо (256x256) и матрицу M    
+    
+    # 7. Возвращаем результат 256x256 И обратную матрицу M_inv
+    return output, M_inv    
 
 def sort_by_order(face, order: str):
     if order == "left-right":
@@ -496,9 +484,9 @@ def swap_face(
     interpolation: str = "Bicubic",
 ):
     # >>>>> РЕШЕНИЕ: Принудительная установка уровня DEBUG <<<<<
-    if logger.getEffectiveLevel() != logging.DEBUG:
-        print("\n--- [ReActor Debug] Принудительная установка уровня логирования на DEBUG (10) в swap_face ---")
-        logger.setLevel(logging.DEBUG)
+#   if logger.getEffectiveLevel() != logging.DEBUG:
+#       print("\n--- [ReActor Debug] Принудительная установка уровня логирования на DEBUG (10) в swap_face ---")
+#       logger.setLevel(logging.DEBUG)
     
     # Проверочное сообщение (теперь оно должно появиться)
     logger.debug("--- ТЕСТ: swap_face запущена, уровень логирования DEBUG активен. ---")